spark_model/layers/ops/
moe_grouped_a2.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! MoE token-routing reduce + CUTLASS grouped ops — extracted from
4//! `moe_grouped_a.rs` during the ≤500-line split. All public items remain
5//! available at `crate::layers::ops::*` via the re-export in `ops.rs`.
6
7#![allow(unused_imports)]
8
9use anyhow::Result;
10use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
11use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
12
13use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
14
15use super::*;
16
17// Counting sort tokens by expert assignment.
18//
19// Produces sorted_token_ids (grouped by expert), expert_offsets (prefix sum),
20// and token_to_perm (reverse map for unpermute).
21//
22// Grid: (1, 1, 1)  Block: (256, 1, 1)
23
24/// Host snapshots of the per-expert pointer/scale tables for the CUTLASS
25/// grouped path, owned by the `MoeLayer` whose device tables they mirror.
26///
27/// The CUTLASS grouped entry needs these on the HOST to build its per-group
28/// problem shapes. They are immutable once `build_cutlass_grouped_sfb` has
29/// run, so re-reading them per call was 6 copies x 2 calls x 47 MoE layers =
30/// 564 pointless D2H transfers per prefill. Only `expert_offsets` genuinely
31/// changes per call (it is produced by the expert sort), so only that one is
32/// still copied at dispatch time.
33///
34/// Why a layer-owned snapshot and not a process-global cache keyed on the
35/// device pointer: an address is not an identity. An in-process model swap
36/// (`model_swap::swap`) tears the outgoing model down — `cuMemFree_v2` on
37/// every table cached here — and the incoming load's near-identical
38/// `cuMemAlloc_v2` sequence reuses those virtual addresses. A pointer-keyed
39/// static then hands the NEW model the OLD model's expert weight pointers,
40/// and the grouped GEMM silently reads whatever now lives there as weights.
41/// Owning the snapshot on the layer makes staleness structurally impossible:
42/// the snapshot dies with the layer, with the model, at teardown.
43pub struct MoeCutlassHostTables {
44    pub gate_packed: Vec<u64>,
45    pub gate_sfb: Vec<u64>,
46    pub gate_scale2: Vec<f32>,
47    pub up_packed: Vec<u64>,
48    pub up_sfb: Vec<u64>,
49    pub up_scale2: Vec<f32>,
50    /// `None` when the checkpoint has no down-projection scale table — the
51    /// grouped down branch is unreachable in that case.
52    pub down: Option<MoeCutlassDownHostTables>,
53}
54
55/// Down-projection third of [`MoeCutlassHostTables`].
56pub struct MoeCutlassDownHostTables {
57    pub packed: Vec<u64>,
58    pub sfb: Vec<u64>,
59    pub scale2: Vec<f32>,
60}
61
62/// One blocking D2H of a device `[n]` u64 pointer table. Load-time only.
63pub fn read_expert_ptrs_u64(gpu: &dyn GpuBackend, p: DevicePtr, n: usize) -> Result<Vec<u64>> {
64    let mut raw = vec![0u8; n * 8];
65    gpu.copy_d2h(p, &mut raw)?;
66    Ok(raw
67        .chunks_exact(8)
68        .map(|x| u64::from_le_bytes(x.try_into().expect("8")))
69        .collect())
70}
71
72/// One blocking D2H of a device `[n]` f32 scale table. Load-time only.
73pub fn read_expert_scales_f32(gpu: &dyn GpuBackend, p: DevicePtr, n: usize) -> Result<Vec<f32>> {
74    let mut raw = vec![0u8; n * 4];
75    gpu.copy_d2h(p, &mut raw)?;
76    Ok(raw
77        .chunks_exact(4)
78        .map(|x| f32::from_le_bytes(x.try_into().expect("4")))
79        .collect())
80}
81
82impl MoeCutlassHostTables {
83    /// Snapshot the grouped-path tables at load. The SFB pointer vectors are
84    /// taken by value because `build_cutlass_grouped_sfb` constructs them on
85    /// the host in the first place — reading them back from the device would
86    /// re-derive data this function's caller already holds. The packed/scale2
87    /// tables exist only on the device (uploaded by the pointer-table build),
88    /// so those are copied down once here.
89    #[allow(clippy::too_many_arguments)]
90    pub fn snapshot(
91        gpu: &dyn GpuBackend,
92        num_experts: usize,
93        gate_packed: DevicePtr,
94        gate_sfb: Vec<u64>,
95        gate_scale2: DevicePtr,
96        up_packed: DevicePtr,
97        up_sfb: Vec<u64>,
98        up_scale2: DevicePtr,
99        down: Option<(DevicePtr, Vec<u64>, DevicePtr)>,
100    ) -> Result<Self> {
101        Ok(Self {
102            gate_packed: read_expert_ptrs_u64(gpu, gate_packed, num_experts)?,
103            gate_sfb,
104            gate_scale2: read_expert_scales_f32(gpu, gate_scale2, num_experts)?,
105            up_packed: read_expert_ptrs_u64(gpu, up_packed, num_experts)?,
106            up_sfb,
107            up_scale2: read_expert_scales_f32(gpu, up_scale2, num_experts)?,
108            down: match down {
109                Some((packed, sfb, scale2)) => Some(MoeCutlassDownHostTables {
110                    packed: read_expert_ptrs_u64(gpu, packed, num_experts)?,
111                    sfb,
112                    scale2: read_expert_scales_f32(gpu, scale2, num_experts)?,
113                }),
114                None => None,
115            },
116        })
117    }
118}
119
120#[allow(clippy::too_many_arguments)]
121pub fn moe_sort_by_expert(
122    gpu: &dyn GpuBackend,
123    kernel: KernelHandle,
124    topk_ids: DevicePtr,
125    sorted_token_ids: DevicePtr,
126    sorted_expert_ids: DevicePtr,
127    expert_offsets: DevicePtr,
128    token_to_perm: DevicePtr,
129    total_expanded: u32,
130    num_experts: u32,
131    topk: u32,
132    stream: u64,
133) -> Result<()> {
134    KernelLaunch::new(gpu, kernel)
135        .grid([1, 1, 1])
136        .block([256, 1, 1])
137        .arg_ptr(topk_ids)
138        .arg_ptr(sorted_token_ids)
139        .arg_ptr(sorted_expert_ids)
140        .arg_ptr(expert_offsets)
141        .arg_ptr(token_to_perm)
142        .arg_u32(total_expanded)
143        .arg_u32(num_experts)
144        .arg_u32(topk)
145        .launch(stream)
146}
147
148/// Unpermute + weighted reduce with pre-built reverse map.
149///
150/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
151#[allow(clippy::too_many_arguments)]
152pub fn moe_unpermute_reduce_indexed(
153    gpu: &dyn GpuBackend,
154    kernel: KernelHandle,
155    expert_output: DevicePtr,
156    output: DevicePtr,
157    token_to_perm: DevicePtr,
158    topk_weights: DevicePtr,
159    hidden_size: u32,
160    num_tokens: u32,
161    topk: u32,
162    stream: u64,
163) -> Result<()> {
164    KernelLaunch::new(gpu, kernel)
165        .grid([num_tokens, 1, 1])
166        .block([256, 1, 1])
167        .arg_ptr(expert_output)
168        .arg_ptr(output)
169        .arg_ptr(token_to_perm)
170        .arg_ptr(topk_weights)
171        .arg_u32(hidden_size)
172        .arg_u32(num_tokens)
173        .arg_u32(topk)
174        .launch(stream)
175}
176
177/// Batched sigmoid blend: output += sigmoid(dot(normed, gate_weight)) * shared_out.
178///
179/// Grid: (num_tokens, 1, 1)  Block: (256, 1, 1)
180pub fn moe_batched_blend(
181    gpu: &dyn GpuBackend,
182    kernel: KernelHandle,
183    output: DevicePtr,
184    shared_out: DevicePtr,
185    normed: DevicePtr,
186    gate_weight: DevicePtr,
187    hidden_size: u32,
188    num_tokens: u32,
189    stream: u64,
190) -> Result<()> {
191    KernelLaunch::new(gpu, kernel)
192        .grid([num_tokens, 1, 1])
193        .block([256, 1, 1])
194        .arg_ptr(output)
195        .arg_ptr(shared_out)
196        .arg_ptr(normed)
197        .arg_ptr(gate_weight)
198        .arg_u32(hidden_size)
199        .arg_u32(num_tokens)
200        .launch(stream)
201}
202
203/// Single-launch CUTLASS grouped NVFP4 fused gate_up GEMM (Phase-2).
204///
205/// Bridges the load-time host snapshot of the per-expert pointer/scale tables
206/// ([`MoeCutlassHostTables`]) to the host-side
207/// [`spark_runtime::cutlass::nvfp4_grouped_gate_up_fused`] entry. `a` is the
208/// expert-contiguous bf16 activation `[total_expanded, k]`; `expert_offsets`
209/// is the device i32 `[num_experts+1]` prefix sum — the only per-call table,
210/// so the only one copied and the only reason for the synchronize (the C
211/// entry indexes offsets on the host before it can launch).
212#[allow(clippy::too_many_arguments)]
213/// Returns the host copy of `expert_offsets` so the paired `down` call can
214/// reuse it instead of repeating the D2H + synchronize. The two calls share the
215/// same offsets (both are driven by one `moe_sort_by_expert`), and each sync
216/// blocks the host until the GPU drains — halving them halves that stall.
217pub fn moe_grouped_gate_up_cutlass(
218    gpu: &dyn GpuBackend,
219    host: &MoeCutlassHostTables,
220    a: DevicePtr,
221    sorted_token_ids: DevicePtr,
222    c_gate: DevicePtr,
223    c_up: DevicePtr,
224    expert_offsets: DevicePtr,
225    inter: u32,
226    hidden: u32,
227    stream: u64,
228) -> Result<Vec<i32>> {
229    let num_experts = host.gate_packed.len();
230    let mut off_raw = vec![0u8; (num_experts + 1) * 4];
231    gpu.copy_d2h_on_stream(expert_offsets, &mut off_raw, stream)?;
232    // The offsets host copy is needed by the C entry before it can launch —
233    // make sure the async D2H has landed.
234    gpu.synchronize(stream)?;
235    let eoff: Vec<i32> = off_raw
236        .chunks_exact(4)
237        .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]))
238        .collect();
239
240    spark_runtime::cutlass::nvfp4_grouped_gate_up_fused(
241        a.0,
242        sorted_token_ids.0,
243        &host.gate_packed,
244        &host.gate_sfb,
245        &host.gate_scale2,
246        &host.up_packed,
247        &host.up_sfb,
248        &host.up_scale2,
249        c_gate.0,
250        c_up.0,
251        &eoff,
252        inter,
253        hidden,
254        stream,
255    )?;
256    Ok(eoff)
257}
258
259/// Single-launch CUTLASS grouped NVFP4 DOWN projection. `a` is the post-SiLU
260/// intermediate `[total_expanded, inter]` (already expert-contiguous — no
261/// gather). `host` is the down third of the load-time snapshot;
262/// `expert_offsets` is the device i32 `[num_experts+1]` prefix sum.
263#[allow(clippy::too_many_arguments)]
264pub fn moe_grouped_down_cutlass(
265    gpu: &dyn GpuBackend,
266    // Host `expert_offsets` from the paired gate_up call. When supplied, the
267    // D2H + synchronize here is skipped entirely — the offsets are identical
268    // (one sort feeds both projections).
269    eoff_cached: Option<&[i32]>,
270    host: &MoeCutlassDownHostTables,
271    a: DevicePtr,
272    c: DevicePtr,
273    expert_offsets: DevicePtr,
274    hidden: u32,
275    inter: u32,
276    stream: u64,
277) -> Result<()> {
278    let num_experts = host.packed.len();
279    // Offsets come from the paired gate_up when available; otherwise fetch them
280    // (D2H + the sync that blocks the host until the GPU drains).
281    let eoff: Vec<i32> = if let Some(e) = eoff_cached {
282        e.to_vec()
283    } else {
284        let mut off_raw = vec![0u8; (num_experts + 1) * 4];
285        gpu.copy_d2h_on_stream(expert_offsets, &mut off_raw, stream)?;
286        gpu.synchronize(stream)?;
287        off_raw
288            .chunks_exact(4)
289            .map(|c| i32::from_le_bytes(c.try_into().expect("4")))
290            .collect()
291    };
292    spark_runtime::cutlass::nvfp4_grouped_down(
293        a.0,
294        &host.packed,
295        &host.sfb,
296        &host.scale2,
297        c.0,
298        &eoff,
299        hidden,
300        inter,
301        stream,
302    )
303}
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308    use spark_runtime::gpu::mock::MockGpuBackend;
309
310    fn upload_u64(gpu: &MockGpuBackend, vals: &[u64]) -> DevicePtr {
311        let bytes: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
312        let p = gpu.alloc(bytes.len()).unwrap();
313        gpu.copy_h2d(&bytes, p).unwrap();
314        p
315    }
316
317    fn upload_f32(gpu: &MockGpuBackend, vals: &[f32]) -> DevicePtr {
318        let bytes: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
319        let p = gpu.alloc(bytes.len()).unwrap();
320        gpu.copy_h2d(&bytes, p).unwrap();
321        p
322    }
323
324    /// The failure this pins: the previous implementation memoized these reads
325    /// in a process-global map keyed on the device ADDRESS. After an
326    /// in-process model swap the incoming load reuses the outgoing model's
327    /// freed virtual addresses, and the pointer-keyed cache handed the new
328    /// model the OLD model's expert weight pointers — the grouped GEMM then
329    /// silently read unrelated tensors as weights. A read must reflect what
330    /// the device holds NOW, so writing new contents to the same address and
331    /// reading again must observe the new contents.
332    #[test]
333    fn read_reflects_current_device_contents_at_a_reused_address() {
334        let gpu = MockGpuBackend::new();
335        let old_model = [0x1111_u64, 0x2222, 0x3333];
336        let p = upload_u64(&gpu, &old_model);
337        assert_eq!(read_expert_ptrs_u64(&gpu, p, 3).unwrap(), old_model);
338
339        // Same address, new contents — the swap's free/realloc collapsed to
340        // its essence (the mock never moves an allocation, which is exactly
341        // the driver's common case for identical alloc sequences).
342        let new_model = [0xaaaa_u64, 0xbbbb, 0xcccc];
343        let bytes: Vec<u8> = new_model.iter().flat_map(|v| v.to_le_bytes()).collect();
344        gpu.copy_h2d(&bytes, p).unwrap();
345        assert_eq!(read_expert_ptrs_u64(&gpu, p, 3).unwrap(), new_model);
346
347        let old_scales = [1.0_f32, 2.0];
348        let ps = upload_f32(&gpu, &old_scales);
349        assert_eq!(read_expert_scales_f32(&gpu, ps, 2).unwrap(), old_scales);
350        let new_scales = [3.0_f32, 4.0];
351        let sbytes: Vec<u8> = new_scales.iter().flat_map(|v| v.to_le_bytes()).collect();
352        gpu.copy_h2d(&sbytes, ps).unwrap();
353        assert_eq!(read_expert_scales_f32(&gpu, ps, 2).unwrap(), new_scales);
354    }
355
356    /// The old cache also ignored the requested length: a hit returned the
357    /// first query's vector whatever `n` the caller asked for. A read of `n`
358    /// elements must return exactly `n` elements.
359    #[test]
360    fn read_honors_the_requested_length() {
361        let gpu = MockGpuBackend::new();
362        let vals = [1_u64, 2, 3, 4];
363        let p = upload_u64(&gpu, &vals);
364        assert_eq!(read_expert_ptrs_u64(&gpu, p, 2).unwrap(), vals[..2]);
365        assert_eq!(read_expert_ptrs_u64(&gpu, p, 4).unwrap(), vals);
366    }
367
368    /// Nine same-typed tables flow into `snapshot`; a transposition would
369    /// type-check and quantize with the wrong expert scales. Pin each field
370    /// to its source.
371    #[test]
372    fn snapshot_maps_every_table_to_its_field() {
373        let gpu = MockGpuBackend::new();
374        let n = 2;
375        let gate_packed = upload_u64(&gpu, &[10, 11]);
376        let gate_scale2 = upload_f32(&gpu, &[0.5, 0.25]);
377        let up_packed = upload_u64(&gpu, &[20, 21]);
378        let up_scale2 = upload_f32(&gpu, &[2.0, 4.0]);
379        let down_packed = upload_u64(&gpu, &[30, 31]);
380        let down_scale2 = upload_f32(&gpu, &[8.0, 16.0]);
381
382        let t = MoeCutlassHostTables::snapshot(
383            &gpu,
384            n,
385            gate_packed,
386            vec![100, 101],
387            gate_scale2,
388            up_packed,
389            vec![200, 201],
390            up_scale2,
391            Some((down_packed, vec![300, 301], down_scale2)),
392        )
393        .unwrap();
394
395        assert_eq!(t.gate_packed, [10, 11]);
396        assert_eq!(t.gate_sfb, [100, 101]);
397        assert_eq!(t.gate_scale2, [0.5, 0.25]);
398        assert_eq!(t.up_packed, [20, 21]);
399        assert_eq!(t.up_sfb, [200, 201]);
400        assert_eq!(t.up_scale2, [2.0, 4.0]);
401        let d = t.down.expect("down tables were supplied");
402        assert_eq!(d.packed, [30, 31]);
403        assert_eq!(d.sfb, [300, 301]);
404        assert_eq!(d.scale2, [8.0, 16.0]);
405    }
406}
spark_model/layers/ops/moe_grouped_a2.rs

spark_model/layers/ops/
moe_grouped_a2.rs
