metrale_model_layers/layers/ops/
moe_grouped_fp4.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launcher for the grouped W4A4 expert up GEMM `moe_w4a4_grouped_gemm_relu2`.
4//!
5//! Owner: model-layers ops.
6//! Invariants: none beyond the types.
7
8#![allow(unused_imports)]
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
13
14use super::*;
15
16/// 2026-09-25: Grouped W4A4 expert up GEMM on the FP4 block-scale MMA, with
17/// relu^2 applied to the accumulator before the BF16 store. A is NVFP4 (packed
18/// E2M1 + per-16 E4M3 scales) in token order, gathered through
19/// `sorted_token_ids`; B is read from the per-expert NVFP4 pointer tables.
20/// `blockIdx.z` is the expert, and `expert_offsets` bounds its rows.
21/// The kernel ships only in
22/// `kernels/gb10/nemotron-labs-3-puzzle-75b-a9b/nvfp4/moe_w4a4_grouped.cu`.
23#[allow(clippy::too_many_arguments)]
24pub fn moe_w4a4_grouped_gemm_relu2(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    a_packed: DevicePtr,
28    a_sf: DevicePtr,
29    b_packed_ptrs: DevicePtr,
30    b_scale_ptrs: DevicePtr,
31    scale2_vals: DevicePtr,
32    output: DevicePtr,
33    expert_offsets: DevicePtr,
34    sorted_token_ids: DevicePtr,
35    num_experts: u32,
36    n_out: u32,
37    k: u32,
38    max_m_tiles: u32,
39    stream: u64,
40) -> Result<()> {
41    KernelLaunch::new(gpu, kernel)
42        .grid([n_tile_blocks(gpu, kernel, n_out)?, max_m_tiles, num_experts])
43        .block([128, 1, 1])
44        .arg_ptr(a_packed)
45        .arg_ptr(a_sf)
46        .arg_ptr(b_packed_ptrs)
47        .arg_ptr(b_scale_ptrs)
48        .arg_ptr(scale2_vals)
49        .arg_ptr(output)
50        .arg_ptr(expert_offsets)
51        .arg_ptr(sorted_token_ids)
52        .arg_u32(num_experts)
53        .arg_u32(n_out)
54        .arg_u32(k)
55        .launch(stream)
56}
metrale_model_layers/layers/ops/moe_grouped_fp4.rs

metrale_model_layers/layers/ops/
moe_grouped_fp4.rs
