spark_model/layers/ops/
moe_grouped_a.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::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
12
13use super::*;
14
15// Explicit `#[path]`: `ops.rs` loads THIS file with one too, so a child
16// module resolves against `ops/` rather than a `moe_grouped_a/` subdirectory.
17#[path = "moe_grouped_a/topk.rs"]
18mod topk;
19// Re-exported, so every existing `ops::moe_topk_*` path still resolves and the
20// split is invisible to callers.
21pub use topk::{moe_topk_sigmoid_batched, moe_topk_softmax_batched, moe_topk_sqrtsoftplus_batched};
22
23/// MoE grouped GEMM: per-expert W4A16 matrix multiply.
24pub fn moe_w4a16_grouped_gemm(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    a: DevicePtr,
28    b_packed: DevicePtr,
29    b_scale: DevicePtr,
30    scale2: f32,
31    c: DevicePtr,
32    expert_offsets: DevicePtr,
33    num_experts: u32,
34    n: u32,
35    k: u32,
36    stream: u64,
37) -> Result<()> {
38    KernelLaunch::new(gpu, kernel)
39        .grid([num_experts, 1, 1])
40        .block([256, 1, 1])
41        .arg_ptr(a)
42        .arg_ptr(b_packed)
43        .arg_ptr(b_scale)
44        .arg_f32(scale2)
45        .arg_ptr(c)
46        .arg_ptr(expert_offsets)
47        .arg_u32(num_experts)
48        .arg_u32(n)
49        .arg_u32(k)
50        .launch(stream)
51}
52
53const PTRTABLE_LEGACY_N_TILE: u32 = 64;
54
55fn ptrtable_legacy_grid_x(n_out: u32) -> u32 {
56    div_ceil(n_out, PTRTABLE_LEGACY_N_TILE)
57}
58
59/// `moe_w4a16_grouped_gemm_ptrtable` with M_TILE=256 (512-thread block, 16
60/// warps). Caller must pass `max_m_tiles` computed against 256, not 64 —
61/// see the `div_ceil(4)` at the call site, mirroring the m128 variant's
62/// `div_ceil(2)`.
63///
64/// DEFAULT-OFF, measured non-win — ~20% slower per call than the base kernel
65/// end-to-end (31.30 vs 26.17 ms avg under nsys). See `launch_grouped_gemm`.
66#[allow(clippy::too_many_arguments)]
67pub fn moe_w4a16_grouped_gemm_ptrtable_m256(
68    gpu: &dyn GpuBackend,
69    kernel: KernelHandle,
70    a: DevicePtr,
71    b_packed_ptrs: DevicePtr,
72    b_scale_ptrs: DevicePtr,
73    scale2_vals: DevicePtr,
74    c: DevicePtr,
75    expert_offsets: DevicePtr,
76    sorted_token_ids: DevicePtr,
77    num_experts: u32,
78    n_out: u32,
79    k: u32,
80    max_m_tiles: u32,
81    stream: u64,
82) -> Result<()> {
83    KernelLaunch::new(gpu, kernel)
84        .grid([ptrtable_legacy_grid_x(n_out), max_m_tiles, num_experts])
85        .block([512, 1, 1])
86        .arg_ptr(a)
87        .arg_ptr(b_packed_ptrs)
88        .arg_ptr(b_scale_ptrs)
89        .arg_ptr(scale2_vals)
90        .arg_ptr(c)
91        .arg_ptr(expert_offsets)
92        .arg_ptr(sorted_token_ids)
93        .arg_u32(num_experts)
94        .arg_u32(n_out)
95        .arg_u32(k)
96        .launch(stream)
97}
98
99/// Pointer-table grouped GEMM: one launch covers all experts.
100///
101/// Grid: (ceil(n_out/64), max_m_tiles, num_experts)  Block: (128, 1, 1)
102#[allow(clippy::too_many_arguments)]
103pub fn moe_w4a16_grouped_gemm_ptrtable(
104    gpu: &dyn GpuBackend,
105    kernel: KernelHandle,
106    a: DevicePtr,
107    b_packed_ptrs: DevicePtr,
108    b_scale_ptrs: DevicePtr,
109    scale2_vals: DevicePtr,
110    c: DevicePtr,
111    expert_offsets: DevicePtr,
112    sorted_token_ids: DevicePtr,
113    num_experts: u32,
114    n_out: u32,
115    k: u32,
116    max_m_tiles: u32,
117    stream: u64,
118) -> Result<()> {
119    KernelLaunch::new(gpu, kernel)
120        .grid([ptrtable_legacy_grid_x(n_out), max_m_tiles, num_experts])
121        .block([128, 1, 1])
122        .arg_ptr(a)
123        .arg_ptr(b_packed_ptrs)
124        .arg_ptr(b_scale_ptrs)
125        .arg_ptr(scale2_vals)
126        .arg_ptr(c)
127        .arg_ptr(expert_offsets)
128        .arg_ptr(sorted_token_ids)
129        .arg_u32(num_experts)
130        .arg_u32(n_out)
131        .arg_u32(k)
132        .launch(stream)
133}
134
135/// Pointer-table grouped GEMM with N_TILE=128 (transposed or wide kernels).
136///
137/// Grid: (ceil(n_out/128), max_m_tiles, num_experts)  Block: (128, 1, 1)
138#[allow(clippy::too_many_arguments)]
139pub fn moe_w4a16_grouped_gemm_ptrtable_n128(
140    gpu: &dyn GpuBackend,
141    kernel: KernelHandle,
142    a: DevicePtr,
143    b_packed_ptrs: DevicePtr,
144    b_scale_ptrs: DevicePtr,
145    scale2_vals: DevicePtr,
146    c: DevicePtr,
147    expert_offsets: DevicePtr,
148    sorted_token_ids: DevicePtr,
149    num_experts: u32,
150    n_out: u32,
151    k: u32,
152    max_m_tiles: u32,
153    stream: u64,
154) -> Result<()> {
155    KernelLaunch::new(gpu, kernel)
156        .grid([div_ceil(n_out, 128), max_m_tiles, num_experts])
157        .block([128, 1, 1])
158        .arg_ptr(a)
159        .arg_ptr(b_packed_ptrs)
160        .arg_ptr(b_scale_ptrs)
161        .arg_ptr(scale2_vals)
162        .arg_ptr(c)
163        .arg_ptr(expert_offsets)
164        .arg_ptr(sorted_token_ids)
165        .arg_u32(num_experts)
166        .arg_u32(n_out)
167        .arg_u32(k)
168        .launch(stream)
169}
170
171/// FP8-A pointer-table grouped GEMM with transposed NVFP4 weights.
172///
173/// A must already be converted to FP8 E4M3. The launch shape mirrors
174/// `moe_w4a16_grouped_gemm_ptrtable_n128`.
175#[allow(clippy::too_many_arguments)]
176pub fn moe_fp8_grouped_gemm_ptrtable_n128(
177    gpu: &dyn GpuBackend,
178    kernel: KernelHandle,
179    a_fp8: DevicePtr,
180    b_packed_ptrs: DevicePtr,
181    b_scale_ptrs: DevicePtr,
182    scale2_vals: DevicePtr,
183    c: DevicePtr,
184    expert_offsets: DevicePtr,
185    sorted_token_ids: DevicePtr,
186    num_experts: u32,
187    n_out: u32,
188    k: u32,
189    max_m_tiles: u32,
190    stream: u64,
191) -> Result<()> {
192    KernelLaunch::new(gpu, kernel)
193        .grid([div_ceil(n_out, 128), max_m_tiles, num_experts])
194        .block([128, 1, 1])
195        .arg_ptr(a_fp8)
196        .arg_ptr(b_packed_ptrs)
197        .arg_ptr(b_scale_ptrs)
198        .arg_ptr(scale2_vals)
199        .arg_ptr(c)
200        .arg_ptr(expert_offsets)
201        .arg_ptr(sorted_token_ids)
202        .arg_u32(num_experts)
203        .arg_u32(n_out)
204        .arg_u32(k)
205        .launch(stream)
206}
207
208/// K64 down GEMM: K_STEP_T=64 eliminates pipeline stall (compute=128 cycles > load ~100 cycles).
209/// Use when K=inter (512 for 35B) — 8 K-steps vs 16 with K32.
210///
211/// Grid: (ceil(n_out/128), max_m_tiles, num_experts)  Block: (128, 1, 1)
212#[allow(clippy::too_many_arguments)]
213pub fn moe_w4a16_grouped_gemm_ptrtable_k64_n128(
214    gpu: &dyn GpuBackend,
215    kernel: KernelHandle,
216    a: DevicePtr,
217    b_packed_ptrs: DevicePtr,
218    b_scale_ptrs: DevicePtr,
219    scale2_vals: DevicePtr,
220    c: DevicePtr,
221    expert_offsets: DevicePtr,
222    sorted_token_ids: DevicePtr,
223    num_experts: u32,
224    n_out: u32,
225    k: u32,
226    max_m_tiles: u32,
227    stream: u64,
228) -> Result<()> {
229    KernelLaunch::new(gpu, kernel)
230        .grid([div_ceil(n_out, 128), max_m_tiles, num_experts])
231        .block([128, 1, 1])
232        .arg_ptr(a)
233        .arg_ptr(b_packed_ptrs)
234        .arg_ptr(b_scale_ptrs)
235        .arg_ptr(scale2_vals)
236        .arg_ptr(c)
237        .arg_ptr(expert_offsets)
238        .arg_ptr(sorted_token_ids)
239        .arg_u32(num_experts)
240        .arg_u32(n_out)
241        .arg_u32(k)
242        .launch(stream)
243}
244
245/// K64 fused gate+up GEMM — zero pipeline stall for K=h (2048 for 35B), 32 K-steps vs 64.
246///
247/// Grid: (ceil(2*n_out/128), max_m_tiles, num_experts)  Block: (128, 1, 1)
248#[allow(clippy::too_many_arguments)]
249pub fn moe_w4a16_fused_gate_up_k64_n128(
250    gpu: &dyn GpuBackend,
251    kernel: KernelHandle,
252    a: DevicePtr,
253    gate_packed_ptrs: DevicePtr,
254    gate_scale_ptrs: DevicePtr,
255    gate_scale2_vals: DevicePtr,
256    up_packed_ptrs: DevicePtr,
257    up_scale_ptrs: DevicePtr,
258    up_scale2_vals: DevicePtr,
259    c_gate: DevicePtr,
260    c_up: DevicePtr,
261    expert_offsets: DevicePtr,
262    sorted_token_ids: DevicePtr,
263    num_experts: u32,
264    n_out: u32,
265    k: u32,
266    max_m_tiles: u32,
267    stream: u64,
268) -> Result<()> {
269    KernelLaunch::new(gpu, kernel)
270        .grid([div_ceil(2 * n_out, 128), max_m_tiles, num_experts])
271        .block([128, 1, 1])
272        .arg_ptr(a)
273        .arg_ptr(gate_packed_ptrs)
274        .arg_ptr(gate_scale_ptrs)
275        .arg_ptr(gate_scale2_vals)
276        .arg_ptr(up_packed_ptrs)
277        .arg_ptr(up_scale_ptrs)
278        .arg_ptr(up_scale2_vals)
279        .arg_ptr(c_gate)
280        .arg_ptr(c_up)
281        .arg_ptr(expert_offsets)
282        .arg_ptr(sorted_token_ids)
283        .arg_u32(num_experts)
284        .arg_u32(n_out)
285        .arg_u32(k)
286        .launch(stream)
287}
288
289/// Gather token rows into expert-sorted order: `permuted[i] = hidden[sorted_token_ids[i]]`.
290/// `permuted` is `[total_expanded, hidden]`. One block per output row, threads
291/// stride over `hidden`. Used by the FP4 grouped gate_up path (the CUTLASS
292/// escape-hatch needs contiguous per-expert rows; the FP8 fused kernel gathers
293/// internally so it doesn't need this).
294///
295/// Retained for the legacy FP4 escape-hatch + potential reuse; the live FP4
296/// path now uses the fused kernel (in-kernel gather), so this is currently
297/// uncalled.
298#[allow(clippy::too_many_arguments, dead_code)]
299pub fn moe_permute_tokens(
300    gpu: &dyn GpuBackend,
301    kernel: KernelHandle,
302    hidden_states: DevicePtr,
303    permuted: DevicePtr,
304    sorted_token_ids: DevicePtr,
305    hidden: u32,
306    total_expanded: u32,
307    stream: u64,
308) -> Result<()> {
309    let threads = hidden.clamp(1, 256);
310    KernelLaunch::new(gpu, kernel)
311        .grid([total_expanded, 1, 1])
312        .block([threads, 1, 1])
313        .arg_ptr(hidden_states)
314        .arg_ptr(permuted)
315        .arg_ptr(sorted_token_ids)
316        .arg_u32(hidden)
317        .arg_u32(total_expanded)
318        .launch(stream)
319}
320
321/// K64 fused gate+up GEMM — M=128 variant (Block D #3 — Avarok pattern).
322///
323/// Doubles M_TILE from 64 → 128. Caller must compute `max_m_tiles_m128`
324/// using divisor 128 (vs 64 for the M=64 variant). Grid covers the same
325/// total work but with half the blocks (and twice the work per block).
326///
327/// Grid: (ceil(2*n_out/128), max_m_tiles_m128, num_experts)  Block: (256, 1, 1)
328#[allow(clippy::too_many_arguments)]
329pub fn moe_w4a16_fused_gate_up_k64_m128(
330    gpu: &dyn GpuBackend,
331    kernel: KernelHandle,
332    a: DevicePtr,
333    gate_packed_ptrs: DevicePtr,
334    gate_scale_ptrs: DevicePtr,
335    gate_scale2_vals: DevicePtr,
336    up_packed_ptrs: DevicePtr,
337    up_scale_ptrs: DevicePtr,
338    up_scale2_vals: DevicePtr,
339    c_gate: DevicePtr,
340    c_up: DevicePtr,
341    expert_offsets: DevicePtr,
342    sorted_token_ids: DevicePtr,
343    num_experts: u32,
344    n_out: u32,
345    k: u32,
346    max_m_tiles_m128: u32,
347    stream: u64,
348) -> Result<()> {
349    KernelLaunch::new(gpu, kernel)
350        .grid([div_ceil(2 * n_out, 128), max_m_tiles_m128, num_experts])
351        .block([256, 1, 1])
352        .arg_ptr(a)
353        .arg_ptr(gate_packed_ptrs)
354        .arg_ptr(gate_scale_ptrs)
355        .arg_ptr(gate_scale2_vals)
356        .arg_ptr(up_packed_ptrs)
357        .arg_ptr(up_scale_ptrs)
358        .arg_ptr(up_scale2_vals)
359        .arg_ptr(c_gate)
360        .arg_ptr(c_up)
361        .arg_ptr(expert_offsets)
362        .arg_ptr(sorted_token_ids)
363        .arg_u32(num_experts)
364        .arg_u32(n_out)
365        .arg_u32(k)
366        .launch(stream)
367}
368
369/// Fused gate+up grouped GEMM — single launch for both projections.
370///
371/// Grid: (ceil(2*n_out/128), max_m_tiles, num_experts)  Block: (128, 1, 1)
372/// First N cols → gate weights/output, last N cols → up weights/output.
373#[allow(clippy::too_many_arguments)]
374pub fn moe_w4a16_fused_gate_up_n128(
375    gpu: &dyn GpuBackend,
376    kernel: KernelHandle,
377    a: DevicePtr,
378    gate_packed_ptrs: DevicePtr,
379    gate_scale_ptrs: DevicePtr,
380    gate_scale2_vals: DevicePtr,
381    up_packed_ptrs: DevicePtr,
382    up_scale_ptrs: DevicePtr,
383    up_scale2_vals: DevicePtr,
384    c_gate: DevicePtr,
385    c_up: DevicePtr,
386    expert_offsets: DevicePtr,
387    sorted_token_ids: DevicePtr,
388    num_experts: u32,
389    n_out: u32,
390    k: u32,
391    max_m_tiles: u32,
392    stream: u64,
393) -> Result<()> {
394    KernelLaunch::new(gpu, kernel)
395        .grid([div_ceil(2 * n_out, 128), max_m_tiles, num_experts])
396        .block([128, 1, 1])
397        .arg_ptr(a)
398        .arg_ptr(gate_packed_ptrs)
399        .arg_ptr(gate_scale_ptrs)
400        .arg_ptr(gate_scale2_vals)
401        .arg_ptr(up_packed_ptrs)
402        .arg_ptr(up_scale_ptrs)
403        .arg_ptr(up_scale2_vals)
404        .arg_ptr(c_gate)
405        .arg_ptr(c_up)
406        .arg_ptr(expert_offsets)
407        .arg_ptr(sorted_token_ids)
408        .arg_u32(num_experts)
409        .arg_u32(n_out)
410        .arg_u32(k)
411        .launch(stream)
412}
413
414/// Element-wise SiLU activation + multiply: `output[i] = silu(gate[i]) * up[i]`.
415///
416/// Grid: (ceil(total_elements/256), 1, 1)  Block: (256, 1, 1)
417pub fn moe_silu_mul(
418    gpu: &dyn GpuBackend,
419    kernel: KernelHandle,
420    gate: DevicePtr,
421    up: DevicePtr,
422    output: DevicePtr,
423    total_elements: u32,
424    stream: u64,
425) -> Result<()> {
426    KernelLaunch::new(gpu, kernel)
427        .grid([div_ceil(total_elements, 256), 1, 1])
428        .block([256, 1, 1])
429        .arg_ptr(gate)
430        .arg_ptr(up)
431        .arg_ptr(output)
432        .arg_u32(total_elements)
433        .launch(stream)
434}
435
436#[cfg(test)]
437mod tests {
438    use super::ptrtable_legacy_grid_x;
439
440    #[test]
441    fn legacy_ptrtable_grid_covers_every_64_column_tile() {
442        assert_eq!(ptrtable_legacy_grid_x(1), 1);
443        assert_eq!(ptrtable_legacy_grid_x(64), 1);
444        assert_eq!(ptrtable_legacy_grid_x(65), 2);
445        assert_eq!(ptrtable_legacy_grid_x(1024), 16);
446        assert_eq!(ptrtable_legacy_grid_x(3072), 48);
447    }
448}
spark_model/layers/ops/moe_grouped_a.rs

spark_model/layers/ops/
moe_grouped_a.rs
