metrale_circuit/
precision_policy.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: [`PolicyPrecision`], the [`EdgePrecision`] the engine's own precision answers
4//! give: the checkpoint's declared plan (`DeclaredPrecisionPlan`) under the served tier
5//! (`WeightQuantPolicy`), plus the formats the engine chooses itself (a recipe's lm_head, the
6//! MTP head), which the checkpoint does not declare.
7//!
8//! Owner: metrale-circuit.
9//! Invariants:
10//! - Engine formats are stated, never inferred: a module they name is answered from them
11//!   first, in order; every other module from the policy.
12//! - Under the `declared` tier a module declared FP8 runs its FP8 weights, with FP8
13//!   activations exactly when `WeightQuantPolicy::fp8_decode_act` says so; one declared FP4
14//!   runs NVFP4 weights, with FP4 activations when the checkpoint declares them. Under
15//!   `nvfp4` every quantized module runs NVFP4 weights and 16-bit activations.
16//! - An unquantized module is BF16 weights and activations.
17
18use metrale_config::precision_plan::{Granularity, Operand};
19use metrale_config::weight_quantization::{ActFormat, KernelCaps};
20use metrale_config::{
21    DeclaredPrecisionPlan, W4a4Downcast, WeightQuantPolicy, WeightQuantTier, WeightQuantization,
22};
23use serde::Deserialize;
24
25use crate::format::{Format, Scale};
26use crate::precision::{EdgePrecision, LinearFormats, PrecisionError, glob};
27
28/// 2026-09-28: The NVFP4 group size the engine's kernels read.
29const NVFP4_GROUP: u32 = 16;
30
31/// 2026-09-28: A checkpoint's `quantization_config`, as the circuit fixtures store it
32/// (`kernels/circuits/checkpoints/<name>.toml`).
33#[derive(Debug, Clone)]
34pub struct CheckpointPlan {
35    /// 2026-09-28: The checkpoint id.
36    pub checkpoint: String,
37    /// 2026-09-28: Its declared plan.
38    pub plan: DeclaredPrecisionPlan,
39}
40
41#[derive(Deserialize)]
42#[serde(deny_unknown_fields)]
43struct FixtureFile {
44    schema: u32,
45    checkpoint: String,
46    quantization_config: String,
47}
48
49impl CheckpointPlan {
50    /// 2026-09-28: Parse a fixture: `checkpoint` and the `quantization_config` object as JSON.
51    pub fn parse(text: &str) -> Result<Self, PrecisionError> {
52        let f: FixtureFile =
53            toml::from_str(text).map_err(|e| PrecisionError::Parse(e.to_string()))?;
54        if f.schema != 1 {
55            return Err(PrecisionError::Parse(format!(
56                "checkpoint fixture schema {} (this build reads 1)",
57                f.schema
58            )));
59        }
60        let qc: serde_json::Value = serde_json::from_str(&f.quantization_config)
61            .map_err(|e| PrecisionError::Parse(format!("quantization_config: {e}")))?;
62        let plan = DeclaredPrecisionPlan::from_quantization_config(&qc)
63            .map_err(|e| PrecisionError::Parse(format!("quantization_config: {e:#}")))?;
64        Ok(CheckpointPlan {
65            checkpoint: f.checkpoint,
66            plan,
67        })
68    }
69}
70
71/// 2026-09-28: The tier named `name` (`declared`, `nvfp4`).
72pub fn tier_named(name: &str) -> Result<WeightQuantTier, PrecisionError> {
73    let tier = WeightQuantization::ALL
74        .into_iter()
75        .find(|t| t.name() == name)
76        .ok_or_else(|| {
77            PrecisionError::Parse(format!("unknown weight-quantization tier `{name}`"))
78        })?;
79    WeightQuantTier::new(tier, W4a4Downcast::Off).map_err(|e| PrecisionError::Parse(e.to_string()))
80}
81
82/// 2026-09-28: The kernel capabilities named `names` (field names of `KernelCaps`).
83pub fn caps_named(names: &[String]) -> Result<KernelCaps, PrecisionError> {
84    let mut caps = KernelCaps::default();
85    for n in names {
86        let bit = match n.as_str() {
87            "w8a8_decode" => &mut caps.w8a8_decode,
88            "w8a8_moe_decode" => &mut caps.w8a8_moe_decode,
89            "w8a8_block_scaled_decode" => &mut caps.w8a8_block_scaled_decode,
90            "fp8_lm_head_batched" => &mut caps.fp8_lm_head_batched,
91            other => {
92                return Err(PrecisionError::Parse(format!(
93                    "unknown kernel cap `{other}`"
94                )));
95            }
96        };
97        *bit = true;
98    }
99    Ok(caps)
100}
101
102/// 2026-09-28: The engine's precision answers as an [`EdgePrecision`].
103pub struct PolicyPrecision<'a> {
104    policy: WeightQuantPolicy<'a>,
105    engine: &'a [(String, LinearFormats)],
106}
107
108impl<'a> PolicyPrecision<'a> {
109    /// 2026-09-28: `engine` lists the engine-chosen formats, first match wins.
110    pub fn new(policy: WeightQuantPolicy<'a>, engine: &'a [(String, LinearFormats)]) -> Self {
111        PolicyPrecision { policy, engine }
112    }
113}
114
115const BF16: LinearFormats = LinearFormats {
116    weight: Format::Bf16,
117    activation: Format::Bf16,
118};
119
120fn fp8_weight(w: Operand) -> Format {
121    Format::Fp8E4m3 {
122        scale: match w.granularity {
123            Granularity::Block(r, c) => Scale::Block(r, c),
124            Granularity::Group(g) => Scale::Group(g),
125            Granularity::Tensor => Scale::PerTensor,
126            _ => Scale::PerChannel,
127        },
128    }
129}
130
131fn fp8_act(a: Option<Operand>) -> Format {
132    Format::Fp8E4m3 {
133        scale: match a.map(|a| a.granularity) {
134            Some(Granularity::Group(g)) => Scale::Group(g),
135            Some(Granularity::Tensor) => Scale::PerTensor,
136            _ => Scale::PerToken,
137        },
138    }
139}
140
141impl EdgePrecision for PolicyPrecision<'_> {
142    fn linear(&self, module: &str) -> LinearFormats {
143        if let Some((_, f)) = self.engine.iter().find(|(p, _)| glob(p, module)) {
144            return *f;
145        }
146        // 2026-10-02: As `DeclaredPrecision`: an expert projection the plan does not name is
147        // declared by its experts module (`precision::expert_container`).
148        let mut declared = self.policy.declared(module);
149        if declared.weight.is_none()
150            && let Some(container) = crate::precision::expert_container(module)
151        {
152            declared = self.policy.declared(container);
153        }
154        let Some(w) = declared.weight else {
155            return BF16;
156        };
157        let nvfp4 = Format::Nvfp4 { group: NVFP4_GROUP };
158        if !self.policy.follows_plan() {
159            return LinearFormats {
160                weight: nvfp4,
161                activation: Format::Bf16,
162            };
163        }
164        if w.is_fp8() {
165            let activation = match self.policy.fp8_decode_act(module) {
166                Some(ActFormat::Fp8) => fp8_act(declared.activation),
167                _ => Format::Bf16,
168            };
169            return LinearFormats {
170                weight: fp8_weight(w),
171                activation,
172            };
173        }
174        LinearFormats {
175            weight: nvfp4,
176            activation: if declared.activation_is_fp4() {
177                nvfp4
178            } else {
179                Format::Bf16
180            },
181        }
182    }
183}
184
185#[cfg(test)]
186#[path = "precision_policy_tests.rs"]
187mod precision_policy_tests;
metrale_circuit/precision_policy.rs

metrale_circuit/
precision_policy.rs
