metrale_circuit/
instances.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: `kernels/circuits/INSTANCES.toml`: which recipes have a circuit, with the arch
4//! shape, precision table, policy and plan matrix each one is shown and checked under.
5//!
6//! The shape is stated here, not read from the checkpoint, so the golden plans are checked
7//! without the checkpoint on disk. `met circuit show` compares it against the checkpoint's
8//! `config.json` when that is in the local cache.
9//!
10//! Owner: metrale-circuit.
11//! Invariants:
12//! - Every policy setting is stated; nothing is filled from an engine default.
13//! - `layer_kinds` spells one letter per layer: `G` a GatedDeltaNet (linear attention) layer,
14//!   `A` a full-attention layer, `M` a Mamba2 layer, `E` a MoE-only layer (the Nemotron-H
15//!   hybrid-pattern letters); whitespace is ignored.
16
17use std::collections::{BTreeMap, BTreeSet};
18
19use serde::Deserialize;
20
21use crate::format::Format;
22use crate::fuser::Policy;
23use crate::ir::{ArchShape, LayerKind};
24use crate::precision::LinearFormats;
25use crate::rules::Mode;
26
27/// 2026-09-28: One recipe's circuit instance.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct Instance {
30    /// 2026-09-28: Recipe id (`qwen3.8/qwen3.8-27b-nvfp4-unsloth`).
31    pub recipe: String,
32    /// 2026-09-28: Checkpoint id.
33    pub checkpoint: String,
34    /// 2026-09-28: Circuit arch: `kernels/circuits/<arch>.toml`.
35    pub arch: String,
36    /// 2026-09-30: What tells this instance's golden plans apart from another golden instance
37    /// of the same arch (the recipe served under another `--weight-quantization` tier); `None`
38    /// for the arch's only golden instance.
39    pub variant: Option<String>,
40    /// 2026-09-28: Where each linear module's formats come from.
41    pub precision: PrecisionSpec,
42    /// 2026-09-28: Kernel target, `hw/model/quant`.
43    pub target: String,
44    /// 2026-09-28: Whether its plans are checked in under `kernels/circuits/plans/`.
45    pub golden: bool,
46    /// 2026-09-28: The arch shape.
47    pub shape: ArchShape,
48    /// 2026-09-28: The policy.
49    pub policy: Policy,
50    /// 2026-09-28: Row counts per mode that have a plan.
51    pub plans: BTreeMap<Mode, Vec<u64>>,
52    /// 2026-09-30: The batched-verify row tables that have a plan. An instance with none is
53    /// not served a batched verify by the executor: each sequence verifies alone.
54    pub verify_batch: Vec<crate::runs::RowTable>,
55}
56
57impl Instance {
58    /// 2026-09-28: The golden file name of one plan: `<stem>-<mode>-n<rows>.txt`.
59    pub fn plan_file(&self, mode: Mode, rows: u64) -> String {
60        format!("{}-{}-n{rows}.txt", self.plan_stem(), mode.name())
61    }
62
63    /// 2026-09-30: The golden file name of one batched-verify plan:
64    /// `<stem>-verify_batch-<table>.txt`, the table's runs joined by `_`, `f` marking a fragmented
65    /// run and a leading `u-` an uncarried table.
66    pub fn table_plan_file(&self, table: &crate::runs::RowTable) -> String {
67        let runs: Vec<String> = table
68            .runs
69            .iter()
70            .map(|r| format!("{}x{}{}", r.k, r.n, if r.contiguous { "" } else { "f" }))
71            .collect();
72        format!(
73            "{}-verify_batch-{}{}.txt",
74            self.plan_stem(),
75            if table.carried { "" } else { "u-" },
76            runs.join("_")
77        )
78    }
79
80    /// 2026-09-30: `<arch>`, or `<arch>.<variant>` for a variant: the prefix of every golden
81    /// file of this instance.
82    pub fn plan_stem(&self) -> String {
83        match &self.variant {
84            Some(v) => format!("{}.{v}", self.arch),
85            None => self.arch.clone(),
86        }
87    }
88}
89
90/// 2026-09-28: An instance's source of linear formats.
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub enum PrecisionSpec {
93    /// 2026-09-28: A stated table, `kernels/circuits/precision/<name>.toml`.
94    Table(String),
95    /// 2026-09-28: The engine's policy over the checkpoint's declared plan
96    /// ([`crate::precision_policy::PolicyPrecision`]).
97    Policy {
98        /// 2026-09-28: The plan fixture, `kernels/circuits/checkpoints/<name>.toml`.
99        checkpoint_plan: String,
100        /// 2026-09-28: The `--weight-quantization` tier served.
101        tier: String,
102        /// 2026-09-28: The kernel capabilities present (`KernelCaps` field names).
103        caps: Vec<String>,
104        /// 2026-09-28: Formats the engine chooses itself, first match wins.
105        engine: Vec<(String, LinearFormats)>,
106    },
107}
108
109impl PrecisionSpec {
110    /// 2026-09-28: The name of the file the spec reads: the table or the plan fixture.
111    pub fn file_name(&self) -> &str {
112        match self {
113            PrecisionSpec::Table(n) => n,
114            PrecisionSpec::Policy {
115                checkpoint_plan, ..
116            } => checkpoint_plan,
117        }
118    }
119}
120
121/// 2026-09-28: Why INSTANCES.toml did not load.
122#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
123pub enum InstanceError {
124    /// 2026-09-28: Not valid TOML or not the shape.
125    #[error("INSTANCES.toml: {0}")]
126    Parse(String),
127    /// 2026-09-28: A bad field in one instance.
128    #[error("instance `{recipe}`: {detail}")]
129    Field {
130        /// 2026-09-28: Recipe id.
131        recipe: String,
132        /// 2026-09-28: What was wrong.
133        detail: String,
134    },
135}
136
137#[derive(Deserialize)]
138#[serde(deny_unknown_fields)]
139struct File {
140    schema: u32,
141    instance: Vec<InstanceFile>,
142}
143
144#[derive(Deserialize)]
145#[serde(deny_unknown_fields)]
146struct InstanceFile {
147    recipe: String,
148    checkpoint: String,
149    arch: String,
150    variant: Option<String>,
151    precision: PrecisionFile,
152    target: String,
153    golden: bool,
154    layer_kinds: String,
155    dims: BTreeMap<String, u64>,
156    policy: PolicyFile,
157    plans: BTreeMap<String, Vec<u64>>,
158    #[serde(default)]
159    verify_batch: Vec<String>,
160}
161
162#[derive(Deserialize)]
163#[serde(untagged)]
164enum PrecisionFile {
165    Table(String),
166    Policy(PolicyPrecisionFile),
167}
168
169#[derive(Deserialize)]
170#[serde(deny_unknown_fields)]
171struct PolicyPrecisionFile {
172    checkpoint_plan: String,
173    tier: String,
174    caps: Vec<String>,
175    engine: Vec<EngineFormatFile>,
176}
177
178#[derive(Deserialize)]
179#[serde(deny_unknown_fields)]
180struct EngineFormatFile {
181    #[serde(rename = "match")]
182    pattern: String,
183    weight: String,
184    activation: String,
185    /// 2026-09-28: Why the engine, not the checkpoint, decides this module's format.
186    why: String,
187}
188
189#[derive(Deserialize)]
190#[serde(deny_unknown_fields)]
191struct PolicyFile {
192    opt_in_levers: Vec<String>,
193    settings: BTreeMap<String, String>,
194}
195
196/// 2026-09-28: Parse INSTANCES.toml text.
197pub fn parse_instances(text: &str) -> Result<Vec<Instance>, InstanceError> {
198    let file: File = toml::from_str(text).map_err(|e| InstanceError::Parse(e.to_string()))?;
199    if file.schema != 1 {
200        return Err(InstanceError::Parse(format!(
201            "schema {} (this build reads 1)",
202            file.schema
203        )));
204    }
205    let mut seen = BTreeSet::new();
206    let mut stems = BTreeSet::new();
207    let mut out = Vec::with_capacity(file.instance.len());
208    for f in file.instance {
209        let field = |detail: String| InstanceError::Field {
210            recipe: f.recipe.clone(),
211            detail,
212        };
213        if !seen.insert(f.recipe.clone()) {
214            return Err(field("listed twice".into()));
215        }
216        if f.variant.as_deref().is_some_and(|v| v.is_empty()) {
217            return Err(field("`variant` is empty".into()));
218        }
219        let stem = (f.arch.clone(), f.variant.clone());
220        if f.golden && !stems.insert(stem) {
221            return Err(field(format!(
222                "a second golden `{}` instance without its own `variant`: the golden plan files \
223                 would collide",
224                f.arch
225            )));
226        }
227        let mut layer_kinds = Vec::new();
228        for c in f.layer_kinds.chars().filter(|c| !c.is_whitespace()) {
229            layer_kinds.push(match c {
230                'G' => LayerKind::LinearAttention,
231                'A' => LayerKind::FullAttention,
232                'M' => LayerKind::Mamba,
233                'E' => LayerKind::Moe,
234                other => return Err(field(format!("layer kind `{other}` is not G, A, M or E"))),
235            });
236        }
237        let mut plans = BTreeMap::new();
238        for (m, rows) in &f.plans {
239            let mode = Mode::parse(m).ok_or_else(|| field(format!("unknown mode `{m}`")))?;
240            if rows.is_empty() || rows.contains(&0) {
241                return Err(field(format!("mode `{m}` needs row counts of at least 1")));
242            }
243            plans.insert(mode, rows.clone());
244        }
245        let verify_batch = f
246            .verify_batch
247            .iter()
248            .map(|t| {
249                crate::runs::RowTable::parse(t).map_err(|e| field(format!("verify_batch: {e}")))
250            })
251            .collect::<Result<Vec<_>, _>>()?;
252        let precision = match f.precision {
253            PrecisionFile::Table(name) => PrecisionSpec::Table(name),
254            PrecisionFile::Policy(p) => {
255                let mut engine = Vec::with_capacity(p.engine.len());
256                for e in p.engine {
257                    if e.why.trim().is_empty() {
258                        return Err(field(format!(
259                            "engine format `{}` states no `why`",
260                            e.pattern
261                        )));
262                    }
263                    let fmt = |s: &str| {
264                        Format::parse(s)
265                            .map_err(|err| field(format!("engine format `{}`: {err}", e.pattern)))
266                    };
267                    let formats = LinearFormats {
268                        weight: fmt(&e.weight)?,
269                        activation: fmt(&e.activation)?,
270                    };
271                    engine.push((e.pattern, formats));
272                }
273                PrecisionSpec::Policy {
274                    checkpoint_plan: p.checkpoint_plan,
275                    tier: p.tier,
276                    caps: p.caps,
277                    engine,
278                }
279            }
280        };
281        out.push(Instance {
282            recipe: f.recipe.clone(),
283            checkpoint: f.checkpoint,
284            arch: f.arch,
285            variant: f.variant,
286            precision,
287            target: f.target,
288            golden: f.golden,
289            shape: ArchShape {
290                layer_kinds,
291                dims: f.dims,
292            },
293            policy: Policy {
294                opt_in_levers: f.policy.opt_in_levers.into_iter().collect(),
295                settings: f.policy.settings,
296            },
297            plans,
298            verify_batch,
299        });
300    }
301    Ok(out)
302}