metrale_model_layers/weight_map/
nvfp4_detect.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Detection of the on-disk weight format (`Nvfp4Variant`) and the variant-dispatching quantized-weight loaders.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants:
7//! - A variant declared by `config.quantization_config` wins over tensor-name sniffing.
8//! - Detection reads only layer 0 and the whole name list, never a layer chosen by EP rank.
9
10#![allow(unused_imports)]
11
12use anyhow::{Context, Result, bail, ensure};
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
14use metrale_model_weights::weights::{WeightDtype, WeightStore};
15
16use super::*;
17
18/// 2026-09-25: Step (1) of [`detect_nvfp4_variant`]: the variant the config
19/// declares, answered without a [`WeightStore`], so the pre-load residency
20/// prediction (`qwen35_dense::predicted_residency`) asks the same question as
21/// the loader. `None` means the config does not say; the sniffing half of
22/// [`detect_nvfp4_variant`] then decides.
23pub fn config_declared_variant(config: &metrale_config::ModelConfig) -> Option<Nvfp4Variant> {
24    let qc = config.quantization_config.as_ref()?;
25    match qc.quant_method.as_str() {
26        "modelopt" if qc.quant_algo.eq_ignore_ascii_case("NVFP4") => Some(Nvfp4Variant::Standard),
27        "modelopt" if qc.quant_algo.eq_ignore_ascii_case("FP8") => Some(Nvfp4Variant::Fp8Dequanted),
28        "compressed-tensors" => {
29            // 2026-09-25: `format` selects: containing "fp8" or "float-quant" is
30            // block-scaled FP8, anything else NVFP4.
31            let fmt = qc.format.to_ascii_lowercase();
32            if fmt.contains("fp8") || fmt.contains("float-quant") {
33                Some(Nvfp4Variant::Fp8Dequanted)
34            } else {
35                Some(Nvfp4Variant::CompressedTensors)
36            }
37        }
38        "fp8" => Some(Nvfp4Variant::Fp8Dequanted),
39        _ => None,
40    }
41}
42
43/// 2026-09-25: Detect the on-disk weight format.
44///
45/// 1. [`config_declared_variant`] wins when it answers. A checkpoint with an
46///    `ignore` list can leave a projection unquantized, which tensor-name
47///    sniffing would misread.
48/// 2. Otherwise the tensor names are sniffed, in this order:
49///    compressed-tensors `weight_packed`, then FP8 `weight_scale_inv` or an
50///    FP8E4M3 projection weight, then `Bf16Raw` when layer 0's MLP
51///    `gate_proj` has no `weight_scale`, and `Standard` last.
52pub fn detect_nvfp4_variant(
53    store: &WeightStore,
54    config: &metrale_config::ModelConfig,
55) -> Nvfp4Variant {
56    if let Some(declared) = config_declared_variant(config) {
57        return declared;
58    }
59
60    let lp = config.layer_prefix(0);
61
62    let local_expert = config.local_expert_range().0;
63    let moe_sehyo_key = format!("{lp}.mlp.experts.{local_expert}.gate_proj.weight_packed");
64    if store.contains(&moe_sehyo_key) {
65        return Nvfp4Variant::CompressedTensors;
66    }
67
68    let dense_sehyo_key = format!("{lp}.mlp.gate_proj.weight_packed");
69    if store.contains(&dense_sehyo_key) {
70        return Nvfp4Variant::CompressedTensors;
71    }
72
73    let mistral_key = format!("layers.0.experts.{local_expert}.w1.weight_packed");
74    if store.contains(&mistral_key) {
75        return Nvfp4Variant::CompressedTensors;
76    }
77
78    if store.names().any(|k| k.ends_with(".weight_packed")) {
79        return Nvfp4Variant::CompressedTensors;
80    }
81
82    // 2026-09-25: Two spellings of layer 0, never a layer index derived from the
83    // EP rank, so every rank reaches the same answer.
84    const ALT_LAYER0_PREFIX: &str = "model.language_model.layers.0";
85    let prefixes_to_check = [lp.clone(), ALT_LAYER0_PREFIX.to_string()];
86    for pfx in &prefixes_to_check {
87        let fp8_key = format!("{pfx}.mlp.experts.{local_expert}.gate_proj.weight_scale_inv");
88        if store.contains(&fp8_key) {
89            return Nvfp4Variant::Fp8Dequanted;
90        }
91        let fp8_dense_key = format!("{pfx}.mlp.gate_proj.weight_scale_inv");
92        if store.contains(&fp8_dense_key) {
93            return Nvfp4Variant::Fp8Dequanted;
94        }
95        let fp8_attn_key = format!("{pfx}.self_attn.q_proj.weight_scale_inv");
96        if store.contains(&fp8_attn_key) {
97            return Nvfp4Variant::Fp8Dequanted;
98        }
99        // 2026-09-25: FP8 with a `weight_scale` instead of `weight_scale_inv` is
100        // told apart from NVFP4 by the FP8E4M3 weight dtype; the `weight_scale`
101        // checks below would otherwise take it for `Standard`.
102        for key in [
103            format!("{pfx}.mlp.experts.{local_expert}.gate_proj.weight"),
104            format!("{pfx}.mlp.gate_proj.weight"),
105            format!("{pfx}.self_attn.q_proj.weight"),
106        ] {
107            if store
108                .get(&key)
109                .map(|w| w.dtype == WeightDtype::FP8E4M3)
110                .unwrap_or(false)
111            {
112                return Nvfp4Variant::Fp8Dequanted;
113            }
114        }
115    }
116    if store.names().any(|k| k.ends_with(".weight_scale_inv")) {
117        return Nvfp4Variant::Fp8Dequanted;
118    }
119
120    let any_standard_scale = store.names().any(|k| k.ends_with(".weight_scale"));
121    if !any_standard_scale {
122        tracing::warn!(
123            "No NVFP4/FP8 quantization metadata found (no .weight_packed / .weight_scale_inv / .weight_scale). \
124             Falling back to runtime BF16→NVFP4 quantization. Quality will be inferior to a calibrated NVFP4 release."
125        );
126        return Nvfp4Variant::Bf16Raw;
127    }
128
129    // 2026-09-25: A `weight_scale` elsewhere (for example on KV scales) is not
130    // enough for `Standard`: layer 0's MLP `gate_proj` must carry one, or
131    // `quantized` would fail on that missing key.
132    let has_mlp_scale = {
133        let k_dense = format!("{lp}.mlp.gate_proj.weight_scale");
134        let k_moe = format!("{lp}.mlp.experts.{local_expert}.gate_proj.weight_scale");
135        store.contains(&k_dense) || store.contains(&k_moe)
136    };
137    if !has_mlp_scale {
138        tracing::warn!(
139            "Partial NVFP4 metadata: `.weight_scale` exists for some tensors (e.g. KV scales) \
140             but not for MLP/MoE projections. Falling back to runtime BF16→NVFP4 quantization. \
141             For best quality use a fully-quantized NVFP4 release (e.g. Sehyo/*-NVFP4)."
142        );
143        return Nvfp4Variant::Bf16Raw;
144    }
145
146    Nvfp4Variant::Standard
147}
148
149/// 2026-09-25: Load a pre-quantized NVFP4 weight by variant: `quantized` for
150/// `Standard`, `quantized_v2` for `CompressedTensors`. Panics for
151/// `Fp8Dequanted` and `Bf16Raw`, which need [`quantized_any`].
152pub fn quantized_auto(
153    store: &WeightStore,
154    prefix: &str,
155    gpu: &dyn GpuBackend,
156    variant: Nvfp4Variant,
157) -> Result<QuantizedWeight> {
158    match variant {
159        Nvfp4Variant::Standard => quantized(store, prefix, gpu),
160        Nvfp4Variant::CompressedTensors => quantized_v2(store, prefix, gpu),
161        Nvfp4Variant::Fp8Dequanted => {
162            unreachable!("Fp8Dequanted must use quantized_auto_fp8 with quant context")
163        }
164        Nvfp4Variant::Bf16Raw => {
165            unreachable!("Bf16Raw must use quantized_any with quant context")
166        }
167    }
168}
169
170/// 2026-09-25: Kernels and stream for load-time NVFP4 quantization (`quantize_to_nvfp4`).
171#[derive(Clone, Copy)]
172pub struct QuantizeCtx {
173    pub absmax_k: metrale_gpu_runtime::gpu::KernelHandle,
174    pub quantize_k: metrale_gpu_runtime::gpu::KernelHandle,
175    pub stream: u64,
176}
177
178/// 2026-09-25: Load a weight `[n, k]` as NVFP4 whatever its on-disk format.
179///
180/// The variant is overridden per key: a key with only a `.weight` is
181/// quantized as `Bf16Raw`, and an FP8E4M3 key with a `weight_scale` or
182/// `weight_scale_inv` but no NVFP4 global scale is dequantized as
183/// `Fp8Dequanted` (unless `variant` is already `Fp8Dequanted` or `Bf16Raw`).
184/// `n` and `k` are used only by those two re-quantizing paths.
185pub fn quantized_any(
186    store: &WeightStore,
187    prefix: &str,
188    n: usize,
189    k: usize,
190    gpu: &dyn GpuBackend,
191    variant: Nvfp4Variant,
192    qctx: QuantizeCtx,
193) -> Result<QuantizedWeight> {
194    let _t_detect = std::time::Instant::now();
195    let has_packed = store.contains(&format!("{prefix}.weight_packed"));
196    let has_scale = store.contains(&format!("{prefix}.weight_scale"));
197    let has_scale_inv = store.contains(&format!("{prefix}.weight_scale_inv"));
198    let has_only_dense =
199        !has_packed && !has_scale && !has_scale_inv && store.contains(&format!("{prefix}.weight"));
200
201    // 2026-09-25: Neither override can take an NVFP4 key: `Standard` keys have
202    // `weight_scale_2`, `CompressedTensors` keys `weight_packed`, and both are
203    // excluded here.
204    let has_fp8_dense = !has_packed
205        && !store.contains(&format!("{prefix}.weight_global_scale"))
206        && !store.contains(&format!("{prefix}.weight_scale_2"))
207        && (has_scale || has_scale_inv)
208        && store
209            .get(&format!("{prefix}.weight"))
210            .map(|w| w.dtype == WeightDtype::FP8E4M3)
211            .unwrap_or(false);
212
213    let effective_variant = if has_only_dense && !matches!(variant, Nvfp4Variant::Bf16Raw) {
214        tracing::debug!("{prefix}: no quantization metadata; falling back to runtime BF16→NVFP4");
215        Nvfp4Variant::Bf16Raw
216    } else if has_fp8_dense
217        && !matches!(variant, Nvfp4Variant::Fp8Dequanted | Nvfp4Variant::Bf16Raw)
218    {
219        // 2026-09-28: KNOWN DEVIATION under `--weight-quantization declared`. A checkpoint that
220        // declares this projection FP8 W8A8 with per-channel scales (unsloth/Qwen3.8-27B-NVFP4:
221        // attention q/k/v/o, GDN in_proj_qkv/in_proj_z/out_proj, MLP of layers 56-63) asks for its
222        // FP8 weights (`WeightQuantPolicy::wants_fp8_weights`), but no decode arm reads a per-row
223        // FP8 scale yet, so it is requantized to NVFP4 here: below the declared weight precision.
224        // The loader stamps the copy `Nvfp4Act::Wide`, so its activations stay 16-bit (no W4A4
225        // below the declared A8). Serving it as declared needs the per-row FP8 decode arms (the
226        // W8A8 decode branch); under `nvfp4` this requantization is the tier's definition.
227        tracing::debug!("{prefix}: FP8 key in an NVFP4 checkpoint; dequant FP8→BF16→NVFP4");
228        Nvfp4Variant::Fp8Dequanted
229    } else {
230        variant
231    };
232
233    let _t_detect_ns = _t_detect.elapsed().as_nanos() as u64;
234    match effective_variant {
235        Nvfp4Variant::Standard => quantized(store, prefix, gpu),
236        Nvfp4Variant::CompressedTensors => quantized_v2(store, prefix, gpu),
237        Nvfp4Variant::Fp8Dequanted => quantized_from_fp8(
238            store,
239            prefix,
240            n,
241            k,
242            gpu,
243            qctx.absmax_k,
244            qctx.quantize_k,
245            qctx.stream,
246        ),
247        Nvfp4Variant::Bf16Raw => {
248            use std::sync::atomic::{AtomicU64, Ordering};
249            static T_DETECT: AtomicU64 = AtomicU64::new(0);
250            static T_GET: AtomicU64 = AtomicU64::new(0);
251            static T_QUANT: AtomicU64 = AtomicU64::new(0);
252            static T_FREE: AtomicU64 = AtomicU64::new(0);
253            static N: AtomicU64 = AtomicU64::new(0);
254            T_DETECT.fetch_add(_t_detect_ns, Ordering::Relaxed);
255            let _t = std::time::Instant::now();
256            let w = store.get(&format!("{prefix}.weight"))?;
257            let bf16 = DenseWeight { weight: w.ptr };
258            T_GET.fetch_add(_t.elapsed().as_nanos() as u64, Ordering::Relaxed);
259            let _t = std::time::Instant::now();
260            let q = quantize_to_nvfp4(
261                &bf16,
262                n,
263                k,
264                gpu,
265                qctx.absmax_k,
266                qctx.quantize_k,
267                qctx.stream,
268            )?;
269            T_QUANT.fetch_add(_t.elapsed().as_nanos() as u64, Ordering::Relaxed);
270            let _t = std::time::Instant::now();
271            // 2026-09-25: The NVFP4 copy is a new allocation, so the store's BF16
272            // weight is freed; the store keeps the freed pointer, and nothing may
273            // read this key afterwards.
274            gpu.free(w.ptr)?;
275            T_FREE.fetch_add(_t.elapsed().as_nanos() as u64, Ordering::Relaxed);
276            let c = N.fetch_add(1, Ordering::Relaxed) + 1;
277            if c.is_multiple_of(512) {
278                let ms = |a: &AtomicU64| a.load(Ordering::Relaxed) as f64 / 1.0e6;
279                tracing::info!(
280                    "quantized_any(Bf16Raw) PROFILE after {c} calls (ms total): detect={:.1} \
281                     store_get={:.1} quantize={:.1} free={:.1} | sum={:.1} per_call={:.3}ms",
282                    ms(&T_DETECT),
283                    ms(&T_GET),
284                    ms(&T_QUANT),
285                    ms(&T_FREE),
286                    ms(&T_DETECT) + ms(&T_GET) + ms(&T_QUANT) + ms(&T_FREE),
287                    (ms(&T_DETECT) + ms(&T_GET) + ms(&T_QUANT) + ms(&T_FREE)) / c as f64,
288                );
289            }
290            Ok(q)
291        }
292    }
293}
294
295/// 2026-09-25: Dequantize a block-scaled FP8 weight `[n, k]` to BF16, quantize
296/// it to NVFP4, and free the BF16 intermediate.
297pub fn quantized_from_fp8(
298    store: &WeightStore,
299    prefix: &str,
300    n: usize,
301    k: usize,
302    gpu: &dyn GpuBackend,
303    absmax_k: metrale_gpu_runtime::gpu::KernelHandle,
304    quantize_k: metrale_gpu_runtime::gpu::KernelHandle,
305    stream: u64,
306) -> Result<QuantizedWeight> {
307    let bf16 = dequant_fp8_blockscaled_to_bf16(store, prefix, gpu)?;
308    let result = quantize_to_nvfp4(&bf16, n, k, gpu, absmax_k, quantize_k, stream)?;
309    gpu.free(bf16.weight)?;
310    Ok(result)
311}
312
313/// 2026-09-25: Dequantize a block-scaled FP8 weight to BF16, without re-quantizing.
314#[allow(dead_code)]
315pub(crate) fn dense_from_fp8(
316    store: &WeightStore,
317    prefix: &str,
318    gpu: &dyn GpuBackend,
319) -> Result<DenseWeight> {
320    dequant_fp8_blockscaled_to_bf16(store, prefix, gpu)
321}
322
323/// 2026-09-25: Load attention weights with q/k/v as the raw `weight_packed`
324/// NVFP4 pointers (typed `DenseWeight`, not dequantized) and o_proj through `quantized_v2`.
325#[allow(dead_code)]
326pub(crate) fn load_attention_qwen35(
327    store: &WeightStore,
328    layer_prefix: &str,
329    gpu: &dyn GpuBackend,
330) -> Result<AttentionWeights> {
331    let p = format!("{layer_prefix}.self_attn");
332    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu)?;
333    Ok(AttentionWeights {
334        q_proj: dense(store, &format!("{p}.q_proj.weight_packed"))?,
335        k_proj: dense(store, &format!("{p}.k_proj.weight_packed"))?,
336        v_proj: dense(store, &format!("{p}.v_proj.weight_packed"))?,
337        o_proj: quantized_v2(store, &format!("{p}.o_proj"), gpu)?,
338        q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
339        k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
340        q_norm_full: None,
341        k_norm_full: None,
342        k_scale,
343        v_scale,
344    })
345}
346
347/// 2026-09-25: `quantized_v2` under another name.
348#[allow(dead_code)]
349pub(crate) fn load_quantized_proj_qwen35(
350    store: &WeightStore,
351    prefix: &str,
352    gpu: &dyn GpuBackend,
353) -> Result<QuantizedWeight> {
354    quantized_v2(store, prefix, gpu)
355}
356
357#[cfg(test)]
358mod ep_detection_tests {
359    use super::*;
360    use metrale_config::ModelConfig;
361    use metrale_model_weights::weights::WeightStore;
362
363    /// 2026-09-25: A store whose tensors are all FP8E4M3 with shape `[1]` and NULL pointers.
364    fn store_with(names: &[String]) -> WeightStore {
365        use std::collections::HashMap;
366        let map: HashMap<String, metrale_model_weights::weights::WeightTensor> = names
367            .iter()
368            .map(|n| {
369                (
370                    n.clone(),
371                    metrale_model_weights::weights::WeightTensor {
372                        ptr: metrale_gpu_runtime::gpu::DevicePtr::NULL,
373                        shape: vec![1],
374                        dtype: metrale_model_weights::weights::WeightDtype::FP8E4M3,
375                    },
376                )
377            })
378            .collect();
379        WeightStore::from_map(map)
380    }
381
382    #[test]
383    fn alternate_layer0_fp8_dtype_is_detected_on_every_ep_rank() {
384        let mut cfg = ModelConfig::qwen3_next_80b_nvfp4();
385        cfg.quantization_config = None;
386        let store =
387            store_with(&["model.language_model.layers.0.self_attn.q_proj.weight".to_string()]);
388
389        cfg.ep_world_size = 2;
390        for ep_rank in 0..2 {
391            cfg.ep_rank = ep_rank;
392            assert_eq!(
393                detect_nvfp4_variant(&store, &cfg),
394                Nvfp4Variant::Fp8Dequanted,
395                "EP rank {ep_rank} must inspect the same layer-zero checkpoint marker"
396            );
397        }
398    }
399
400    #[test]
401    fn scale_inv_suffix_fallback_detects_an_unexpected_prefix() {
402        let mut cfg = ModelConfig::qwen3_next_80b_nvfp4();
403        cfg.quantization_config = None;
404        let store =
405            store_with(&["third_party.transformer.blocks.17.attn.q.weight_scale_inv".to_string()]);
406        assert_eq!(
407            detect_nvfp4_variant(&store, &cfg),
408            Nvfp4Variant::Fp8Dequanted
409        );
410    }
411}
metrale_model_layers/weight_map/nvfp4_detect.rs

metrale_model_layers/weight_map/
nvfp4_detect.rs
