metrale_config/
precision_plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The declared precision plan: what a checkpoint's `quantization_config`
4//! declares, per linear layer, for the weight AND the input activations.
5//!
6//! This is the one place that reads the precision half of `quantization_config`
7//! (compressed-tensors `config_groups`, ModelOpt `quantized_layers` / `quant_algo`, and the
8//! HF `fp8` block scheme). Loaders and kernel dispatch ask [`DeclaredPrecisionPlan::resolve`]
9//! for a module and run what it says by default. A lever may only go BELOW what the plan
10//! declares; running above it is not a silent default either.
11//!
12//! Owner: config (quantization).
13//! Invariants:
14//! - Pure: no environment, no I/O. A malformed block is an error, never a guess.
15//! - A module the plan does not cover (ignored, unmatched, or no quantization at all)
16//!   resolves to [`LayerPrecision::UNQUANTIZED`]: 16-bit weights and activations.
17
18use anyhow::Result;
19
20#[path = "precision_plan/parse.rs"]
21mod parse;
22#[path = "precision_plan/target.rs"]
23mod target;
24
25pub use target::{IgnoreDialect, Target};
26
27/// 2026-09-28: Element type of a quantized operand.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum NumKind {
30    /// 2026-09-28: A float format (E4M3 at 8 bits, E2M1 at 4 bits).
31    Float,
32    /// 2026-09-28: An integer format.
33    Int,
34}
35
36/// 2026-09-28: How many values share one scale.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum Granularity {
39    /// 2026-09-28: One scale for the whole tensor.
40    Tensor,
41    /// 2026-09-28: One scale per output channel (weights).
42    Channel,
43    /// 2026-09-28: One scale per token (activations).
44    Token,
45    /// 2026-09-28: One scale per `n` consecutive values along K.
46    Group(u32),
47    /// 2026-09-28: One scale per `n` values plus a per-tensor global scale
48    /// (compressed-tensors `tensor_group`, ModelOpt NVFP4).
49    TensorGroup(u32),
50    /// 2026-09-28: One scale per `[rows, cols]` block (weights).
51    Block(u32, u32),
52    /// 2026-09-28: Not stated by the checkpoint.
53    Unstated,
54}
55
56/// 2026-09-28: When an operand's scales are computed.
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum ScaleTiming {
59    /// 2026-09-28: Calibrated and stored in the checkpoint.
60    Static,
61    /// 2026-09-28: Computed at run time.
62    Dynamic,
63    /// 2026-09-28: Local scales at run time under a stored global scale
64    /// (compressed-tensors `dynamic: "local"`).
65    Local,
66}
67
68/// 2026-09-28: One quantized operand as declared.
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub struct Operand {
71    /// 2026-09-28: Float or int.
72    pub kind: NumKind,
73    /// 2026-09-28: Bits per element.
74    pub bits: u8,
75    /// 2026-09-28: Scale sharing.
76    pub granularity: Granularity,
77    /// 2026-09-28: Scale timing.
78    pub timing: ScaleTiming,
79}
80
81impl Operand {
82    /// 2026-09-28: NVFP4 as ModelOpt and compressed-tensors declare it: E2M1, group 16
83    /// under a per-tensor scale.
84    pub const NVFP4: Operand = Operand {
85        kind: NumKind::Float,
86        bits: 4,
87        granularity: Granularity::TensorGroup(16),
88        timing: ScaleTiming::Static,
89    };
90
91    /// 2026-09-28: A 4-bit float (NVFP4 or MXFP4).
92    pub fn is_fp4(&self) -> bool {
93        self.kind == NumKind::Float && self.bits == 4
94    }
95
96    /// 2026-09-28: An 8-bit float (E4M3).
97    pub fn is_fp8(&self) -> bool {
98        self.kind == NumKind::Float && self.bits == 8
99    }
100}
101
102/// 2026-09-28: What one linear layer declares. `None` means that operand is not quantized
103/// (16-bit).
104#[derive(Debug, Clone, Copy, PartialEq, Eq)]
105pub struct LayerPrecision {
106    /// 2026-09-28: The weight's declared format.
107    pub weight: Option<Operand>,
108    /// 2026-09-28: The input activations' declared format.
109    pub activation: Option<Operand>,
110}
111
112impl LayerPrecision {
113    /// 2026-09-28: 16-bit weights and activations.
114    pub const UNQUANTIZED: LayerPrecision = LayerPrecision {
115        weight: None,
116        activation: None,
117    };
118
119    /// 2026-09-28: The activations are declared 4-bit float (W4A4 on an FP4 weight).
120    pub fn activation_is_fp4(&self) -> bool {
121        self.activation.is_some_and(|a| a.is_fp4())
122    }
123
124    /// 2026-09-28: The activations are declared 8-bit float.
125    pub fn activation_is_fp8(&self) -> bool {
126        self.activation.is_some_and(|a| a.is_fp8())
127    }
128
129    /// 2026-09-28: A short label, e.g. `W4A4`, `W8A8`, `W4A16`, `W16A16`.
130    pub fn label(&self) -> String {
131        let bits = |o: Option<Operand>| o.map_or(16, |o| o.bits);
132        format!("W{}A{}", bits(self.weight), bits(self.activation))
133    }
134}
135
136/// 2026-09-28: Which `quantization_config` dialect the plan came from.
137#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
138pub enum PlanSource {
139    /// 2026-09-28: No precision declared (no block, or a block with no scheme).
140    #[default]
141    Undeclared,
142    /// 2026-09-28: compressed-tensors `config_groups`.
143    CompressedTensors,
144    /// 2026-09-28: ModelOpt (`quantized_layers`, `config_groups` or a global `quant_algo`).
145    ModelOpt,
146    /// 2026-09-28: The HF `fp8` method (`weight_block_size`, `activation_scheme`).
147    Fp8,
148}
149
150/// 2026-09-28: One scheme and the modules it targets.
151#[derive(Debug, Clone)]
152pub struct Rule {
153    /// 2026-09-28: Targets, in the order the checkpoint lists them.
154    pub targets: Vec<Target>,
155    /// 2026-09-28: What the targeted modules declare.
156    pub precision: LayerPrecision,
157}
158
159/// 2026-09-28: The per-layer precision a checkpoint declares.
160#[derive(Debug, Clone, Default)]
161pub struct DeclaredPrecisionPlan {
162    /// 2026-09-28: Dialect.
163    pub source: PlanSource,
164    /// 2026-09-28: Schemes.
165    pub rules: Vec<Rule>,
166    /// 2026-09-28: Modules left unquantized.
167    pub ignore: Vec<Target>,
168}
169
170impl DeclaredPrecisionPlan {
171    /// 2026-09-28: A plan that declares nothing (no `quantization_config`).
172    pub const UNDECLARED: DeclaredPrecisionPlan = DeclaredPrecisionPlan {
173        source: PlanSource::Undeclared,
174        rules: Vec::new(),
175        ignore: Vec::new(),
176    };
177
178    /// 2026-09-28: Parse the `quantization_config` object (the value, not the whole
179    /// config.json). A ModelOpt `hf_quant_config.json` nested under `"quantization"` is
180    /// accepted as is.
181    pub fn from_quantization_config(qc: &serde_json::Value) -> Result<Self> {
182        parse::plan(qc)
183    }
184
185    /// 2026-09-28: What `module` (a checkpoint module path such as
186    /// `model.language_model.layers.3.mlp.gate_proj`, no `.weight` suffix) declares.
187    ///
188    /// Ignored modules are unquantized. Otherwise the compressed-tensors precedence
189    /// applies: exact names first, then patterns ordered by their text, then class
190    /// targets (`Linear`), and the first match wins.
191    pub fn resolve(&self, module: &str) -> LayerPrecision {
192        if self.ignore.iter().any(|t| t.matches_name(module)) {
193            return LayerPrecision::UNQUANTIZED;
194        }
195        let mut best: Option<(u8, &str, LayerPrecision)> = None;
196        for rule in &self.rules {
197            for t in &rule.targets {
198                let Some(tier) = t.tier_for(module) else {
199                    continue;
200                };
201                let key = (tier, t.text());
202                if best.is_none_or(|(bt, bs, _)| key < (bt, bs)) {
203                    best = Some((tier, t.text(), rule.precision));
204                }
205            }
206        }
207        best.map_or(LayerPrecision::UNQUANTIZED, |(_, _, p)| p)
208    }
209
210    /// 2026-09-28: True when no module declares anything.
211    pub fn is_undeclared(&self) -> bool {
212        self.source == PlanSource::Undeclared
213    }
214}
215
216#[cfg(test)]
217#[path = "precision_plan_tests.rs"]
218mod precision_plan_tests;
metrale_config/precision_plan.rs

metrale_config/
precision_plan.rs
