metrale_model_arch/
nemotron_mamba2.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Nemotron-H Mamba-2 SSM layer (no FFN), implementing
4//! `TransformerLayer`.
5//!
6//! Decode, per token (`trait_impl.rs`):
7//!   1. RMS norm, the layer input copied to `residual`
8//!   2. in_proj GEMV into one row `[z (d_inner) | xBC (d_xbc) | dt (num_heads)]`
9//!   3. causal conv1d update on xBC, with bias and a fused SiLU
10//!   4. xBC split into x, B, C
11//!   5. Mamba-2 SSM decode (state update and output y)
12//!   6. gated RMS norm of y with gate z, in the order of the resolved
13//!      `gated_rms_norm` kernel: the gb10 Nemotron trees' kernels gate first,
14//!      `y * silu(z)`, then normalise each `d_inner / n_groups` slice
15//!   7. out_proj GEMV
16//!   8. residual add into `hidden`
17//!
18//! Prefill runs the same steps over all tokens (`prefill.rs`, projection arms
19//! in `prefill_proj.rs`).
20//!
21//! Owner: model-arch (Nemotron-H).
22//! Invariants: none beyond the types.
23
24use anyhow::Result;
25use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
26use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
27
28use metrale_model_layers::weight_map::{
29    DenseWeight, Fp8Weight, NemotronSsmWeights, QuantizedWeight,
30};
31
32mod prefill;
33mod prefill_proj;
34mod trait_impl;
35
36#[allow(dead_code)]
37pub struct NemotronMamba2Layer {
38    input_norm: DenseWeight,
39    ssm: NemotronSsmWeights,
40    // 2026-09-25: The checkpoint's own block-scaled FP8 projections. Without
41    // them an FP8 checkpoint's projections are dequantized to BF16 and
42    // requantized to NVFP4 at load.
43    in_proj_fp8: Option<Fp8Weight>,
44    out_proj_fp8: Option<Fp8Weight>,
45    // 2026-09-25: Whether prefill may use the native FP8 weights above. False
46    // in the `METRALE_NEMOTRON_NATIVE_FP8_SSM=decode` mode, where the native
47    // weights are installed for decode only and the requantized NVFP4 copies
48    // are still built and used by prefill. Prefill keys off this flag, not off
49    // `in_proj_fp8.is_some()`.
50    native_fp8_prefill: bool,
51    // 2026-09-25: Transposed NVFP4 projections for the `w4a16_gemm_t` prefill arms.
52    in_proj_t: Option<QuantizedWeight>,
53    out_proj_t: Option<QuantizedWeight>,
54    // 2026-09-25: Pre-dequantized FP8 E4M3 copies of the two SSM projections,
55    // [N, K], for the prefill `fp8_gemm_t_k` / `fp8_fp8_gemm_t_k` arms.
56    in_proj_pd_fp8: Option<DevicePtr>,
57    out_proj_pd_fp8: Option<DevicePtr>,
58    // 2026-09-25: The checkpoint's own BF16 projections, for a layer the
59    // checkpoint left unquantized; used instead of requantizing them to NVFP4.
60    in_proj_bf16: Option<DenseWeight>,
61    out_proj_bf16: Option<DenseWeight>,
62    rms_norm_residual_k: KernelHandle,
63    w4a16_gemv_k: KernelHandle,
64    /// 2026-09-25: Single-warp `w4a16_gemv_sw`; `KernelHandle(0)` when the
65    /// kernel is missing, and then decode uses the base GEMV.
66    w4a16_gemv_sw_k: KernelHandle,
67    w8a16_gemv_k: KernelHandle,
68    conv1d_update_k: KernelHandle,
69    mamba2_ssm_k: KernelHandle,
70    gated_rms_norm_k: KernelHandle,
71    residual_add_k: KernelHandle,
72    w4a16_gemm_k: KernelHandle,
73    // 2026-09-25: Native FP8 block-scaled prefill GEMM (paired with in_proj_fp8/out_proj_fp8).
74    w8a16_gemm_k: KernelHandle,
75    w8a16_gemm_pipelined_k: KernelHandle,
76    w4a16_gemm_t_k: KernelHandle,
77    w4a16_gemm_t_m128_k: KernelHandle,
78    fp8_gemm_t_k: KernelHandle,
79    fp8_fp8_gemm_t_k: KernelHandle,
80    // 2026-09-25: Native-BF16 projection kernels (paired with in_proj_bf16/out_proj_bf16).
81    dense_gemm_bf16_k: KernelHandle,
82    dense_gemv_bf16_k: KernelHandle,
83    bf16_to_fp8_k: KernelHandle,
84    w4a4_gemm_k: KernelHandle,
85    quantize_nvfp4_k: KernelHandle,
86    conv1d_prefill_k: KernelHandle,
87    conv1d_prefill_tp_k: KernelHandle,
88    mamba2_ssm_prefill_k: KernelHandle,
89    mamba2_ssm_prefill_persistent_k: KernelHandle,
90    // 2026-09-25: SSD chunked prefill scan (`ops::ssm_ssd`, chunks of `SSD_L` tokens).
91    ssd_cumsum_k: KernelHandle,
92    ssd_bmm_k: KernelHandle,
93    ssd_scan_k: KernelHandle,
94    d_inner: usize,
95    d_xbc: usize,
96    in_proj_size: usize,
97    num_heads: usize,
98    head_dim: usize,
99    state_size: usize,
100    n_groups: usize,
101    d_conv: usize,
102    h_state_bytes: usize,
103    conv_state_bytes: usize,
104    layer_idx: usize,
105}
106
107impl NemotronMamba2Layer {
108    pub fn new(
109        input_norm: DenseWeight,
110        ssm: NemotronSsmWeights,
111        config: &metrale_config::ModelConfig,
112        gpu: &dyn GpuBackend,
113        layer_idx: usize,
114    ) -> Result<Self> {
115        let num_heads = config.mamba_num_heads;
116        let head_dim = config.mamba_head_dim;
117        let state_size = config.ssm_state_size;
118        let n_groups = config.n_groups;
119        let d_conv = config.linear_conv_kernel_dim;
120        let d_inner = config.mamba2_d_inner();
121        let d_xbc = config.mamba2_d_xbc();
122        let in_proj_size = config.mamba2_in_proj_size();
123
124        Ok(Self {
125            input_norm,
126            ssm,
127            in_proj_fp8: None,
128            out_proj_fp8: None,
129            native_fp8_prefill: false,
130            in_proj_t: None,
131            out_proj_t: None,
132            in_proj_pd_fp8: None,
133            out_proj_pd_fp8: None,
134            in_proj_bf16: None,
135            out_proj_bf16: None,
136            rms_norm_residual_k: gpu.kernel("norm", "rms_norm_residual")?,
137            w4a16_gemv_k: gpu.kernel("w4a16_gemv", "w4a16_gemv")?,
138            w4a16_gemv_sw_k: metrale_model_layers::layers::try_kernel(
139                gpu,
140                "w4a16_gemv",
141                "w4a16_gemv_sw",
142            ),
143            w8a16_gemv_k: metrale_model_layers::layers::try_kernel(gpu, "w8a16_gemv", "w8a16_gemv"),
144            conv1d_update_k: gpu.kernel("causal_conv1d", "causal_conv1d_update")?,
145            mamba2_ssm_k: gpu.kernel("mamba2_ssm", "mamba2_ssm_decode")?,
146            gated_rms_norm_k: gpu.kernel("norm", "gated_rms_norm")?,
147            residual_add_k: gpu.kernel("residual_add", "bf16_residual_add")?,
148            w4a16_gemm_k: gpu.kernel("w4a16", "w4a16_gemm")?,
149            w8a16_gemm_k: metrale_model_layers::layers::try_kernel(gpu, "w8a16_gemm", "w8a16_gemm"),
150            w8a16_gemm_pipelined_k: metrale_model_layers::layers::try_kernel(
151                gpu,
152                "w8a16_gemm_pipelined",
153                "w8a16_gemm_pipelined",
154            ),
155            w4a16_gemm_t_k: metrale_model_layers::layers::try_kernel(gpu, "w4a16", "w4a16_gemm_t"),
156            w4a16_gemm_t_m128_k: metrale_model_layers::layers::try_kernel(
157                gpu,
158                "w4a16",
159                "w4a16_gemm_t_m128",
160            ),
161            fp8_gemm_t_k: metrale_model_layers::layers::try_kernel(
162                gpu,
163                "w4a16",
164                "fp8_gemm_t_m128_mfast",
165            ),
166            fp8_fp8_gemm_t_k: metrale_model_layers::layers::try_kernel(
167                gpu,
168                "w4a16",
169                "fp8_fp8_gemm_t_m128_mfast",
170            ),
171            dense_gemm_bf16_k: metrale_model_layers::layers::try_kernel(
172                gpu,
173                "gemm",
174                "dense_gemm_bf16_pipelined",
175            ),
176            dense_gemv_bf16_k: metrale_model_layers::layers::try_kernel(
177                gpu,
178                "gemv",
179                "dense_gemv_bf16",
180            ),
181            bf16_to_fp8_k: metrale_model_layers::layers::try_kernel(gpu, "w4a16", "bf16_to_fp8"),
182            w4a4_gemm_k: metrale_model_layers::layers::try_kernel(gpu, "w4a4", "w4a4_gemm_mfast"),
183            quantize_nvfp4_k: metrale_model_layers::layers::try_kernel(
184                gpu,
185                "quantize_nvfp4",
186                "quantize_bf16_to_nvfp4",
187            ),
188            conv1d_prefill_k: gpu.kernel("causal_conv1d", "causal_conv1d_update_prefill")?,
189            conv1d_prefill_tp_k: metrale_model_layers::layers::try_kernel(
190                gpu,
191                "causal_conv1d",
192                "causal_conv1d_update_prefill_tp",
193            ),
194            mamba2_ssm_prefill_k: gpu.kernel("mamba2_ssm", "mamba2_ssm_prefill")?,
195            ssd_cumsum_k: metrale_model_layers::layers::try_kernel(
196                gpu,
197                "mamba2_ssd_chunk",
198                "mamba2_ssd_cumsum",
199            ),
200            ssd_bmm_k: metrale_model_layers::layers::try_kernel(
201                gpu,
202                "mamba2_ssd_chunk",
203                "mamba2_ssd_bmm",
204            ),
205            ssd_scan_k: metrale_model_layers::layers::try_kernel(
206                gpu,
207                "mamba2_ssd_chunk",
208                "mamba2_ssd_scan",
209            ),
210            mamba2_ssm_prefill_persistent_k: metrale_model_layers::layers::try_kernel(
211                gpu,
212                "mamba2_ssm",
213                "mamba2_ssm_prefill_persistent",
214            ),
215            d_inner,
216            d_xbc,
217            in_proj_size,
218            num_heads,
219            head_dim,
220            state_size,
221            n_groups,
222            d_conv,
223            h_state_bytes: num_heads * head_dim * state_size * 4,
224            conv_state_bytes: d_xbc * d_conv * 4,
225            layer_idx,
226        })
227    }
228
229    /// 2026-09-25: Install the checkpoint's native FP8 projections. Decode then
230    /// uses `w8a16_gemv`, and when `prefill` is true the prefill GEMMs use
231    /// `w8a16_gemm_pipelined` (or `w8a16_gemm`) ahead of the NVFP4 / W4A4 arms.
232    ///
233    /// Inputs must be tagged `WeightQuantFormat::Fp8BlockScaled`: the w8a16
234    /// kernels index `block_scale[n_block * k_blocks + k_block]` with 128-wide
235    /// blocks, so a per-row `[N]` scale or a scalar `weight_scale` would be read
236    /// past its end. A wrong tag panics (`expect`); a missing `w8a16_gemv`, or
237    /// with `prefill` a missing prefill GEMM, is an error here at load.
238    ///
239    /// When `prefill` is false (`METRALE_NEMOTRON_NATIVE_FP8_SSM=decode`) only
240    /// `w8a16_gemv` reads these weights and prefill uses the NVFP4 /
241    /// pre-dequantized copies, which the loader builds in that mode.
242    pub fn set_fp8_weights(
243        &mut self,
244        in_proj: Option<Fp8Weight>,
245        out_proj: Option<Fp8Weight>,
246        prefill: bool,
247    ) -> Result<()> {
248        use metrale_model_layers::weight_map::WeightQuantFormat;
249        if let Some(ref w) = in_proj {
250            w.scale_format.expect(
251                WeightQuantFormat::Fp8BlockScaled,
252                "nemotron mamba2 in_proj (w8a16 expects [ceil(N/128),ceil(K/128)] FP32 block scales)",
253            );
254        }
255        if let Some(ref w) = out_proj {
256            w.scale_format.expect(
257                WeightQuantFormat::Fp8BlockScaled,
258                "nemotron mamba2 out_proj (w8a16 expects [ceil(N/128),ceil(K/128)] FP32 block scales)",
259            );
260        }
261        anyhow::ensure!(
262            self.w8a16_gemv_k.0 != 0,
263            "native FP8 SSM requires the w8a16_gemv kernel (decode)"
264        );
265        anyhow::ensure!(
266            !prefill || self.w8a16_gemm_pipelined_k.0 != 0 || self.w8a16_gemm_k.0 != 0,
267            "native FP8 SSM requires w8a16_gemm[_pipelined] (prefill)"
268        );
269        self.in_proj_fp8 = in_proj;
270        self.out_proj_fp8 = out_proj;
271        self.native_fp8_prefill = prefill;
272        Ok(())
273    }
274
275    /// 2026-09-25: The layer's Mamba-2 mixer weights.
276    pub fn ssm_weights(&self) -> &NemotronSsmWeights {
277        &self.ssm
278    }
279
280    /// 2026-09-25: Install transposed NVFP4 projections: prefill then uses
281    /// `w4a16_gemm_t` (or, above 128 tokens and when it resolved,
282    /// `w4a16_gemm_t_m128`) instead of the plain `w4a16_gemm`, unless an
283    /// earlier arm in `prefill_proj.rs` applies.
284    pub fn set_prefill_weights(
285        &mut self,
286        in_proj_t: Option<QuantizedWeight>,
287        out_proj_t: Option<QuantizedWeight>,
288    ) {
289        self.in_proj_t = in_proj_t;
290        self.out_proj_t = out_proj_t;
291    }
292
293    /// 2026-09-25: Install the checkpoint's own BF16 projections; decode and
294    /// prefill then use the dense BF16 kernels ahead of every other arm. Valid
295    /// only when both projections are BF16 in the checkpoint and the dense
296    /// kernels resolved; the caller checks that (`bf16_native_ready`).
297    pub fn set_bf16_weights(&mut self, in_proj: DenseWeight, out_proj: DenseWeight) {
298        self.in_proj_bf16 = Some(in_proj);
299        self.out_proj_bf16 = Some(out_proj);
300    }
301
302    /// 2026-09-25: Whether this layer can run natively BF16 (weights installed
303    /// and both dense kernels present).
304    pub fn bf16_native_ready(&self) -> bool {
305        self.in_proj_bf16.is_some()
306            && self.out_proj_bf16.is_some()
307            && self.dense_gemm_bf16_k.0 != 0
308            && self.dense_gemv_bf16_k.0 != 0
309    }
310
311    /// 2026-09-25: Install pre-dequantized FP8 E4M3 copies of in_proj /
312    /// out_proj for prefill; their arms come after the native and W4A4 arms
313    /// and before the transposed NVFP4 ones (`prefill_proj.rs`).
314    pub fn set_fp8_prefill_weights(&mut self, in_proj: DevicePtr, out_proj: DevicePtr) {
315        self.in_proj_pd_fp8 = Some(in_proj);
316        self.out_proj_pd_fp8 = Some(out_proj);
317    }
318
319    /// 2026-09-25: Conv1d update with the layer's conv1d bias.
320    ///
321    /// Kernel: `causal_conv1d_update(conv_state, input, weight, bias, output,
322    ///          batch, dim, d_conv)`
323    fn conv1d_update_biased(
324        &self,
325        gpu: &dyn GpuBackend,
326        conv_state: DevicePtr,
327        input: DevicePtr,
328        output: DevicePtr,
329        d_inner: u32,
330        d_conv: u32,
331        batch_size: u32,
332        stream: u64,
333    ) -> Result<()> {
334        KernelLaunch::new(gpu, self.conv1d_update_k)
335            .grid([div_ceil(d_inner, 256), batch_size, 1])
336            .block([256, 1, 1])
337            .arg_ptr(conv_state)
338            .arg_ptr(input)
339            .arg_ptr(self.ssm.conv1d_weight.weight)
340            .arg_ptr(self.ssm.conv1d_bias.weight)
341            .arg_ptr(output)
342            .arg_u32(batch_size)
343            .arg_u32(d_inner)
344            .arg_u32(d_conv)
345            .launch(stream)
346    }
347
348    /// 2026-09-25: Launch the Mamba-2 SSM decode kernel.
349    ///
350    /// Grid: (num_heads, batch, 1)  Block: (state_size, 1, 1)
351    #[allow(clippy::too_many_arguments)]
352    fn ssm_decode(
353        &self,
354        gpu: &dyn GpuBackend,
355        h_state: DevicePtr,
356        x: DevicePtr,
357        b_proj: DevicePtr,
358        c_proj: DevicePtr,
359        dt_raw: DevicePtr,
360        output: DevicePtr,
361        batch_size: u32,
362        stream: u64,
363    ) -> Result<()> {
364        KernelLaunch::new(gpu, self.mamba2_ssm_k)
365            .grid([self.num_heads as u32, batch_size, 1])
366            .block([self.state_size as u32, 1, 1])
367            .arg_ptr(h_state)
368            .arg_ptr(x)
369            .arg_ptr(b_proj)
370            .arg_ptr(c_proj)
371            .arg_ptr(dt_raw)
372            .arg_ptr(self.ssm.a_log.weight)
373            .arg_ptr(self.ssm.d_param.weight)
374            .arg_ptr(self.ssm.dt_bias.weight)
375            .arg_ptr(output)
376            .arg_u32(batch_size)
377            .arg_u32(self.num_heads as u32)
378            .arg_u32(self.head_dim as u32)
379            .arg_u32(self.state_size as u32)
380            .arg_u32(self.n_groups as u32)
381            .arg_f32(1e-9)
382            .arg_f32(1e9)
383            .launch(stream)
384    }
385}
metrale_model_arch/nemotron_mamba2.rs

metrale_model_arch/
nemotron_mamba2.rs
