metrale_model_arch/nemotron_mamba2/
trait_impl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `impl TransformerLayer for NemotronMamba2Layer`: the decode
4//! steps listed in the parent module, prefill (`prefill.rs`) and the per-layer
5//! SSM state.
6//!
7//! Owner: model-arch (Nemotron-H).
8//! Invariants:
9//! - `alloc_state` returns zeroed FP32 `h_state` and `conv_state` buffers.
10
11use anyhow::Result;
12use metrale_cache::kv_cache::PagedKvCache;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
14
15use super::NemotronMamba2Layer;
16use metrale_model_layers::layer::{ForwardContext, LayerState, SsmLayerState, TransformerLayer};
17use metrale_model_layers::layer::{
18    LayerAuxState, LayerCapabilities, LayerGraphHooks, LayerSplitPrefill, LayerWeightSetup,
19    LayerWriteOnAccept,
20};
21use metrale_model_layers::layers::ops;
22
23impl TransformerLayer for NemotronMamba2Layer {
24    fn decode(
25        &self,
26        hidden: DevicePtr,
27        residual: DevicePtr,
28        state: &mut dyn LayerState,
29        _kv_cache: &mut metrale_cache::kv_cache::PagedKvCache,
30        _seq_len: usize,
31        _block_table: &mut Vec<u32>,
32        _disk_block_ids: &mut Vec<u32>,
33        _disk_last_offloaded_per_layer: &mut Vec<u32>,
34        ctx: &ForwardContext,
35        stream: u64,
36    ) -> Result<()> {
37        let h = ctx.config.hidden_size;
38        let eps = ctx.config.rms_norm_eps as f32;
39
40        let ssm_state = state
41            .as_any_mut()
42            .downcast_mut::<SsmLayerState>()
43            .ok_or_else(|| anyhow::anyhow!("Expected SsmLayerState"))?;
44
45        let normed = ctx.buffers.norm_output();
46        ops::rms_norm_residual(
47            ctx.gpu,
48            self.rms_norm_residual_k,
49            hidden,
50            &self.input_norm,
51            normed,
52            residual,
53            1,
54            h as u32,
55            eps,
56            stream,
57        )?;
58
59        let proj = ctx.buffers.ssm_qkvz();
60        // 2026-09-25: Projection precedence, here and for out_proj: native
61        // BF16, then native FP8 (`w8a16_gemv`), then the NVFP4 GEMV.
62        if let Some(ref w) = self.in_proj_bf16 {
63            // 2026-09-25: Native BF16: `ssm.in_proj` is NULL here (never quantized).
64            ops::dense_gemv(
65                ctx.gpu,
66                self.dense_gemv_bf16_k,
67                normed,
68                w,
69                proj,
70                self.in_proj_size as u32,
71                h as u32,
72                stream,
73            )?;
74        } else if let Some(ref fp8w) = self.in_proj_fp8 {
75            ops::w8a16_gemv(
76                ctx.gpu,
77                self.w8a16_gemv_k,
78                normed,
79                fp8w.weight,
80                fp8w.row_scale,
81                proj,
82                self.in_proj_size as u32,
83                h as u32,
84                stream,
85            )?;
86        } else {
87            ops::w4a16_decode_gemv(
88                ctx.gpu,
89                self.w4a16_gemv_k,
90                self.w4a16_gemv_sw_k,
91                ctx.levers.gemv_sw,
92                normed,
93                &self.ssm.in_proj,
94                proj,
95                self.in_proj_size as u32,
96                h as u32,
97                stream,
98            )?;
99        }
100
101        let z_ptr = proj;
102        let xbc_ptr = proj.offset(self.d_inner * 2);
103        let dt_ptr = proj.offset((self.d_inner + self.d_xbc) * 2);
104
105        let xbc_out = ctx.buffers.ssm_deinterleaved();
106        self.conv1d_update_biased(
107            ctx.gpu,
108            ssm_state.conv_state,
109            xbc_ptr,
110            xbc_out,
111            self.d_xbc as u32,
112            self.d_conv as u32,
113            1,
114            stream,
115        )?;
116
117        let x_ptr = xbc_out;
118        let gs = self.n_groups * self.state_size;
119        let b_ptr = xbc_out.offset(self.d_inner * 2);
120        let c_ptr = xbc_out.offset((self.d_inner + gs) * 2);
121
122        let y_ptr = ctx.buffers.attn_output();
123        self.ssm_decode(
124            ctx.gpu,
125            ssm_state.h_state,
126            x_ptr,
127            b_ptr,
128            c_ptr,
129            dt_ptr,
130            y_ptr,
131            1,
132            stream,
133        )?;
134
135        let gated_out = ctx.buffers.norm_output();
136        let group_size = (self.d_inner / self.n_groups) as u32;
137        ops::gated_rms_norm(
138            ctx.gpu,
139            self.gated_rms_norm_k,
140            y_ptr,
141            z_ptr,
142            &self.ssm.ssm_norm,
143            gated_out,
144            1,
145            self.d_inner as u32,
146            self.in_proj_size as u32,
147            eps,
148            group_size,
149            stream,
150        )?;
151
152        let out = ctx.buffers.qkv_output();
153        if let Some(ref w) = self.out_proj_bf16 {
154            ops::dense_gemv(
155                ctx.gpu,
156                self.dense_gemv_bf16_k,
157                gated_out,
158                w,
159                out,
160                h as u32,
161                self.d_inner as u32,
162                stream,
163            )?;
164        } else if let Some(ref fp8w) = self.out_proj_fp8 {
165            ops::w8a16_gemv(
166                ctx.gpu,
167                self.w8a16_gemv_k,
168                gated_out,
169                fp8w.weight,
170                fp8w.row_scale,
171                out,
172                h as u32,
173                self.d_inner as u32,
174                stream,
175            )?;
176        } else {
177            ops::w4a16_decode_gemv(
178                ctx.gpu,
179                self.w4a16_gemv_k,
180                self.w4a16_gemv_sw_k,
181                ctx.levers.gemv_sw,
182                gated_out,
183                &self.ssm.out_proj,
184                out,
185                h as u32,
186                self.d_inner as u32,
187                stream,
188            )?;
189        }
190
191        // 2026-09-25: `hidden` still holds the layer input (`rms_norm_residual`
192        // only reads it), so this adds the mixer output to it.
193        ops::residual_add(ctx.gpu, self.residual_add_k, hidden, out, h as u32, stream)?;
194
195        Ok(())
196    }
197
198    fn prefill(
199        &self,
200        hidden: DevicePtr,
201        residual: DevicePtr,
202        num_tokens: usize,
203        state: &mut dyn LayerState,
204        _kv_cache: &mut PagedKvCache,
205        _seq_len_start: usize,
206        _block_table: &mut Vec<u32>,
207        _disk_block_ids: &mut Vec<u32>,
208        _disk_last_offloaded_per_layer: &mut Vec<u32>,
209        _kv_write_start: usize,
210        ctx: &ForwardContext,
211        stream: u64,
212    ) -> Result<()> {
213        self.prefill_ssm(hidden, residual, num_tokens, state, ctx, stream)
214    }
215
216    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
217        let h_state = gpu.alloc(self.h_state_bytes)?;
218        gpu.memset(h_state, 0, self.h_state_bytes)?;
219        let conv_state = gpu.alloc(self.conv_state_bytes)?;
220        gpu.memset(conv_state, 0, self.conv_state_bytes)?;
221        Ok(Box::new(SsmLayerState {
222            h_state,
223            conv_state,
224            h_state_checkpoint: None,
225            conv_state_checkpoint: None,
226            h_state_intermediates: Vec::new(),
227            conv_state_intermediates: Vec::new(),
228            h_is_f16: false,
229            h_prefill_stage: None,
230            ple: None,
231        }))
232    }
233}
234
235impl LayerCapabilities for NemotronMamba2Layer {}
236impl LayerWeightSetup for NemotronMamba2Layer {}
237impl LayerWriteOnAccept for NemotronMamba2Layer {}
238impl LayerGraphHooks for NemotronMamba2Layer {}
239impl LayerAuxState for NemotronMamba2Layer {}
240impl LayerSplitPrefill for NemotronMamba2Layer {}
241impl metrale_model_layers::circuit_exec::CircuitBindings for NemotronMamba2Layer {}
metrale_model_arch/nemotron_mamba2/trait_impl.rs

metrale_model_arch/nemotron_mamba2/
trait_impl.rs
