metrale_model_layers/layers/ops/
moe_grouped_a2.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the MoE expert sort, the unpermute reduce and the shared-expert blend, and the host side of the CUTLASS grouped NVFP4 gate_up and down GEMMs.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - The two CUTLASS calls take the weight tables from [`MoeCutlassHostTables`]; per call they
8//!   copy at most `expert_offsets` from the device.
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
17
18use super::*;
19
20/// 2026-09-25: Host copies of the per-expert pointer and scale tables for the CUTLASS grouped
21/// path, which takes them on the host. `MoeLayer::build_cutlass_grouped_sfb` builds them at load
22/// into the layer's `cutlass_grouped_host`, so they are dropped with the layer.
23pub struct MoeCutlassHostTables {
24    pub gate_packed: Vec<u64>,
25    pub gate_sfb: Vec<u64>,
26    pub gate_scale2: Vec<f32>,
27    pub up_packed: Vec<u64>,
28    pub up_sfb: Vec<u64>,
29    pub up_scale2: Vec<f32>,
30    /// 2026-09-25: `None` when the layer has no down-projection scale table; the grouped down
31    /// branch of `forward_prefill_routed` is then not taken.
32    pub down: Option<MoeCutlassDownHostTables>,
33}
34
35/// 2026-09-25: The down-projection tables of [`MoeCutlassHostTables`].
36pub struct MoeCutlassDownHostTables {
37    pub packed: Vec<u64>,
38    pub sfb: Vec<u64>,
39    pub scale2: Vec<f32>,
40}
41
42/// 2026-09-25: One blocking device-to-host copy of an `[n]` u64 pointer table.
43pub fn read_expert_ptrs_u64(gpu: &dyn GpuBackend, p: DevicePtr, n: usize) -> Result<Vec<u64>> {
44    let mut raw = vec![0u8; n * 8];
45    gpu.copy_d2h(p, &mut raw)?;
46    Ok(raw
47        .chunks_exact(8)
48        .map(|x| u64::from_le_bytes(x.try_into().expect("8")))
49        .collect())
50}
51
52/// 2026-09-25: One blocking device-to-host copy of an `[n]` f32 scale table.
53pub fn read_expert_scales_f32(gpu: &dyn GpuBackend, p: DevicePtr, n: usize) -> Result<Vec<f32>> {
54    let mut raw = vec![0u8; n * 4];
55    gpu.copy_d2h(p, &mut raw)?;
56    Ok(raw
57        .chunks_exact(4)
58        .map(|x| f32::from_le_bytes(x.try_into().expect("4")))
59        .collect())
60}
61
62impl MoeCutlassHostTables {
63    /// 2026-09-25: Snapshot the grouped-path tables. The SFB pointer vectors are taken by value
64    /// because `build_cutlass_grouped_sfb` builds them on the host; the packed and scale2 tables
65    /// are copied from the device here.
66    #[allow(clippy::too_many_arguments)]
67    pub fn snapshot(
68        gpu: &dyn GpuBackend,
69        num_experts: usize,
70        gate_packed: DevicePtr,
71        gate_sfb: Vec<u64>,
72        gate_scale2: DevicePtr,
73        up_packed: DevicePtr,
74        up_sfb: Vec<u64>,
75        up_scale2: DevicePtr,
76        down: Option<(DevicePtr, Vec<u64>, DevicePtr)>,
77    ) -> Result<Self> {
78        Ok(Self {
79            gate_packed: read_expert_ptrs_u64(gpu, gate_packed, num_experts)?,
80            gate_sfb,
81            gate_scale2: read_expert_scales_f32(gpu, gate_scale2, num_experts)?,
82            up_packed: read_expert_ptrs_u64(gpu, up_packed, num_experts)?,
83            up_sfb,
84            up_scale2: read_expert_scales_f32(gpu, up_scale2, num_experts)?,
85            down: match down {
86                Some((packed, sfb, scale2)) => Some(MoeCutlassDownHostTables {
87                    packed: read_expert_ptrs_u64(gpu, packed, num_experts)?,
88                    sfb,
89                    scale2: read_expert_scales_f32(gpu, scale2, num_experts)?,
90                }),
91                None => None,
92            },
93        })
94    }
95}
96
97/// 2026-09-25: Counting sort of the `total_expanded` (token, slot) routes by expert, in one block:
98/// writes `sorted_token_ids` and `sorted_expert_ids` in expert order, `expert_offsets` (the
99/// `[num_experts + 1]` prefix sum) and `token_to_perm` (each route's sorted position).
100#[allow(clippy::too_many_arguments)]
101pub fn moe_sort_by_expert(
102    gpu: &dyn GpuBackend,
103    kernel: KernelHandle,
104    topk_ids: DevicePtr,
105    sorted_token_ids: DevicePtr,
106    sorted_expert_ids: DevicePtr,
107    expert_offsets: DevicePtr,
108    token_to_perm: DevicePtr,
109    total_expanded: u32,
110    num_experts: u32,
111    topk: u32,
112    stream: u64,
113) -> Result<()> {
114    KernelLaunch::new(gpu, kernel)
115        .grid([1, 1, 1])
116        .block([256, 1, 1])
117        .arg_ptr(topk_ids)
118        .arg_ptr(sorted_token_ids)
119        .arg_ptr(sorted_expert_ids)
120        .arg_ptr(expert_offsets)
121        .arg_ptr(token_to_perm)
122        .arg_u32(total_expanded)
123        .arg_u32(num_experts)
124        .arg_u32(topk)
125        .launch(stream)
126}
127
128/// 2026-09-25: For each token, `output[t] = sum_k topk_weights[t, k] * expert_output[token_to_perm[t, k]]`,
129/// one block per token.
130#[allow(clippy::too_many_arguments)]
131pub fn moe_unpermute_reduce_indexed(
132    gpu: &dyn GpuBackend,
133    kernel: KernelHandle,
134    expert_output: DevicePtr,
135    output: DevicePtr,
136    token_to_perm: DevicePtr,
137    topk_weights: DevicePtr,
138    hidden_size: u32,
139    num_tokens: u32,
140    topk: u32,
141    stream: u64,
142) -> Result<()> {
143    KernelLaunch::new(gpu, kernel)
144        .grid([num_tokens, 1, 1])
145        .block([256, 1, 1])
146        .arg_ptr(expert_output)
147        .arg_ptr(output)
148        .arg_ptr(token_to_perm)
149        .arg_ptr(topk_weights)
150        .arg_u32(hidden_size)
151        .arg_u32(num_tokens)
152        .arg_u32(topk)
153        .launch(stream)
154}
155
156/// 2026-09-28: [`moe_unpermute_reduce_indexed`] then [`moe_batched_blend`] in one launch
157/// (kernel `moe_unpermute_blend`, `kernels/gb10/common/moe_unpermute_blend.cu`), with the same
158/// output bits; one 256-thread block per token. Needs `hidden_size % 8 == 0`.
159#[allow(clippy::too_many_arguments)]
160pub fn moe_unpermute_blend(
161    gpu: &dyn GpuBackend,
162    kernel: KernelHandle,
163    expert_output: DevicePtr,
164    output: DevicePtr,
165    token_to_perm: DevicePtr,
166    topk_weights: DevicePtr,
167    shared_out: DevicePtr,
168    normed: DevicePtr,
169    gate_weight: DevicePtr,
170    hidden_size: u32,
171    num_tokens: u32,
172    topk: u32,
173    stream: u64,
174) -> Result<()> {
175    anyhow::ensure!(
176        hidden_size.is_multiple_of(8),
177        "moe_unpermute_blend: hidden {hidden_size} is not a multiple of 8"
178    );
179    KernelLaunch::new(gpu, kernel)
180        .grid([num_tokens, 1, 1])
181        .block([256, 1, 1])
182        .arg_ptr(expert_output)
183        .arg_ptr(output)
184        .arg_ptr(token_to_perm)
185        .arg_ptr(topk_weights)
186        .arg_ptr(shared_out)
187        .arg_ptr(normed)
188        .arg_ptr(gate_weight)
189        .arg_u32(hidden_size)
190        .arg_u32(num_tokens)
191        .arg_u32(topk)
192        .launch(stream)
193}
194
195/// 2026-09-25: `output[t] += sigmoid(dot(normed[t], gate_weight)) * shared_out[t]`, one block per
196/// token. A null `gate_weight` blends at weight 1.
197pub fn moe_batched_blend(
198    gpu: &dyn GpuBackend,
199    kernel: KernelHandle,
200    output: DevicePtr,
201    shared_out: DevicePtr,
202    normed: DevicePtr,
203    gate_weight: DevicePtr,
204    hidden_size: u32,
205    num_tokens: u32,
206    stream: u64,
207) -> Result<()> {
208    KernelLaunch::new(gpu, kernel)
209        .grid([num_tokens, 1, 1])
210        .block([256, 1, 1])
211        .arg_ptr(output)
212        .arg_ptr(shared_out)
213        .arg_ptr(normed)
214        .arg_ptr(gate_weight)
215        .arg_u32(hidden_size)
216        .arg_u32(num_tokens)
217        .launch(stream)
218}
219
220/// 2026-09-25: Single-launch CUTLASS grouped NVFP4 gate_up GEMM through
221/// [`metrale_gpu_runtime::cutlass::nvfp4_grouped_gate_up_fused`], with the weights from `host`.
222/// `expert_offsets` is the device i32 `[num_experts + 1]` prefix sum; the host entry needs it on
223/// the host, so it is copied and the stream synchronised first. Returns that host copy, which
224/// [`moe_grouped_down_cutlass`] accepts to skip its own copy and synchronise.
225#[allow(clippy::too_many_arguments)]
226pub fn moe_grouped_gate_up_cutlass(
227    gpu: &dyn GpuBackend,
228    host: &MoeCutlassHostTables,
229    a: DevicePtr,
230    sorted_token_ids: DevicePtr,
231    c_gate: DevicePtr,
232    c_up: DevicePtr,
233    expert_offsets: DevicePtr,
234    inter: u32,
235    hidden: u32,
236    stream: u64,
237) -> Result<Vec<i32>> {
238    let num_experts = host.gate_packed.len();
239    let mut off_raw = vec![0u8; (num_experts + 1) * 4];
240    gpu.copy_d2h_on_stream(expert_offsets, &mut off_raw, stream)?;
241    gpu.synchronize(stream)?;
242    let eoff: Vec<i32> = off_raw
243        .chunks_exact(4)
244        .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]))
245        .collect();
246
247    metrale_gpu_runtime::cutlass::nvfp4_grouped_gate_up_fused(
248        a.0,
249        sorted_token_ids.0,
250        &host.gate_packed,
251        &host.gate_sfb,
252        &host.gate_scale2,
253        &host.up_packed,
254        &host.up_sfb,
255        &host.up_scale2,
256        c_gate.0,
257        c_up.0,
258        &eoff,
259        inter,
260        hidden,
261        stream,
262    )?;
263    Ok(eoff)
264}
265
266/// 2026-09-25: Single-launch CUTLASS grouped NVFP4 down GEMM through
267/// [`metrale_gpu_runtime::cutlass::nvfp4_grouped_down`]. `a` is the expert-sorted post-SiLU
268/// intermediate `[total_expanded, inter]`; `expert_offsets` is the device prefix sum.
269#[allow(clippy::too_many_arguments)]
270pub fn moe_grouped_down_cutlass(
271    gpu: &dyn GpuBackend,
272    // 2026-09-25: The host `expert_offsets` returned by the paired gate_up call. When supplied,
273    // no device copy or synchronise happens here.
274    eoff_cached: Option<&[i32]>,
275    host: &MoeCutlassDownHostTables,
276    a: DevicePtr,
277    c: DevicePtr,
278    expert_offsets: DevicePtr,
279    hidden: u32,
280    inter: u32,
281    stream: u64,
282) -> Result<()> {
283    let num_experts = host.packed.len();
284    let eoff: Vec<i32> = if let Some(e) = eoff_cached {
285        e.to_vec()
286    } else {
287        let mut off_raw = vec![0u8; (num_experts + 1) * 4];
288        gpu.copy_d2h_on_stream(expert_offsets, &mut off_raw, stream)?;
289        gpu.synchronize(stream)?;
290        off_raw
291            .chunks_exact(4)
292            .map(|c| i32::from_le_bytes(c.try_into().expect("4")))
293            .collect()
294    };
295    metrale_gpu_runtime::cutlass::nvfp4_grouped_down(
296        a.0,
297        &host.packed,
298        &host.sfb,
299        &host.scale2,
300        c.0,
301        &eoff,
302        hidden,
303        inter,
304        stream,
305    )
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311    use metrale_gpu_runtime::gpu::mock::MockGpuBackend;
312
313    fn upload_u64(gpu: &MockGpuBackend, vals: &[u64]) -> DevicePtr {
314        let bytes: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
315        let p = gpu.alloc(bytes.len()).unwrap();
316        gpu.copy_h2d(&bytes, p).unwrap();
317        p
318    }
319
320    fn upload_f32(gpu: &MockGpuBackend, vals: &[f32]) -> DevicePtr {
321        let bytes: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
322        let p = gpu.alloc(bytes.len()).unwrap();
323        gpu.copy_h2d(&bytes, p).unwrap();
324        p
325    }
326
327    /// 2026-09-25: A read returns what the device holds at the time of the call: new contents
328    /// written at the same address are observed, so nothing is cached by address.
329    #[test]
330    fn read_reflects_current_device_contents_at_a_reused_address() {
331        let gpu = MockGpuBackend::new();
332        let old_model = [0x1111_u64, 0x2222, 0x3333];
333        let p = upload_u64(&gpu, &old_model);
334        assert_eq!(read_expert_ptrs_u64(&gpu, p, 3).unwrap(), old_model);
335
336        let new_model = [0xaaaa_u64, 0xbbbb, 0xcccc];
337        let bytes: Vec<u8> = new_model.iter().flat_map(|v| v.to_le_bytes()).collect();
338        gpu.copy_h2d(&bytes, p).unwrap();
339        assert_eq!(read_expert_ptrs_u64(&gpu, p, 3).unwrap(), new_model);
340
341        let old_scales = [1.0_f32, 2.0];
342        let ps = upload_f32(&gpu, &old_scales);
343        assert_eq!(read_expert_scales_f32(&gpu, ps, 2).unwrap(), old_scales);
344        let new_scales = [3.0_f32, 4.0];
345        let sbytes: Vec<u8> = new_scales.iter().flat_map(|v| v.to_le_bytes()).collect();
346        gpu.copy_h2d(&sbytes, ps).unwrap();
347        assert_eq!(read_expert_scales_f32(&gpu, ps, 2).unwrap(), new_scales);
348    }
349
350    /// 2026-09-25: A read of `n` elements returns exactly `n` elements.
351    #[test]
352    fn read_honors_the_requested_length() {
353        let gpu = MockGpuBackend::new();
354        let vals = [1_u64, 2, 3, 4];
355        let p = upload_u64(&gpu, &vals);
356        assert_eq!(read_expert_ptrs_u64(&gpu, p, 2).unwrap(), vals[..2]);
357        assert_eq!(read_expert_ptrs_u64(&gpu, p, 4).unwrap(), vals);
358    }
359
360    /// 2026-09-25: Each of the nine tables `snapshot` takes lands in its own field; a swapped
361    /// argument of the same type would still compile.
362    #[test]
363    fn snapshot_maps_every_table_to_its_field() {
364        let gpu = MockGpuBackend::new();
365        let n = 2;
366        let gate_packed = upload_u64(&gpu, &[10, 11]);
367        let gate_scale2 = upload_f32(&gpu, &[0.5, 0.25]);
368        let up_packed = upload_u64(&gpu, &[20, 21]);
369        let up_scale2 = upload_f32(&gpu, &[2.0, 4.0]);
370        let down_packed = upload_u64(&gpu, &[30, 31]);
371        let down_scale2 = upload_f32(&gpu, &[8.0, 16.0]);
372
373        let t = MoeCutlassHostTables::snapshot(
374            &gpu,
375            n,
376            gate_packed,
377            vec![100, 101],
378            gate_scale2,
379            up_packed,
380            vec![200, 201],
381            up_scale2,
382            Some((down_packed, vec![300, 301], down_scale2)),
383        )
384        .unwrap();
385
386        assert_eq!(t.gate_packed, [10, 11]);
387        assert_eq!(t.gate_sfb, [100, 101]);
388        assert_eq!(t.gate_scale2, [0.5, 0.25]);
389        assert_eq!(t.up_packed, [20, 21]);
390        assert_eq!(t.up_sfb, [200, 201]);
391        assert_eq!(t.up_scale2, [2.0, 4.0]);
392        let d = t.down.expect("down tables were supplied");
393        assert_eq!(d.packed, [30, 31]);
394        assert_eq!(d.sfb, [300, 301]);
395        assert_eq!(d.scale2, [8.0, 16.0]);
396    }
397}
metrale_model_layers/layers/ops/moe_grouped_a2.rs

metrale_model_layers/layers/ops/
moe_grouped_a2.rs
