metrale_circuit/
format.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Tensor formats of edges and weights, and their one text spelling.
4//!
5//! Owner: metrale-circuit.
6//! Invariants:
7//! - `Format::parse(f.name()) == Ok(f)` for every format: the spelling is canonical, so the
8//!   plan digest and the golden files see one text per format.
9//! - Edge formats carry activation scales only (`token`, `tensor`, `g<n>`); `channel` and
10//!   `block` are weight layouts, refused on an edge by the circuit loader.
11
12use std::fmt;
13
14/// 2026-09-28: How the scales of an 8-bit float tensor are shared.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
16pub enum Scale {
17    /// 2026-09-28: One scale for the whole tensor.
18    PerTensor,
19    /// 2026-09-28: One scale per row (activations).
20    PerToken,
21    /// 2026-09-28: One scale per output channel (weights).
22    PerChannel,
23    /// 2026-09-28: One scale per `n` consecutive values along K.
24    Group(u32),
25    /// 2026-09-28: One scale per `[rows, cols]` block (weights).
26    Block(u32, u32),
27}
28
29/// 2026-09-28: A tensor format.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
31pub enum Format {
32    /// 2026-09-28: bfloat16.
33    Bf16,
34    /// 2026-09-28: float32.
35    F32,
36    /// 2026-09-28: int32 (token and expert ids).
37    I32,
38    /// 2026-09-28: FP8 E4M3 with the given scale layout.
39    Fp8E4m3 {
40        /// 2026-09-28: Scale sharing.
41        scale: Scale,
42    },
43    /// 2026-09-28: NVFP4: E2M1 values, one E4M3 scale per `group` values, one F32 global.
44    Nvfp4 {
45        /// 2026-09-28: Values per scale.
46        group: u32,
47    },
48}
49
50/// 2026-09-28: A format string that is not one of the spellings [`Format::parse`] accepts.
51#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
52#[error(
53    "unknown format `{0}` (expected bf16, f32, i32, fp8/<token|tensor|channel|g<n>|block<r>x<c>> or nvfp4/g<n>)"
54)]
55pub struct FormatError(pub String);
56
57impl Format {
58    /// 2026-09-28: Parse the canonical spelling, e.g. `bf16`, `fp8/token`, `fp8/block128x128`,
59    /// `nvfp4/g16`.
60    pub fn parse(s: &str) -> Result<Self, FormatError> {
61        let bad = || FormatError(s.to_string());
62        match s {
63            "bf16" => return Ok(Format::Bf16),
64            "f32" => return Ok(Format::F32),
65            "i32" => return Ok(Format::I32),
66            _ => {}
67        }
68        let (head, tail) = s.split_once('/').ok_or_else(bad)?;
69        match head {
70            "fp8" => Ok(Format::Fp8E4m3 {
71                scale: parse_scale(tail).ok_or_else(bad)?,
72            }),
73            "nvfp4" => {
74                let group = tail.strip_prefix('g').and_then(positive).ok_or_else(bad)?;
75                Ok(Format::Nvfp4 { group })
76            }
77            _ => Err(bad()),
78        }
79    }
80
81    /// 2026-09-28: The canonical spelling.
82    pub fn name(&self) -> String {
83        match self {
84            Format::Bf16 => "bf16".into(),
85            Format::F32 => "f32".into(),
86            Format::I32 => "i32".into(),
87            Format::Fp8E4m3 { scale } => format!("fp8/{}", scale_name(*scale)),
88            Format::Nvfp4 { group } => format!("nvfp4/g{group}"),
89        }
90    }
91
92    /// 2026-09-28: True for an unquantized format (16/32-bit floats and int32).
93    pub fn is_plain(&self) -> bool {
94        matches!(self, Format::Bf16 | Format::F32 | Format::I32)
95    }
96
97    /// 2026-09-28: True when the scale layout is one an activation edge can carry.
98    pub fn is_edge_format(&self) -> bool {
99        match self {
100            Format::Bf16 | Format::F32 | Format::I32 | Format::Nvfp4 { .. } => true,
101            Format::Fp8E4m3 { scale } => {
102                matches!(scale, Scale::PerTensor | Scale::PerToken | Scale::Group(_))
103            }
104        }
105    }
106
107    /// 2026-09-30: Bytes of a `[n, k]` weight in this format, scales included: as
108    /// [`Format::bytes`], except that an NVFP4 weight has one F32 global scale for the whole
109    /// tensor (`weight_scale_2`), where an activation has one per row.
110    pub fn weight_bytes(&self, n: u64, k: u64) -> Option<u64> {
111        match self {
112            Format::Nvfp4 { .. } => self
113                .bytes(n, k)?
114                .checked_sub(n.checked_mul(4)?)?
115                .checked_add(4),
116            _ => self.bytes(n, k),
117        }
118    }
119
120    /// 2026-09-28: Bytes of a `rows x dim` tensor in this format, scales included. Used by
121    /// the buffer planner. `None` when `dim` is not a multiple of the scale group.
122    pub fn bytes(&self, rows: u64, dim: u64) -> Option<u64> {
123        let elems = rows.checked_mul(dim)?;
124        match self {
125            Format::Bf16 => elems.checked_mul(2),
126            Format::F32 | Format::I32 => elems.checked_mul(4),
127            Format::Fp8E4m3 { scale } => {
128                let scales = match scale {
129                    Scale::PerTensor => 1,
130                    Scale::PerToken => rows,
131                    Scale::PerChannel => dim,
132                    Scale::Group(g) => {
133                        let g = u64::from(*g);
134                        if !dim.is_multiple_of(g) {
135                            return None;
136                        }
137                        rows.checked_mul(dim / g)?
138                    }
139                    Scale::Block(r, c) => {
140                        let (r, c) = (u64::from(*r), u64::from(*c));
141                        rows.div_ceil(r).checked_mul(dim.div_ceil(c))?
142                    }
143                };
144                elems.checked_add(scales.checked_mul(4)?)
145            }
146            Format::Nvfp4 { group } => {
147                let g = u64::from(*group);
148                if !dim.is_multiple_of(g) || !elems.is_multiple_of(2) {
149                    return None;
150                }
151                // 2026-09-28: Packed E2M1 pairs, one E4M3 byte per group. 2026-09-30: One F32
152                // global per row, as the activation quantizer writes it (`w4a4_quant_rows`);
153                // an edge is always an activation.
154                (elems / 2)
155                    .checked_add(elems / g)?
156                    .checked_add(rows.checked_mul(4)?)
157            }
158        }
159    }
160}
161
162impl fmt::Display for Format {
163    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
164        f.write_str(&self.name())
165    }
166}
167
168fn positive(s: &str) -> Option<u32> {
169    s.parse::<u32>().ok().filter(|&n| n > 0)
170}
171
172fn parse_scale(s: &str) -> Option<Scale> {
173    match s {
174        "tensor" => Some(Scale::PerTensor),
175        "token" => Some(Scale::PerToken),
176        "channel" => Some(Scale::PerChannel),
177        _ => {
178            if let Some(g) = s.strip_prefix('g') {
179                return positive(g).map(Scale::Group);
180            }
181            let (r, c) = s.strip_prefix("block")?.split_once('x')?;
182            Some(Scale::Block(positive(r)?, positive(c)?))
183        }
184    }
185}
186
187fn scale_name(s: Scale) -> String {
188    match s {
189        Scale::PerTensor => "tensor".into(),
190        Scale::PerToken => "token".into(),
191        Scale::PerChannel => "channel".into(),
192        Scale::Group(g) => format!("g{g}"),
193        Scale::Block(r, c) => format!("block{r}x{c}"),
194    }
195}
196
197#[cfg(test)]
198#[path = "format_tests.rs"]
199mod format_tests;
metrale_circuit/format.rs

metrale_circuit/
format.rs
