metrale_circuit/
precision.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The precision a linear node runs at: the [`EdgePrecision`] resolver, and
4//! [`PrecisionTable`], a table-driven implementation read from
5//! `kernels/circuits/precision/<checkpoint>.toml`.
6//!
7//! The trait is shaped after `DeclaredPrecisionPlan::resolve(module) -> LayerPrecision`
8//! (branch perf/weight-quantization, `crates/config/src/precision_plan.rs`): a module path
9//! in, the weight and input-activation formats out. Once that plan and `WeightQuantPolicy`
10//! are on main, an adapter implements this trait from them and the tables go.
11//!
12//! Owner: metrale-circuit.
13//! Invariants:
14//! - A table ends with a catch-all `match = "*"` entry, so every module resolves to a
15//!   stated format; there is no implicit default.
16//! - The first matching entry wins, in file order.
17
18use serde::Deserialize;
19
20use crate::format::{Format, FormatError};
21
22/// 2026-09-28: The formats one linear module runs at.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub struct LinearFormats {
25    /// 2026-09-28: Weight format as loaded.
26    pub weight: Format,
27    /// 2026-09-28: Input activation format the kernel reads.
28    pub activation: Format,
29}
30
31/// 2026-09-28: Answers the formats of a linear module, e.g.
32/// `layers.3.mlp.down_proj`. Module paths are relative to the circuit's `module_prefix`.
33pub trait EdgePrecision {
34    /// 2026-09-28: The formats `module` runs at.
35    fn linear(&self, module: &str) -> LinearFormats;
36}
37
38/// 2026-09-28: Why a precision table did not load.
39#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
40pub enum PrecisionError {
41    /// 2026-09-28: Not valid TOML, or not the table shape.
42    #[error("precision table: {0}")]
43    Parse(String),
44    /// 2026-09-28: A format string no spelling matches.
45    #[error("precision table entry `{pattern}`: {source}")]
46    Format {
47        /// 2026-09-28: The entry's `match`.
48        pattern: String,
49        /// 2026-09-28: The format error.
50        source: FormatError,
51    },
52    /// 2026-09-28: The last entry is not `match = "*"`.
53    #[error("precision table must end with a catch-all `match = \"*\"` entry")]
54    NoCatchAll,
55}
56
57#[derive(Deserialize)]
58#[serde(deny_unknown_fields)]
59struct TableFile {
60    schema: u32,
61    checkpoint: String,
62    tier: String,
63    linear: Vec<EntryFile>,
64}
65
66#[derive(Deserialize)]
67#[serde(deny_unknown_fields)]
68struct EntryFile {
69    #[serde(rename = "match")]
70    pattern: String,
71    weight: String,
72    activation: String,
73}
74
75/// 2026-09-28: A glob-matched table of linear formats.
76#[derive(Debug, Clone, PartialEq, Eq)]
77pub struct PrecisionTable {
78    /// 2026-09-28: The checkpoint the table describes.
79    pub checkpoint: String,
80    /// 2026-09-28: The `--weight-quantization` tier it describes.
81    pub tier: String,
82    entries: Vec<(String, LinearFormats)>,
83    catch_all: LinearFormats,
84}
85
86impl PrecisionTable {
87    /// 2026-09-28: Parse a table from its TOML text.
88    pub fn parse(text: &str) -> Result<Self, PrecisionError> {
89        let file: TableFile =
90            toml::from_str(text).map_err(|e| PrecisionError::Parse(e.to_string()))?;
91        if file.schema != 1 {
92            return Err(PrecisionError::Parse(format!(
93                "schema {} (this build reads 1)",
94                file.schema
95            )));
96        }
97        let mut entries = Vec::with_capacity(file.linear.len());
98        for e in file.linear {
99            let fmt = |s: &str| {
100                Format::parse(s).map_err(|source| PrecisionError::Format {
101                    pattern: e.pattern.clone(),
102                    source,
103                })
104            };
105            let formats = LinearFormats {
106                weight: fmt(&e.weight)?,
107                activation: fmt(&e.activation)?,
108            };
109            entries.push((e.pattern, formats));
110        }
111        let catch_all = match entries.pop() {
112            Some((p, f)) if p == "*" => f,
113            _ => return Err(PrecisionError::NoCatchAll),
114        };
115        Ok(PrecisionTable {
116            checkpoint: file.checkpoint,
117            tier: file.tier,
118            entries,
119            catch_all,
120        })
121    }
122}
123
124impl EdgePrecision for PrecisionTable {
125    fn linear(&self, module: &str) -> LinearFormats {
126        self.entries
127            .iter()
128            .find(|(p, _)| glob(p, module))
129            .map_or(self.catch_all, |(_, f)| *f)
130    }
131}
132
133/// 2026-10-02: The fused-experts module an expert projection belongs to: `<p>.experts` for
134/// `<p>.experts.<i>.<proj>`. A checkpoint that quantizes the routed experts as one module (a
135/// ModelOpt `quantized_layers` entry naming `...mlp.experts`, nvidia/Qwen3.6-35B-A3B-NVFP4)
136/// declares every expert projection through it; the circuit asks expert 0's projection.
137pub(crate) fn expert_container(module: &str) -> Option<&str> {
138    let at = module.rfind(".experts.")?;
139    let rest = &module[at + ".experts.".len()..];
140    let (index, proj) = rest.split_once('.')?;
141    (!index.is_empty() && index.bytes().all(|b| b.is_ascii_digit()) && !proj.is_empty())
142        .then(|| &module[..at + ".experts".len()])
143}
144
145/// 2026-09-28: `*` matches any run of characters, dots included; everything else is literal.
146pub(crate) fn glob(pattern: &str, text: &str) -> bool {
147    let parts: Vec<&str> = pattern.split('*').collect();
148    if parts.len() == 1 {
149        return pattern == text;
150    }
151    let (first, last) = (parts[0], parts[parts.len() - 1]);
152    if !text.starts_with(first) || text.len() < first.len() + last.len() || !text.ends_with(last) {
153        return false;
154    }
155    let mut rest = &text[first.len()..text.len() - last.len()];
156    for mid in &parts[1..parts.len() - 1] {
157        match rest.find(mid) {
158            Some(at) => rest = &rest[at + mid.len()..],
159            None => return false,
160        }
161    }
162    true
163}
164
165#[cfg(test)]
166#[path = "precision_tests.rs"]
167mod precision_tests;
metrale_circuit/precision.rs

metrale_circuit/
precision.rs
