metrale_config/
activation_quantization.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: `--activation-quantization`: which activation precision each decode projection
4//! runs at, as a function of the projection family and of the number of rows in the launch
5//! (decode batch rows, MTP verify rows included).
6//!
7//! A value is either one format for every row count, or a ladder that maps row ranges to
8//! formats, optionally overridden per projection family:
9//!
10//! ```text
11//! spec   := ladder ("," family ":" ladder)*
12//! ladder := format | rung (";" rung)*
13//! rung   := range "=" format
14//! range  := N | N "-" | N "-" M              (1-based rows, inclusive)
15//! format := "bf16" | "fp8" | "nvfp4" | "declared" | "adaptive"
16//! family := "gdn" | "attn" | "ffn" | "moe" | "lm_head"
17//! ```
18//!
19//! - `bf16`, `fp8` and `nvfp4` run that activation format through one fixed-order kernel family
20//!   per weight format, so a row's output bits do not depend on the row count or on the other
21//!   rows. `declared` is the checkpoint's declared input-activation format per layer
22//!   (`quantization_config`), run the same way.
23//! - `adaptive` is today's routing: each dispatch site picks its kernel by row count as it did
24//!   before this flag, byte for byte. The named value `adaptive` is the ladder `1-=adaptive`; a
25//!   rung may also hand only its range to it (`1=bf16;2-=adaptive`).
26//! - A ladder's rungs start at row 1, leave no gap and do not overlap, and the last one is
27//!   open-ended. A family override replaces the base ladder for that family only.
28//!
29//! The same shape, a rule per (family, row range), is what the auto-fuser's circuit rules
30//! express per row range (kernels/circuits); a later step makes the circuit read this value.
31//!
32//! Owner: config (quantization).
33//! Invariants:
34//! - Pure: no environment, no I/O, no process state. The serve parses the flag here once and
35//!   publishes the result; every dispatch site asks [`ActivationQuantization::route`].
36//! - Its `Display` form is canonical and parses back to an equal value.
37
38use std::collections::BTreeMap;
39use std::fmt;
40
41use anyhow::{Result, bail, ensure};
42
43/// 2026-09-30: The activation format of one rung.
44#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
45pub enum ActQuantFormat {
46    /// 2026-09-30: 16-bit activations.
47    Bf16,
48    /// 2026-09-30: FP8 E4M3 activations, one scale per row (and per 128-column group where
49    /// the weights are block-scaled).
50    Fp8,
51    /// 2026-09-30: NVFP4 activations (E2M1, one E4M3 scale per 16 columns of a row).
52    Nvfp4,
53    /// 2026-09-30: The checkpoint's declared input-activation format of the layer.
54    Declared,
55    /// 2026-09-30: Today's per-row-count routing.
56    Adaptive,
57}
58
59impl ActQuantFormat {
60    /// 2026-09-30: Every format, in the order the flag lists them.
61    pub const ALL: [Self; 5] = [
62        Self::Bf16,
63        Self::Fp8,
64        Self::Nvfp4,
65        Self::Declared,
66        Self::Adaptive,
67    ];
68
69    /// 2026-09-30: The flag value, recipe value and record value.
70    pub fn name(self) -> &'static str {
71        match self {
72            Self::Bf16 => "bf16",
73            Self::Fp8 => "fp8",
74            Self::Nvfp4 => "nvfp4",
75            Self::Declared => "declared",
76            Self::Adaptive => "adaptive",
77        }
78    }
79
80    fn parse(s: &str) -> Result<Self> {
81        Self::ALL
82            .into_iter()
83            .find(|f| f.name() == s)
84            .ok_or_else(|| {
85                anyhow::anyhow!(
86                    "unknown activation format {s:?} (one of {})",
87                    Self::ALL.map(Self::name).join(", ")
88                )
89            })
90    }
91}
92
93/// 2026-09-30: The projection families a family override can name.
94#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
95pub enum ProjFamily {
96    /// 2026-09-30: GDN (linear attention) in_proj_qkvz, in_proj_ba and out_proj.
97    Gdn,
98    /// 2026-09-30: Full attention q (+gate), k, v and o.
99    Attn,
100    /// 2026-09-30: Dense MLP gate, up and down.
101    Ffn,
102    /// 2026-09-30: MoE router, routed experts and shared expert.
103    Moe,
104    /// 2026-09-30: The LM head.
105    LmHead,
106}
107
108impl ProjFamily {
109    /// 2026-09-30: Every family, in the order records list them.
110    pub const ALL: [Self; 5] = [Self::Gdn, Self::Attn, Self::Ffn, Self::Moe, Self::LmHead];
111
112    pub fn name(self) -> &'static str {
113        match self {
114            Self::Gdn => "gdn",
115            Self::Attn => "attn",
116            Self::Ffn => "ffn",
117            Self::Moe => "moe",
118            Self::LmHead => "lm_head",
119        }
120    }
121
122    fn parse(s: &str) -> Result<Self> {
123        Self::ALL
124            .into_iter()
125            .find(|f| f.name() == s)
126            .ok_or_else(|| {
127                anyhow::anyhow!(
128                    "unknown projection family {s:?} (one of {})",
129                    Self::ALL.map(Self::name).join(", ")
130                )
131            })
132    }
133}
134
135/// 2026-09-30: One rung: rows `lo..=hi` (`hi` `None`: no upper bound) run `format`.
136#[derive(Clone, Copy, Debug, PartialEq, Eq)]
137pub struct Rung {
138    pub lo: u32,
139    pub hi: Option<u32>,
140    pub format: ActQuantFormat,
141}
142
143/// 2026-09-30: Rungs covering every row count from 1, in order, without gaps or overlaps.
144#[derive(Clone, Debug, PartialEq, Eq)]
145pub struct Ladder(Vec<Rung>);
146
147impl Ladder {
148    /// 2026-09-30: One format at every row count.
149    pub fn uniform(format: ActQuantFormat) -> Self {
150        Self(vec![Rung {
151            lo: 1,
152            hi: None,
153            format,
154        }])
155    }
156
157    pub fn rungs(&self) -> &[Rung] {
158        &self.0
159    }
160
161    /// 2026-09-30: The format of `rows` (at least 1).
162    pub fn format(&self, rows: u32) -> ActQuantFormat {
163        let rows = rows.max(1);
164        self.0
165            .iter()
166            .find(|r| rows >= r.lo && r.hi.is_none_or(|hi| rows <= hi))
167            .map(|r| r.format)
168            .expect("a validated ladder covers every row count from 1")
169    }
170
171    /// 2026-09-30: Whether every rung runs one format (no `adaptive` rung and one rung).
172    pub fn is_uniform_invariant(&self) -> bool {
173        self.0.len() == 1 && self.0[0].format != ActQuantFormat::Adaptive
174    }
175
176    fn parse(s: &str) -> Result<Self> {
177        let s = s.trim();
178        ensure!(!s.is_empty(), "empty ladder");
179        if !s.contains('=') {
180            return Ok(Self::uniform(ActQuantFormat::parse(s)?));
181        }
182        let mut rungs = Vec::new();
183        for part in s.split(';') {
184            let (range, format) = part
185                .split_once('=')
186                .ok_or_else(|| anyhow::anyhow!("rung {part:?} is not RANGE=FORMAT"))?;
187            let (lo, hi) = parse_range(range.trim())?;
188            rungs.push(Rung {
189                lo,
190                hi,
191                format: ActQuantFormat::parse(format.trim())?,
192            });
193        }
194        Self::validated(rungs)
195    }
196
197    fn validated(rungs: Vec<Rung>) -> Result<Self> {
198        let mut next = 1u32;
199        for (i, r) in rungs.iter().enumerate() {
200            ensure!(
201                r.lo == next,
202                "rung {} starts at row {} but must start at row {next} (rungs run from row 1 \
203                 in order, without gaps or overlaps)",
204                i + 1,
205                r.lo
206            );
207            match r.hi {
208                Some(hi) => {
209                    ensure!(hi >= r.lo, "rung {} ends at {hi}, before it starts", i + 1);
210                    next = hi
211                        .checked_add(1)
212                        .ok_or_else(|| anyhow::anyhow!("rung {} ends past u32", i + 1))?;
213                }
214                None => {
215                    ensure!(
216                        i + 1 == rungs.len(),
217                        "rung {} is open-ended but rungs follow it",
218                        i + 1
219                    );
220                    return Ok(Self(merge_equal(rungs)));
221                }
222            }
223        }
224        bail!("the last rung must be open-ended (N-), so every row count has a format")
225    }
226}
227
228/// 2026-09-30: Adjacent rungs of one format merged, so equal ladders compare and print equal.
229fn merge_equal(rungs: Vec<Rung>) -> Vec<Rung> {
230    let mut out: Vec<Rung> = Vec::with_capacity(rungs.len());
231    for r in rungs {
232        match out.last_mut() {
233            Some(prev) if prev.format == r.format => prev.hi = r.hi,
234            _ => out.push(r),
235        }
236    }
237    out
238}
239
240fn parse_row(s: &str) -> Result<u32> {
241    let n: u32 = s
242        .parse()
243        .map_err(|_| anyhow::anyhow!("row bound {s:?} is not a whole number"))?;
244    ensure!(n >= 1, "rows are counted from 1");
245    Ok(n)
246}
247
248fn parse_range(s: &str) -> Result<(u32, Option<u32>)> {
249    match s.split_once('-') {
250        None => {
251            let n = parse_row(s)?;
252            Ok((n, Some(n)))
253        }
254        Some((lo, "")) => Ok((parse_row(lo)?, None)),
255        Some((lo, hi)) => Ok((parse_row(lo)?, Some(parse_row(hi)?))),
256    }
257}
258
259impl fmt::Display for Ladder {
260    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
261        if self.0.len() == 1 {
262            return f.write_str(self.0[0].format.name());
263        }
264        for (i, r) in self.0.iter().enumerate() {
265            if i > 0 {
266                f.write_str(";")?;
267            }
268            match r.hi {
269                Some(hi) if hi == r.lo => write!(f, "{}", r.lo)?,
270                Some(hi) => write!(f, "{}-{hi}", r.lo)?,
271                None => write!(f, "{}-", r.lo)?,
272            }
273            write!(f, "={}", r.format.name())?;
274        }
275        Ok(())
276    }
277}
278
279/// 2026-09-30: A parsed `--activation-quantization` value.
280#[derive(Clone, Debug, PartialEq, Eq)]
281pub struct ActivationQuantization {
282    base: Ladder,
283    families: BTreeMap<ProjFamily, Ladder>,
284}
285
286impl Default for ActivationQuantization {
287    /// 2026-09-30: The flag's default: the checkpoint's declared formats, row-invariant.
288    fn default() -> Self {
289        Self::uniform(ActQuantFormat::Declared)
290    }
291}
292
293impl ActivationQuantization {
294    /// 2026-09-30: One format for every family and row count.
295    pub fn uniform(format: ActQuantFormat) -> Self {
296        Self {
297            base: Ladder::uniform(format),
298            families: BTreeMap::new(),
299        }
300    }
301
302    /// 2026-09-30: Today's routing everywhere, as recipes pin it.
303    pub fn adaptive() -> Self {
304        Self::uniform(ActQuantFormat::Adaptive)
305    }
306
307    pub fn parse(s: &str) -> Result<Self> {
308        let mut parts = s.split(',');
309        let base = Ladder::parse(parts.next().unwrap_or(""))
310            .map_err(|e| anyhow::anyhow!("--activation-quantization {s:?}: {e}"))?;
311        let mut families = BTreeMap::new();
312        for part in parts {
313            let (fam, ladder) = part.split_once(':').ok_or_else(|| {
314                anyhow::anyhow!(
315                    "--activation-quantization {s:?}: override {part:?} is not FAMILY:LADDER"
316                )
317            })?;
318            let fam = ProjFamily::parse(fam.trim())
319                .map_err(|e| anyhow::anyhow!("--activation-quantization {s:?}: {e}"))?;
320            let ladder = Ladder::parse(ladder).map_err(|e| {
321                anyhow::anyhow!("--activation-quantization {s:?}, {}: {e}", fam.name())
322            })?;
323            ensure!(
324                families.insert(fam, ladder).is_none(),
325                "--activation-quantization {s:?}: family {} is overridden twice",
326                fam.name()
327            );
328        }
329        // 2026-09-30: An override equal to the base says nothing; drop it so the canonical
330        // form is unique.
331        families.retain(|_, l| *l != base);
332        Ok(Self { base, families })
333    }
334
335    /// 2026-09-30: The ladder `family` runs.
336    pub fn ladder(&self, family: ProjFamily) -> &Ladder {
337        self.families.get(&family).unwrap_or(&self.base)
338    }
339
340    /// 2026-09-30: The format a `family` projection of `rows` rows runs.
341    pub fn route(&self, family: ProjFamily, rows: u32) -> ActQuantFormat {
342        self.ladder(family).format(rows)
343    }
344
345    /// 2026-09-30: Whether every family runs `adaptive` at every row count (today's routing).
346    pub fn is_adaptive(&self) -> bool {
347        ProjFamily::ALL.iter().all(|&f| {
348            self.ladder(f)
349                .rungs()
350                .iter()
351                .all(|r| r.format == ActQuantFormat::Adaptive)
352        })
353    }
354
355    /// 2026-09-30: Whether `family` is row-invariant: one non-adaptive format at every row
356    /// count.
357    pub fn is_invariant(&self, family: ProjFamily) -> bool {
358        self.ladder(family).is_uniform_invariant()
359    }
360}
361
362impl fmt::Display for ActivationQuantization {
363    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
364        write!(f, "{}", self.base)?;
365        for (fam, l) in &self.families {
366            write!(f, ",{}:{l}", fam.name())?;
367        }
368        Ok(())
369    }
370}
371
372#[cfg(test)]
373#[path = "activation_quantization_tests.rs"]
374mod activation_quantization_tests;
metrale_config/activation_quantization.rs

metrale_config/
activation_quantization.rs
