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}