metrale_circuit/
declared_precision.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: [`DeclaredPrecision`], the [`EdgePrecision`] that gives every projection the
4//! formats its checkpoint DECLARES (`DeclaredPrecisionPlan`), before any serving policy: the
5//! weight as stored, and the input activation as the checkpoint declares it (FP8 per token,
6//! per tensor with a static scale, NVFP4 group 16, or 16-bit).
7//!
8//! Owner: metrale-circuit.
9//! Invariants:
10//! - A declared format the circuit has no edge format for (an integer scheme, an FP8 weight
11//!   with an unstated granularity, ...) is recorded, never guessed; [`DeclaredPrecision::refusals`]
12//!   lists them and the caller refuses the checkpoint.
13//! - A module the plan does not cover is 16-bit (`LayerPrecision::UNQUANTIZED`); an expert
14//!   projection is covered by its fused-experts module's declaration (2026-10-02).
15
16use std::cell::RefCell;
17
18use metrale_config::DeclaredPrecisionPlan;
19use metrale_config::precision_plan::{Granularity, NumKind, Operand};
20
21use crate::format::{Format, Scale};
22use crate::precision::{EdgePrecision, LinearFormats};
23
24/// 2026-09-30: The checkpoint's declared formats as an [`EdgePrecision`].
25pub struct DeclaredPrecision<'a> {
26    plan: &'a DeclaredPrecisionPlan,
27    refusals: RefCell<Vec<String>>,
28}
29
30impl<'a> DeclaredPrecision<'a> {
31    /// 2026-09-30: Answer from `plan`.
32    pub fn new(plan: &'a DeclaredPrecisionPlan) -> Self {
33        DeclaredPrecision {
34            plan,
35            refusals: RefCell::new(Vec::new()),
36        }
37    }
38
39    /// 2026-09-30: Every module asked whose declaration has no circuit format, with why.
40    pub fn refusals(&self) -> Vec<String> {
41        self.refusals.borrow().clone()
42    }
43
44    fn refuse(&self, module: &str, what: &str, o: Operand) {
45        self.refusals.borrow_mut().push(format!(
46            "`{module}` declares a {what} the circuit has no format for: {o:?}"
47        ));
48    }
49}
50
51/// 2026-09-30: The weight format `o` declares, if the circuit has one.
52fn weight_format(o: Operand) -> Option<Format> {
53    if o.kind != NumKind::Float {
54        return None;
55    }
56    match (o.bits, o.granularity) {
57        (4, Granularity::TensorGroup(g) | Granularity::Group(g)) => {
58            Some(Format::Nvfp4 { group: g })
59        }
60        (8, g) => Some(Format::Fp8E4m3 {
61            scale: match g {
62                Granularity::Tensor => Scale::PerTensor,
63                Granularity::Channel => Scale::PerChannel,
64                Granularity::Group(n) => Scale::Group(n),
65                Granularity::Block(r, c) => Scale::Block(r, c),
66                Granularity::Token | Granularity::TensorGroup(_) | Granularity::Unstated => {
67                    return None;
68                }
69            },
70        }),
71        _ => None,
72    }
73}
74
75/// 2026-09-30: The activation (edge) format `o` declares, if the circuit has one.
76fn activation_format(o: Operand) -> Option<Format> {
77    if o.kind != NumKind::Float {
78        return None;
79    }
80    match (o.bits, o.granularity) {
81        (4, Granularity::TensorGroup(g) | Granularity::Group(g)) => {
82            Some(Format::Nvfp4 { group: g })
83        }
84        (8, Granularity::Tensor) => Some(Format::Fp8E4m3 {
85            scale: Scale::PerTensor,
86        }),
87        (8, Granularity::Token | Granularity::Unstated) => Some(Format::Fp8E4m3 {
88            scale: Scale::PerToken,
89        }),
90        (8, Granularity::Group(g)) => Some(Format::Fp8E4m3 {
91            scale: Scale::Group(g),
92        }),
93        _ => None,
94    }
95}
96
97impl EdgePrecision for DeclaredPrecision<'_> {
98    fn linear(&self, module: &str) -> LinearFormats {
99        // 2026-10-02: An expert projection the plan does not name is declared by its experts
100        // module, where the checkpoint quantizes that as one (`precision::expert_container`); a
101        // projection the checkpoint ignores keeps its own (16-bit) answer.
102        let mut declared = self.plan.resolve(module);
103        if declared.weight.is_none()
104            && !self.plan.ignore.iter().any(|t| t.matches_name(module))
105            && let Some(container) = crate::precision::expert_container(module)
106        {
107            declared = self.plan.resolve(container);
108        }
109        let weight = match declared.weight {
110            None => Format::Bf16,
111            Some(o) => weight_format(o).unwrap_or_else(|| {
112                self.refuse(module, "weight", o);
113                Format::Bf16
114            }),
115        };
116        let activation = match declared.activation {
117            None => Format::Bf16,
118            Some(o) => activation_format(o).unwrap_or_else(|| {
119                self.refuse(module, "activation", o);
120                Format::Bf16
121            }),
122        };
123        LinearFormats { weight, activation }
124    }
125}
126
127#[cfg(test)]
128#[path = "declared_precision_tests.rs"]
129mod declared_precision_tests;
metrale_circuit/declared_precision.rs

metrale_circuit/
declared_precision.rs
