metrale_model_layers/layers/ops/
moe_expert_shared.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the routed+shared expert GEMVs of one token (NVFP4, FP8, BF16) and of two tokens (BF16).
4//!
5//! Owner: model-layers ops.
6//! Invariants: none beyond the types.
7//!
8//! Single-token kernels: `blockIdx.y < top_k` is a routed expert, looked up through the device
9//! pointer tables; `blockIdx.y == top_k` is the shared expert, with direct weight pointers.
10
11use super::*;
12
13/// 2026-09-25: NVFP4 gate and up GEMVs of the routed experts and the shared expert in one launch.
14#[allow(clippy::too_many_arguments)]
15pub fn moe_expert_gate_up_shared(
16    gpu: &dyn GpuBackend,
17    kernel: KernelHandle,
18    input: DevicePtr,
19    gate_packed_ptrs: DevicePtr,
20    gate_scale_ptrs: DevicePtr,
21    gate_scale2_vals: DevicePtr,
22    gate_out: DevicePtr,
23    up_packed_ptrs: DevicePtr,
24    up_scale_ptrs: DevicePtr,
25    up_scale2_vals: DevicePtr,
26    up_out: DevicePtr,
27    expert_indices: DevicePtr,
28    sh_gate: &QuantizedWeight,
29    sh_gate_out: DevicePtr,
30    sh_up: &QuantizedWeight,
31    sh_up_out: DevicePtr,
32    n: u32,
33    k: u32,
34    top_k: u32,
35    stream: u64,
36) -> Result<()> {
37    KernelLaunch::new(gpu, kernel)
38        .grid([div_ceil(n, 8), top_k + 1, 2])
39        .block([128, 1, 1])
40        .arg_ptr(input)
41        .arg_ptr(gate_packed_ptrs)
42        .arg_ptr(gate_scale_ptrs)
43        .arg_ptr(gate_scale2_vals)
44        .arg_ptr(gate_out)
45        .arg_ptr(up_packed_ptrs)
46        .arg_ptr(up_scale_ptrs)
47        .arg_ptr(up_scale2_vals)
48        .arg_ptr(up_out)
49        .arg_ptr(expert_indices)
50        .arg_ptr(sh_gate.weight)
51        .arg_ptr(sh_gate.weight_scale)
52        .arg_f32(sh_gate.weight_scale_2)
53        .arg_ptr(sh_gate_out)
54        .arg_ptr(sh_up.weight)
55        .arg_ptr(sh_up.weight_scale)
56        .arg_f32(sh_up.weight_scale_2)
57        .arg_ptr(sh_up_out)
58        .arg_u32(n)
59        .arg_u32(k)
60        .arg_u32(top_k)
61        .launch(stream)
62}
63
64/// 2026-09-25: NVFP4 SiLU+down GEMVs in one launch. Routed experts read `gate_out`/`up_out`;
65/// the shared expert reads `sh_gate_in`/`sh_up_in`. The kernel's `s_act` is dynamic shared
66/// memory of K floats.
67#[allow(clippy::too_many_arguments)]
68pub fn moe_expert_silu_down_shared(
69    gpu: &dyn GpuBackend,
70    kernel: KernelHandle,
71    gate_out: DevicePtr,
72    up_out: DevicePtr,
73    packed_ptrs: DevicePtr,
74    scale_ptrs: DevicePtr,
75    scale2_vals: DevicePtr,
76    output: DevicePtr,
77    expert_indices: DevicePtr,
78    sh_gate_in: DevicePtr,
79    sh_up_in: DevicePtr,
80    sh_down: &QuantizedWeight,
81    sh_down_out: DevicePtr,
82    n: u32,
83    k: u32,
84    top_k: u32,
85    stream: u64,
86) -> Result<()> {
87    KernelLaunch::new(gpu, kernel)
88        .grid([div_ceil(n, 8), top_k + 1, 1])
89        .block([128, 1, 1])
90        .shared_mem(k * 4)
91        .arg_ptr(gate_out)
92        .arg_ptr(up_out)
93        .arg_ptr(packed_ptrs)
94        .arg_ptr(scale_ptrs)
95        .arg_ptr(scale2_vals)
96        .arg_ptr(output)
97        .arg_ptr(expert_indices)
98        .arg_ptr(sh_gate_in)
99        .arg_ptr(sh_up_in)
100        .arg_ptr(sh_down.weight)
101        .arg_ptr(sh_down.weight_scale)
102        .arg_f32(sh_down.weight_scale_2)
103        .arg_ptr(sh_down_out)
104        .arg_u32(n)
105        .arg_u32(k)
106        .arg_u32(top_k)
107        .launch(stream)
108}
109
110/// 2026-09-25: FP8 variant of [`moe_expert_gate_up_shared`]: two pointer tables per projection
111/// (weights, scales) and the shared expert as direct `Fp8Weight` pointers.
112#[allow(clippy::too_many_arguments)]
113pub fn moe_expert_gate_up_shared_fp8(
114    gpu: &dyn GpuBackend,
115    kernel: KernelHandle,
116    input: DevicePtr,
117    gate_weight_ptrs: DevicePtr,
118    gate_scale_ptrs: DevicePtr,
119    gate_out: DevicePtr,
120    up_weight_ptrs: DevicePtr,
121    up_scale_ptrs: DevicePtr,
122    up_out: DevicePtr,
123    expert_indices: DevicePtr,
124    sh_gate: &Fp8Weight,
125    sh_gate_out: DevicePtr,
126    sh_up: &Fp8Weight,
127    sh_up_out: DevicePtr,
128    n: u32,
129    k: u32,
130    top_k: u32,
131    stream: u64,
132) -> Result<()> {
133    KernelLaunch::new(gpu, kernel)
134        .grid([div_ceil(n, 8), top_k + 1, 2])
135        .block([128, 1, 1])
136        .arg_ptr(input)
137        .arg_ptr(gate_weight_ptrs)
138        .arg_ptr(gate_scale_ptrs)
139        .arg_ptr(gate_out)
140        .arg_ptr(up_weight_ptrs)
141        .arg_ptr(up_scale_ptrs)
142        .arg_ptr(up_out)
143        .arg_ptr(expert_indices)
144        .arg_ptr(sh_gate.weight)
145        .arg_ptr(sh_gate.row_scale)
146        .arg_ptr(sh_gate_out)
147        .arg_ptr(sh_up.weight)
148        .arg_ptr(sh_up.row_scale)
149        .arg_ptr(sh_up_out)
150        .arg_u32(n)
151        .arg_u32(k)
152        .arg_u32(top_k)
153        .launch(stream)
154}
155
156/// 2026-09-25: BF16 variant of [`moe_expert_gate_up_shared_fp8`], with no scale tables, for
157/// experts installed by `MoeLayer::set_bf16_experts`.
158#[allow(clippy::too_many_arguments)]
159pub fn moe_expert_gate_up_shared_bf16(
160    gpu: &dyn GpuBackend,
161    kernel: KernelHandle,
162    input: DevicePtr,
163    gate_weight_ptrs: DevicePtr,
164    gate_out: DevicePtr,
165    up_weight_ptrs: DevicePtr,
166    up_out: DevicePtr,
167    expert_indices: DevicePtr,
168    sh_gate_weight: DevicePtr,
169    sh_gate_out: DevicePtr,
170    sh_up_weight: DevicePtr,
171    sh_up_out: DevicePtr,
172    n: u32,
173    k: u32,
174    top_k: u32,
175    stream: u64,
176) -> Result<()> {
177    KernelLaunch::new(gpu, kernel)
178        .grid([div_ceil(n, 8), top_k + 1, 2])
179        .block([128, 1, 1])
180        .arg_ptr(input)
181        .arg_ptr(gate_weight_ptrs)
182        .arg_ptr(gate_out)
183        .arg_ptr(up_weight_ptrs)
184        .arg_ptr(up_out)
185        .arg_ptr(expert_indices)
186        .arg_ptr(sh_gate_weight)
187        .arg_ptr(sh_gate_out)
188        .arg_ptr(sh_up_weight)
189        .arg_ptr(sh_up_out)
190        .arg_u32(n)
191        .arg_u32(k)
192        .arg_u32(top_k)
193        .launch(stream)
194}
195
196/// 2026-09-25: BF16 variant of [`moe_expert_silu_down_shared_fp8`], with no scale tables.
197#[allow(clippy::too_many_arguments)]
198pub fn moe_expert_silu_down_shared_bf16(
199    gpu: &dyn GpuBackend,
200    kernel: KernelHandle,
201    gate_out: DevicePtr,
202    up_out: DevicePtr,
203    down_weight_ptrs: DevicePtr,
204    output: DevicePtr,
205    expert_indices: DevicePtr,
206    sh_gate_in: DevicePtr,
207    sh_up_in: DevicePtr,
208    sh_down_weight: DevicePtr,
209    sh_down_out: DevicePtr,
210    n: u32,
211    k: u32,
212    top_k: u32,
213    stream: u64,
214) -> Result<()> {
215    KernelLaunch::new(gpu, kernel)
216        .grid([div_ceil(n, 8), top_k + 1, 1])
217        .block([128, 1, 1])
218        .arg_ptr(gate_out)
219        .arg_ptr(up_out)
220        .arg_ptr(down_weight_ptrs)
221        .arg_ptr(output)
222        .arg_ptr(expert_indices)
223        .arg_ptr(sh_gate_in)
224        .arg_ptr(sh_up_in)
225        .arg_ptr(sh_down_weight)
226        .arg_ptr(sh_down_out)
227        .arg_u32(n)
228        .arg_u32(k)
229        .arg_u32(top_k)
230        .launch(stream)
231}
232
233/// 2026-09-25: [`moe_expert_gate_up_shared_bf16`] for two tokens in one launch (called from
234/// `MoeLayer::forward_k2`). `blockIdx.y` in `[0, 2*top_k)` is a routed expert (token
235/// `y / top_k`, slot `y % top_k`, output row `token * top_k + slot`); `blockIdx.y == 2*top_k` is
236/// the shared expert, computed for both tokens into rows 0 and 1.
237#[allow(clippy::too_many_arguments)]
238pub fn moe_expert_gate_up_shared_bf16_batch2(
239    gpu: &dyn GpuBackend,
240    kernel: KernelHandle,
241    input: DevicePtr,
242    gate_weight_ptrs: DevicePtr,
243    gate_out: DevicePtr,
244    up_weight_ptrs: DevicePtr,
245    up_out: DevicePtr,
246    expert_indices: DevicePtr,
247    sh_gate_weight: DevicePtr,
248    sh_gate_out: DevicePtr,
249    sh_up_weight: DevicePtr,
250    sh_up_out: DevicePtr,
251    n: u32,
252    k: u32,
253    top_k: u32,
254    stream: u64,
255) -> Result<()> {
256    KernelLaunch::new(gpu, kernel)
257        .grid([div_ceil(n, 8), 2 * top_k + 1, 2])
258        .block([128, 1, 1])
259        .arg_ptr(input)
260        .arg_ptr(gate_weight_ptrs)
261        .arg_ptr(gate_out)
262        .arg_ptr(up_weight_ptrs)
263        .arg_ptr(up_out)
264        .arg_ptr(expert_indices)
265        .arg_ptr(sh_gate_weight)
266        .arg_ptr(sh_gate_out)
267        .arg_ptr(sh_up_weight)
268        .arg_ptr(sh_up_out)
269        .arg_u32(n)
270        .arg_u32(k)
271        .arg_u32(top_k)
272        .launch(stream)
273}
274
275/// 2026-09-25: [`moe_expert_silu_down_shared_bf16`] for two tokens, with the
276/// [`moe_expert_gate_up_shared_bf16_batch2`] row layout.
277#[allow(clippy::too_many_arguments)]
278pub fn moe_expert_silu_down_shared_bf16_batch2(
279    gpu: &dyn GpuBackend,
280    kernel: KernelHandle,
281    gate_out: DevicePtr,
282    up_out: DevicePtr,
283    down_weight_ptrs: DevicePtr,
284    output: DevicePtr,
285    expert_indices: DevicePtr,
286    sh_gate_in: DevicePtr,
287    sh_up_in: DevicePtr,
288    sh_down_weight: DevicePtr,
289    sh_down_out: DevicePtr,
290    n: u32,
291    k: u32,
292    top_k: u32,
293    stream: u64,
294) -> Result<()> {
295    KernelLaunch::new(gpu, kernel)
296        .grid([div_ceil(n, 8), 2 * top_k + 1, 1])
297        .block([128, 1, 1])
298        .arg_ptr(gate_out)
299        .arg_ptr(up_out)
300        .arg_ptr(down_weight_ptrs)
301        .arg_ptr(output)
302        .arg_ptr(expert_indices)
303        .arg_ptr(sh_gate_in)
304        .arg_ptr(sh_up_in)
305        .arg_ptr(sh_down_weight)
306        .arg_ptr(sh_down_out)
307        .arg_u32(n)
308        .arg_u32(k)
309        .arg_u32(top_k)
310        .launch(stream)
311}
312
313/// 2026-09-25: FP8 variant of [`moe_expert_silu_down_shared`]: two pointer tables (weights,
314/// scales) and the shared down weight as a direct `Fp8Weight` pointer.
315#[allow(clippy::too_many_arguments)]
316pub fn moe_expert_silu_down_shared_fp8(
317    gpu: &dyn GpuBackend,
318    kernel: KernelHandle,
319    gate_out: DevicePtr,
320    up_out: DevicePtr,
321    down_weight_ptrs: DevicePtr,
322    down_scale_ptrs: DevicePtr,
323    output: DevicePtr,
324    expert_indices: DevicePtr,
325    sh_gate_in: DevicePtr,
326    sh_up_in: DevicePtr,
327    sh_down: &Fp8Weight,
328    sh_down_out: DevicePtr,
329    n: u32,
330    k: u32,
331    top_k: u32,
332    stream: u64,
333) -> Result<()> {
334    KernelLaunch::new(gpu, kernel)
335        .grid([div_ceil(n, 8), top_k + 1, 1])
336        .block([128, 1, 1])
337        .arg_ptr(gate_out)
338        .arg_ptr(up_out)
339        .arg_ptr(down_weight_ptrs)
340        .arg_ptr(down_scale_ptrs)
341        .arg_ptr(output)
342        .arg_ptr(expert_indices)
343        .arg_ptr(sh_gate_in)
344        .arg_ptr(sh_up_in)
345        .arg_ptr(sh_down.weight)
346        .arg_ptr(sh_down.row_scale)
347        .arg_ptr(sh_down_out)
348        .arg_u32(n)
349        .arg_u32(k)
350        .arg_u32(top_k)
351        .launch(stream)
352}
metrale_model_layers/layers/ops/moe_expert_shared.rs

metrale_model_layers/layers/ops/
moe_expert_shared.rs
