spark_model/layers/ops/
moe_expert.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `ops.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11use crate::layers::moe;
12use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
13
14use super::*;
15
16/// Batched MoE expert W4A16 GEMV: runs top_k expert GEMVs in one launch.
17///
18/// Uses device-side pointer tables for weight indirection.
19/// expert_indices come from GPU top-K (device memory).
20///
21/// input_stride: 0 = shared input (gate/up), K = per-expert input (down).
22///
23/// Kernel: `moe_expert_gemv(A, packed_ptrs, scale_ptrs, scale2_vals,
24///          C, expert_indices, N, K, top_k, input_stride)`
25/// Grid: (ceil(N/4), top_k, 1)  Block: (128, 1, 1)
26pub fn moe_expert_gemv(
27    gpu: &dyn GpuBackend,
28    kernel: KernelHandle,
29    input: DevicePtr,
30    packed_ptrs: DevicePtr,
31    scale_ptrs: DevicePtr,
32    scale2_vals: DevicePtr,
33    output: DevicePtr,
34    expert_indices: DevicePtr,
35    n: u32,
36    k: u32,
37    top_k: u32,
38    input_stride: u32,
39    stream: u64,
40) -> Result<()> {
41    KernelLaunch::new(gpu, kernel)
42        .grid([div_ceil(n, 4), top_k, 1])
43        .block([128, 1, 1])
44        .arg_ptr(input)
45        .arg_ptr(packed_ptrs)
46        .arg_ptr(scale_ptrs)
47        .arg_ptr(scale2_vals)
48        .arg_ptr(output)
49        .arg_ptr(expert_indices)
50        .arg_u32(n)
51        .arg_u32(k)
52        .arg_u32(top_k)
53        .arg_u32(input_stride)
54        .launch(stream)
55}
56
57/// Fused gate+up expert GEMV: both projections in one kernel launch.
58///
59/// blockIdx.z selects gate (0) vs up (1). Saves 48 launches per decode step.
60///
61/// Grid: (ceil(N/4), top_k, 2)  Block: (128, 1, 1)
62#[allow(clippy::too_many_arguments)]
63pub fn moe_expert_gemv_gate_up(
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    n: u32,
77    k: u32,
78    top_k: u32,
79    stream: u64,
80) -> Result<()> {
81    KernelLaunch::new(gpu, kernel)
82        .grid([div_ceil(n, 4), top_k, 2])
83        .block([128, 1, 1])
84        .arg_ptr(input)
85        .arg_ptr(gate_packed_ptrs)
86        .arg_ptr(gate_scale_ptrs)
87        .arg_ptr(gate_scale2_vals)
88        .arg_ptr(gate_out)
89        .arg_ptr(up_packed_ptrs)
90        .arg_ptr(up_scale_ptrs)
91        .arg_ptr(up_scale2_vals)
92        .arg_ptr(up_out)
93        .arg_ptr(expert_indices)
94        .arg_u32(n)
95        .arg_u32(k)
96        .arg_u32(top_k)
97        .launch(stream)
98}
99
100/// Register-tiled fused gate+up expert GEMV: 2 output rows per thread.
101///
102/// Same as gate_up but each thread computes 2 adjacent output rows,
103/// reusing the input vector from registers. Doubles weight reads per
104/// iteration for better LPDDR5X bandwidth utilization.
105///
106/// Grid: (ceil(N/8), top_k, 2)  Block: (128, 1, 1)
107#[allow(clippy::too_many_arguments)]
108pub fn moe_expert_gemv_gate_up_2x(
109    gpu: &dyn GpuBackend,
110    kernel: KernelHandle,
111    input: DevicePtr,
112    gate_packed_ptrs: DevicePtr,
113    gate_scale_ptrs: DevicePtr,
114    gate_scale2_vals: DevicePtr,
115    gate_out: DevicePtr,
116    up_packed_ptrs: DevicePtr,
117    up_scale_ptrs: DevicePtr,
118    up_scale2_vals: DevicePtr,
119    up_out: DevicePtr,
120    expert_indices: DevicePtr,
121    n: u32,
122    k: u32,
123    top_k: u32,
124    stream: u64,
125) -> Result<()> {
126    KernelLaunch::new(gpu, kernel)
127        .grid([div_ceil(n, 8), top_k, 2])
128        .block([128, 1, 1])
129        .arg_ptr(input)
130        .arg_ptr(gate_packed_ptrs)
131        .arg_ptr(gate_scale_ptrs)
132        .arg_ptr(gate_scale2_vals)
133        .arg_ptr(gate_out)
134        .arg_ptr(up_packed_ptrs)
135        .arg_ptr(up_scale_ptrs)
136        .arg_ptr(up_scale2_vals)
137        .arg_ptr(up_out)
138        .arg_ptr(expert_indices)
139        .arg_u32(n)
140        .arg_u32(k)
141        .arg_u32(top_k)
142        .launch(stream)
143}
144
145/// Fused SiLU+down expert GEMV: computes silu(gate)*up inline as activation.
146///
147/// Eliminates separate silu_mul kernel. Reads both gate_out and up_out,
148/// computes silu(gate)*up per element, then GEMV with down weights.
149///
150/// Grid: (ceil(N/4), top_k, 1)  Block: (128, 1, 1)
151#[allow(clippy::too_many_arguments)]
152pub fn moe_expert_gemv_silu_down(
153    gpu: &dyn GpuBackend,
154    kernel: KernelHandle,
155    gate_out: DevicePtr,
156    up_out: DevicePtr,
157    packed_ptrs: DevicePtr,
158    scale_ptrs: DevicePtr,
159    scale2_vals: DevicePtr,
160    output: DevicePtr,
161    expert_indices: DevicePtr,
162    n: u32,
163    k: u32,
164    top_k: u32,
165    stream: u64,
166) -> Result<()> {
167    KernelLaunch::new(gpu, kernel)
168        .grid([div_ceil(n, 4), top_k, 1])
169        .block([128, 1, 1])
170        .arg_ptr(gate_out)
171        .arg_ptr(up_out)
172        .arg_ptr(packed_ptrs)
173        .arg_ptr(scale_ptrs)
174        .arg_ptr(scale2_vals)
175        .arg_ptr(output)
176        .arg_ptr(expert_indices)
177        .arg_u32(n)
178        .arg_u32(k)
179        .arg_u32(top_k)
180        .launch(stream)
181}
182
183/// Register-tiled fused SiLU+down expert GEMV: 2 output rows per thread.
184///
185/// Same as silu_down but each thread computes 2 adjacent output rows,
186/// reusing the SiLU(gate)*up activation from registers. Doubles weight
187/// reads per iteration for better LPDDR5X bandwidth utilization.
188///
189/// Grid: (ceil(N/8), top_k, 1)  Block: (128, 1, 1)
190#[allow(clippy::too_many_arguments)]
191pub fn moe_expert_gemv_silu_down_2x(
192    gpu: &dyn GpuBackend,
193    kernel: KernelHandle,
194    gate_out: DevicePtr,
195    up_out: DevicePtr,
196    packed_ptrs: DevicePtr,
197    scale_ptrs: DevicePtr,
198    scale2_vals: DevicePtr,
199    output: DevicePtr,
200    expert_indices: DevicePtr,
201    n: u32,
202    k: u32,
203    top_k: u32,
204    stream: u64,
205) -> Result<()> {
206    KernelLaunch::new(gpu, kernel)
207        .grid([div_ceil(n, 8), top_k, 1])
208        .block([128, 1, 1])
209        .arg_ptr(gate_out)
210        .arg_ptr(up_out)
211        .arg_ptr(packed_ptrs)
212        .arg_ptr(scale_ptrs)
213        .arg_ptr(scale2_vals)
214        .arg_ptr(output)
215        .arg_ptr(expert_indices)
216        .arg_u32(n)
217        .arg_u32(k)
218        .arg_u32(top_k)
219        .launch(stream)
220}
221
222/// Fused gate+up expert GEMV with shared expert as extra blockIdx.y slot.
223///
224/// blockIdx.y < top_k: routed expert (pointer table lookup).
225/// blockIdx.y == top_k: shared expert (direct weight pointers).
226/// Eliminates separate shared expert gate+up kernel launch.
227///
228/// Grid: (ceil(N/8), top_k+1, 2)  Block: (128, 1, 1)
229#[allow(clippy::too_many_arguments)]
230pub fn moe_expert_gate_up_shared(
231    gpu: &dyn GpuBackend,
232    kernel: KernelHandle,
233    input: DevicePtr,
234    gate_packed_ptrs: DevicePtr,
235    gate_scale_ptrs: DevicePtr,
236    gate_scale2_vals: DevicePtr,
237    gate_out: DevicePtr,
238    up_packed_ptrs: DevicePtr,
239    up_scale_ptrs: DevicePtr,
240    up_scale2_vals: DevicePtr,
241    up_out: DevicePtr,
242    expert_indices: DevicePtr,
243    sh_gate: &QuantizedWeight,
244    sh_gate_out: DevicePtr,
245    sh_up: &QuantizedWeight,
246    sh_up_out: DevicePtr,
247    n: u32,
248    k: u32,
249    top_k: u32,
250    stream: u64,
251) -> Result<()> {
252    KernelLaunch::new(gpu, kernel)
253        .grid([div_ceil(n, 8), top_k + 1, 2])
254        .block([128, 1, 1])
255        .arg_ptr(input)
256        .arg_ptr(gate_packed_ptrs)
257        .arg_ptr(gate_scale_ptrs)
258        .arg_ptr(gate_scale2_vals)
259        .arg_ptr(gate_out)
260        .arg_ptr(up_packed_ptrs)
261        .arg_ptr(up_scale_ptrs)
262        .arg_ptr(up_scale2_vals)
263        .arg_ptr(up_out)
264        .arg_ptr(expert_indices)
265        .arg_ptr(sh_gate.weight)
266        .arg_ptr(sh_gate.weight_scale)
267        .arg_f32(sh_gate.weight_scale_2)
268        .arg_ptr(sh_gate_out)
269        .arg_ptr(sh_up.weight)
270        .arg_ptr(sh_up.weight_scale)
271        .arg_f32(sh_up.weight_scale_2)
272        .arg_ptr(sh_up_out)
273        .arg_u32(n)
274        .arg_u32(k)
275        .arg_u32(top_k)
276        .launch(stream)
277}
278
279/// Fused SiLU+down expert GEMV with shared expert as extra blockIdx.y slot.
280///
281/// blockIdx.y < top_k: routed expert (pointer table + expert gate/up buffers).
282/// blockIdx.y == top_k: shared expert (direct pointers + sh_gate_in/up_in).
283/// Eliminates separate shared expert silu+down kernel launch.
284///
285/// Grid: (ceil(N/8), top_k+1, 1)  Block: (128, 1, 1)
286#[allow(clippy::too_many_arguments)]
287pub fn moe_expert_silu_down_shared(
288    gpu: &dyn GpuBackend,
289    kernel: KernelHandle,
290    gate_out: DevicePtr,
291    up_out: DevicePtr,
292    packed_ptrs: DevicePtr,
293    scale_ptrs: DevicePtr,
294    scale2_vals: DevicePtr,
295    output: DevicePtr,
296    expert_indices: DevicePtr,
297    sh_gate_in: DevicePtr,
298    sh_up_in: DevicePtr,
299    sh_down: &QuantizedWeight,
300    sh_down_out: DevicePtr,
301    n: u32,
302    k: u32,
303    top_k: u32,
304    stream: u64,
305) -> Result<()> {
306    KernelLaunch::new(gpu, kernel)
307        .grid([div_ceil(n, 8), top_k + 1, 1])
308        .block([128, 1, 1])
309        .shared_mem(k * 4) // s_act[K] for precomputed SiLU(gate)*up activation
310        .arg_ptr(gate_out)
311        .arg_ptr(up_out)
312        .arg_ptr(packed_ptrs)
313        .arg_ptr(scale_ptrs)
314        .arg_ptr(scale2_vals)
315        .arg_ptr(output)
316        .arg_ptr(expert_indices)
317        .arg_ptr(sh_gate_in)
318        .arg_ptr(sh_up_in)
319        .arg_ptr(sh_down.weight)
320        .arg_ptr(sh_down.weight_scale)
321        .arg_f32(sh_down.weight_scale_2)
322        .arg_ptr(sh_down_out)
323        .arg_u32(n)
324        .arg_u32(k)
325        .arg_u32(top_k)
326        .launch(stream)
327}
328
329/// Fused gate+up expert GEMV with shared expert for FP8 weights.
330///
331/// FP8 variant of `moe_expert_gate_up_shared`: uses 2 pointer tables per
332/// projection (weight_ptrs + scale_ptrs) instead of NVFP4's 3 (packed +
333/// scale + scale2). Shared expert weights are passed as direct Fp8Weight
334/// pointers.
335///
336/// Grid: (ceil(N/8), top_k+1, 2)  Block: (128, 1, 1)
337#[allow(clippy::too_many_arguments)]
338pub fn moe_expert_gate_up_shared_fp8(
339    gpu: &dyn GpuBackend,
340    kernel: KernelHandle,
341    input: DevicePtr,
342    gate_weight_ptrs: DevicePtr,
343    gate_scale_ptrs: DevicePtr,
344    gate_out: DevicePtr,
345    up_weight_ptrs: DevicePtr,
346    up_scale_ptrs: DevicePtr,
347    up_out: DevicePtr,
348    expert_indices: DevicePtr,
349    sh_gate: &Fp8Weight,
350    sh_gate_out: DevicePtr,
351    sh_up: &Fp8Weight,
352    sh_up_out: DevicePtr,
353    n: u32,
354    k: u32,
355    top_k: u32,
356    stream: u64,
357) -> Result<()> {
358    KernelLaunch::new(gpu, kernel)
359        .grid([div_ceil(n, 8), top_k + 1, 2])
360        .block([128, 1, 1])
361        .arg_ptr(input)
362        .arg_ptr(gate_weight_ptrs)
363        .arg_ptr(gate_scale_ptrs)
364        .arg_ptr(gate_out)
365        .arg_ptr(up_weight_ptrs)
366        .arg_ptr(up_scale_ptrs)
367        .arg_ptr(up_out)
368        .arg_ptr(expert_indices)
369        .arg_ptr(sh_gate.weight)
370        .arg_ptr(sh_gate.row_scale)
371        .arg_ptr(sh_gate_out)
372        .arg_ptr(sh_up.weight)
373        .arg_ptr(sh_up.row_scale)
374        .arg_ptr(sh_up_out)
375        .arg_u32(n)
376        .arg_u32(k)
377        .arg_u32(top_k)
378        .launch(stream)
379}
380
381/// Fused gate+up expert GEMV with shared expert for BF16 weights.
382///
383/// BF16 variant of `moe_expert_gate_up_shared_fp8`: no scale tables, direct
384/// BF16 weight pointers. For models loaded via the FP8-dequant-on-load path.
385///
386/// Grid: (ceil(N/8), top_k+1, 2)  Block: (128, 1, 1)
387#[allow(clippy::too_many_arguments)]
388pub fn moe_expert_gate_up_shared_bf16(
389    gpu: &dyn GpuBackend,
390    kernel: KernelHandle,
391    input: DevicePtr,
392    gate_weight_ptrs: DevicePtr,
393    gate_out: DevicePtr,
394    up_weight_ptrs: DevicePtr,
395    up_out: DevicePtr,
396    expert_indices: DevicePtr,
397    sh_gate_weight: DevicePtr,
398    sh_gate_out: DevicePtr,
399    sh_up_weight: DevicePtr,
400    sh_up_out: DevicePtr,
401    n: u32,
402    k: u32,
403    top_k: u32,
404    stream: u64,
405) -> Result<()> {
406    KernelLaunch::new(gpu, kernel)
407        .grid([div_ceil(n, 8), top_k + 1, 2])
408        .block([128, 1, 1])
409        .arg_ptr(input)
410        .arg_ptr(gate_weight_ptrs)
411        .arg_ptr(gate_out)
412        .arg_ptr(up_weight_ptrs)
413        .arg_ptr(up_out)
414        .arg_ptr(expert_indices)
415        .arg_ptr(sh_gate_weight)
416        .arg_ptr(sh_gate_out)
417        .arg_ptr(sh_up_weight)
418        .arg_ptr(sh_up_out)
419        .arg_u32(n)
420        .arg_u32(k)
421        .arg_u32(top_k)
422        .launch(stream)
423}
424
425/// Fused SiLU+down expert GEMV with shared expert for BF16 weights.
426///
427/// BF16 variant of `moe_expert_silu_down_shared_fp8`: no scale tables.
428///
429/// Grid: (ceil(N/8), top_k+1, 1)  Block: (128, 1, 1)
430#[allow(clippy::too_many_arguments)]
431pub fn moe_expert_silu_down_shared_bf16(
432    gpu: &dyn GpuBackend,
433    kernel: KernelHandle,
434    gate_out: DevicePtr,
435    up_out: DevicePtr,
436    down_weight_ptrs: DevicePtr,
437    output: DevicePtr,
438    expert_indices: DevicePtr,
439    sh_gate_in: DevicePtr,
440    sh_up_in: DevicePtr,
441    sh_down_weight: DevicePtr,
442    sh_down_out: DevicePtr,
443    n: u32,
444    k: u32,
445    top_k: u32,
446    stream: u64,
447) -> Result<()> {
448    KernelLaunch::new(gpu, kernel)
449        .grid([div_ceil(n, 8), top_k + 1, 1])
450        .block([128, 1, 1])
451        .arg_ptr(gate_out)
452        .arg_ptr(up_out)
453        .arg_ptr(down_weight_ptrs)
454        .arg_ptr(output)
455        .arg_ptr(expert_indices)
456        .arg_ptr(sh_gate_in)
457        .arg_ptr(sh_up_in)
458        .arg_ptr(sh_down_weight)
459        .arg_ptr(sh_down_out)
460        .arg_u32(n)
461        .arg_u32(k)
462        .arg_u32(top_k)
463        .launch(stream)
464}
465
466/// Fused gate+up expert GEMV with shared expert for BF16 weights — K=2 batch.
467///
468/// BF16 K=2 variant of `moe_expert_gate_up_shared_fp8_batch2`: processes 2
469/// tokens (MTP verify) in one launch. Direct BF16 weight pointers, no scale.
470/// Output layout matches the FP8 batch2 path (routed at flat_slot=token*top_k+
471/// slot, shared at token). For models loaded via the FP8-dequant-on-load path.
472///
473/// Grid: (ceil(N/8), 2*top_k+1, 2)  Block: (128, 1, 1)
474/// y in [0,2*top_k) = routed (per token); y==2*top_k = shared (both tokens).
475#[allow(clippy::too_many_arguments)]
476pub fn moe_expert_gate_up_shared_bf16_batch2(
477    gpu: &dyn GpuBackend,
478    kernel: KernelHandle,
479    input: DevicePtr,
480    gate_weight_ptrs: DevicePtr,
481    gate_out: DevicePtr,
482    up_weight_ptrs: DevicePtr,
483    up_out: DevicePtr,
484    expert_indices: DevicePtr,
485    sh_gate_weight: DevicePtr,
486    sh_gate_out: DevicePtr,
487    sh_up_weight: DevicePtr,
488    sh_up_out: DevicePtr,
489    n: u32,
490    k: u32,
491    top_k: u32,
492    stream: u64,
493) -> Result<()> {
494    KernelLaunch::new(gpu, kernel)
495        .grid([div_ceil(n, 8), 2 * top_k + 1, 2])
496        .block([128, 1, 1])
497        .arg_ptr(input)
498        .arg_ptr(gate_weight_ptrs)
499        .arg_ptr(gate_out)
500        .arg_ptr(up_weight_ptrs)
501        .arg_ptr(up_out)
502        .arg_ptr(expert_indices)
503        .arg_ptr(sh_gate_weight)
504        .arg_ptr(sh_gate_out)
505        .arg_ptr(sh_up_weight)
506        .arg_ptr(sh_up_out)
507        .arg_u32(n)
508        .arg_u32(k)
509        .arg_u32(top_k)
510        .launch(stream)
511}
512
513/// Fused SiLU+down expert GEMV with shared expert for BF16 weights — K=2 batch.
514///
515/// BF16 K=2 variant of `moe_expert_silu_down_shared_fp8_batch2`.
516///
517/// Grid: (ceil(N/8), 2*top_k+1, 1)  Block: (128, 1, 1)
518/// y in [0,2*top_k) = routed (per token); y==2*top_k = shared (both tokens).
519#[allow(clippy::too_many_arguments)]
520pub fn moe_expert_silu_down_shared_bf16_batch2(
521    gpu: &dyn GpuBackend,
522    kernel: KernelHandle,
523    gate_out: DevicePtr,
524    up_out: DevicePtr,
525    down_weight_ptrs: DevicePtr,
526    output: DevicePtr,
527    expert_indices: DevicePtr,
528    sh_gate_in: DevicePtr,
529    sh_up_in: DevicePtr,
530    sh_down_weight: DevicePtr,
531    sh_down_out: DevicePtr,
532    n: u32,
533    k: u32,
534    top_k: u32,
535    stream: u64,
536) -> Result<()> {
537    KernelLaunch::new(gpu, kernel)
538        .grid([div_ceil(n, 8), 2 * top_k + 1, 1])
539        .block([128, 1, 1])
540        .arg_ptr(gate_out)
541        .arg_ptr(up_out)
542        .arg_ptr(down_weight_ptrs)
543        .arg_ptr(output)
544        .arg_ptr(expert_indices)
545        .arg_ptr(sh_gate_in)
546        .arg_ptr(sh_up_in)
547        .arg_ptr(sh_down_weight)
548        .arg_ptr(sh_down_out)
549        .arg_u32(n)
550        .arg_u32(k)
551        .arg_u32(top_k)
552        .launch(stream)
553}
554
555/// Fused SiLU+down expert GEMV with shared expert for FP8 weights.
556///
557/// FP8 variant of `moe_expert_silu_down_shared`: uses 2 pointer tables
558/// (weight_ptrs + scale_ptrs) instead of NVFP4's 3. Shared expert down
559/// weight passed as direct Fp8Weight pointer.
560///
561/// Grid: (ceil(N/8), top_k+1, 1)  Block: (128, 1, 1)
562#[allow(clippy::too_many_arguments)]
563pub fn moe_expert_silu_down_shared_fp8(
564    gpu: &dyn GpuBackend,
565    kernel: KernelHandle,
566    gate_out: DevicePtr,
567    up_out: DevicePtr,
568    down_weight_ptrs: DevicePtr,
569    down_scale_ptrs: DevicePtr,
570    output: DevicePtr,
571    expert_indices: DevicePtr,
572    sh_gate_in: DevicePtr,
573    sh_up_in: DevicePtr,
574    sh_down: &Fp8Weight,
575    sh_down_out: DevicePtr,
576    n: u32,
577    k: u32,
578    top_k: u32,
579    stream: u64,
580) -> Result<()> {
581    KernelLaunch::new(gpu, kernel)
582        .grid([div_ceil(n, 8), top_k + 1, 1])
583        .block([128, 1, 1])
584        .arg_ptr(gate_out)
585        .arg_ptr(up_out)
586        .arg_ptr(down_weight_ptrs)
587        .arg_ptr(down_scale_ptrs)
588        .arg_ptr(output)
589        .arg_ptr(expert_indices)
590        .arg_ptr(sh_gate_in)
591        .arg_ptr(sh_up_in)
592        .arg_ptr(sh_down.weight)
593        .arg_ptr(sh_down.row_scale)
594        .arg_ptr(sh_down_out)
595        .arg_u32(n)
596        .arg_u32(k)
597        .arg_u32(top_k)
598        .launch(stream)
599}
spark_model/layers/ops/moe_expert.rs

spark_model/layers/ops/
moe_expert.rs
