spark_model/weight_map/
nemotron.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `weight_map.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::{Context, Result, bail, ensure};
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9use spark_runtime::weights::{WeightDtype, WeightStore};
10
11use super::*;
12
13/// Nemotron-H Mamba-2 SSM weights.
14///
15/// in_proj produces [z(d_inner), x(d_inner), B(n_groups*state), C(n_groups*state), dt(num_heads)].
16pub struct NemotronSsmWeights {
17    /// in_proj: [in_proj_size, hidden_size] NVFP4.
18    pub in_proj: QuantizedWeight,
19    /// out_proj: `[hidden_size, d_inner]` NVFP4.
20    pub out_proj: QuantizedWeight,
21    /// conv1d weight: `[d_xBC, 1, conv_kernel]` BF16.
22    pub conv1d_weight: DenseWeight,
23    /// conv1d bias: `[d_xBC]` BF16.
24    pub conv1d_bias: DenseWeight,
25    /// A_log: `[mamba_num_heads]` BF16 (cast to FP32 at runtime).
26    pub a_log: DenseWeight,
27    /// D skip-connection: `[mamba_num_heads]` BF16.
28    pub d_param: DenseWeight,
29    /// dt_bias: `[mamba_num_heads]` BF16.
30    pub dt_bias: DenseWeight,
31    /// SSM internal norm: `[d_inner]` BF16 (applied to y before gating with z).
32    pub ssm_norm: DenseWeight,
33}
34
35/// Nemotron-H 2-projection expert (up_proj + relu² + down_proj, no gate_proj).
36#[derive(Debug, Clone, Copy)]
37pub struct NemotronExpertWeight {
38    pub up_proj: QuantizedWeight,
39    pub down_proj: QuantizedWeight,
40}
41
42impl NemotronExpertWeight {
43    pub fn null() -> Self {
44        Self {
45            up_proj: QuantizedWeight::null(),
46            down_proj: QuantizedWeight::null(),
47        }
48    }
49}
50
51/// Nemotron-H MoE layer weights.
52pub struct NemotronMoeWeights {
53    /// Router gate: [num_experts, hidden_size] F32→BF16.
54    pub gate: DenseWeight,
55    /// Expert score correction bias: `[num_experts]` F32.
56    pub e_score_correction_bias: DenseWeight,
57    /// Per-expert weights (routed): NVFP4.
58    pub experts: Vec<NemotronExpertWeight>,
59    /// Shared expert up_proj: [shared_inter, hidden_size] NVFP4.
60    pub shared_up: QuantizedWeight,
61    /// Shared expert up_proj kept as NATIVE FP8 when the checkpoint ships it that
62    /// way (ModelOpt MIXED_PRECISION), instead of the FP8→BF16→NVFP4 requant.
63    /// `Some` only under `ATLAS_NEMOTRON_NATIVE_FP8_SSM`; decode prefers it via
64    /// `w8a16_gemv`. Measured on Puzzle-75B: with the SSM projections already
65    /// native, a 977-token story went from calling the dog "Rover"/"Rex" to
66    /// using the given name "Rufus" 8 times with no substitutions — proper-noun
67    /// retrieval is what the requant was destroying.
68    pub shared_up_fp8: Option<Fp8Weight>,
69    /// Shared expert down_proj: [hidden_size, shared_inter] NVFP4.
70    pub shared_down: QuantizedWeight,
71    /// Shared expert down_proj kept as NATIVE FP8, same rationale as
72    /// `shared_up_fp8`. Consumed by relu_squared_inplace + `w8a16_gemv` instead
73    /// of the fused `moe_expert_relu2_down_shared` kernel, which only speaks
74    /// NVFP4 — the fused launch simply drops its shared slot (grid.y = top_k)
75    /// when this is present.
76    pub shared_down_fp8: Option<Fp8Weight>,
77    /// LatentMoE: fc1 [moe_latent_size, hidden_size] BF16 (dequant from FP8 at load).
78    /// Present only for Super 120B (moe_latent_size > 0).
79    pub fc1_latent_proj: Option<DenseWeight>,
80    /// LatentMoE: fc2 [hidden_size, moe_latent_size] BF16.
81    /// Present only for Super 120B (moe_latent_size > 0).
82    pub fc2_latent_proj: Option<DenseWeight>,
83}
84
85/// SSM weight quantization format detected at load time.
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
87pub enum NemotronSsmQuant {
88    /// NVFP4: has weight_scale + weight_scale_2 (Nano NVFP4 model).
89    Nvfp4,
90    /// FP8 E4M3: has weight_scale but no weight_scale_2 (Super 120B).
91    Fp8,
92    /// BF16: no weight_scale at all (BF16 layers adjacent to attention).
93    Bf16,
94}
95
96/// Load Nemotron-H Mamba-2 SSM weights.
97///
98/// Mixed quantization: NVFP4, FP8, or BF16 depending on layer and model.
99/// Non-NVFP4 projections use QuantizedWeight::null() — caller runtime-quantizes.
100pub(crate) fn load_nemotron_ssm(
101    store: &WeightStore,
102    _layer: usize,
103    gpu: &dyn GpuBackend,
104    layer_prefix: &str,
105) -> Result<(NemotronSsmWeights, NemotronSsmQuant)> {
106    let p = format!("{layer_prefix}.mixer");
107    let has_scale = store.contains(&format!("{p}.in_proj.weight_scale"));
108    let has_scale2 = store.contains(&format!("{p}.in_proj.weight_scale_2"));
109    let quant = if has_scale && has_scale2 {
110        NemotronSsmQuant::Nvfp4
111    } else if has_scale {
112        NemotronSsmQuant::Fp8
113    } else {
114        NemotronSsmQuant::Bf16
115    };
116    let in_proj = if quant == NemotronSsmQuant::Nvfp4 {
117        quantized(store, &format!("{p}.in_proj"), gpu)?
118    } else {
119        QuantizedWeight::null()
120    };
121    let out_proj = if quant == NemotronSsmQuant::Nvfp4 {
122        quantized(store, &format!("{p}.out_proj"), gpu)?
123    } else {
124        QuantizedWeight::null()
125    };
126    // A_log, D, dt_bias, conv1d.bias are BF16 in safetensors but consumed as FP32 by kernels.
127    Ok((
128        NemotronSsmWeights {
129            in_proj,
130            out_proj,
131            conv1d_weight: dense(store, &format!("{p}.conv1d.weight"))?,
132            conv1d_bias: dense_bf16_as_f32(store, &format!("{p}.conv1d.bias"), gpu)?,
133            a_log: dense_bf16_as_f32(store, &format!("{p}.A_log"), gpu)?,
134            d_param: dense_bf16_as_f32(store, &format!("{p}.D"), gpu)?,
135            dt_bias: dense_bf16_as_f32(store, &format!("{p}.dt_bias"), gpu)?,
136            ssm_norm: dense(store, &format!("{p}.norm.weight"))?,
137        },
138        quant,
139    ))
140}
141
142/// Load Nemotron-H attention weights.
143///
144/// Standard GQA: 32 Q heads, 2 KV heads, head_dim=128.
145/// Mixed quantization: some attention layers are NVFP4, some BF16.
146/// BF16 layers return dense Q/K/V in AttentionWeights + null QuantizedWeights.
147pub(crate) 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    // Distinguish NVFP4 (has weight_scale_2) from FP8 (has weight_scale only) from BF16 (neither).
162    let is_nvfp4 = store.contains(&format!("{p}.q_proj.weight_scale_2"));
163    let is_fp8 = !is_nvfp4 && store.contains(&format!("{p}.q_proj.weight_scale"));
164    let dummy = DenseWeight {
165        weight: DevicePtr::NULL,
166    };
167
168    let (q_dense, k_dense, v_dense, o_dense, o_proj, q_nvfp4, k_nvfp4, v_nvfp4) = if is_nvfp4 {
169        let q = quantized(store, &format!("{p}.q_proj"), gpu)?;
170        let k = quantized(store, &format!("{p}.k_proj"), gpu)?;
171        let v = quantized(store, &format!("{p}.v_proj"), gpu)?;
172        let o = quantized(store, &format!("{p}.o_proj"), gpu)?;
173        (dummy, dummy, dummy, dummy, o, Some(q), Some(k), Some(v))
174    } else {
175        // FP8 or BF16 — dequant to BF16 dense, caller will quantize to NVFP4.
176        let load_proj = |name: &str| -> Result<DenseWeight> {
177            let prefix = format!("{p}.{name}");
178            if store.contains(&format!("{prefix}.weight_scale")) {
179                dequant_fp8_to_bf16(store, &prefix, gpu)
180            } else {
181                dense(store, &format!("{prefix}.weight"))
182            }
183        };
184        let q = load_proj("q_proj")?;
185        let k = load_proj("k_proj")?;
186        let v = load_proj("v_proj")?;
187        let o = load_proj("o_proj")?;
188        if is_fp8 && layer < 2 {
189            tracing::info!("L{layer} Attention: FP8 → BF16 (runtime quantization to NVFP4)");
190        }
191        (q, k, v, o, QuantizedWeight::null(), None, None, None)
192    };
193
194    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);
195    let attn = AttentionWeights {
196        q_proj: q_dense,
197        k_proj: k_dense,
198        v_proj: v_dense,
199        o_proj,
200        q_norm: dummy,
201        k_norm: dummy,
202        q_norm_full: None,
203        k_norm_full: None,
204        k_scale,
205        v_scale,
206    };
207    Ok((attn, q_nvfp4, k_nvfp4, v_nvfp4, o_dense, is_nvfp4))
208}
209
210/// Load Nemotron-H MoE weights.
211///
212/// Handles both Nano (all NVFP4) and Super 120B (mixed FP8/NVFP4/BF16 + LatentMoE).
213/// For FP8 shared_up and fc1_latent_proj, dequants to BF16 then runtime-quantizes to NVFP4.
214pub(crate) fn load_nemotron_moe(
215    store: &WeightStore,
216    layer: usize,
217    num_experts: usize,
218    gpu: &dyn GpuBackend,
219    config: &atlas_core::config::ModelConfig,
220    absmax_k: Option<spark_runtime::gpu::KernelHandle>,
221    quantize_k: Option<spark_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    // Gate weight: F32 in Nano 30B, BF16 in Super 120B. Convert to BF16 if needed.
228    // The gate GEMV kernel expects BF16 input.
229    let gate_name = format!("{p}.gate.weight");
230    let gate_w = store.get(&gate_name)?;
231    let gate = if gate_w.dtype == WeightDtype::FP32 {
232        dense_f32_as_bf16(store, &gate_name, gpu)?
233    } else {
234        DenseWeight { weight: gate_w.ptr }
235    };
236    // e_score_correction_bias stays F32 — bias_add_bf16_f32 kernel consumes F32 bias.
237    let e_score_correction_bias = dense(store, &format!("{p}.gate.e_score_correction_bias"))?;
238
239    // Shared expert: detect FP8 vs NVFP4 by presence of weight_scale_2.
240    let shared_up_prefix = format!("{p}.shared_experts.up_proj");
241    let shared_up_has_s2 = store.contains(&format!("{shared_up_prefix}.weight_scale_2"));
242    let shared_up_has_s = store.contains(&format!("{shared_up_prefix}.weight_scale"));
243    // Native FP8 for the shared-expert up_proj: skip the FP8→BF16→NVFP4 requant
244    // and hand `w8a16_gemv` the checkpoint's own bytes. Same gate as the SSM
245    // projections — both are ModelOpt MIXED_PRECISION tensors in an otherwise
246    // NVFP4 repo, and both were being needlessly re-quantized.
247    //
248    // `shared_down` is NOT converted here: it is consumed inside the fused
249    // `relu2_down_shared` kernel, which takes NVFP4 (packed + scale + scale_2)
250    // arguments, so making it native needs CUDA work rather than a loader change.
251    let native_fp8_mode =
252        std::env::var("ATLAS_NEMOTRON_NATIVE_FP8_SSM").unwrap_or_else(|_| "1".to_string());
253    let want_native_fp8 = matches!(native_fp8_mode.as_str(), "1" | "both" | "decode");
254    let shared_up_fp8 = if want_native_fp8 && !shared_up_has_s2 && shared_up_has_s {
255        match load_fp8_block_scaled_as_fp8weight(store, &shared_up_prefix, gpu) {
256            Ok(mut w) => {
257                // OWN the bytes — the helper aliases the WeightStore pointer, which
258                // is only safe for a loader that consumes them during load. This one
259                // dereferences per-token forever. Aliasing here produced garbage
260                // weights that still ran (see ssm_layer.rs for the same fix).
261                let bytes = (w.n as usize) * (w.k as usize);
262                let owned = gpu.alloc(bytes)?;
263                gpu.copy_d2d(w.weight, owned, bytes)?;
264                w.weight = owned;
265                Some(w)
266            }
267            Err(e) => {
268                tracing::warn!("shared_up native FP8 unavailable ({e}) — using NVFP4 requant");
269                None
270            }
271        }
272    } else {
273        None
274    };
275    // Under native FP8 nothing reads the NVFP4 copy, so do not spend the load
276    // time or the ~2.9 GB building it. Every consumer is gated on
277    // `shared_up_fp8`/`shared_down_fp8` — including the LatentMoE decode path,
278    // which is the live one on Puzzle (`hybrid_mamba2_latent_moe_...`) and was
279    // the site that faulted with CUDA 700 when it was missed.
280    let shared_up = if shared_up_fp8.is_some() {
281        QuantizedWeight::null()
282    } else if shared_up_has_s2 {
283        quantized(store, &shared_up_prefix, gpu)?
284    } else {
285        // FP8 or BF16 — dequant to scratch, then quantize to NVFP4
286        let bf16 = if shared_up_has_s {
287            if let Some(s) = scratch {
288                dequant_fp8_to_bf16_into(store, &shared_up_prefix, gpu, s)?
289            } else {
290                dequant_fp8_to_bf16(store, &shared_up_prefix, gpu)?
291            }
292        } else {
293            dense(store, &format!("{shared_up_prefix}.weight"))?
294        };
295        quantize_to_nvfp4(
296            &bf16,
297            config.shared_expert_intermediate_size,
298            config.hidden_size,
299            gpu,
300            absmax_k.unwrap(),
301            quantize_k.unwrap(),
302            stream,
303        )?
304    };
305
306    let shared_down_prefix = format!("{p}.shared_experts.down_proj");
307    let shared_down_has_s2 = store.contains(&format!("{shared_down_prefix}.weight_scale_2"));
308    let shared_down_has_s = store.contains(&format!("{shared_down_prefix}.weight_scale"));
309    let shared_down_fp8 = if want_native_fp8 && !shared_down_has_s2 && shared_down_has_s {
310        match load_fp8_block_scaled_as_fp8weight(store, &shared_down_prefix, gpu) {
311            Ok(mut w) => {
312                // Own the bytes (the helper aliases the WeightStore pointer).
313                let bytes = (w.n as usize) * (w.k as usize);
314                let owned = gpu.alloc(bytes)?;
315                gpu.copy_d2d(w.weight, owned, bytes)?;
316                w.weight = owned;
317                Some(w)
318            }
319            Err(e) => {
320                tracing::warn!("shared_down native FP8 unavailable ({e}) — using NVFP4 requant");
321                None
322            }
323        }
324    } else {
325        None
326    };
327    let shared_down = if shared_down_fp8.is_some() {
328        QuantizedWeight::null()
329    } else if shared_down_has_s2 {
330        quantized(store, &shared_down_prefix, gpu)?
331    } else {
332        let bf16 = if shared_down_has_s {
333            if let Some(s) = scratch {
334                dequant_fp8_to_bf16_into(store, &shared_down_prefix, gpu, s)?
335            } else {
336                dequant_fp8_to_bf16(store, &shared_down_prefix, gpu)?
337            }
338        } else {
339            dense(store, &format!("{shared_down_prefix}.weight"))?
340        };
341        quantize_to_nvfp4(
342            &bf16,
343            config.hidden_size,
344            config.shared_expert_intermediate_size,
345            gpu,
346            absmax_k.unwrap(),
347            quantize_k.unwrap(),
348            stream,
349        )?
350    };
351
352    // LatentMoE projections (Super 120B only).
353    // fc1 persists as BF16 — must allocate (not scratch).
354    let (fc1_latent_proj, fc2_latent_proj) = if config.moe_latent_size > 0 {
355        let fc1_prefix = format!("{p}.fc1_latent_proj");
356        let fc1 = if store.contains(&format!("{fc1_prefix}.weight_scale")) {
357            dequant_fp8_to_bf16(store, &fc1_prefix, gpu)?
358        } else {
359            dense(store, &format!("{fc1_prefix}.weight"))?
360        };
361        let fc2_prefix = format!("{p}.fc2_latent_proj");
362        let fc2 = if store.contains(&format!("{fc2_prefix}.weight_scale")) {
363            dequant_fp8_to_bf16(store, &fc2_prefix, gpu)?
364        } else {
365            dense(store, &format!("{fc2_prefix}.weight"))?
366        };
367        (Some(fc1), Some(fc2))
368    } else {
369        (None, None)
370    };
371
372    // Routed experts: detect NVFP4 vs FP8 vs BF16 from first local expert.
373    // Puzzle: per-layer intermediate size from block_configs.
374    let moe_input = config.moe_input_size();
375    let moe_inter = config.moe_intermediate_size_for(layer);
376    let first_local = (0..num_experts).find(|e| config.is_local_expert(*e));
377    let experts_are_nvfp4 = first_local
378        .is_none_or(|e| store.contains(&format!("{p}.experts.{e}.up_proj.weight_scale_2")));
379    let experts_are_fp8 = !experts_are_nvfp4
380        && first_local
381            .is_some_and(|e| store.contains(&format!("{p}.experts.{e}.up_proj.weight_scale")));
382    if !experts_are_nvfp4 && layer < 2 {
383        tracing::info!(
384            "L{layer} MoE experts: {} → NVFP4 (runtime quantization, {} experts)",
385            if experts_are_fp8 { "FP8" } else { "BF16" },
386            num_experts,
387        );
388    }
389
390    let mut experts = Vec::with_capacity(num_experts);
391    for e in 0..num_experts {
392        if config.is_local_expert(e) {
393            let up_prefix = format!("{p}.experts.{e}.up_proj");
394            let down_prefix = format!("{p}.experts.{e}.down_proj");
395            let (up_proj, down_proj) = if experts_are_nvfp4 {
396                (
397                    quantized(store, &up_prefix, gpu)?,
398                    quantized(store, &down_prefix, gpu)?,
399                )
400            } else {
401                let up_bf16 = if experts_are_fp8 {
402                    if let Some(s) = scratch {
403                        dequant_fp8_to_bf16_into(store, &up_prefix, gpu, s)?
404                    } else {
405                        dequant_fp8_to_bf16(store, &up_prefix, gpu)?
406                    }
407                } else {
408                    dense(store, &format!("{up_prefix}.weight"))?
409                };
410                let up = quantize_to_nvfp4(
411                    &up_bf16,
412                    moe_inter,
413                    moe_input,
414                    gpu,
415                    absmax_k.unwrap(),
416                    quantize_k.unwrap(),
417                    stream,
418                )?;
419                let down_bf16 = if experts_are_fp8 {
420                    if let Some(s) = scratch {
421                        dequant_fp8_to_bf16_into(store, &down_prefix, gpu, s)?
422                    } else {
423                        dequant_fp8_to_bf16(store, &down_prefix, gpu)?
424                    }
425                } else {
426                    dense(store, &format!("{down_prefix}.weight"))?
427                };
428                let down = quantize_to_nvfp4(
429                    &down_bf16,
430                    moe_input,
431                    moe_inter,
432                    gpu,
433                    absmax_k.unwrap(),
434                    quantize_k.unwrap(),
435                    stream,
436                )?;
437                (up, down)
438            };
439            experts.push(NemotronExpertWeight { up_proj, down_proj });
440        } else {
441            experts.push(NemotronExpertWeight::null());
442        }
443    }
444
445    Ok(NemotronMoeWeights {
446        gate,
447        e_score_correction_bias,
448        experts,
449        shared_up,
450        shared_up_fp8,
451        shared_down_fp8,
452        shared_down,
453        fc1_latent_proj,
454        fc2_latent_proj,
455    })
456}
spark_model/weight_map/nemotron.rs

spark_model/weight_map/
nemotron.rs
