metrale_model_layers/layers/ops/
nvfp4_moe_grouped.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-27: Launchers for the grouped NVFP4 MoE decode kernels
4//! (`kernels/gb10/common/moe_nvfp4_grouped.cu`).
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8//!
9//! The layout is that of the grouped FP8 decode (`fp8_moe_grouped.rs`): rows grouped by
10//! expert and the active experts listed by `moe_fp8_grouped_sort`, intermediates laid out
11//! by sorted position, and the blend
12//! `moe_weighted_sum_blend_fp8_grouped`. Only the weight format differs: row-major
13//! NVFP4 (packed E2M1, E4M3 block scales of 16, per-tensor scale 2).
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
18
19use super::Fp8GroupedGeometry;
20use crate::weight_map::QuantizedWeight;
21
22/// 2026-09-27: Output columns per gate+up CTA. Must equal `NG_GU_COLS_PER_CTA` in the `.cu`.
23pub const NVFP4_GROUPED_GATE_UP_COLS_PER_CTA: u32 = 16;
24
25/// 2026-09-27: Rows per gate+up pass. Must equal `NG_GU_ROWS` in the `.cu`.
26pub const NVFP4_GROUPED_GATE_UP_ROWS_PER_PASS: u32 = 4;
27
28/// 2026-09-27: Output columns per down CTA. Must equal `NG_DOWN_COLS_PER_CTA` in the `.cu`.
29pub const NVFP4_GROUPED_DOWN_COLS_PER_CTA: u32 = 128;
30
31/// 2026-09-27: Rows per down pass. Must equal `NG_DOWN_ROWS` in the `.cu`.
32pub const NVFP4_GROUPED_DOWN_ROWS_PER_PASS: u32 = 4;
33
34/// 2026-10-02: `moe_expert_gate_up_act_nvfp4_grouped` (CUDA cores, FP32 SiLU products).
35pub const NVFP4_GROUPED_GATE_UP_SCALAR: Fp8GroupedGeometry = Fp8GroupedGeometry {
36    cols_per_cta: NVFP4_GROUPED_GATE_UP_COLS_PER_CTA,
37    rows_per_pass: NVFP4_GROUPED_GATE_UP_ROWS_PER_PASS,
38    threads: 128,
39};
40
41/// 2026-10-02: `moe_expert_down_act_nvfp4_grouped` (reads FP32 SiLU products).
42pub const NVFP4_GROUPED_DOWN_SCALAR: Fp8GroupedGeometry = Fp8GroupedGeometry {
43    cols_per_cta: NVFP4_GROUPED_DOWN_COLS_PER_CTA,
44    rows_per_pass: NVFP4_GROUPED_DOWN_ROWS_PER_PASS,
45    threads: 256,
46};
47
48/// 2026-10-02: Rows per pass of the tensor-core kernels (`moe_nvfp4_grouped_tc.cu`,
49/// `TC_ROWS`), gate+up and down alike.
50pub const NVFP4_GROUPED_TC_ROWS_PER_PASS: u32 = 8;
51
52/// 2026-10-02: `moe_expert_gate_up_act_nvfp4_grouped_tc` (`NTC_GU_COLS` columns per CTA; writes
53/// the SiLU product as BF16 hi + lo).
54pub const NVFP4_GROUPED_GATE_UP_TC: Fp8GroupedGeometry = Fp8GroupedGeometry {
55    cols_per_cta: 64,
56    rows_per_pass: NVFP4_GROUPED_TC_ROWS_PER_PASS,
57    threads: 128,
58};
59
60/// 2026-10-02: `moe_expert_down_act_nvfp4_grouped_tc` (`NTC_DOWN_COLS` columns per CTA; reads
61/// the hi + lo SiLU rows).
62pub const NVFP4_GROUPED_DOWN_TC: Fp8GroupedGeometry = Fp8GroupedGeometry {
63    cols_per_cta: 128,
64    rows_per_pass: NVFP4_GROUPED_TC_ROWS_PER_PASS,
65    threads: 128,
66};
67
68/// 2026-10-02: Whether the tensor-core kernels take an `n`-column, `k`-deep projection: `k` in
69/// whole load groups of two 128-wide chunks, `n` in whole CTAs of `geometry`.
70pub fn nvfp4_grouped_tc_shape_ok(n: u32, k: u32, geometry: Fp8GroupedGeometry) -> bool {
71    n > 0 && k > 0 && k.is_multiple_of(256) && n.is_multiple_of(geometry.cols_per_cta)
72}
73
74/// 2026-09-27: One NVFP4 projection across the routed experts: device tables of each
75/// expert's packed and block-scale pointers, and its per-tensor scales as f32.
76#[derive(Clone, Copy, Debug)]
77pub struct Nvfp4ExpertTables {
78    pub packed_ptrs: DevicePtr,
79    pub scale_ptrs: DevicePtr,
80    pub scale2_vals: DevicePtr,
81}
82
83/// 2026-09-27: Grouped NVFP4 gate+up and SiLU over `k`-wide input rows. `cap` is
84/// `fp8_grouped_active_cap`; the first `ceil(num_tokens / geometry.rows_per_pass)` block rows
85/// are the shared expert. Writes the FP32 product `silu(bf16(gate)) * bf16(up)`,
86/// `[positions, n]` for the routed experts into `act` and `[num_tokens, n]` for the shared
87/// expert into `sh_act` (2026-10-02: the tensor-core kernel writes it as BF16 hi + lo in the
88/// same bytes, which only its down kernel reads).
89#[allow(clippy::too_many_arguments)]
90pub fn moe_expert_gate_up_act_nvfp4_grouped(
91    gpu: &dyn GpuBackend,
92    kernel: KernelHandle,
93    geometry: Fp8GroupedGeometry,
94    input: DevicePtr,
95    gate: Nvfp4ExpertTables,
96    up: Nvfp4ExpertTables,
97    act: DevicePtr,
98    expert_offsets: DevicePtr,
99    sorted_token_ids: DevicePtr,
100    active_experts: DevicePtr,
101    active_count: DevicePtr,
102    sh_gate: &QuantizedWeight,
103    sh_up: &QuantizedWeight,
104    sh_act: DevicePtr,
105    n: u32,
106    k: u32,
107    cap: u32,
108    num_tokens: u32,
109    stream: u64,
110) -> Result<()> {
111    KernelLaunch::new(gpu, kernel)
112        .grid([
113            div_ceil(n, geometry.cols_per_cta),
114            cap + div_ceil(num_tokens, geometry.rows_per_pass),
115            1,
116        ])
117        .block([geometry.threads, 1, 1])
118        .arg_ptr(input)
119        .arg_ptr(gate.packed_ptrs)
120        .arg_ptr(gate.scale_ptrs)
121        .arg_ptr(gate.scale2_vals)
122        .arg_ptr(up.packed_ptrs)
123        .arg_ptr(up.scale_ptrs)
124        .arg_ptr(up.scale2_vals)
125        .arg_ptr(act)
126        .arg_ptr(expert_offsets)
127        .arg_ptr(sorted_token_ids)
128        .arg_ptr(active_experts)
129        .arg_ptr(active_count)
130        .arg_ptr(sh_gate.weight)
131        .arg_ptr(sh_gate.weight_scale)
132        .arg_f32(sh_gate.weight_scale_2)
133        .arg_ptr(sh_up.weight)
134        .arg_ptr(sh_up.weight_scale)
135        .arg_f32(sh_up.weight_scale_2)
136        .arg_ptr(sh_act)
137        .arg_u32(n)
138        .arg_u32(k)
139        .arg_u32(cap)
140        .arg_u32(num_tokens)
141        .launch(stream)
142}
143
144/// 2026-09-27: Grouped NVFP4 down over the SiLU product. `k` is the intermediate width; rows
145/// of `act` and `output` are sorted positions, rows of `sh_act` and `sh_down_out` tokens.
146#[allow(clippy::too_many_arguments)]
147/// 2026-10-04: Repacks one row-major NVFP4 matrix `[n, k]` (packed E2M1 `[n, k / 2]` at `packed`,
148/// E4M3 block scales `[n, k / 16]` at `scale`) in place into the lean layout the `_lean`
149/// tensor-core expert kernels read (`nvfp4_tc_lean_repack`, moe_nvfp4_grouped_tc.cu): same bytes
150/// and size, outputs byte-identical to the row-major point. Sets the device u32 at `bad` (zeroed
151/// by the caller) on a negative or NaN scale, after which the matrix is unusable. Needs
152/// [`nvfp4_lean_repack_shape_ok`].
153pub fn nvfp4_tc_lean_repack(
154    gpu: &dyn GpuBackend,
155    kernel: KernelHandle,
156    packed: DevicePtr,
157    scale: DevicePtr,
158    n: u32,
159    k: u32,
160    bad: DevicePtr,
161    stream: u64,
162) -> Result<()> {
163    anyhow::ensure!(
164        nvfp4_lean_repack_shape_ok(n, k),
165        "nvfp4_tc_lean_repack: shape [{n}, {k}] not admitted"
166    );
167    KernelLaunch::new(gpu, kernel)
168        .grid([n / 16, 1, 1])
169        .block([256, 1, 1])
170        .shared_mem(NVFP4_LEAN_REPACK_SMEM_PER_K * k)
171        .arg_ptr(packed)
172        .arg_ptr(scale)
173        .arg_u32(k)
174        .arg_ptr(bad)
175        .launch(stream)
176}
177
178/// 2026-10-04: Shared bytes `nvfp4_tc_lean_repack` stages per unit of K: a 16-row tile's packed
179/// (16 * K / 2) and scale (16 * K / 16) bytes.
180pub const NVFP4_LEAN_REPACK_SMEM_PER_K: u32 = 9;
181
182/// 2026-10-04: `nvfp4_tc_lean_repack` admits whole 16-row tiles, whole 128-K chunks, and a tile
183/// that fits the default 48 KiB of shared memory.
184pub fn nvfp4_lean_repack_shape_ok(n: u32, k: u32) -> bool {
185    n > 0
186        && n.is_multiple_of(16)
187        && k > 0
188        && k.is_multiple_of(128)
189        && NVFP4_LEAN_REPACK_SMEM_PER_K * k <= 48 * 1024
190}
191
192pub fn moe_expert_down_act_nvfp4_grouped(
193    gpu: &dyn GpuBackend,
194    kernel: KernelHandle,
195    geometry: Fp8GroupedGeometry,
196    act: DevicePtr,
197    down: Nvfp4ExpertTables,
198    output: DevicePtr,
199    expert_offsets: DevicePtr,
200    active_experts: DevicePtr,
201    active_count: DevicePtr,
202    sh_act: DevicePtr,
203    sh_down: &QuantizedWeight,
204    sh_down_out: DevicePtr,
205    n: u32,
206    k: u32,
207    cap: u32,
208    num_tokens: u32,
209    stream: u64,
210) -> Result<()> {
211    KernelLaunch::new(gpu, kernel)
212        .grid([
213            div_ceil(n, geometry.cols_per_cta),
214            cap + div_ceil(num_tokens, geometry.rows_per_pass),
215            1,
216        ])
217        .block([geometry.threads, 1, 1])
218        .arg_ptr(act)
219        .arg_ptr(down.packed_ptrs)
220        .arg_ptr(down.scale_ptrs)
221        .arg_ptr(down.scale2_vals)
222        .arg_ptr(output)
223        .arg_ptr(expert_offsets)
224        .arg_ptr(active_experts)
225        .arg_ptr(active_count)
226        .arg_ptr(sh_act)
227        .arg_ptr(sh_down.weight)
228        .arg_ptr(sh_down.weight_scale)
229        .arg_f32(sh_down.weight_scale_2)
230        .arg_ptr(sh_down_out)
231        .arg_u32(n)
232        .arg_u32(k)
233        .arg_u32(cap)
234        .arg_u32(num_tokens)
235        .launch(stream)
236}
237
238#[cfg(test)]
239#[path = "nvfp4_moe_grouped_tests.rs"]
240mod tests;
metrale_model_layers/layers/ops/nvfp4_moe_grouped.rs

metrale_model_layers/layers/ops/
nvfp4_moe_grouped.rs
