metrale_gpu_runtime/cutlass/
grouped.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Grouped (per-expert) CUTLASS NVFP4 MoE GEMM wrappers.
4//!
5//! Owner: gpu-runtime.
6//! Invariants:
7//! - Before any call into C, every wrapper checks that each per-expert host
8//!   array has `num_experts` entries and that `expert_offsets` has
9//!   `num_experts + 1` non-negative, non-decreasing entries.
10
11use anyhow::{Result, bail};
12
13#[cfg(metrale_cutlass)]
14use std::ffi::c_void;
15
16#[cfg(metrale_cutlass)]
17use super::*;
18
19/// 2026-09-25: Check the host arrays a grouped launch passes to C as bare
20/// pointers. The C side reads `num_experts` entries from each per-expert array
21/// and `num_experts + 1` from `expert_offsets`, and the lengths do not cross the
22/// FFI, so a short slice would be read past its end.
23///
24/// Outside the `cfg(metrale_cutlass)` arms, so it runs, and is tested, without
25/// CUTLASS. `offsets` must also be non-negative and non-decreasing. Its upper
26/// bound (`offsets[num_experts] <= M_total`) is the caller's to check: no
27/// wrapper takes `M_total`.
28fn ensure_group_arrays(
29    who: &str,
30    num_experts: usize,
31    per_expert: &[(&str, usize)],
32    offsets: &[i32],
33) -> Result<()> {
34    if offsets.len() != num_experts + 1 {
35        bail!(
36            "{who}: expert_offsets len {} != num_experts+1 {}",
37            offsets.len(),
38            num_experts + 1
39        );
40    }
41    for (name, len) in per_expert {
42        if *len != num_experts {
43            bail!("{who}: {name} len {len} != num_experts {num_experts}");
44        }
45    }
46    if offsets[0] < 0 {
47        bail!("{who}: expert_offsets[0] = {} is negative", offsets[0]);
48    }
49    for w in offsets.windows(2) {
50        if w[1] < w[0] {
51            bail!(
52                "{who}: expert_offsets is not non-decreasing ({} then {}) — a group would \
53                 have a negative row count",
54                w[0],
55                w[1]
56            );
57        }
58    }
59    Ok(())
60}
61
62/// 2026-09-25: Per-expert NVFP4 gate and up GEMMs: for each expert with rows,
63/// the C side calls `metrale_cutlass_nvfp4_gemm_bf16_act_weight_t` once for gate
64/// and once for up, so the result matches `nvfp4_gemm_bf16_act_weight_t` on each
65/// expert's rows.
66///
67/// `a` is bf16 `[M_total, K]`; expert `e` owns rows
68/// `[expert_offsets[e], expert_offsets[e+1])`. `*_packed_ptrs` and `*_scale_ptrs`
69/// hold one device pointer per expert, in the `pack_bf16_weight_to_nvfp4_t`
70/// layout (`[N,K/2]` and `[K/16,N]`). The `*_scale2_vals` and `expert_offsets`
71/// slices are host arrays.
72#[allow(clippy::too_many_arguments)]
73pub fn nvfp4_grouped_gate_up(
74    a: u64,
75    gate_packed_ptrs: &[u64],
76    gate_scale_ptrs: &[u64],
77    gate_scale2_vals: &[f32],
78    up_packed_ptrs: &[u64],
79    up_scale_ptrs: &[u64],
80    up_scale2_vals: &[f32],
81    c_gate: u64,
82    c_up: u64,
83    expert_offsets: &[i32],
84    n: u32,
85    k: u32,
86    stream: u64,
87) -> Result<()> {
88    let num_experts = gate_packed_ptrs.len();
89    ensure_group_arrays(
90        "nvfp4_grouped_gate_up",
91        num_experts,
92        &[
93            ("gate_scale_ptrs", gate_scale_ptrs.len()),
94            ("gate_scale2_vals", gate_scale2_vals.len()),
95            ("up_packed_ptrs", up_packed_ptrs.len()),
96            ("up_scale_ptrs", up_scale_ptrs.len()),
97            ("up_scale2_vals", up_scale2_vals.len()),
98        ],
99        expert_offsets,
100    )?;
101    #[cfg(metrale_cutlass)]
102    {
103        let ctx = ctx()?;
104        let status = unsafe {
105            metrale_cutlass_nvfp4_grouped_gate_up(
106                a as *const c_void,
107                gate_packed_ptrs.as_ptr(),
108                gate_scale_ptrs.as_ptr(),
109                gate_scale2_vals.as_ptr(),
110                up_packed_ptrs.as_ptr(),
111                up_scale_ptrs.as_ptr(),
112                up_scale2_vals.as_ptr(),
113                c_gate as *mut c_void,
114                c_up as *mut c_void,
115                expert_offsets.as_ptr(),
116                num_experts as i32,
117                n as i32,
118                k as i32,
119                ctx.workspace as *mut c_void,
120                ctx.ws_size,
121                stream as *mut c_void,
122            )
123        };
124        if status != 0 {
125            bail!("CUTLASS nvfp4 grouped gate_up failed: status {status}");
126        }
127        Ok(())
128    }
129    #[cfg(not(metrale_cutlass))]
130    {
131        let _ = (
132            a,
133            gate_packed_ptrs,
134            gate_scale_ptrs,
135            gate_scale2_vals,
136            up_packed_ptrs,
137            up_scale_ptrs,
138            up_scale2_vals,
139            c_gate,
140            c_up,
141            expert_offsets,
142            n,
143            k,
144            stream,
145        );
146        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
147    }
148}
149
150/// 2026-09-25: Grouped NVFP4 gate and up GEMMs, one `GemmUniversalMode::kGrouped`
151/// launch per projection over the active experts. `a` is packed once and
152/// shared by both.
153///
154/// `a` is bf16 and token-major. Group row `i` reads token `sorted_token_ids[i]`,
155/// or token `i` when `sorted_token_ids` is 0 (null); expert `e` owns group rows
156/// `[expert_offsets_host[e], expert_offsets_host[e+1])`, and `c_gate`/`c_up` are
157/// written in that order. `*_packed_ptrs` hold one device pointer per expert to
158/// a CUTLASS `[N,K/2]` weight; `*_sfb_ptrs` to its swizzled SFB scales (see
159/// `pack_weight_sfb`). `*_scale2_vals` and `expert_offsets_host` are host arrays.
160#[allow(clippy::too_many_arguments)]
161pub fn nvfp4_grouped_gate_up_fused(
162    a: u64,
163    sorted_token_ids: u64,
164    gate_packed_ptrs: &[u64],
165    gate_sfb_ptrs: &[u64],
166    gate_scale2_vals: &[f32],
167    up_packed_ptrs: &[u64],
168    up_sfb_ptrs: &[u64],
169    up_scale2_vals: &[f32],
170    c_gate: u64,
171    c_up: u64,
172    expert_offsets_host: &[i32],
173    n: u32,
174    k: u32,
175    stream: u64,
176) -> Result<()> {
177    let num_experts = gate_packed_ptrs.len();
178    ensure_group_arrays(
179        "nvfp4_grouped_gate_up_fused",
180        num_experts,
181        &[
182            ("gate_sfb_ptrs", gate_sfb_ptrs.len()),
183            ("gate_scale2_vals", gate_scale2_vals.len()),
184            ("up_packed_ptrs", up_packed_ptrs.len()),
185            ("up_sfb_ptrs", up_sfb_ptrs.len()),
186            ("up_scale2_vals", up_scale2_vals.len()),
187        ],
188        expert_offsets_host,
189    )?;
190    #[cfg(metrale_cutlass)]
191    {
192        let ctx = ctx()?;
193        let status = unsafe {
194            metrale_cutlass_nvfp4_grouped_gate_up_fused(
195                a as *const c_void,
196                sorted_token_ids as *const i32,
197                gate_packed_ptrs.as_ptr(),
198                gate_sfb_ptrs.as_ptr(),
199                gate_scale2_vals.as_ptr(),
200                up_packed_ptrs.as_ptr(),
201                up_sfb_ptrs.as_ptr(),
202                up_scale2_vals.as_ptr(),
203                c_gate as *mut c_void,
204                c_up as *mut c_void,
205                expert_offsets_host.as_ptr(),
206                num_experts as i32,
207                n as i32,
208                k as i32,
209                ctx.workspace as *mut c_void,
210                ctx.ws_size,
211                stream as *mut c_void,
212            )
213        };
214        if status != 0 {
215            bail!("CUTLASS nvfp4 grouped(fused) gate_up failed: status {status}");
216        }
217        Ok(())
218    }
219    #[cfg(not(metrale_cutlass))]
220    {
221        let _ = (
222            a,
223            sorted_token_ids,
224            gate_packed_ptrs,
225            gate_sfb_ptrs,
226            gate_scale2_vals,
227            up_packed_ptrs,
228            up_sfb_ptrs,
229            up_scale2_vals,
230            c_gate,
231            c_up,
232            expert_offsets_host,
233            n,
234            k,
235            stream,
236        );
237        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
238    }
239}
240
241/// 2026-09-25: Grouped NVFP4 down projection, one `kGrouped` launch. `a` is the
242/// bf16 intermediate `[M_total, K]`, already in expert order (no gather).
243/// `packed_ptrs` and `sfb_ptrs` hold one device pointer per expert to a
244/// `[N,K/2]` weight and its swizzled SFB scales; `scale2_vals` and
245/// `expert_offsets_host` are host arrays. Writes `c` `[M_total, N]`.
246#[allow(clippy::too_many_arguments)]
247pub fn nvfp4_grouped_down(
248    a: u64,
249    packed_ptrs: &[u64],
250    sfb_ptrs: &[u64],
251    scale2_vals: &[f32],
252    c: u64,
253    expert_offsets_host: &[i32],
254    n: u32,
255    k: u32,
256    stream: u64,
257) -> Result<()> {
258    let num_experts = packed_ptrs.len();
259    ensure_group_arrays(
260        "nvfp4_grouped_down",
261        num_experts,
262        &[
263            ("sfb_ptrs", sfb_ptrs.len()),
264            ("scale2_vals", scale2_vals.len()),
265        ],
266        expert_offsets_host,
267    )?;
268    #[cfg(metrale_cutlass)]
269    {
270        let ctx = ctx()?;
271        let status = unsafe {
272            metrale_cutlass_nvfp4_grouped_down(
273                a as *const c_void,
274                packed_ptrs.as_ptr(),
275                sfb_ptrs.as_ptr(),
276                scale2_vals.as_ptr(),
277                c as *mut c_void,
278                expert_offsets_host.as_ptr(),
279                num_experts as i32,
280                n as i32,
281                k as i32,
282                ctx.workspace as *mut c_void,
283                ctx.ws_size,
284                stream as *mut c_void,
285            )
286        };
287        if status != 0 {
288            bail!("CUTLASS nvfp4 grouped down failed: status {status}");
289        }
290        Ok(())
291    }
292    #[cfg(not(metrale_cutlass))]
293    {
294        let _ = (
295            a,
296            packed_ptrs,
297            sfb_ptrs,
298            scale2_vals,
299            c,
300            expert_offsets_host,
301            n,
302            k,
303            stream,
304        );
305        bail!("CUTLASS support was not built; set CUTLASS_HOME when building")
306    }
307}
308
309#[cfg(test)]
310mod tests {
311    use super::ensure_group_arrays;
312
313    /// 2026-09-25: Every per-expert array is checked, not only the first.
314    #[test]
315    fn rejects_any_short_per_expert_array() {
316        let offsets = [0i32, 4, 9];
317        ensure_group_arrays("t", 2, &[("a", 2), ("b", 2)], &offsets).unwrap();
318        // 2026-09-25: Each array one entry short is rejected, and the message
319        // names the array.
320        let e = ensure_group_arrays("t", 2, &[("a", 1), ("b", 2)], &offsets).unwrap_err();
321        assert!(e.to_string().contains("a len 1"), "{e}");
322        let e = ensure_group_arrays("t", 2, &[("a", 2), ("b", 1)], &offsets).unwrap_err();
323        assert!(e.to_string().contains("b len 1"), "{e}");
324        // 2026-09-25: Too long is rejected too: the caller disagrees with
325        // `num_experts`.
326        assert!(ensure_group_arrays("t", 2, &[("a", 3), ("b", 2)], &offsets).is_err());
327    }
328
329    #[test]
330    fn rejects_bad_expert_offsets() {
331        assert!(
332            ensure_group_arrays("t", 2, &[], &[0i32, 4]).is_err(),
333            "short"
334        );
335        assert!(
336            ensure_group_arrays("t", 2, &[], &[0i32, 4, 9, 12]).is_err(),
337            "long"
338        );
339        assert!(
340            ensure_group_arrays("t", 2, &[], &[-1i32, 4, 9]).is_err(),
341            "negative base"
342        );
343        assert!(
344            ensure_group_arrays("t", 2, &[], &[0i32, 9, 4]).is_err(),
345            "decreasing => negative group row count"
346        );
347        // 2026-09-25: Equal consecutive offsets are an expert with no rows, which
348        // is accepted.
349        ensure_group_arrays("t", 2, &[], &[0i32, 4, 4]).unwrap();
350    }
351
352    /// 2026-09-25: Zero experts with a one-entry `offsets` passes the check
353    /// without a panic.
354    #[test]
355    fn zero_experts_is_not_a_panic() {
356        ensure_group_arrays("t", 0, &[], &[0i32]).unwrap();
357    }
358}
metrale_gpu_runtime/cutlass/grouped.rs

metrale_gpu_runtime/cutlass/
grouped.rs
