spark_runtime/cutlass/
grouped.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Grouped (per-expert) CUTLASS NVFP4 MoE GEMM host wrappers.
3
4use anyhow::{Result, bail};
5
6#[cfg(atlas_cutlass)]
7use std::ffi::c_void;
8
9#[cfg(atlas_cutlass)]
10use super::*;
11
12/// Validate every host array a grouped launch hands to C++ as a bare pointer.
13///
14/// The C++ side reads `num_experts` elements from EACH per-expert host array and
15/// `num_experts + 1` from `expert_offsets` — the lengths never cross the FFI
16/// boundary, so a slice that is one entry short becomes a host heap
17/// over-read whose result is then dereferenced as a DEVICE pointer. Each of the
18/// three wrappers below used to check `expert_offsets` alone and pass the other
19/// five or six arrays unchecked; the shape held only because the single caller
20/// happened to build them all from one `num_experts`. This function is where
21/// that rule lives now, so a wrapper cannot check some of its arrays.
22///
23/// Deliberately OUTSIDE the `#[cfg(atlas_cutlass)]` arms: the guard has to run
24/// (and be testable) on a build without CUTLASS, which is what CI builds.
25///
26/// `offsets` is additionally checked for the one property that is complete
27/// without more parameters — non-negative and non-decreasing, so no group gets a
28/// negative row count. The UPPER bound (`offsets[num_experts] <= M_total`)
29/// cannot be checked here: `M_total` is not a parameter of any of these
30/// wrappers, only the device pointer `a` is. That check belongs to the caller
31/// that owns the activation buffer.
32fn ensure_group_arrays(
33    who: &str,
34    num_experts: usize,
35    per_expert: &[(&str, usize)],
36    offsets: &[i32],
37) -> Result<()> {
38    if offsets.len() != num_experts + 1 {
39        bail!(
40            "{who}: expert_offsets len {} != num_experts+1 {}",
41            offsets.len(),
42            num_experts + 1
43        );
44    }
45    for (name, len) in per_expert {
46        if *len != num_experts {
47            bail!("{who}: {name} len {len} != num_experts {num_experts}");
48        }
49    }
50    if offsets[0] < 0 {
51        bail!("{who}: expert_offsets[0] = {} is negative", offsets[0]);
52    }
53    for w in offsets.windows(2) {
54        if w[1] < w[0] {
55            bail!(
56                "{who}: expert_offsets is not non-decreasing ({} then {}) — a group would \
57                 have a negative row count",
58                w[0],
59                w[1]
60            );
61        }
62    }
63    Ok(())
64}
65
66/// Grouped (per-expert) NVFP4 fused gate_up GEMM — Holo MoE Phase-1
67/// escape-hatch path. Dispatches the proven Sm120 NVFP4 collective once per
68/// active expert over its token slice; bit-faithful to
69/// `nvfp4_gemm_bf16_act_weight_t` (it IS that collective), at one launch per
70/// expert. Used to validate that the FP4 math integrates correctly in grouped
71/// form before the hand-rolled block-scaled mma (Phase 2).
72///
73/// `a` is bf16 `[M_total, K]`; expert `e` owns rows
74/// `[expert_offsets[e], expert_offsets[e+1])`. `*_packed_ptrs`/`*_scale_ptrs`
75/// are device-pointer arrays (one per expert) in the
76/// `pack_bf16_weight_to_nvfp4_t` layout (`[N,K/2]` + `[K/16,N]`); the
77/// `*_scale2_vals` and `expert_offsets` slices are HOST arrays.
78#[allow(clippy::too_many_arguments)]
79pub fn nvfp4_grouped_gate_up(
80    a: u64,
81    gate_packed_ptrs: &[u64],
82    gate_scale_ptrs: &[u64],
83    gate_scale2_vals: &[f32],
84    up_packed_ptrs: &[u64],
85    up_scale_ptrs: &[u64],
86    up_scale2_vals: &[f32],
87    c_gate: u64,
88    c_up: u64,
89    expert_offsets: &[i32],
90    n: u32,
91    k: u32,
92    stream: u64,
93) -> Result<()> {
94    let num_experts = gate_packed_ptrs.len();
95    ensure_group_arrays(
96        "nvfp4_grouped_gate_up",
97        num_experts,
98        &[
99            ("gate_scale_ptrs", gate_scale_ptrs.len()),
100            ("gate_scale2_vals", gate_scale2_vals.len()),
101            ("up_packed_ptrs", up_packed_ptrs.len()),
102            ("up_scale_ptrs", up_scale_ptrs.len()),
103            ("up_scale2_vals", up_scale2_vals.len()),
104        ],
105        expert_offsets,
106    )?;
107    #[cfg(atlas_cutlass)]
108    {
109        let ctx = ctx()?;
110        let status = unsafe {
111            atlas_cutlass_nvfp4_grouped_gate_up(
112                a as *const c_void,
113                gate_packed_ptrs.as_ptr(),
114                gate_scale_ptrs.as_ptr(),
115                gate_scale2_vals.as_ptr(),
116                up_packed_ptrs.as_ptr(),
117                up_scale_ptrs.as_ptr(),
118                up_scale2_vals.as_ptr(),
119                c_gate as *mut c_void,
120                c_up as *mut c_void,
121                expert_offsets.as_ptr(),
122                num_experts as i32,
123                n as i32,
124                k as i32,
125                ctx.workspace as *mut c_void,
126                ctx.ws_size,
127                stream as *mut c_void,
128            )
129        };
130        if status != 0 {
131            bail!("CUTLASS nvfp4 grouped gate_up failed: status {status}");
132        }
133        Ok(())
134    }
135    #[cfg(not(atlas_cutlass))]
136    {
137        let _ = (
138            a,
139            gate_packed_ptrs,
140            gate_scale_ptrs,
141            gate_scale2_vals,
142            up_packed_ptrs,
143            up_scale_ptrs,
144            up_scale2_vals,
145            c_gate,
146            c_up,
147            expert_offsets,
148            n,
149            k,
150            stream,
151        );
152        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
153    }
154}
155
156/// Single-launch grouped (`GemmUniversalMode::kGrouped`) NVFP4 fused gate_up
157/// GEMM — the Phase-2 successor to [`nvfp4_grouped_gate_up`]. Replaces the
158/// per-expert collective loop with ONE grouped launch over all active experts,
159/// eliminating the N-launch overhead.
160///
161/// `a` is bf16 `[M_total, K]`, expert-contiguous (caller permuted so expert `e`
162/// owns rows `[expert_offsets_host[e], expert_offsets_host[e+1])`).
163/// `*_packed_ptrs` are device-pointer arrays (one per expert) into the CUTLASS
164/// `[N,K/2]` packed weight tables; `*_sfb_ptrs` are device-pointer arrays into
165/// the swizzled SFB (ue4m3) scale tables (see `pack_weight_sfb`).
166/// `*_scale2_vals` and `expert_offsets_host` are HOST arrays.
167#[allow(clippy::too_many_arguments)]
168pub fn nvfp4_grouped_gate_up_fused(
169    a: u64,
170    sorted_token_ids: u64,
171    gate_packed_ptrs: &[u64],
172    gate_sfb_ptrs: &[u64],
173    gate_scale2_vals: &[f32],
174    up_packed_ptrs: &[u64],
175    up_sfb_ptrs: &[u64],
176    up_scale2_vals: &[f32],
177    c_gate: u64,
178    c_up: u64,
179    expert_offsets_host: &[i32],
180    n: u32,
181    k: u32,
182    stream: u64,
183) -> Result<()> {
184    let num_experts = gate_packed_ptrs.len();
185    ensure_group_arrays(
186        "nvfp4_grouped_gate_up_fused",
187        num_experts,
188        &[
189            ("gate_sfb_ptrs", gate_sfb_ptrs.len()),
190            ("gate_scale2_vals", gate_scale2_vals.len()),
191            ("up_packed_ptrs", up_packed_ptrs.len()),
192            ("up_sfb_ptrs", up_sfb_ptrs.len()),
193            ("up_scale2_vals", up_scale2_vals.len()),
194        ],
195        expert_offsets_host,
196    )?;
197    #[cfg(atlas_cutlass)]
198    {
199        let ctx = ctx()?;
200        let status = unsafe {
201            atlas_cutlass_nvfp4_grouped_gate_up_fused(
202                a as *const c_void,
203                sorted_token_ids as *const i32,
204                gate_packed_ptrs.as_ptr(),
205                gate_sfb_ptrs.as_ptr(),
206                gate_scale2_vals.as_ptr(),
207                up_packed_ptrs.as_ptr(),
208                up_sfb_ptrs.as_ptr(),
209                up_scale2_vals.as_ptr(),
210                c_gate as *mut c_void,
211                c_up as *mut c_void,
212                expert_offsets_host.as_ptr(),
213                num_experts as i32,
214                n as i32,
215                k as i32,
216                ctx.workspace as *mut c_void,
217                ctx.ws_size,
218                stream as *mut c_void,
219            )
220        };
221        if status != 0 {
222            bail!("CUTLASS nvfp4 grouped(fused) gate_up failed: status {status}");
223        }
224        Ok(())
225    }
226    #[cfg(not(atlas_cutlass))]
227    {
228        let _ = (
229            a,
230            sorted_token_ids,
231            gate_packed_ptrs,
232            gate_sfb_ptrs,
233            gate_scale2_vals,
234            up_packed_ptrs,
235            up_sfb_ptrs,
236            up_scale2_vals,
237            c_gate,
238            c_up,
239            expert_offsets_host,
240            n,
241            k,
242            stream,
243        );
244        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
245    }
246}
247
248/// Single-launch grouped NVFP4 DOWN projection (`atlas_cutlass_nvfp4_grouped_down`).
249/// `a` is the post-SiLU bf16 intermediate `[M_total, K=inter]`, ALREADY
250/// expert-contiguous (no gather). `packed_ptrs`/`sfb_ptrs` are device-pointer
251/// arrays into the `[N=hidden,K/2]` packed + swizzled-SFB down tables; `scale2_vals`
252/// and `expert_offsets_host` are HOST arrays. Writes `c` `[M_total, N=hidden]`.
253#[allow(clippy::too_many_arguments)]
254pub fn nvfp4_grouped_down(
255    a: u64,
256    packed_ptrs: &[u64],
257    sfb_ptrs: &[u64],
258    scale2_vals: &[f32],
259    c: u64,
260    expert_offsets_host: &[i32],
261    n: u32,
262    k: u32,
263    stream: u64,
264) -> Result<()> {
265    let num_experts = packed_ptrs.len();
266    ensure_group_arrays(
267        "nvfp4_grouped_down",
268        num_experts,
269        &[
270            ("sfb_ptrs", sfb_ptrs.len()),
271            ("scale2_vals", scale2_vals.len()),
272        ],
273        expert_offsets_host,
274    )?;
275    #[cfg(atlas_cutlass)]
276    {
277        let ctx = ctx()?;
278        let status = unsafe {
279            atlas_cutlass_nvfp4_grouped_down(
280                a as *const c_void,
281                packed_ptrs.as_ptr(),
282                sfb_ptrs.as_ptr(),
283                scale2_vals.as_ptr(),
284                c as *mut c_void,
285                expert_offsets_host.as_ptr(),
286                num_experts as i32,
287                n as i32,
288                k as i32,
289                ctx.workspace as *mut c_void,
290                ctx.ws_size,
291                stream as *mut c_void,
292            )
293        };
294        if status != 0 {
295            bail!("CUTLASS nvfp4 grouped down failed: status {status}");
296        }
297        Ok(())
298    }
299    #[cfg(not(atlas_cutlass))]
300    {
301        let _ = (
302            a,
303            packed_ptrs,
304            sfb_ptrs,
305            scale2_vals,
306            c,
307            expert_offsets_host,
308            n,
309            k,
310            stream,
311        );
312        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use super::ensure_group_arrays;
319
320    /// The rule these wrappers exist to enforce: EVERY per-expert host array is
321    /// read `num_experts` deep by C++, not just the one that used to be checked.
322    #[test]
323    fn rejects_any_short_per_expert_array() {
324        let offsets = [0i32, 4, 9];
325        // All lengths correct → accepted.
326        ensure_group_arrays("t", 2, &[("a", 2), ("b", 2)], &offsets).unwrap();
327        // Each array in turn, one entry short → rejected, and the message names
328        // the array so the failure is actionable.
329        let e = ensure_group_arrays("t", 2, &[("a", 1), ("b", 2)], &offsets).unwrap_err();
330        assert!(e.to_string().contains("a len 1"), "{e}");
331        let e = ensure_group_arrays("t", 2, &[("a", 2), ("b", 1)], &offsets).unwrap_err();
332        assert!(e.to_string().contains("b len 1"), "{e}");
333        // Over-long is rejected too: it means the caller disagrees with
334        // `num_experts`, and C++ would silently use only a prefix.
335        assert!(ensure_group_arrays("t", 2, &[("a", 3), ("b", 2)], &offsets).is_err());
336    }
337
338    #[test]
339    fn rejects_bad_expert_offsets() {
340        assert!(
341            ensure_group_arrays("t", 2, &[], &[0i32, 4]).is_err(),
342            "short"
343        );
344        assert!(
345            ensure_group_arrays("t", 2, &[], &[0i32, 4, 9, 12]).is_err(),
346            "long"
347        );
348        assert!(
349            ensure_group_arrays("t", 2, &[], &[-1i32, 4, 9]).is_err(),
350            "negative base"
351        );
352        assert!(
353            ensure_group_arrays("t", 2, &[], &[0i32, 9, 4]).is_err(),
354            "decreasing => negative group row count"
355        );
356        // Empty groups (equal consecutive offsets) are legitimate: an expert
357        // that no token routed to.
358        ensure_group_arrays("t", 2, &[], &[0i32, 4, 4]).unwrap();
359    }
360
361    /// A zero-expert launch is a valid no-op shape, and the check must not
362    /// index `offsets[0]` out of bounds when it happens.
363    #[test]
364    fn zero_experts_is_not_a_panic() {
365        ensure_group_arrays("t", 0, &[], &[0i32]).unwrap();
366    }
367}
spark_runtime/cutlass/grouped.rs

spark_runtime/cutlass/
grouped.rs
