metrale_circuit/
circuit_toml.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The `kernels/circuits/<arch>.toml` schema: block templates, the layout rule
4//! that maps a model's layer kinds to templates, and the typed load errors.
5//!
6//! Owner: metrale-circuit.
7//! Invariants:
8//! - Every edge states its format and its shape (`"n x hidden"`); nothing defaults.
9//! - Loading is fail-fast: the first unknown op, dangling or duplicate edge, format no
10//!   consumer accepts, or template/layout mismatch is returned as a [`CircuitError`].
11
12use std::collections::BTreeMap;
13
14use serde::Deserialize;
15
16use crate::dims::DimError;
17use crate::format::FormatError;
18use crate::ir::OpParseError;
19
20/// 2026-09-28: Why a circuit did not load or instantiate.
21#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
22pub enum CircuitError {
23    /// 2026-09-28: Not valid TOML, or not the schema.
24    #[error("circuit TOML: {0}")]
25    Parse(String),
26    /// 2026-09-28: An op the vocabulary does not have, or a bad role or format qualifier.
27    #[error("block `{block}` node `{node}`: {source}")]
28    Op {
29        /// 2026-09-28: Template.
30        block: String,
31        /// 2026-09-28: Node id.
32        node: String,
33        /// 2026-09-28: The op error.
34        source: OpParseError,
35    },
36    /// 2026-09-28: A format spelling no format matches.
37    #[error("block `{block}` edge `{edge}`: {source}")]
38    Format {
39        /// 2026-09-28: Template.
40        block: String,
41        /// 2026-09-28: Edge.
42        edge: String,
43        /// 2026-09-28: The format error.
44        source: FormatError,
45    },
46    /// 2026-09-28: A weight-only scale layout (`channel`, `block`) on an edge.
47    #[error("block `{block}` edge `{edge}`: `{format}` is a weight layout, not an edge format")]
48    WeightFormatOnEdge {
49        /// 2026-09-28: Template.
50        block: String,
51        /// 2026-09-28: Edge.
52        edge: String,
53        /// 2026-09-28: The format.
54        format: String,
55    },
56    /// 2026-09-28: A shape that is not `<rows> x <dim>`, or a bad expression in it.
57    #[error("block `{block}` edge `{edge}`: {detail}")]
58    Shape {
59        /// 2026-09-28: Template.
60        block: String,
61        /// 2026-09-28: Edge.
62        edge: String,
63        /// 2026-09-28: What was wrong.
64        detail: String,
65    },
66    /// 2026-09-30: A state declaration that does not parse or evaluate.
67    #[error("block `{block}` state `{state}`: {detail}")]
68    State {
69        /// 2026-09-30: Template.
70        block: String,
71        /// 2026-09-30: State id.
72        state: String,
73        /// 2026-09-30: What was wrong.
74        detail: String,
75    },
76    /// 2026-09-28: A dimension the arch shape does not define, or an overflow.
77    #[error("block `{block}`: {source}")]
78    Dim {
79        /// 2026-09-28: Template.
80        block: String,
81        /// 2026-09-28: The dimension error.
82        source: DimError,
83    },
84    /// 2026-09-28: A node reads an edge nothing produced.
85    #[error("block `{block}` node `{node}` reads `{edge}`, which no earlier node produces")]
86    DanglingInput {
87        /// 2026-09-28: Template.
88        block: String,
89        /// 2026-09-28: Node id.
90        node: String,
91        /// 2026-09-28: Edge.
92        edge: String,
93    },
94    /// 2026-09-28: An edge nothing reads that is not a declared output.
95    #[error("edge `{0}` is produced but never read, and is not a declared output")]
96    DanglingOutput(String),
97    /// 2026-09-28: Two producers of one edge, or two nodes with one id.
98    #[error("block `{block}`: `{name}` is defined twice")]
99    Duplicate {
100        /// 2026-09-28: Template.
101        block: String,
102        /// 2026-09-28: Edge or node id.
103        name: String,
104    },
105    /// 2026-09-28: An edge whose format a reader does not accept.
106    #[error(
107        "node `{node}` ({op}) does not accept `{edge}` as {format}{}",
108        expected.as_ref().map(|e| format!(" (expects {e})")).unwrap_or_default()
109    )]
110    FormatMismatch {
111        /// 2026-09-28: The reading node.
112        node: String,
113        /// 2026-09-28: Its op.
114        op: String,
115        /// 2026-09-28: The edge.
116        edge: String,
117        /// 2026-09-28: The edge's format.
118        format: String,
119        /// 2026-09-28: What the reader expects, when it is one format.
120        expected: Option<String>,
121    },
122    /// 2026-09-28: The modules of one node resolve to different formats.
123    #[error("node `{node}`: bound modules resolve to different formats ({detail})")]
124    MixedPrecision {
125        /// 2026-09-28: Node id.
126        node: String,
127        /// 2026-09-28: The two resolutions.
128        detail: String,
129    },
130    /// 2026-09-28: A weight-reading node without a binding, or a binding on a node that
131    /// reads no weight but is not a norm.
132    #[error("node `{node}`: {detail}")]
133    Binding {
134        /// 2026-09-28: Node id.
135        node: String,
136        /// 2026-09-28: What was wrong.
137        detail: String,
138    },
139    /// 2026-09-28: A layout that names a missing template, or a model whose layers the
140    /// layout does not describe.
141    #[error("layout: {0}")]
142    Layout(String),
143    /// 2026-09-28: An included block file that was not supplied, does not parse, or defines a
144    /// block another file also defines.
145    #[error("include `{name}`: {detail}")]
146    Include {
147        /// 2026-09-28: The include's name.
148        name: String,
149        /// 2026-09-28: What was wrong.
150        detail: String,
151    },
152    /// 2026-09-28: A dim the circuit requires that the arch shape does not give, or one it
153    /// gives that the circuit does not declare.
154    #[error("arch shape: {0}")]
155    ShapeMismatch(String),
156}
157
158#[derive(Deserialize)]
159#[serde(deny_unknown_fields)]
160pub(crate) struct CircuitFile {
161    pub schema: u32,
162    pub arch: String,
163    pub description: String,
164    pub layer_module: String,
165    pub include: Vec<String>,
166    pub dims: Vec<String>,
167    pub layout: LayoutFile,
168    pub prologue: Vec<String>,
169    pub epilogue: Vec<String>,
170    pub draft: Vec<String>,
171    pub draft_module: Option<String>,
172    /// 2026-09-30: The draft blocks are instantiated only when this switch holds
173    /// ([`when_holds`]); absent, always.
174    pub draft_when: Option<String>,
175    /// 2026-09-30: Per switch, a draft block list replacing `draft` while it holds. A draft
176    /// entry `template@module` instantiates `template` with `{L}` = `module` (one template
177    /// serves a layer and a draft module).
178    #[serde(default)]
179    pub draft_variant: BTreeMap<String, Vec<String>>,
180    pub block: BTreeMap<String, BlockFile>,
181    /// 2026-09-30: The blocks the circuit file itself defines (not its libraries'); only these
182    /// must all be used.
183    #[serde(skip)]
184    pub local_blocks: std::collections::BTreeSet<String>,
185}
186
187#[derive(Deserialize)]
188#[serde(deny_unknown_fields)]
189pub(crate) struct LayoutFile {
190    pub kind: String,
191    pub period: Option<usize>,
192    pub blocks: BTreeMap<String, Vec<String>>,
193    /// 2026-09-30: Per switch ([`when_holds`]), layer kinds whose blocks are replaced while it
194    /// holds (`moe_latent = { moe = ["moe_latent"] }`).
195    #[serde(default)]
196    pub when: BTreeMap<String, BTreeMap<String, Vec<String>>>,
197}
198
199#[derive(Deserialize)]
200#[serde(deny_unknown_fields)]
201pub(crate) struct BlockFile {
202    pub stream_in: Option<String>,
203    pub stream_out: Option<String>,
204    #[serde(default)]
205    pub outputs: Vec<OutputFile>,
206    pub node: Vec<NodeFile>,
207    /// 2026-09-30: The state the block keeps between steps (`crate::state`).
208    #[serde(default)]
209    pub state: Vec<StateFile>,
210}
211
212/// 2026-09-30: A declared output: the edge, and (in a served circuit) the model buffer it binds
213/// to (`crate::model_buffer::ModelBuffer`). A bare edge name declares an output with no buffer.
214#[derive(Deserialize)]
215#[serde(untagged)]
216pub(crate) enum OutputFile {
217    Bare(String),
218    Bound { edge: String, buffer: String },
219}
220
221impl OutputFile {
222    /// 2026-09-30: The edge it names.
223    pub fn edge(&self) -> &str {
224        match self {
225            Self::Bare(e) | Self::Bound { edge: e, .. } => e,
226        }
227    }
228}
229
230/// 2026-09-30: One `[[block.<name>.state]]`: `kind` (`recurrent` | `paged_kv`, 2026-10-02: or a
231/// cache kind, `crate::state::StateKind`), `format` (a dtype or a `{key}` the plan inputs give),
232/// `shape` (axes of dim expressions joined by ` x `: one slot of a recurrent state, one token of
233/// a KV side, one unit of a cache), and for a recurrent state the verify intermediates it keeps
234/// (`h_steps` | `conv_steps`). 2026-10-02: A snapshot kind names the state of its block it copies
235/// (`of`) instead of a shape, so the copy cannot be sized apart from the original.
236#[derive(Deserialize)]
237#[serde(deny_unknown_fields)]
238pub(crate) struct StateFile {
239    pub id: String,
240    pub kind: String,
241    /// 2026-09-30: `sequence` (recurrent) or `model` (a KV pool); see `crate::state::Lifetime`.
242    pub lifetime: String,
243    pub format: String,
244    pub shape: Option<String>,
245    pub verify: Option<String>,
246    pub of: Option<String>,
247}
248
249#[derive(Deserialize)]
250#[serde(deny_unknown_fields)]
251pub(crate) struct NodeFile {
252    pub id: String,
253    pub op: String,
254    /// 2026-09-30: The block states the node touches, by state id (`crate::state::StateAccess`).
255    #[serde(default)]
256    pub state: BTreeMap<String, String>,
257    pub role: Option<String>,
258    pub format: Option<String>,
259    #[serde(default, rename = "in")]
260    pub inputs: Vec<String>,
261    pub out: Vec<EdgeFile>,
262    #[serde(default)]
263    pub binding: Vec<String>,
264    #[serde(default)]
265    pub params: BTreeMap<String, String>,
266    /// 2026-09-30: The node is present only when this switch holds ([`when_holds`]). An absent
267    /// node has one input and one output of the same format and width, and its output IS its
268    /// input (the identity it degenerates to: a norm, bias or projection the checkpoint lacks).
269    pub when: Option<String>,
270    /// 2026-09-30: Params merged in when their switch holds (`attn_bias = { bias = "true" }`).
271    #[serde(default)]
272    pub params_when: BTreeMap<String, BTreeMap<String, String>>,
273    /// 2026-10-02: The weight is no quantization target (a parameter, not a linear layer), so
274    /// no checkpoint-wide quantization config reaches it: the precision is not asked; the
275    /// weight is 16-bit and the node reads its input edge as is.
276    #[serde(default)]
277    pub unquantized: bool,
278}
279
280#[derive(Deserialize)]
281#[serde(deny_unknown_fields)]
282pub(crate) struct EdgeFile {
283    pub edge: String,
284    pub format: String,
285    pub shape: String,
286}
287
288/// 2026-09-28: The layer-kind rule a circuit declares.
289#[derive(Debug, Clone, PartialEq, Eq)]
290pub enum LayoutRule {
291    /// 2026-09-28: `full_attention` on every layer `i` with `(i + 1) % period == 0`,
292    /// `linear_attention` elsewhere; the HF `full_attention_interval` convention.
293    Interval {
294        /// 2026-09-28: The interval.
295        period: usize,
296    },
297    /// 2026-09-28: Any sequence of the kinds `blocks` maps.
298    List,
299}
300
301/// 2026-09-28: A block library: `kernels/circuits/blocks/<name>.toml`, holding only block
302/// templates that several circuits share.
303#[derive(Deserialize)]
304#[serde(deny_unknown_fields)]
305struct BlocksFile {
306    schema: u32,
307    #[allow(dead_code)]
308    description: String,
309    block: BTreeMap<String, BlockFile>,
310}
311
312/// 2026-09-28: The block libraries a circuit TOML includes, in its order, so the caller knows
313/// which files to read.
314pub fn includes_of(text: &str) -> Result<Vec<String>, CircuitError> {
315    Ok(parse_one(text)?.include)
316}
317
318fn parse_one(text: &str) -> Result<CircuitFile, CircuitError> {
319    let file: CircuitFile = toml::from_str(text).map_err(|e| CircuitError::Parse(e.to_string()))?;
320    if file.schema != 1 {
321        return Err(CircuitError::Parse(format!(
322            "schema {} (this build reads 1)",
323            file.schema
324        )));
325    }
326    Ok(file)
327}
328
329/// 2026-09-28: Parse a circuit and merge in its includes. `includes` maps an include name to
330/// its text; one not supplied, a library that does not parse, and a block defined twice are
331/// errors, never an override.
332pub(crate) fn parse_file(
333    text: &str,
334    includes: &[(&str, &str)],
335) -> Result<CircuitFile, CircuitError> {
336    let mut file = parse_one(text)?;
337    file.local_blocks = file.block.keys().cloned().collect();
338    for name in file.include.clone() {
339        let err = |detail: String| CircuitError::Include {
340            name: name.clone(),
341            detail,
342        };
343        let lib_text = includes
344            .iter()
345            .find(|(n, _)| *n == name)
346            .map(|(_, t)| *t)
347            .ok_or_else(|| err("not supplied".into()))?;
348        let lib: BlocksFile = toml::from_str(lib_text).map_err(|e| err(e.to_string()))?;
349        if lib.schema != 1 {
350            return Err(err(format!("schema {} (this build reads 1)", lib.schema)));
351        }
352        for (block, tpl) in lib.block {
353            if file.block.contains_key(&block) {
354                return Err(err(format!("block `{block}` is defined twice")));
355            }
356            file.block.insert(block, tpl);
357        }
358    }
359    Ok(file)
360}
361
362pub(crate) fn layout_rule(l: &LayoutFile) -> Result<LayoutRule, CircuitError> {
363    match (l.kind.as_str(), l.period) {
364        ("interval", Some(p)) if p > 0 => Ok(LayoutRule::Interval { period: p }),
365        ("interval", _) => Err(CircuitError::Layout(
366            "`interval` needs a `period` of at least 1".into(),
367        )),
368        ("list", None) => Ok(LayoutRule::List),
369        ("list", Some(_)) => Err(CircuitError::Layout("`list` takes no `period`".into())),
370        (other, _) => Err(CircuitError::Layout(format!(
371            "unknown layout kind `{other}` (interval | list)"
372        ))),
373    }
374}
375
376/// 2026-09-30: Whether switch `expr` holds: `name` when the dim `name` is non-zero, `!name`
377/// when it is zero. The dim must be in the circuit's `dims` list.
378pub(crate) fn when_holds(
379    expr: &str,
380    declared: &[String],
381    dims: &BTreeMap<String, u64>,
382) -> Result<bool, CircuitError> {
383    let (negate, name) = match expr.trim().strip_prefix('!') {
384        Some(n) => (true, n.trim()),
385        None => (false, expr.trim()),
386    };
387    if !declared.iter().any(|d| d == name) {
388        return Err(CircuitError::ShapeMismatch(format!(
389            "switch `{expr}` names `{name}`, which is not in the circuit's `dims` list"
390        )));
391    }
392    let v = dims.get(name).copied().ok_or_else(|| {
393        CircuitError::ShapeMismatch(format!("switch `{expr}`: the arch shape lacks `{name}`"))
394    })?;
395    Ok((v != 0) != negate)
396}
397
398/// 2026-09-28: Split `"<rows> x <dim>"`.
399pub(crate) fn split_shape(s: &str) -> Option<(&str, &str)> {
400    let (rows, dim) = s.split_once(" x ")?;
401    Some((rows.trim(), dim.trim()))
402}
metrale_circuit/circuit_toml.rs

metrale_circuit/
circuit_toml.rs
