metrale_model_layers/weight_map/
nemotron.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Nemotron-H weight structs and loaders for the Mamba-2 mixer, attention and MoE layers.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants: none beyond the types.
7
8#![allow(unused_imports)]
9
10use anyhow::{Context, Result, bail, ensure};
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12use metrale_model_weights::weights::{WeightDtype, WeightStore};
13
14use super::*;
15
16/// 2026-09-25: Nemotron-H Mamba-2 mixer weights.
17///
18/// in_proj produces `[z(d_inner), x(d_inner), B(n_groups*state), C(n_groups*state), dt(num_heads)]`.
19pub struct NemotronSsmWeights {
20    /// 2026-09-25: NVFP4, or `QuantizedWeight::null()` when the layer is not NVFP4.
21    pub in_proj: QuantizedWeight,
22    /// 2026-09-25: NVFP4, or `QuantizedWeight::null()` when the layer is not NVFP4.
23    pub out_proj: QuantizedWeight,
24    pub conv1d_weight: DenseWeight,
25    /// 2026-09-25: F32, widened from the BF16 checkpoint tensor at load.
26    pub conv1d_bias: DenseWeight,
27    /// 2026-09-25: F32, widened from the BF16 checkpoint tensor at load.
28    pub a_log: DenseWeight,
29    /// 2026-09-25: The D skip connection; F32, widened from BF16 at load.
30    pub d_param: DenseWeight,
31    /// 2026-09-25: F32, widened from the BF16 checkpoint tensor at load.
32    pub dt_bias: DenseWeight,
33    pub ssm_norm: DenseWeight,
34}
35
36/// 2026-09-25: Nemotron-H expert: `up_proj` and `down_proj`, no `gate_proj`.
37#[derive(Debug, Clone, Copy)]
38pub struct NemotronExpertWeight {
39    pub up_proj: QuantizedWeight,
40    pub down_proj: QuantizedWeight,
41}
42
43impl NemotronExpertWeight {
44    pub fn null() -> Self {
45        Self {
46            up_proj: QuantizedWeight::null(),
47            down_proj: QuantizedWeight::null(),
48        }
49    }
50}
51
52/// 2026-09-25: Nemotron-H MoE layer weights.
53pub struct NemotronMoeWeights {
54    /// 2026-09-25: Router gate, BF16 (converted when the checkpoint stores FP32).
55    pub gate: DenseWeight,
56    /// 2026-09-25: The store's `gate.e_score_correction_bias` pointer, unconverted;
57    /// `moe_topk_sigmoid` reads it as F32.
58    pub e_score_correction_bias: DenseWeight,
59    /// 2026-09-25: Routed experts as NVFP4; experts this rank does not hold are null.
60    pub experts: Vec<NemotronExpertWeight>,
61    /// 2026-09-25: Shared-expert up_proj as NVFP4; NULL when `shared_up_fp8` is `Some`.
62    pub shared_up: QuantizedWeight,
63    /// 2026-09-25: Shared-expert up_proj as the checkpoint's own FP8 bytes, in an
64    /// owned copy. `Some` when `METRALE_NEMOTRON_NATIVE_FP8_SSM` is unset, `1`,
65    /// `both` or `decode`, the tensor has `weight_scale` but no `weight_scale_2`,
66    /// and it loads.
67    pub shared_up_fp8: Option<Fp8Weight>,
68    /// 2026-09-25: Shared-expert down_proj as NVFP4; NULL when `shared_down_fp8` is `Some`.
69    pub shared_down: QuantizedWeight,
70    /// 2026-09-25: Shared-expert down_proj as native FP8, under the same condition
71    /// as `shared_up_fp8`. When its kernels are available, decode applies relu²
72    /// in place and `w8a16_gemv` to it, and launches the fused NVFP4 kernel
73    /// without its shared slot.
74    pub shared_down_fp8: Option<Fp8Weight>,
75    /// 2026-09-25: LatentMoE fc1, BF16 (FP8 is dequantized at load). `Some` when
76    /// `config.moe_latent_size > 0`.
77    pub fc1_latent_proj: Option<DenseWeight>,
78    /// 2026-09-25: LatentMoE fc2, BF16 (FP8 is dequantized at load). `Some` when
79    /// `config.moe_latent_size > 0`.
80    pub fc2_latent_proj: Option<DenseWeight>,
81}
82
83/// 2026-09-25: Format of a Mamba-2 layer's `in_proj`, detected at load.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum NemotronSsmQuant {
86    /// 2026-09-25: `weight_scale` and `weight_scale_2` present.
87    Nvfp4,
88    /// 2026-09-25: `weight_scale` without `weight_scale_2`.
89    Fp8,
90    /// 2026-09-25: No `weight_scale`.
91    Bf16,
92}
93
94/// 2026-09-25: Load a Nemotron-H Mamba-2 mixer (`{layer_prefix}.mixer`).
95///
96/// The format is read from `in_proj` alone. Unless it is NVFP4, `in_proj` and
97/// `out_proj` come back as `QuantizedWeight::null()`, and the returned
98/// `NemotronSsmQuant` tells the caller how to load them.
99pub fn load_nemotron_ssm(
100    store: &WeightStore,
101    _layer: usize,
102    gpu: &dyn GpuBackend,
103    layer_prefix: &str,
104) -> Result<(NemotronSsmWeights, NemotronSsmQuant)> {
105    let p = format!("{layer_prefix}.mixer");
106    let has_scale = store.contains(&format!("{p}.in_proj.weight_scale"));
107    let has_scale2 = store.contains(&format!("{p}.in_proj.weight_scale_2"));
108    let quant = if has_scale && has_scale2 {
109        NemotronSsmQuant::Nvfp4
110    } else if has_scale {
111        NemotronSsmQuant::Fp8
112    } else {
113        NemotronSsmQuant::Bf16
114    };
115    let in_proj = if quant == NemotronSsmQuant::Nvfp4 {
116        quantized(store, &format!("{p}.in_proj"), gpu)?
117    } else {
118        QuantizedWeight::null()
119    };
120    let out_proj = if quant == NemotronSsmQuant::Nvfp4 {
121        quantized(store, &format!("{p}.out_proj"), gpu)?
122    } else {
123        QuantizedWeight::null()
124    };
125    Ok((
126        NemotronSsmWeights {
127            in_proj,
128            out_proj,
129            conv1d_weight: dense(store, &format!("{p}.conv1d.weight"))?,
130            conv1d_bias: dense_bf16_as_f32(store, &format!("{p}.conv1d.bias"), gpu)?,
131            a_log: dense_bf16_as_f32(store, &format!("{p}.A_log"), gpu)?,
132            d_param: dense_bf16_as_f32(store, &format!("{p}.D"), gpu)?,
133            dt_bias: dense_bf16_as_f32(store, &format!("{p}.dt_bias"), gpu)?,
134            ssm_norm: dense(store, &format!("{p}.norm.weight"))?,
135        },
136        quant,
137    ))
138}
139
140/// 2026-09-25: Load a Nemotron-H attention mixer.
141///
142/// Returns `(attn, q, k, v, o_dense, is_nvfp4)`. When `q_proj` has a
143/// `weight_scale_2`, q/k/v are `Some` NVFP4 weights, `attn.o_proj` is NVFP4
144/// and the dense fields are NULL. Otherwise q/k/v/o are BF16 in `attn` and
145/// `o_dense` (an FP8 projection with `weight_scale` is dequantized), q/k/v are
146/// `None`, and `attn.o_proj` is null. The q/k norms are always NULL.
147pub fn load_nemotron_attention(
148    store: &WeightStore,
149    layer: usize,
150    gpu: &dyn GpuBackend,
151    layer_prefix: &str,
152) -> Result<(
153    AttentionWeights,
154    Option<QuantizedWeight>,
155    Option<QuantizedWeight>,
156    Option<QuantizedWeight>,
157    DenseWeight,
158    bool,
159)> {
160    let p = format!("{layer_prefix}.mixer");
161    let is_nvfp4 = store.contains(&format!("{p}.q_proj.weight_scale_2"));
162    let is_fp8 = !is_nvfp4 && store.contains(&format!("{p}.q_proj.weight_scale"));
163    let dummy = DenseWeight {
164        weight: DevicePtr::NULL,
165    };
166
167    let (q_dense, k_dense, v_dense, o_dense, o_proj, q_nvfp4, k_nvfp4, v_nvfp4) = if is_nvfp4 {
168        let q = quantized(store, &format!("{p}.q_proj"), gpu)?;
169        let k = quantized(store, &format!("{p}.k_proj"), gpu)?;
170        let v = quantized(store, &format!("{p}.v_proj"), gpu)?;
171        let o = quantized(store, &format!("{p}.o_proj"), gpu)?;
172        (dummy, dummy, dummy, dummy, o, Some(q), Some(k), Some(v))
173    } else {
174        let load_proj = |name: &str| -> Result<DenseWeight> {
175            let prefix = format!("{p}.{name}");
176            if store.contains(&format!("{prefix}.weight_scale")) {
177                dequant_fp8_to_bf16(store, &prefix, gpu)
178            } else {
179                dense(store, &format!("{prefix}.weight"))
180            }
181        };
182        let q = load_proj("q_proj")?;
183        let k = load_proj("k_proj")?;
184        let v = load_proj("v_proj")?;
185        let o = load_proj("o_proj")?;
186        if is_fp8 && layer < 2 {
187            tracing::info!("L{layer} Attention: FP8 → BF16 (runtime quantization to NVFP4)");
188        }
189        (q, k, v, o, QuantizedWeight::null(), None, None, None)
190    };
191
192    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu)?;
193    let attn = AttentionWeights {
194        q_proj: q_dense,
195        k_proj: k_dense,
196        v_proj: v_dense,
197        o_proj,
198        q_norm: dummy,
199        k_norm: dummy,
200        q_norm_full: None,
201        k_norm_full: None,
202        k_scale,
203        v_scale,
204    };
205    Ok((attn, q_nvfp4, k_nvfp4, v_nvfp4, o_dense, is_nvfp4))
206}
207
208/// 2026-09-25: Load a Nemotron-H MoE mixer.
209///
210/// Expert projections that are not NVFP4 on disk (FP8 or BF16) are quantized
211/// to NVFP4 at load, except shared-expert projections kept as native FP8.
212/// That needs `absmax_k` and `quantize_k` (it panics without them). Expert FP8
213/// dequants go into `scratch` when it is given. The LatentMoE projections stay BF16.
214pub fn load_nemotron_moe(
215    store: &WeightStore,
216    layer: usize,
217    num_experts: usize,
218    gpu: &dyn GpuBackend,
219    config: &metrale_config::ModelConfig,
220    absmax_k: Option<metrale_gpu_runtime::gpu::KernelHandle>,
221    quantize_k: Option<metrale_gpu_runtime::gpu::KernelHandle>,
222    stream: u64,
223    scratch: Option<DevicePtr>,
224    layer_prefix: &str,
225) -> Result<NemotronMoeWeights> {
226    let p = format!("{layer_prefix}.mixer");
227    let gate_name = format!("{p}.gate.weight");
228    let gate_w = store.get(&gate_name)?;
229    let gate = if gate_w.dtype == WeightDtype::FP32 {
230        dense_f32_as_bf16(store, &gate_name, gpu)?
231    } else {
232        DenseWeight { weight: gate_w.ptr }
233    };
234    let e_score_correction_bias = dense(store, &format!("{p}.gate.e_score_correction_bias"))?;
235
236    // 2026-09-25: `weight_scale_2` marks NVFP4; `weight_scale` alone marks FP8.
237    let shared_up_prefix = format!("{p}.shared_experts.up_proj");
238    let shared_up_has_s2 = store.contains(&format!("{shared_up_prefix}.weight_scale_2"));
239    let shared_up_has_s = store.contains(&format!("{shared_up_prefix}.weight_scale"));
240    // 2026-09-25: The same `METRALE_NEMOTRON_NATIVE_FP8_SSM` modes as the Mamba-2
241    // loader (`weight_loader/nemotron/ssm_layer.rs`): unset, `1`, `both` or
242    // `decode` keep an FP8 shared expert as FP8.
243    let native_fp8_mode =
244        std::env::var("METRALE_NEMOTRON_NATIVE_FP8_SSM").unwrap_or_else(|_| "1".to_string());
245    let want_native_fp8 = matches!(native_fp8_mode.as_str(), "1" | "both" | "decode");
246    let shared_up_fp8 = if want_native_fp8 && !shared_up_has_s2 && shared_up_has_s {
247        match load_fp8_block_scaled_as_fp8weight(store, &shared_up_prefix, gpu) {
248            Ok(mut w) => {
249                let bytes = (w.n as usize) * (w.k as usize);
250                let owned = gpu.alloc(bytes)?;
251                gpu.copy_d2d(w.weight, owned, bytes)?;
252                w.weight = owned;
253                Some(w)
254            }
255            Err(e) => {
256                tracing::warn!("shared_up native FP8 unavailable ({e}) — using NVFP4 requant");
257                None
258            }
259        }
260    } else {
261        None
262    };
263    let shared_up = if shared_up_fp8.is_some() {
264        QuantizedWeight::null()
265    } else if shared_up_has_s2 {
266        quantized(store, &shared_up_prefix, gpu)?
267    } else {
268        let bf16 = if shared_up_has_s {
269            if let Some(s) = scratch {
270                dequant_fp8_to_bf16_into(store, &shared_up_prefix, gpu, s)?
271            } else {
272                dequant_fp8_to_bf16(store, &shared_up_prefix, gpu)?
273            }
274        } else {
275            dense(store, &format!("{shared_up_prefix}.weight"))?
276        };
277        quantize_to_nvfp4(
278            &bf16,
279            config.shared_expert_intermediate_size,
280            config.hidden_size,
281            gpu,
282            absmax_k.unwrap(),
283            quantize_k.unwrap(),
284            stream,
285        )?
286    };
287
288    let shared_down_prefix = format!("{p}.shared_experts.down_proj");
289    let shared_down_has_s2 = store.contains(&format!("{shared_down_prefix}.weight_scale_2"));
290    let shared_down_has_s = store.contains(&format!("{shared_down_prefix}.weight_scale"));
291    let shared_down_fp8 = if want_native_fp8 && !shared_down_has_s2 && shared_down_has_s {
292        match load_fp8_block_scaled_as_fp8weight(store, &shared_down_prefix, gpu) {
293            Ok(mut w) => {
294                let bytes = (w.n as usize) * (w.k as usize);
295                let owned = gpu.alloc(bytes)?;
296                gpu.copy_d2d(w.weight, owned, bytes)?;
297                w.weight = owned;
298                Some(w)
299            }
300            Err(e) => {
301                tracing::warn!("shared_down native FP8 unavailable ({e}) — using NVFP4 requant");
302                None
303            }
304        }
305    } else {
306        None
307    };
308    let shared_down = if shared_down_fp8.is_some() {
309        QuantizedWeight::null()
310    } else if shared_down_has_s2 {
311        quantized(store, &shared_down_prefix, gpu)?
312    } else {
313        let bf16 = if shared_down_has_s {
314            if let Some(s) = scratch {
315                dequant_fp8_to_bf16_into(store, &shared_down_prefix, gpu, s)?
316            } else {
317                dequant_fp8_to_bf16(store, &shared_down_prefix, gpu)?
318            }
319        } else {
320            dense(store, &format!("{shared_down_prefix}.weight"))?
321        };
322        quantize_to_nvfp4(
323            &bf16,
324            config.hidden_size,
325            config.shared_expert_intermediate_size,
326            gpu,
327            absmax_k.unwrap(),
328            quantize_k.unwrap(),
329            stream,
330        )?
331    };
332
333    // 2026-09-25: fc1 and fc2 stay BF16 for the model's lifetime, so their
334    // dequants get their own allocations, never `scratch`.
335    let (fc1_latent_proj, fc2_latent_proj) = if config.moe_latent_size > 0 {
336        let fc1_prefix = format!("{p}.fc1_latent_proj");
337        let fc1 = if store.contains(&format!("{fc1_prefix}.weight_scale")) {
338            dequant_fp8_to_bf16(store, &fc1_prefix, gpu)?
339        } else {
340            dense(store, &format!("{fc1_prefix}.weight"))?
341        };
342        let fc2_prefix = format!("{p}.fc2_latent_proj");
343        let fc2 = if store.contains(&format!("{fc2_prefix}.weight_scale")) {
344            dequant_fp8_to_bf16(store, &fc2_prefix, gpu)?
345        } else {
346            dense(store, &format!("{fc2_prefix}.weight"))?
347        };
348        (Some(fc1), Some(fc2))
349    } else {
350        (None, None)
351    };
352
353    // 2026-09-25: The first local expert's `up_proj` decides the format of every
354    // routed expert. The intermediate size is per layer
355    // (`moe_intermediate_size_for`).
356    let moe_input = config.moe_input_size();
357    let moe_inter = config.moe_intermediate_size_for(layer);
358    let first_local = (0..num_experts).find(|e| config.is_local_expert(*e));
359    let experts_are_nvfp4 = first_local
360        .is_none_or(|e| store.contains(&format!("{p}.experts.{e}.up_proj.weight_scale_2")));
361    let experts_are_fp8 = !experts_are_nvfp4
362        && first_local
363            .is_some_and(|e| store.contains(&format!("{p}.experts.{e}.up_proj.weight_scale")));
364    if !experts_are_nvfp4 && layer < 2 {
365        tracing::info!(
366            "L{layer} MoE experts: {} → NVFP4 (runtime quantization, {} experts)",
367            if experts_are_fp8 { "FP8" } else { "BF16" },
368            num_experts,
369        );
370    }
371
372    let mut experts = Vec::with_capacity(num_experts);
373    for e in 0..num_experts {
374        if config.is_local_expert(e) {
375            let up_prefix = format!("{p}.experts.{e}.up_proj");
376            let down_prefix = format!("{p}.experts.{e}.down_proj");
377            let (up_proj, down_proj) = if experts_are_nvfp4 {
378                (
379                    quantized(store, &up_prefix, gpu)?,
380                    quantized(store, &down_prefix, gpu)?,
381                )
382            } else {
383                let up_bf16 = if experts_are_fp8 {
384                    if let Some(s) = scratch {
385                        dequant_fp8_to_bf16_into(store, &up_prefix, gpu, s)?
386                    } else {
387                        dequant_fp8_to_bf16(store, &up_prefix, gpu)?
388                    }
389                } else {
390                    dense(store, &format!("{up_prefix}.weight"))?
391                };
392                let up = quantize_to_nvfp4(
393                    &up_bf16,
394                    moe_inter,
395                    moe_input,
396                    gpu,
397                    absmax_k.unwrap(),
398                    quantize_k.unwrap(),
399                    stream,
400                )?;
401                let down_bf16 = if experts_are_fp8 {
402                    if let Some(s) = scratch {
403                        dequant_fp8_to_bf16_into(store, &down_prefix, gpu, s)?
404                    } else {
405                        dequant_fp8_to_bf16(store, &down_prefix, gpu)?
406                    }
407                } else {
408                    dense(store, &format!("{down_prefix}.weight"))?
409                };
410                let down = quantize_to_nvfp4(
411                    &down_bf16,
412                    moe_input,
413                    moe_inter,
414                    gpu,
415                    absmax_k.unwrap(),
416                    quantize_k.unwrap(),
417                    stream,
418                )?;
419                (up, down)
420            };
421            experts.push(NemotronExpertWeight { up_proj, down_proj });
422        } else {
423            experts.push(NemotronExpertWeight::null());
424        }
425    }
426
427    Ok(NemotronMoeWeights {
428        gate,
429        e_score_correction_bias,
430        experts,
431        shared_up,
432        shared_up_fp8,
433        shared_down_fp8,
434        shared_down,
435        fc1_latent_proj,
436        fc2_latent_proj,
437    })
438}
metrale_model_layers/weight_map/nemotron.rs

metrale_model_layers/weight_map/
nemotron.rs
