spark_model/weight_map/
nvfp4_detect.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/// Detect the weight quantization variant from the weight store.
14///
15/// Dispatch order matches vLLM / TRT-LLM / SGLang:
16///   1. **Config-declared scheme** (`config.quantization_config.quant_method`)
17///      wins outright. This is the authoritative signal and the only one
18///      that correctly handles checkpoints with an `ignore` list (e.g.
19///      `lukealonso/MiniMax-M2.7-NVFP4`, whose MLP `gate_proj` is
20///      intentionally unquantized and therefore has no `.weight_scale`
21///      tensor — sniffing would mis-detect the whole checkpoint as
22///      `Bf16Raw` and then read uint8-packed FP4 as BF16, which is the
23///      4× byte overrun that surfaces as `CUDA_ERROR_ILLEGAL_ADDRESS`
24///      ten seconds into load).
25///   2. **Tensor-name sniffing** for the many checkpoints in the wild
26///      that ship without a `quantization_config` block.
27pub fn detect_nvfp4_variant(
28    store: &WeightStore,
29    config: &atlas_core::config::ModelConfig,
30) -> Nvfp4Variant {
31    // (1) Config-first dispatch. See module docs on `quant_format` for
32    // the full rationale — this is the fix for the Discord 2026-04-17
33    // `CUDA_ERROR_ILLEGAL_ADDRESS` bug.
34    if let Some(qc) = &config.quantization_config {
35        match qc.quant_method.as_str() {
36            "modelopt" if qc.quant_algo.eq_ignore_ascii_case("NVFP4") => {
37                return Nvfp4Variant::Standard;
38            }
39            "modelopt" if qc.quant_algo.eq_ignore_ascii_case("FP8") => {
40                return Nvfp4Variant::Fp8Dequanted;
41            }
42            "compressed-tensors" => {
43                // `format` is the sub-selector here. Block-scaled FP8 is tagged
44                // either with a literal "fp8" OR with compressed-tensors'
45                // `"float-quantized"` (8-bit float = FP8 E4M3, e.g.
46                // Hcompany/Holo-3.1-*-FP8); the rest ("nvfp4-pack-quantized",
47                // "pack-quantized") are NVFP4.
48                let fmt = qc.format.to_ascii_lowercase();
49                if fmt.contains("fp8") || fmt.contains("float-quant") {
50                    return Nvfp4Variant::Fp8Dequanted;
51                }
52                return Nvfp4Variant::CompressedTensors;
53            }
54            "fp8" => {
55                return Nvfp4Variant::Fp8Dequanted;
56            }
57            _ => {
58                // Unknown method with non-empty ignore list — fall
59                // through to heuristic detection. A warning was already
60                // emitted by `quant_format::detect_quant_format`.
61            }
62        }
63    }
64
65    let lp = config.layer_prefix(0);
66
67    // Check MoE expert key first (most models are MoE).
68    let local_expert = config.local_expert_range().0;
69    let moe_sehyo_key = format!("{lp}.mlp.experts.{local_expert}.gate_proj.weight_packed");
70    if store.contains(&moe_sehyo_key) {
71        return Nvfp4Variant::CompressedTensors;
72    }
73
74    // Check dense FFN key (non-MoE models like Qwen3.5-27B).
75    let dense_sehyo_key = format!("{lp}.mlp.gate_proj.weight_packed");
76    if store.contains(&dense_sehyo_key) {
77        return Nvfp4Variant::CompressedTensors;
78    }
79
80    // Mistral uses "layers.{i}.experts.{e}.w1" naming (no "model." prefix, no ".mlp.").
81    let mistral_key = format!("layers.0.experts.{local_expert}.w1.weight_packed");
82    if store.contains(&mistral_key) {
83        return Nvfp4Variant::CompressedTensors;
84    }
85
86    // Fallback: scan any tensor name for `.weight_packed` suffix.
87    // Catches compressed-tensors checkpoints with unexpected naming conventions.
88    if store.names().any(|k| k.ends_with(".weight_packed")) {
89        return Nvfp4Variant::CompressedTensors;
90    }
91
92    // Check for FP8 block-scaled weights (e.g. Qwen/Qwen3.5-35B-A3B-FP8):
93    // FP8 models have `weight_scale_inv` alongside FP8E4M3 weights.
94    //
95    // Two SPELLINGS of the same layer, not two layers. The second entry used to
96    // be indexed by `local_expert_range().0` — an EXPERT index used as a LAYER
97    // index. `local_expert_range` returns global expert ids (see its sibling
98    // `is_local_expert`), so on a single node or EP rank 0 it is 0 and layer 0
99    // is probed by accident; on rank >= 1 it is `ep_rank * num_experts /
100    // ep_world_size` and each rank probes a DIFFERENT layer. With 256 experts
101    // over 2 ranks, rank 1 probed layer 39 — a full-attention layer whose FP8
102    // `q_proj` trips the attention sniff below — so two ranks loading one
103    // checkpoint could disagree about its quantisation variant.
104    //
105    // Detection must not depend on EP rank: every rank sees the same file.
106    const ALT_LAYER0_PREFIX: &str = "model.language_model.layers.0";
107    let prefixes_to_check = [lp.clone(), ALT_LAYER0_PREFIX.to_string()];
108    for pfx in &prefixes_to_check {
109        let fp8_key = format!("{pfx}.mlp.experts.{local_expert}.gate_proj.weight_scale_inv");
110        if store.contains(&fp8_key) {
111            return Nvfp4Variant::Fp8Dequanted;
112        }
113        let fp8_dense_key = format!("{pfx}.mlp.gate_proj.weight_scale_inv");
114        if store.contains(&fp8_dense_key) {
115            return Nvfp4Variant::Fp8Dequanted;
116        }
117        let fp8_attn_key = format!("{pfx}.self_attn.q_proj.weight_scale_inv");
118        if store.contains(&fp8_attn_key) {
119            return Nvfp4Variant::Fp8Dequanted;
120        }
121        // compressed-tensors `float-quantized` FP8 (e.g. Hcompany/Holo-3.1-*-FP8)
122        // ships block-FP8 as an FP8E4M3 `.weight` + 2D `.weight_scale` — NO
123        // `.weight_packed` (that's NVFP4) and NO `.weight_scale_inv` (that's
124        // DeepSeek/Qwen-native FP8). The `.weight_scale` name alias-collides
125        // with compressed-tensors NVFP4, so the `.weight_scale` checks below
126        // would misroute it to an NVFP4 variant. Disambiguate by the
127        // unambiguous FP8E4M3 weight dtype: an FP8E4M3 projection weight is
128        // always block-FP8 (Fp8Dequanted; the FP8→BF16→NVFP4 requant path in
129        // `quantized_from_fp8` reads the 2D `.weight_scale`).
130        for key in [
131            format!("{pfx}.mlp.experts.{local_expert}.gate_proj.weight"),
132            format!("{pfx}.mlp.gate_proj.weight"),
133            format!("{pfx}.self_attn.q_proj.weight"),
134        ] {
135            if store
136                .get(&key)
137                .map(|w| w.dtype == WeightDtype::FP8E4M3)
138                .unwrap_or(false)
139            {
140                return Nvfp4Variant::Fp8Dequanted;
141            }
142        }
143    }
144    // Fallback: scan any tensor name for `.weight_scale_inv` suffix.
145    // Catches FP8 checkpoints where the layer prefix hasn't been resolved yet.
146    if store.names().any(|k| k.ends_with(".weight_scale_inv")) {
147        return Nvfp4Variant::Fp8Dequanted;
148    }
149
150    // BF16/FP16 fine-tune detection: no quantization markers at all.
151    // If even `.weight_scale` is absent (i.e., not a Standard NVFP4 model
152    // either), fall through to runtime quantization from raw BF16/FP16.
153    // Catches third-party fine-tunes like samuelcardillo/Carnice-MoE-35B-A3B
154    // that ship only `.weight` tensors with no per-channel scales.
155    let any_standard_scale = store.names().any(|k| k.ends_with(".weight_scale"));
156    if !any_standard_scale {
157        tracing::warn!(
158            "No NVFP4/FP8 quantization metadata found (no .weight_packed / .weight_scale_inv / .weight_scale). \
159             Falling back to runtime BF16→NVFP4 quantization. Quality will be inferior to a calibrated NVFP4 release."
160        );
161        return Nvfp4Variant::Bf16Raw;
162    }
163
164    // Partial-NVFP4 guard: some upstream checkpoints (notably google/gemma-4-26B-A4B-it)
165    // ship `.weight_scale` on KV-cache scale tensors but NOT on the MLP/MoE
166    // projections Atlas actually consumes. If we claim Standard here the
167    // loader will then fail with a cryptic `Weight '...mlp.gate_proj.weight_scale'
168    // not found in store` half-way through load (logged against #bugs 2026-04-15
169    // by kiiv6565). Sniff the canonical L0 MLP gate_proj — if its `.weight_scale`
170    // is missing, the right answer is BF16 runtime quantization, not Standard.
171    let has_mlp_scale = {
172        let k_dense = format!("{lp}.mlp.gate_proj.weight_scale");
173        let k_moe = format!("{lp}.mlp.experts.{local_expert}.gate_proj.weight_scale");
174        store.contains(&k_dense) || store.contains(&k_moe)
175    };
176    if !has_mlp_scale {
177        tracing::warn!(
178            "Partial NVFP4 metadata: `.weight_scale` exists for some tensors (e.g. KV scales) \
179             but not for MLP/MoE projections. Falling back to runtime BF16→NVFP4 quantization. \
180             For best quality use a fully-quantized NVFP4 release (e.g. Sehyo/*-NVFP4)."
181        );
182        return Nvfp4Variant::Bf16Raw;
183    }
184
185    Nvfp4Variant::Standard
186}
187
188/// Load a quantized weight using the appropriate naming convention.
189///
190/// For `Fp8Dequanted`, requires `quant_ctx` (absmax_k, quantize_k, stream)
191/// to runtime-quantize the dequanted BF16 to NVFP4.
192pub(crate) fn quantized_auto(
193    store: &WeightStore,
194    prefix: &str,
195    gpu: &dyn GpuBackend,
196    variant: Nvfp4Variant,
197) -> Result<QuantizedWeight> {
198    match variant {
199        Nvfp4Variant::Standard => quantized(store, prefix, gpu),
200        Nvfp4Variant::CompressedTensors => quantized_v2(store, prefix, gpu),
201        Nvfp4Variant::Fp8Dequanted => {
202            unreachable!("Fp8Dequanted must use quantized_auto_fp8 with quant context")
203        }
204        Nvfp4Variant::Bf16Raw => {
205            unreachable!("Bf16Raw must use quantized_any with quant context")
206        }
207    }
208}
209
210/// Quantize context for FP8→BF16→NVFP4 runtime conversion.
211#[derive(Clone, Copy)]
212pub(crate) struct QuantizeCtx {
213    pub absmax_k: spark_runtime::gpu::KernelHandle,
214    pub quantize_k: spark_runtime::gpu::KernelHandle,
215    pub stream: u64,
216}
217
218/// Load a quantized weight, dispatching by variant. Handles all three on-disk formats
219/// including FP8 block-scaled (requires dimensions for FP8→BF16→NVFP4 conversion).
220pub(crate) fn quantized_any(
221    store: &WeightStore,
222    prefix: &str,
223    n: usize,
224    k: usize,
225    gpu: &dyn GpuBackend,
226    variant: Nvfp4Variant,
227    qctx: QuantizeCtx,
228) -> Result<QuantizedWeight> {
229    let _t_detect = std::time::Instant::now();
230    // Per-key fallback (B8 #bugs RedHatAI/Qwen3-Coder-Next-NVFP4): some
231    // models that are CompressedTensors overall keep certain projections
232    // (e.g. `linear_attn.out_proj`) as raw BF16 with no quantization
233    // metadata. Detect that case here and runtime-quantize, instead of
234    // failing the whole load with "weight_global_scale not found".
235    let has_packed = store.contains(&format!("{prefix}.weight_packed"));
236    let has_scale = store.contains(&format!("{prefix}.weight_scale"));
237    let has_scale_inv = store.contains(&format!("{prefix}.weight_scale_inv"));
238    let has_only_dense =
239        !has_packed && !has_scale && !has_scale_inv && store.contains(&format!("{prefix}.weight"));
240
241    // Per-key fallback #2 (unsloth/Qwen3.6-{27B,35B-A3B}-NVFP4, re-quantized
242    // 2026-07-10): mixed-precision checkpoints that are NVFP4 for most of the
243    // net but leave a tail of layers — and, in the MoE, the shared experts —
244    // as FP8 E4M3 with a per-row `weight_scale` ([N,1] BF16). Those keys carry
245    // no NVFP4 metadata at all (no `weight_packed`, no `weight_global_scale`,
246    // no `weight_scale_2`), so the declared NVFP4 variant cannot load them and
247    // the whole model dies on `weight_global_scale not found in store`.
248    // Detect the FP8 layout per key and dequant→runtime-quantize instead.
249    //
250    // The three NVFP4 layouts are all excluded by construction, so this can
251    // never steal a key that IS NVFP4:
252    //   Standard (ModelOpt/nvidia) -> has `weight_scale_2`
253    //   CompressedTensors (Sehyo)  -> has `weight_packed` + `weight_global_scale`
254    //   this FP8 case              -> has neither, and `.weight` is FP8E4M3
255    let has_fp8_dense = !has_packed
256        && !store.contains(&format!("{prefix}.weight_global_scale"))
257        && !store.contains(&format!("{prefix}.weight_scale_2"))
258        && (has_scale || has_scale_inv)
259        && store
260            .get(&format!("{prefix}.weight"))
261            .map(|w| w.dtype == WeightDtype::FP8E4M3)
262            .unwrap_or(false);
263
264    let effective_variant = if has_only_dense && !matches!(variant, Nvfp4Variant::Bf16Raw) {
265        tracing::debug!("{prefix}: no quantization metadata; falling back to runtime BF16→NVFP4");
266        Nvfp4Variant::Bf16Raw
267    } else if has_fp8_dense
268        && !matches!(variant, Nvfp4Variant::Fp8Dequanted | Nvfp4Variant::Bf16Raw)
269    {
270        tracing::debug!("{prefix}: FP8 key in an NVFP4 checkpoint; dequant FP8→BF16→NVFP4");
271        Nvfp4Variant::Fp8Dequanted
272    } else {
273        variant
274    };
275
276    let _t_detect_ns = _t_detect.elapsed().as_nanos() as u64;
277    match effective_variant {
278        Nvfp4Variant::Standard => quantized(store, prefix, gpu),
279        Nvfp4Variant::CompressedTensors => quantized_v2(store, prefix, gpu),
280        Nvfp4Variant::Fp8Dequanted => quantized_from_fp8(
281            store,
282            prefix,
283            n,
284            k,
285            gpu,
286            qctx.absmax_k,
287            qctx.quantize_k,
288            qctx.stream,
289        ),
290        Nvfp4Variant::Bf16Raw => {
291            use std::sync::atomic::{AtomicU64, Ordering};
292            static T_DETECT: AtomicU64 = AtomicU64::new(0);
293            static T_GET: AtomicU64 = AtomicU64::new(0);
294            static T_QUANT: AtomicU64 = AtomicU64::new(0);
295            static T_FREE: AtomicU64 = AtomicU64::new(0);
296            static N: AtomicU64 = AtomicU64::new(0);
297            T_DETECT.fetch_add(_t_detect_ns, Ordering::Relaxed);
298            // Raw BF16/FP16 fine-tune: load the dense weight then runtime-quantize.
299            let _t = std::time::Instant::now();
300            let w = store.get(&format!("{prefix}.weight"))?;
301            let bf16 = DenseWeight { weight: w.ptr };
302            T_GET.fetch_add(_t.elapsed().as_nanos() as u64, Ordering::Relaxed);
303            let _t = std::time::Instant::now();
304            let q = quantize_to_nvfp4(
305                &bf16,
306                n,
307                k,
308                gpu,
309                qctx.absmax_k,
310                qctx.quantize_k,
311                qctx.stream,
312            )?;
313            T_QUANT.fetch_add(_t.elapsed().as_nanos() as u64, Ordering::Relaxed);
314            let _t = std::time::Instant::now();
315            // Free the BF16 source: the NVFP4 buffer is a fresh allocation, so the
316            // on-disk BF16 weight is now redundant. Without this a 35B BF16 MoE
317            // (Bf16Raw, SEPARATE per-expert layout routed through here by #200's
318            // `quantized_any`) holds BOTH the ~60GB BF16 experts AND the ~22GB
319            // NVFP4 copies → ~109GB pre-KV, no room for KV. Safe + mirrors
320            // `quantized_from_fp8` which frees its BF16 intermediate the same way.
321            gpu.free(w.ptr)?;
322            T_FREE.fetch_add(_t.elapsed().as_nanos() as u64, Ordering::Relaxed);
323            let c = N.fetch_add(1, Ordering::Relaxed) + 1;
324            if c.is_multiple_of(512) {
325                let ms = |a: &AtomicU64| a.load(Ordering::Relaxed) as f64 / 1.0e6;
326                tracing::info!(
327                    "quantized_any(Bf16Raw) PROFILE after {c} calls (ms total): detect={:.1} \
328                     store_get={:.1} quantize={:.1} free={:.1} | sum={:.1} per_call={:.3}ms",
329                    ms(&T_DETECT),
330                    ms(&T_GET),
331                    ms(&T_QUANT),
332                    ms(&T_FREE),
333                    ms(&T_DETECT) + ms(&T_GET) + ms(&T_QUANT) + ms(&T_FREE),
334                    (ms(&T_DETECT) + ms(&T_GET) + ms(&T_QUANT) + ms(&T_FREE)) / c as f64,
335                );
336            }
337            Ok(q)
338        }
339    }
340}
341
342/// Load a quantized weight from FP8 block-scaled data: FP8→BF16→NVFP4.
343///
344/// `n` and `k` are the logical weight dimensions (e.g. [inter, hidden] for gate_proj).
345pub(crate) fn quantized_from_fp8(
346    store: &WeightStore,
347    prefix: &str,
348    n: usize,
349    k: usize,
350    gpu: &dyn GpuBackend,
351    absmax_k: spark_runtime::gpu::KernelHandle,
352    quantize_k: spark_runtime::gpu::KernelHandle,
353    stream: u64,
354) -> Result<QuantizedWeight> {
355    let bf16 = dequant_fp8_blockscaled_to_bf16(store, prefix, gpu)?;
356    let result = quantize_to_nvfp4(&bf16, n, k, gpu, absmax_k, quantize_k, stream)?;
357    // Free the BF16 intermediate — only the NVFP4 result is needed.
358    gpu.free(bf16.weight)?;
359    Ok(result)
360}
361
362/// Load FP8 block-scaled weight as BF16 dense (no NVFP4 re-quantization).
363///
364/// Use this when the runtime NVFP4 quantization produces degenerate weights
365/// (e.g., FP8 checkpoints where double-quantization degrades quality).
366/// The weight stays in BF16 and uses `dense_gemv`/`dense_gemm` kernels.
367#[allow(dead_code)]
368pub(crate) fn dense_from_fp8(
369    store: &WeightStore,
370    prefix: &str,
371    gpu: &dyn GpuBackend,
372) -> Result<DenseWeight> {
373    dequant_fp8_blockscaled_to_bf16(store, prefix, gpu)
374}
375
376/// Load full attention weights for Qwen3.5 (all Q/K/V/O are NVFP4 on disk).
377#[allow(dead_code)]
378pub(crate) fn load_attention_qwen35(
379    store: &WeightStore,
380    layer_prefix: &str,
381    gpu: &dyn GpuBackend,
382) -> Result<AttentionWeights> {
383    let p = format!("{layer_prefix}.self_attn");
384    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);
385    Ok(AttentionWeights {
386        // Q/K/V are NVFP4 quantized — load packed, return as dense (the weight_packed data)
387        // The weight_loader will handle creating QuantizedWeight from these
388        q_proj: dense(store, &format!("{p}.q_proj.weight_packed"))?,
389        k_proj: dense(store, &format!("{p}.k_proj.weight_packed"))?,
390        v_proj: dense(store, &format!("{p}.v_proj.weight_packed"))?,
391        o_proj: quantized_v2(store, &format!("{p}.o_proj"), gpu)?,
392        q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
393        k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
394        q_norm_full: None,
395        k_norm_full: None,
396        k_scale,
397        v_scale,
398    })
399}
400
401/// Load NVFP4 quantized projection for Qwen3.5 full attention layer.
402#[allow(dead_code)]
403pub(crate) fn load_quantized_proj_qwen35(
404    store: &WeightStore,
405    prefix: &str,
406    gpu: &dyn GpuBackend,
407) -> Result<QuantizedWeight> {
408    quantized_v2(store, prefix, gpu)
409}
410
411#[cfg(test)]
412mod ep_detection_tests {
413    use super::*;
414    use atlas_core::config::ModelConfig;
415    use spark_runtime::weights::WeightStore;
416
417    /// A store holding only the FP8 attention marker at a given layer, which is
418    /// what the detector sniffs for. Names are all the detector reads.
419    fn store_with(names: &[String]) -> WeightStore {
420        use std::collections::HashMap;
421        let map: HashMap<String, spark_runtime::weights::WeightTensor> = names
422            .iter()
423            .map(|n| {
424                (
425                    n.clone(),
426                    spark_runtime::weights::WeightTensor {
427                        ptr: spark_runtime::gpu::DevicePtr::NULL,
428                        shape: vec![1],
429                        dtype: spark_runtime::weights::WeightDtype::FP8E4M3,
430                    },
431                )
432            })
433            .collect();
434        WeightStore::from_map(map)
435    }
436
437    #[test]
438    fn alternate_layer0_fp8_dtype_is_detected_on_every_ep_rank() {
439        let mut cfg = ModelConfig::qwen3_next_80b_nvfp4();
440        cfg.quantization_config = None;
441        let store =
442            store_with(&["model.language_model.layers.0.self_attn.q_proj.weight".to_string()]);
443
444        cfg.ep_world_size = 2;
445        for ep_rank in 0..2 {
446            cfg.ep_rank = ep_rank;
447            assert_eq!(
448                detect_nvfp4_variant(&store, &cfg),
449                Nvfp4Variant::Fp8Dequanted,
450                "EP rank {ep_rank} must inspect the same layer-zero checkpoint marker"
451            );
452        }
453    }
454
455    #[test]
456    fn scale_inv_suffix_fallback_detects_an_unexpected_prefix() {
457        let mut cfg = ModelConfig::qwen3_next_80b_nvfp4();
458        cfg.quantization_config = None;
459        let store =
460            store_with(&["third_party.transformer.blocks.17.attn.q.weight_scale_inv".to_string()]);
461        assert_eq!(
462            detect_nvfp4_variant(&store, &cfg),
463            Nvfp4Variant::Fp8Dequanted
464        );
465    }
466}
spark_model/weight_map/nvfp4_detect.rs

spark_model/weight_map/
nvfp4_detect.rs
