metrale_model_layers/layers/ops/
moe_expert_more.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the MoE blend (`moe_weighted_sum_blend`) and for the two- and three-token routed+shared expert kernels (`moe_shared_expert_fused_batch2.cu`, `_batch3.cu`).
4//!
5//! Owner: model-layers ops.
6//! Invariants: none beyond the types.
7//!
8//! Row layout of the batch kernels, for `T` tokens: `blockIdx.y` in `[0, T*top_k)` is a routed
9//! expert (token `y / top_k`, slot `y % top_k`), and the last `T` values of `blockIdx.y` are the
10//! shared expert of token `y - T*top_k`. Routed outputs are rows `token * top_k + slot`; shared
11//! outputs are rows `token`.
12
13#![allow(unused_imports)]
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
18
19use crate::layers::moe;
20use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
21
22use super::*;
23
24/// 2026-09-25: The MoE combine for one token:
25/// `output[j] = sum_e weights[e] * expert_out[e, j] + sigmoid(dot(input, gate_weight)) * shared_out[j]`.
26/// Every block computes the gate dot product itself. A null `gate_weight` means an ungated
27/// shared expert, blended at weight 1.
28#[allow(clippy::too_many_arguments)]
29pub fn moe_weighted_sum_blend(
30    gpu: &dyn GpuBackend,
31    kernel: KernelHandle,
32    output: DevicePtr,
33    expert_out: DevicePtr,
34    expert_weights: DevicePtr,
35    shared_out: DevicePtr,
36    input: DevicePtr,
37    gate_weight: DevicePtr,
38    hidden: u32,
39    top_k: u32,
40    k: u32,
41    stream: u64,
42) -> Result<()> {
43    KernelLaunch::new(gpu, kernel)
44        .grid([div_ceil(hidden, 256), 1, 1])
45        .block([256, 1, 1])
46        .arg_ptr(output)
47        .arg_ptr(expert_out)
48        .arg_ptr(expert_weights)
49        .arg_ptr(shared_out)
50        .arg_ptr(input)
51        .arg_ptr(gate_weight)
52        .arg_u32(hidden)
53        .arg_u32(top_k)
54        .arg_u32(k)
55        .launch(stream)
56}
57
58/// 2026-09-25: Gate and up GEMVs of the routed and shared experts for two tokens in one launch
59/// (row layout in the module doc). The shared expert uses direct weight pointers, the routed
60/// experts the pointer tables. A `block_size` of 256 selects the kernel's wide decomposition
61/// (`blockDim.x == 256`); any other width runs the 128-thread one.
62#[allow(clippy::too_many_arguments)]
63pub fn moe_expert_gate_up_shared_batch2(
64    gpu: &dyn GpuBackend,
65    kernel: KernelHandle,
66    input: DevicePtr,
67    gate_packed_ptrs: DevicePtr,
68    gate_scale_ptrs: DevicePtr,
69    gate_scale2_vals: DevicePtr,
70    gate_out: DevicePtr,
71    up_packed_ptrs: DevicePtr,
72    up_scale_ptrs: DevicePtr,
73    up_scale2_vals: DevicePtr,
74    up_out: DevicePtr,
75    expert_indices: DevicePtr,
76    sh_gate: &QuantizedWeight,
77    sh_gate_out: DevicePtr,
78    sh_up: &QuantizedWeight,
79    sh_up_out: DevicePtr,
80    n: u32,
81    k: u32,
82    top_k: u32,
83    block_size: u32,
84    stream: u64,
85) -> Result<()> {
86    KernelLaunch::new(gpu, kernel)
87        .grid([div_ceil(n, 8), 2 * (top_k + 1), 2])
88        .block([block_size, 1, 1])
89        .arg_ptr(input)
90        .arg_ptr(gate_packed_ptrs)
91        .arg_ptr(gate_scale_ptrs)
92        .arg_ptr(gate_scale2_vals)
93        .arg_ptr(gate_out)
94        .arg_ptr(up_packed_ptrs)
95        .arg_ptr(up_scale_ptrs)
96        .arg_ptr(up_scale2_vals)
97        .arg_ptr(up_out)
98        .arg_ptr(expert_indices)
99        .arg_ptr(sh_gate.weight)
100        .arg_ptr(sh_gate.weight_scale)
101        .arg_f32(sh_gate.weight_scale_2)
102        .arg_ptr(sh_gate_out)
103        .arg_ptr(sh_up.weight)
104        .arg_ptr(sh_up.weight_scale)
105        .arg_f32(sh_up.weight_scale_2)
106        .arg_ptr(sh_up_out)
107        .arg_u32(n)
108        .arg_u32(k)
109        .arg_u32(top_k)
110        .launch(stream)
111}
112
113/// 2026-09-25: SiLU+down GEMVs of the routed and shared experts for two tokens in one launch.
114/// `block_size` as in [`moe_expert_gate_up_shared_batch2`].
115#[allow(clippy::too_many_arguments)]
116pub fn moe_expert_silu_down_shared_batch2(
117    gpu: &dyn GpuBackend,
118    kernel: KernelHandle,
119    gate_out: DevicePtr,
120    up_out: DevicePtr,
121    packed_ptrs: DevicePtr,
122    scale_ptrs: DevicePtr,
123    scale2_vals: DevicePtr,
124    output: DevicePtr,
125    expert_indices: DevicePtr,
126    sh_gate_in: DevicePtr,
127    sh_up_in: DevicePtr,
128    sh_down: &QuantizedWeight,
129    sh_down_out: DevicePtr,
130    n: u32,
131    k: u32,
132    top_k: u32,
133    block_size: u32,
134    stream: u64,
135) -> Result<()> {
136    // 2026-09-25: The kernel's `s_act` is dynamic shared memory of K floats.
137    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
138    KernelLaunch::new(gpu, kernel)
139        .grid([div_ceil(n, 8), 2 * (top_k + 1), 1])
140        .block([block_size, 1, 1])
141        .shared_mem(smem_bytes)
142        .arg_ptr(gate_out)
143        .arg_ptr(up_out)
144        .arg_ptr(packed_ptrs)
145        .arg_ptr(scale_ptrs)
146        .arg_ptr(scale2_vals)
147        .arg_ptr(output)
148        .arg_ptr(expert_indices)
149        .arg_ptr(sh_gate_in)
150        .arg_ptr(sh_up_in)
151        .arg_ptr(sh_down.weight)
152        .arg_ptr(sh_down.weight_scale)
153        .arg_f32(sh_down.weight_scale_2)
154        .arg_ptr(sh_down_out)
155        .arg_u32(n)
156        .arg_u32(k)
157        .arg_u32(top_k)
158        .launch(stream)
159}
160
161/// 2026-09-25: [`moe_weighted_sum_blend`] for two tokens; `blockIdx.y` is the token.
162#[allow(clippy::too_many_arguments)]
163pub fn moe_weighted_sum_blend_batch2(
164    gpu: &dyn GpuBackend,
165    kernel: KernelHandle,
166    output: DevicePtr,
167    expert_out: DevicePtr,
168    expert_weights: DevicePtr,
169    shared_out: DevicePtr,
170    input: DevicePtr,
171    gate_weight: DevicePtr,
172    hidden: u32,
173    top_k: u32,
174    k: u32,
175    stream: u64,
176) -> Result<()> {
177    KernelLaunch::new(gpu, kernel)
178        .grid([div_ceil(hidden, 256), 2, 1])
179        .block([256, 1, 1])
180        .arg_ptr(output)
181        .arg_ptr(expert_out)
182        .arg_ptr(expert_weights)
183        .arg_ptr(shared_out)
184        .arg_ptr(input)
185        .arg_ptr(gate_weight)
186        .arg_u32(hidden)
187        .arg_u32(top_k)
188        .arg_u32(k)
189        .launch(stream)
190}
191
192/// 2026-09-25: [`moe_expert_gate_up_shared_batch2`] for three tokens.
193#[allow(clippy::too_many_arguments)]
194pub fn moe_expert_gate_up_shared_batch3(
195    gpu: &dyn GpuBackend,
196    kernel: KernelHandle,
197    input: DevicePtr,
198    gate_packed_ptrs: DevicePtr,
199    gate_scale_ptrs: DevicePtr,
200    gate_scale2_vals: DevicePtr,
201    gate_out: DevicePtr,
202    up_packed_ptrs: DevicePtr,
203    up_scale_ptrs: DevicePtr,
204    up_scale2_vals: DevicePtr,
205    up_out: DevicePtr,
206    expert_indices: DevicePtr,
207    sh_gate: &QuantizedWeight,
208    sh_gate_out: DevicePtr,
209    sh_up: &QuantizedWeight,
210    sh_up_out: DevicePtr,
211    n: u32,
212    k: u32,
213    top_k: u32,
214    stream: u64,
215) -> Result<()> {
216    KernelLaunch::new(gpu, kernel)
217        .grid([div_ceil(n, 8), 3 * (top_k + 1), 2])
218        .block([128, 1, 1])
219        .arg_ptr(input)
220        .arg_ptr(gate_packed_ptrs)
221        .arg_ptr(gate_scale_ptrs)
222        .arg_ptr(gate_scale2_vals)
223        .arg_ptr(gate_out)
224        .arg_ptr(up_packed_ptrs)
225        .arg_ptr(up_scale_ptrs)
226        .arg_ptr(up_scale2_vals)
227        .arg_ptr(up_out)
228        .arg_ptr(expert_indices)
229        .arg_ptr(sh_gate.weight)
230        .arg_ptr(sh_gate.weight_scale)
231        .arg_f32(sh_gate.weight_scale_2)
232        .arg_ptr(sh_gate_out)
233        .arg_ptr(sh_up.weight)
234        .arg_ptr(sh_up.weight_scale)
235        .arg_f32(sh_up.weight_scale_2)
236        .arg_ptr(sh_up_out)
237        .arg_u32(n)
238        .arg_u32(k)
239        .arg_u32(top_k)
240        .launch(stream)
241}
242
243/// 2026-09-25: [`moe_expert_silu_down_shared_batch2`] for three tokens.
244#[allow(clippy::too_many_arguments)]
245pub fn moe_expert_silu_down_shared_batch3(
246    gpu: &dyn GpuBackend,
247    kernel: KernelHandle,
248    gate_out: DevicePtr,
249    up_out: DevicePtr,
250    packed_ptrs: DevicePtr,
251    scale_ptrs: DevicePtr,
252    scale2_vals: DevicePtr,
253    output: DevicePtr,
254    expert_indices: DevicePtr,
255    sh_gate_in: DevicePtr,
256    sh_up_in: DevicePtr,
257    sh_down: &QuantizedWeight,
258    sh_down_out: DevicePtr,
259    n: u32,
260    k: u32,
261    top_k: u32,
262    stream: u64,
263) -> Result<()> {
264    // 2026-09-25: The kernel's `s_act` is dynamic shared memory of K floats.
265    let smem_bytes = (k as usize * std::mem::size_of::<f32>()) as u32;
266    KernelLaunch::new(gpu, kernel)
267        .grid([div_ceil(n, 8), 3 * (top_k + 1), 1])
268        .block([128, 1, 1])
269        .shared_mem(smem_bytes)
270        .arg_ptr(gate_out)
271        .arg_ptr(up_out)
272        .arg_ptr(packed_ptrs)
273        .arg_ptr(scale_ptrs)
274        .arg_ptr(scale2_vals)
275        .arg_ptr(output)
276        .arg_ptr(expert_indices)
277        .arg_ptr(sh_gate_in)
278        .arg_ptr(sh_up_in)
279        .arg_ptr(sh_down.weight)
280        .arg_ptr(sh_down.weight_scale)
281        .arg_f32(sh_down.weight_scale_2)
282        .arg_ptr(sh_down_out)
283        .arg_u32(n)
284        .arg_u32(k)
285        .arg_u32(top_k)
286        .launch(stream)
287}
288
289/// 2026-09-25: [`moe_weighted_sum_blend`] for three tokens.
290#[allow(clippy::too_many_arguments)]
291pub fn moe_weighted_sum_blend_batch3(
292    gpu: &dyn GpuBackend,
293    kernel: KernelHandle,
294    output: DevicePtr,
295    expert_out: DevicePtr,
296    expert_weights: DevicePtr,
297    shared_out: DevicePtr,
298    input: DevicePtr,
299    gate_weight: DevicePtr,
300    hidden: u32,
301    top_k: u32,
302    k: u32,
303    stream: u64,
304) -> Result<()> {
305    KernelLaunch::new(gpu, kernel)
306        .grid([div_ceil(hidden, 256), 3, 1])
307        .block([256, 1, 1])
308        .arg_ptr(output)
309        .arg_ptr(expert_out)
310        .arg_ptr(expert_weights)
311        .arg_ptr(shared_out)
312        .arg_ptr(input)
313        .arg_ptr(gate_weight)
314        .arg_u32(hidden)
315        .arg_u32(top_k)
316        .arg_u32(k)
317        .launch(stream)
318}
metrale_model_layers/layers/ops/moe_expert_more.rs

metrale_model_layers/layers/ops/
moe_expert_more.rs
