spark_model/layers/
nemotron_mamba2.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Nemotron-H Mamba-2 SSM layer implementing TransformerLayer.
4//!
5//! Standalone SSM layer (no FFN component). Forward pass:
6//!   1. RMS norm (standard weight*x scaling)
7//!   2. in_proj GEMV → [z, xBC, dt]
8//!   3. Conv1d update on xBC (WITH bias, fused SiLU)
9//!   4. Split xBC_out → x, B, C
10//!   5. Mamba-2 SSM decode (state update + output)
11//!   6. Gated RMS norm: rms_norm(y, ssm_norm) * silu(z)
12//!   7. out_proj GEMV
13//!   8. Residual add
14
15use anyhow::Result;
16use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
18
19use crate::weight_map::{DenseWeight, Fp8Weight, NemotronSsmWeights, QuantizedWeight};
20
21mod prefill;
22mod prefill_proj;
23mod trait_impl;
24
25#[allow(dead_code)]
26pub struct NemotronMamba2Layer {
27    input_norm: DenseWeight,
28    ssm: NemotronSsmWeights,
29    // FP8 native weights (skip double-quantization FP8→BF16→NVFP4)
30    in_proj_fp8: Option<Fp8Weight>,
31    out_proj_fp8: Option<Fp8Weight>,
32    // Whether PREFILL may use the native FP8 weights above. False in the
33    // `ATLAS_NEMOTRON_NATIVE_FP8_SSM=decode` bisect mode, where the native
34    // weights are installed for decode only and the legacy NVFP4 copies are
35    // still built and used by prefill. Prefill must key off this flag, not off
36    // `in_proj_fp8.is_some()`.
37    native_fp8_prefill: bool,
38    // Transposed NVFP4 weights for fast prefill GEMM (FP8 MMA, N128, cp.async)
39    in_proj_t: Option<QuantizedWeight>,
40    out_proj_t: Option<QuantizedWeight>,
41    // Pre-dequantized FP8 E4M3 copies of the two SSM projections, [N, K].
42    // Consumed by `fp8_gemm_t`, which has NO dequant phase at all.
43    in_proj_pd_fp8: Option<DevicePtr>,
44    out_proj_pd_fp8: Option<DevicePtr>,
45    // NATIVE BF16 projections. A mixed-precision checkpoint can leave some
46    // Mamba layers unquantized (Nano-30B: 6 of 23 in_proj/out_proj are BF16
47    // while 17 are NVFP4). Requantizing those to NVFP4 under ONE global scale
48    // is what the FP8 arm above already documents as destroying context
49    // retrieval; keeping them BF16 avoids inventing a quantization the
50    // checkpoint never asked for.
51    in_proj_bf16: Option<DenseWeight>,
52    out_proj_bf16: Option<DenseWeight>,
53    // Kernel handles — decode
54    rms_norm_residual_k: KernelHandle,
55    w4a16_gemv_k: KernelHandle,
56    /// Single-warp `w4a16_gemv_sw`. `KernelHandle(0)` on miss → base GEMV.
57    w4a16_gemv_sw_k: KernelHandle,
58    w8a16_gemv_k: KernelHandle,
59    conv1d_update_k: KernelHandle,
60    mamba2_ssm_k: KernelHandle,
61    gated_rms_norm_k: KernelHandle,
62    residual_add_k: KernelHandle,
63    // Kernel handles — prefill (GEMM + batched kernels)
64    w4a16_gemm_k: KernelHandle,
65    // Native FP8 block-scaled prefill GEMM (paired with in_proj_fp8/out_proj_fp8).
66    w8a16_gemm_k: KernelHandle,
67    w8a16_gemm_pipelined_k: KernelHandle,
68    w4a16_gemm_t_k: KernelHandle,
69    w4a16_gemm_t_m128_k: KernelHandle,
70    fp8_gemm_t_k: KernelHandle,
71    fp8_fp8_gemm_t_k: KernelHandle,
72    // Native-BF16 projection kernels (paired with in_proj_bf16/out_proj_bf16).
73    dense_gemm_bf16_k: KernelHandle,
74    dense_gemv_bf16_k: KernelHandle,
75    bf16_to_fp8_k: KernelHandle,
76    w4a4_gemm_k: KernelHandle,
77    quantize_nvfp4_k: KernelHandle,
78    conv1d_prefill_k: KernelHandle,
79    conv1d_prefill_tp_k: KernelHandle,
80    mamba2_ssm_prefill_k: KernelHandle,
81    mamba2_ssm_prefill_persistent_k: KernelHandle,
82    // SSD chunked prefill scan (tensor-core; ceil(T/64) serial links instead of T).
83    ssd_cumsum_k: KernelHandle,
84    ssd_bmm_k: KernelHandle,
85    ssd_scan_k: KernelHandle,
86    // Pre-computed dimensions
87    d_inner: usize,
88    d_xbc: usize,
89    in_proj_size: usize,
90    num_heads: usize,
91    head_dim: usize,
92    state_size: usize,
93    n_groups: usize,
94    d_conv: usize,
95    h_state_bytes: usize,
96    conv_state_bytes: usize,
97    layer_idx: usize,
98}
99
100impl NemotronMamba2Layer {
101    pub fn new(
102        input_norm: DenseWeight,
103        ssm: NemotronSsmWeights,
104        config: &atlas_core::config::ModelConfig,
105        gpu: &dyn GpuBackend,
106        layer_idx: usize,
107    ) -> Result<Self> {
108        let num_heads = config.mamba_num_heads;
109        let head_dim = config.mamba_head_dim;
110        let state_size = config.ssm_state_size;
111        let n_groups = config.n_groups;
112        let d_conv = config.linear_conv_kernel_dim;
113        let d_inner = config.mamba2_d_inner();
114        let d_xbc = config.mamba2_d_xbc();
115        let in_proj_size = config.mamba2_in_proj_size();
116
117        Ok(Self {
118            input_norm,
119            ssm,
120            in_proj_fp8: None,
121            out_proj_fp8: None,
122            native_fp8_prefill: false,
123            in_proj_t: None,
124            out_proj_t: None,
125            in_proj_pd_fp8: None,
126            out_proj_pd_fp8: None,
127            in_proj_bf16: None,
128            out_proj_bf16: None,
129            rms_norm_residual_k: gpu.kernel("norm", "rms_norm_residual")?,
130            w4a16_gemv_k: gpu.kernel("w4a16_gemv", "w4a16_gemv")?,
131            w4a16_gemv_sw_k: super::try_kernel(gpu, "w4a16_gemv", "w4a16_gemv_sw"),
132            w8a16_gemv_k: super::try_kernel(gpu, "w8a16_gemv", "w8a16_gemv"),
133            conv1d_update_k: gpu.kernel("causal_conv1d", "causal_conv1d_update")?,
134            mamba2_ssm_k: gpu.kernel("mamba2_ssm", "mamba2_ssm_decode")?,
135            gated_rms_norm_k: gpu.kernel("norm", "gated_rms_norm")?,
136            residual_add_k: gpu.kernel("residual_add", "bf16_residual_add")?,
137            w4a16_gemm_k: gpu.kernel("w4a16", "w4a16_gemm")?,
138            w8a16_gemm_k: super::try_kernel(gpu, "w8a16_gemm", "w8a16_gemm"),
139            w8a16_gemm_pipelined_k: super::try_kernel(
140                gpu,
141                "w8a16_gemm_pipelined",
142                "w8a16_gemm_pipelined",
143            ),
144            w4a16_gemm_t_k: super::try_kernel(gpu, "w4a16", "w4a16_gemm_t"),
145            w4a16_gemm_t_m128_k: super::try_kernel(gpu, "w4a16", "w4a16_gemm_t_m128"),
146            fp8_gemm_t_k: super::try_kernel(gpu, "w4a16", "fp8_gemm_t_m128_mfast"),
147            fp8_fp8_gemm_t_k: super::try_kernel(gpu, "w4a16", "fp8_fp8_gemm_t_m128_mfast"),
148            dense_gemm_bf16_k: super::try_kernel(gpu, "gemm", "dense_gemm_bf16_pipelined"),
149            dense_gemv_bf16_k: super::try_kernel(gpu, "gemv", "dense_gemv_bf16"),
150            bf16_to_fp8_k: super::try_kernel(gpu, "w4a16", "bf16_to_fp8"),
151            w4a4_gemm_k: super::try_kernel(gpu, "w4a4", "w4a4_gemm_mfast"),
152            quantize_nvfp4_k: super::try_kernel(gpu, "quantize_nvfp4", "quantize_bf16_to_nvfp4"),
153            conv1d_prefill_k: gpu.kernel("causal_conv1d", "causal_conv1d_update_prefill")?,
154            conv1d_prefill_tp_k: super::try_kernel(
155                gpu,
156                "causal_conv1d",
157                "causal_conv1d_update_prefill_tp",
158            ),
159            mamba2_ssm_prefill_k: gpu.kernel("mamba2_ssm", "mamba2_ssm_prefill")?,
160            ssd_cumsum_k: super::try_kernel(gpu, "mamba2_ssd_chunk", "mamba2_ssd_cumsum"),
161            ssd_bmm_k: super::try_kernel(gpu, "mamba2_ssd_chunk", "mamba2_ssd_bmm"),
162            ssd_scan_k: super::try_kernel(gpu, "mamba2_ssd_chunk", "mamba2_ssd_scan"),
163            mamba2_ssm_prefill_persistent_k: super::try_kernel(
164                gpu,
165                "mamba2_ssm",
166                "mamba2_ssm_prefill_persistent",
167            ),
168            d_inner,
169            d_xbc,
170            in_proj_size,
171            num_heads,
172            head_dim,
173            state_size,
174            n_groups,
175            d_conv,
176            h_state_bytes: num_heads * head_dim * state_size * 4, // FP32
177            conv_state_bytes: d_xbc * d_conv * 4,                 // FP32
178            layer_idx,
179        })
180    }
181
182    /// Set native FP8 weights to skip double-quantization (FP8→BF16→NVFP4).
183    /// When set, decode uses `w8a16_gemv` and prefill uses `w8a16_gemm` /
184    /// `w8a16_gemm_pipelined` instead of the NVFP4/W4A4 arms.
185    ///
186    /// Inputs MUST be tagged `WeightQuantFormat::Fp8BlockScaled`: every w8a16
187    /// kernel indexes `block_scale[n/128 * k_blocks + k/128]`, so a per-row `[N]`
188    /// scale — or the checkpoint's raw 4-byte scalar `weight_scale` — reads far
189    /// past the end of its allocation (illegal address, not wrong numbers). The
190    /// kernel-handle checks are the same contract: once the FP8 weights are
191    /// installed the NVFP4 fallbacks are NULL, so a missing kernel must fail
192    /// here at load, not deref NULL on the first token.
193    ///
194    /// `prefill` selects whether the prefill GEMMs may use these weights. When
195    /// false (`ATLAS_NEMOTRON_NATIVE_FP8_SSM=decode`) only `w8a16_gemv` reads
196    /// them and prefill stays on the legacy NVFP4 / pre-dequantized copies,
197    /// which the loader still builds in that mode.
198    pub fn set_fp8_weights(
199        &mut self,
200        in_proj: Option<Fp8Weight>,
201        out_proj: Option<Fp8Weight>,
202        prefill: bool,
203    ) -> Result<()> {
204        use crate::weight_map::WeightQuantFormat;
205        if let Some(ref w) = in_proj {
206            w.scale_format.expect(
207                WeightQuantFormat::Fp8BlockScaled,
208                "nemotron mamba2 in_proj (w8a16 expects [ceil(N/128),ceil(K/128)] FP32 block scales)",
209            );
210        }
211        if let Some(ref w) = out_proj {
212            w.scale_format.expect(
213                WeightQuantFormat::Fp8BlockScaled,
214                "nemotron mamba2 out_proj (w8a16 expects [ceil(N/128),ceil(K/128)] FP32 block scales)",
215            );
216        }
217        anyhow::ensure!(
218            self.w8a16_gemv_k.0 != 0,
219            "native FP8 SSM requires the w8a16_gemv kernel (decode)"
220        );
221        anyhow::ensure!(
222            !prefill || self.w8a16_gemm_pipelined_k.0 != 0 || self.w8a16_gemm_k.0 != 0,
223            "native FP8 SSM requires w8a16_gemm[_pipelined] (prefill)"
224        );
225        self.in_proj_fp8 = in_proj;
226        self.out_proj_fp8 = out_proj;
227        self.native_fp8_prefill = prefill;
228        Ok(())
229    }
230
231    /// Access SSM weights (needed by weight loader for transpose).
232    pub fn ssm_weights(&self) -> &NemotronSsmWeights {
233        &self.ssm
234    }
235
236    /// Set transposed NVFP4 weights for fast prefill GEMM (FP8 MMA, N128, cp.async).
237    /// Switches prefill from w4a16_gemm (M64,N64,K16 BF16) to w4a16_gemm_t
238    /// (M64,N128,K32 FP8 MMA) — est. 3-4x TTFT improvement for SSM layers.
239    pub fn set_prefill_weights(
240        &mut self,
241        in_proj_t: Option<QuantizedWeight>,
242        out_proj_t: Option<QuantizedWeight>,
243    ) {
244        self.in_proj_t = in_proj_t;
245        self.out_proj_t = out_proj_t;
246    }
247
248    /// Set pre-dequantized FP8 E4M3 copies of in_proj/out_proj for prefill.
249    ///
250    /// `w4a16_gemm_t_m128` dequantizes its NVFP4 B tile from FP4 to FP8 in
251    /// shared memory on every K step, and that work is redone by every M-block:
252    /// the cost is N*K*(M/M_TILE), so a 1k-token prefill pays for it 8x over.
253    /// Measured on Puzzle: ablating just that dequant ALU cut a 1k prefill from
254    /// 557 ms to 424 ms. Converting the weights once at load time removes it
255    /// entirely and lets prefill use `fp8_gemm_t`, which has no dequant phase.
256    /// Install the checkpoint's own BF16 projections, bypassing the NVFP4
257    /// requant entirely. Only valid when BOTH projections are BF16 in the
258    /// checkpoint and the dense kernels resolved; the caller checks that.
259    pub fn set_bf16_weights(&mut self, in_proj: DenseWeight, out_proj: DenseWeight) {
260        self.in_proj_bf16 = Some(in_proj);
261        self.out_proj_bf16 = Some(out_proj);
262    }
263
264    /// Whether this layer can run natively BF16 (weights installed AND both
265    /// dense kernels present).
266    pub fn bf16_native_ready(&self) -> bool {
267        self.in_proj_bf16.is_some()
268            && self.out_proj_bf16.is_some()
269            && self.dense_gemm_bf16_k.0 != 0
270            && self.dense_gemv_bf16_k.0 != 0
271    }
272
273    pub fn set_fp8_prefill_weights(&mut self, in_proj: DevicePtr, out_proj: DevicePtr) {
274        self.in_proj_pd_fp8 = Some(in_proj);
275        self.out_proj_pd_fp8 = Some(out_proj);
276    }
277
278    /// Conv1d update with bias (Nemotron conv1d has learned bias, unlike Qwen3).
279    ///
280    /// Kernel: `causal_conv1d_update(conv_state, input, weight, bias, output,
281    ///          batch, dim, d_conv)`
282    fn conv1d_update_biased(
283        &self,
284        gpu: &dyn GpuBackend,
285        conv_state: DevicePtr,
286        input: DevicePtr,
287        output: DevicePtr,
288        d_inner: u32,
289        d_conv: u32,
290        batch_size: u32,
291        stream: u64,
292    ) -> Result<()> {
293        KernelLaunch::new(gpu, self.conv1d_update_k)
294            .grid([div_ceil(d_inner, 256), batch_size, 1])
295            .block([256, 1, 1])
296            .arg_ptr(conv_state)
297            .arg_ptr(input)
298            .arg_ptr(self.ssm.conv1d_weight.weight)
299            .arg_ptr(self.ssm.conv1d_bias.weight)
300            .arg_ptr(output)
301            .arg_u32(batch_size)
302            .arg_u32(d_inner)
303            .arg_u32(d_conv)
304            .launch(stream)
305    }
306
307    /// Launch Mamba-2 SSM decode kernel.
308    ///
309    /// Grid: (num_heads, batch, 1)  Block: (state_size, 1, 1)
310    #[allow(clippy::too_many_arguments)]
311    fn ssm_decode(
312        &self,
313        gpu: &dyn GpuBackend,
314        h_state: DevicePtr,
315        x: DevicePtr,
316        b_proj: DevicePtr,
317        c_proj: DevicePtr,
318        dt_raw: DevicePtr,
319        output: DevicePtr,
320        batch_size: u32,
321        stream: u64,
322    ) -> Result<()> {
323        KernelLaunch::new(gpu, self.mamba2_ssm_k)
324            .grid([self.num_heads as u32, batch_size, 1])
325            .block([self.state_size as u32, 1, 1])
326            .arg_ptr(h_state)
327            .arg_ptr(x)
328            .arg_ptr(b_proj)
329            .arg_ptr(c_proj)
330            .arg_ptr(dt_raw)
331            .arg_ptr(self.ssm.a_log.weight)
332            .arg_ptr(self.ssm.d_param.weight)
333            .arg_ptr(self.ssm.dt_bias.weight)
334            .arg_ptr(output)
335            .arg_u32(batch_size)
336            .arg_u32(self.num_heads as u32)
337            .arg_u32(self.head_dim as u32)
338            .arg_u32(self.state_size as u32)
339            .arg_u32(self.n_groups as u32)
340            .arg_f32(1e-9) // dt_min (no effective clamp — reference uses no clamping)
341            .arg_f32(1e9) // dt_max (no effective clamp — reference uses no clamping)
342            .launch(stream)
343    }
344}
spark_model/layers/nemotron_mamba2.rs

spark_model/layers/
nemotron_mamba2.rs
