metrale_model_layers/circuit_exec/
sources.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The circuit texts this binary was built from: INSTANCES.toml, the circuits, their
4//! block libraries, the precision tables and each hardware's FUSIONS.toml, embedded so a plan
5//! matches the kernels compiled in. `met circuit` and the executor both read them here.
6//!
7//! Owner: model-layers circuit executor.
8//! Invariants:
9//! - One copy of every text: the CLI and the executor never embed their own.
10//! - A lookup of a name the tables lack is an error, never a fallback to another entry.
11
12use anyhow::{Result, anyhow, bail};
13use metrale_circuit::{ArchShape, Instance, PrecisionSpec, Sources};
14
15/// 2026-09-28: kernels/circuits/INSTANCES.toml as built.
16pub const INSTANCES: &str = include_str!("../../../../kernels/circuits/INSTANCES.toml");
17
18/// 2026-09-28: Every circuit an instance can name, by arch.
19pub const CIRCUITS: [(&str, &str); 3] = [
20    (
21        "qwen3_5",
22        include_str!("../../../../kernels/circuits/qwen3_5.toml"),
23    ),
24    (
25        "qwen3_6_moe",
26        include_str!("../../../../kernels/circuits/qwen3_6_moe.toml"),
27    ),
28    (
29        "nemotron_h",
30        include_str!("../../../../kernels/circuits/nemotron_h.toml"),
31    ),
32];
33
34/// 2026-09-28: Every precision table an instance can name.
35pub const PRECISION: [(&str, &str); 3] = [
36    (
37        "qwen3.6-35b-a3b-fp8-bf16head",
38        include_str!("../../../../kernels/circuits/precision/qwen3.6-35b-a3b-fp8-bf16head.toml"),
39    ),
40    (
41        "qwen3.6-35b-a3b-nvfp4-declared",
42        include_str!("../../../../kernels/circuits/precision/qwen3.6-35b-a3b-nvfp4-declared.toml"),
43    ),
44    (
45        "nemotron-3.5-lightning-30b-a3b-nvfp4",
46        include_str!(
47            "../../../../kernels/circuits/precision/nemotron-3.5-lightning-30b-a3b-nvfp4.toml"
48        ),
49    ),
50];
51
52/// 2026-09-28: Every checkpoint plan fixture an instance can name.
53pub const CHECKPOINTS: [(&str, &str); 1] = [(
54    "unsloth--Qwen3.8-27B-NVFP4",
55    include_str!("../../../../kernels/circuits/checkpoints/unsloth--Qwen3.8-27B-NVFP4.toml"),
56)];
57
58/// 2026-09-28: Every block library a circuit can include.
59pub const BLOCKS: [(&str, &str); 1] = [(
60    "qwen3_hybrid",
61    include_str!("../../../../kernels/circuits/blocks/qwen3_hybrid.toml"),
62)];
63
64/// 2026-09-28: FUSIONS.toml per hardware.
65pub const FUSIONS: [(&str, &str); 1] = [(
66    "gb10",
67    include_str!("../../../../kernels/gb10/common/FUSIONS.toml"),
68)];
69
70/// 2026-09-28: The entry `key` of `table`; `what` names the table in the error.
71pub fn lookup<'a>(table: &[(&str, &'a str)], key: &str, what: &str) -> Result<&'a str> {
72    table
73        .iter()
74        .find(|(k, _)| *k == key)
75        .map(|(_, v)| *v)
76        .ok_or_else(|| anyhow!("{what} `{key}` is not built into this binary"))
77}
78
79/// 2026-09-28: The embedded texts `instance` is built from.
80pub fn sources(instance: &Instance) -> Result<Sources<'static>> {
81    let hw = instance.target.split('/').next().unwrap_or_default();
82    Ok(Sources {
83        circuit: lookup(&CIRCUITS, &instance.arch, "circuit")?,
84        precision: match &instance.precision {
85            PrecisionSpec::Table(n) => lookup(&PRECISION, n, "precision table")?,
86            PrecisionSpec::Policy {
87                checkpoint_plan, ..
88            } => lookup(&CHECKPOINTS, checkpoint_plan, "checkpoint plan")?,
89        },
90        rules: lookup(&FUSIONS, hw, "FUSIONS.toml for hardware")?,
91        blocks: &BLOCKS,
92    })
93}
94
95/// 2026-09-28: The instance serving `recipe`.
96pub fn instance(recipe: &str) -> Result<Instance> {
97    let all = metrale_circuit::parse_instances(INSTANCES)?;
98    let known: Vec<String> = all.iter().map(|i| i.recipe.clone()).collect();
99    all.into_iter().find(|i| i.recipe == recipe).ok_or_else(|| {
100        anyhow!(
101            "no circuit instance for recipe `{recipe}`; kernels/circuits/INSTANCES.toml has: {}",
102            known.join(", ")
103        )
104    })
105}
106
107/// 2026-09-28: The instance whose checkpoint and kernel target are these. Several recipes may
108/// serve one checkpoint on one target; they must then agree on the circuit and the precision
109/// source, the only parts the executor takes from an instance (the policy, and 2026-09-30 a
110/// policy precision's tier and kernel capabilities, are read live: `CircuitExec::build`).
111pub fn instance_for(checkpoint: &str, target: &str) -> Result<Instance> {
112    let all = metrale_circuit::parse_instances(INSTANCES)?;
113    let hits: Vec<Instance> = all
114        .into_iter()
115        .filter(|i| i.checkpoint == checkpoint && i.target == target)
116        .collect();
117    let Some(first) = hits.first() else {
118        bail!(
119            "no circuit instance serves checkpoint `{checkpoint}` on target `{target}` \
120             (kernels/circuits/INSTANCES.toml)"
121        );
122    };
123    if let Some(other) = hits
124        .iter()
125        .find(|i| i.arch != first.arch || !same_source(&i.precision, &first.precision))
126    {
127        bail!(
128            "recipes `{}` and `{}` both serve `{checkpoint}` on `{target}` with different \
129             circuits or precision tables",
130            first.recipe,
131            other.recipe
132        );
133    }
134    Ok(first.clone())
135}
136
137/// 2026-09-30: `a` and `b` read the same precision table, or the same checkpoint plan with the
138/// same engine formats; the tier and capabilities of a policy are the process's own.
139fn same_source(a: &PrecisionSpec, b: &PrecisionSpec) -> bool {
140    match (a, b) {
141        (
142            PrecisionSpec::Policy {
143                checkpoint_plan: pa,
144                engine: ea,
145                ..
146            },
147            PrecisionSpec::Policy {
148                checkpoint_plan: pb,
149                engine: eb,
150                ..
151            },
152        ) => pa == pb && ea == eb,
153        _ => a == b,
154    }
155}
156
157/// 2026-09-30: The shape a served checkpoint runs: its `config.json` through the circuit's config
158/// map (`metrale_circuit::map_checkpoint`, the one config-to-shape mapping), with `vocab` at the
159/// served vocabulary: `ModelConfig::vocab_size` after the server caps it to the tokenizer
160/// (`cap_vocab_size_to_tokenizer`). [`served_shape`] then checks it against INSTANCES.toml.
161pub fn checkpoint_shape(config_json: &str, served_vocab: u64) -> Result<ArchShape> {
162    let mut shape = metrale_circuit::map_checkpoint(config_json)?.shape;
163    shape.dims.insert("vocab".to_string(), served_vocab);
164    Ok(shape)
165}
166
167/// 2026-09-28: Every way `from_config` disagrees with the instance's stated shape.
168pub fn shape_drift(stated: &ArchShape, from_config: &ArchShape) -> Vec<String> {
169    let mut out = Vec::new();
170    if stated.layer_kinds != from_config.layer_kinds {
171        out.push(format!(
172            "layer kinds: INSTANCES.toml has {} layers, config.json {} (or the kinds differ)",
173            stated.layer_kinds.len(),
174            from_config.layer_kinds.len()
175        ));
176    }
177    for (k, v) in &stated.dims {
178        match from_config.dims.get(k) {
179            Some(c) if c == v => {}
180            other => out.push(format!("{k}: INSTANCES.toml {v}, config.json {other:?}")),
181        }
182    }
183    out
184}
185
186/// 2026-09-28: The shape a served model runs, checked against the instance's stated shape. Every
187/// dim must agree except `vocab`, which the server caps to the tokenizer's vocabulary
188/// (`cap_vocab_size_to_tokenizer`) and so may be smaller; the served value is the one the logits
189/// are sized by, so the plan is instantiated with it.
190pub fn served_shape(stated: &ArchShape, served: &ArchShape) -> Result<ArchShape> {
191    let capped = |k: &str| {
192        k == "vocab"
193            && matches!(
194                (stated.dims.get(k), served.dims.get(k)),
195                (Some(s), Some(v)) if v <= s
196            )
197    };
198    let drift: Vec<String> = shape_drift(stated, served)
199        .into_iter()
200        .filter(|line| !line.split(':').next().is_some_and(capped))
201        .collect();
202    if !drift.is_empty() {
203        bail!(
204            "kernels/circuits/INSTANCES.toml disagrees with the loaded model:\n  {}",
205            drift.join("\n  ")
206        );
207    }
208    Ok(served.clone())
209}
210
211#[cfg(test)]
212mod served_shape_tests {
213    use super::*;
214    use metrale_circuit::LayerKind;
215
216    fn shape(vocab: u64, hidden: u64) -> ArchShape {
217        ArchShape {
218            layer_kinds: vec![LayerKind::LinearAttention, LayerKind::FullAttention],
219            dims: [("vocab".to_string(), vocab), ("hidden".to_string(), hidden)].into(),
220        }
221    }
222
223    #[test]
224    fn a_vocab_capped_to_the_tokenizer_is_served_and_nothing_else_may_move() {
225        let stated = shape(248_320, 5120);
226        assert_eq!(
227            served_shape(&stated, &shape(248_077, 5120)).unwrap().dims["vocab"],
228            248_077
229        );
230        assert_eq!(served_shape(&stated, &stated).unwrap(), stated);
231        let wider = served_shape(&stated, &shape(248_321, 5120))
232            .unwrap_err()
233            .to_string();
234        assert!(wider.contains("vocab"), "{wider}");
235        let hidden = served_shape(&stated, &shape(248_077, 4096))
236            .unwrap_err()
237            .to_string();
238        assert!(
239            hidden.contains("hidden") && !hidden.contains("vocab"),
240            "{hidden}"
241        );
242        let mut kinds = shape(248_320, 5120);
243        kinds.layer_kinds.pop();
244        assert!(served_shape(&stated, &kinds).is_err());
245    }
246}
247
248#[cfg(test)]
249mod instance_for_tests {
250    use super::*;
251
252    /// 2026-09-30: The dense checkpoint has an instance per tier; the executor takes the tier
253    /// from the process, so either serves it. Before `same_source` the pair was refused.
254    #[test]
255    fn two_tiers_of_one_checkpoint_plan_serve_one_checkpoint() {
256        let got = instance_for("unsloth/Qwen3.8-27B-NVFP4", "gb10/qwen3.8-27b/nvfp4").unwrap();
257        assert_eq!(got.recipe, "qwen3.8/qwen3.8-27b-nvfp4-unsloth");
258    }
259
260    #[test]
261    fn another_checkpoint_plan_or_engine_format_is_another_source() {
262        let all = metrale_circuit::parse_instances(INSTANCES).unwrap();
263        let find = |r: &str| {
264            all.iter()
265                .find(|i| i.recipe == r)
266                .unwrap()
267                .precision
268                .clone()
269        };
270        let (nvfp4, declared) = (
271            find("qwen3.8/qwen3.8-27b-nvfp4-unsloth"),
272            find("qwen3.8/qwen3.8-27b-nvfp4-unsloth-declared"),
273        );
274        assert!(nvfp4 != declared && same_source(&nvfp4, &declared));
275        let edit = |f: &dyn Fn(&mut String, &mut Vec<(String, metrale_circuit::LinearFormats)>)| {
276            let mut p = declared.clone();
277            if let PrecisionSpec::Policy {
278                checkpoint_plan,
279                engine,
280                ..
281            } = &mut p
282            {
283                f(checkpoint_plan, engine);
284            }
285            p
286        };
287        assert!(!same_source(&nvfp4, &edit(&|plan, _| plan.push('x'))));
288        assert!(!same_source(
289            &nvfp4,
290            &edit(&|_, engine| {
291                engine.pop();
292            })
293        ));
294        assert!(!same_source(
295            &nvfp4,
296            &PrecisionSpec::Table("unsloth--Qwen3.8-27B-NVFP4".into())
297        ));
298    }
299}
metrale_model_layers/circuit_exec/sources.rs

metrale_model_layers/circuit_exec/
sources.rs
