spark_model/quant_format/
compressed_tensors.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Neural Magic `llm-compressor` / compressed-tensors NVFP4 serialization.
4//!
5//! The dominant community NVFP4 format, shipped by `Sehyo/*`, `RedHatAI/*`,
6//! `nm-testing/*`, and most third-party re-quants. Tensor-name convention:
7//!
8//! | field                | tensor name            | dtype         |
9//! | -------------------- | ---------------------- | ------------- |
10//! | packed FP4 payload   | `.weight_packed`       | uint8 packed  |
11//! | per-group FP8 scales | `.weight_scale`        | float8_e4m3   |
12//! | per-tensor scalar    | `.weight_global_scale` | f32 scalar    |
13//! | activation scale     | `.input_global_scale`  | f32 scalar    |
14//!
15//! Scale convention: `weight_global_scale` is the RECIPROCAL of ModelOpt's
16//! `weight_scale_2` (verified empirically — see `quantized_v2` comment in
17//! `weight_map.rs`).
18//!
19//! Unquantized modules are declared either in the top-level `ignore`
20//! array or in `config_groups.group_N.targets` / `exclude_modules`
21//! (vLLM style). Both are folded into a single list during config parse.
22//!
23//! Maps to the existing [`Nvfp4Variant::CompressedTensors`] dispatch.
24
25use crate::quant_format::{QuantFormat, module_matches_pattern};
26use crate::weight_map::Nvfp4Variant;
27
28/// compressed-tensors NVFP4 checkpoint.
29#[derive(Debug)]
30pub struct CompressedTensorsFormat {
31    /// `format` string from config (e.g. `"nvfp4-pack-quantized"`). Log only.
32    pub format: String,
33    /// Module-path globs that stay BF16 rather than NVFP4.
34    pub ignore_modules: Vec<String>,
35}
36
37impl CompressedTensorsFormat {
38    pub fn new(format: String, ignore_modules: Vec<String>) -> Self {
39        Self {
40            format,
41            ignore_modules,
42        }
43    }
44}
45
46impl QuantFormat for CompressedTensorsFormat {
47    fn name(&self) -> &'static str {
48        "compressed-tensors"
49    }
50
51    fn base_variant(&self) -> Nvfp4Variant {
52        Nvfp4Variant::CompressedTensors
53    }
54
55    fn is_ignored(&self, module_path: &str) -> bool {
56        self.ignore_modules
57            .iter()
58            .any(|pat| module_matches_pattern(module_path, pat))
59    }
60}