metrale_circuit/config_map.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The declarative config mapping of one architecture
4//! (`kernels/circuits/<arch>.config.toml`): how a checkpoint's `config.json` becomes the arch
5//! shape a circuit instantiates under (layer kinds, dims) and the math parameters the kernels
6//! read, with every key the config carries classified.
7//!
8//! Owner: metrale-circuit.
9//! Invariants:
10//! - PCND: every key of the mapped object is classified (a dim, a param, a required value, a
11//! nested object, or ignorable with a stated reason). An unclassified key is refused with
12//! its path; a math-changing key the circuit does not model is refused by its rule, never
13//! silently dropped.
14//! - An absent key takes the default its rule states (the Hugging Face class default, cited
15//! in the TOML); a requirement applies to a present (or defaulted) value, and an absent key
16//! without a default is simply absent. A dim with no default is required.
17//! - JSON `null` is absent.
18
19use std::collections::BTreeMap;
20
21use serde::Deserialize;
22
23#[path = "config_map/apply.rs"]
24mod apply;
25
26pub use apply::{MappedConfig, map_config};
27
28/// 2026-09-30: Why a config did not map.
29#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
30pub enum ConfigMapError {
31 /// 2026-09-30: The mapping TOML is malformed.
32 #[error("config map: {0}")]
33 Schema(String),
34 /// 2026-09-30: `config.json` is not a JSON object, or a nested value is the wrong type.
35 #[error("config.json: {0}")]
36 Json(String),
37 /// 2026-09-30: A key no rule classifies.
38 #[error(
39 "config.json key `{key}` is not mapped for arch `{arch}`: it may change the math, so it \
40 is refused until the config map classifies it"
41 )]
42 UnmappedKey {
43 /// 2026-09-30: The arch.
44 arch: String,
45 /// 2026-09-30: Dotted path of the key.
46 key: String,
47 },
48 /// 2026-09-30: A value the circuit does not model.
49 #[error("config.json `{key}` = {value}: {why}")]
50 Refused {
51 /// 2026-09-30: Dotted path of the key.
52 key: String,
53 /// 2026-09-30: The value, as JSON.
54 value: String,
55 /// 2026-09-30: Why it is refused.
56 why: String,
57 },
58 /// 2026-09-30: A required key is absent and its rule states no default.
59 #[error("config.json lacks `{key}`, which the `{arch}` mapping requires")]
60 MissingKey {
61 /// 2026-09-30: The arch.
62 arch: String,
63 /// 2026-09-30: Dotted path of the key.
64 key: String,
65 },
66}
67
68/// 2026-09-30: A parsed config map.
69#[derive(Debug, Clone)]
70pub struct ConfigMap {
71 /// 2026-09-30: The circuit arch it produces a shape for (`kernels/circuits/<arch>.toml`).
72 pub arch: String,
73 pub(crate) file: MapFile,
74}
75
76impl ConfigMap {
77 /// 2026-09-30: Parse a mapping TOML.
78 pub fn parse(text: &str) -> Result<Self, ConfigMapError> {
79 let file: MapFile =
80 toml::from_str(text).map_err(|e| ConfigMapError::Schema(e.to_string()))?;
81 if file.schema != 1 {
82 return Err(ConfigMapError::Schema(format!(
83 "schema {} (this build reads 1)",
84 file.schema
85 )));
86 }
87 if file.model_types.is_empty() {
88 return Err(ConfigMapError::Schema("`model_types` is empty".into()));
89 }
90 for v in file.variant.keys() {
91 if !file.model_types.contains(v) {
92 return Err(ConfigMapError::Schema(format!(
93 "variant `{v}` is not one of `model_types`"
94 )));
95 }
96 }
97 if file.nest.is_some() == file.root.is_empty() {
98 return Err(ConfigMapError::Schema(
99 "`[root]` classifies the top-level keys of a nested config; it is required \
100 exactly when `nest` is set"
101 .into(),
102 ));
103 }
104 Ok(ConfigMap {
105 arch: file.arch.clone(),
106 file,
107 })
108 }
109
110 /// 2026-09-30: The top-level `model_type` values this map serves.
111 pub fn model_types(&self) -> &[String] {
112 &self.file.model_types
113 }
114
115 /// 2026-09-30: Whether this map's circuit serves an engine-configured `model_type`: one of
116 /// `model_types`, or a name the engine's parser gives such a checkpoint.
117 pub fn serves_engine_model_type(&self, model_type: &str) -> bool {
118 self.file
119 .model_types
120 .iter()
121 .chain(&self.file.engine_model_types)
122 .any(|t| t == model_type)
123 }
124}
125
126#[derive(Debug, Clone, Deserialize)]
127#[serde(deny_unknown_fields)]
128pub(crate) struct MapFile {
129 pub schema: u32,
130 pub arch: String,
131 pub model_types: Vec<String>,
132 /// 2026-09-30: The `ModelConfig::model_type` names the engine's config parser gives these
133 /// checkpoints where it renames them (`qwen3_6_moe`, `holo3_1_moe` for a `qwen3_5_moe`
134 /// config); a model configured by the engine finds its circuit by these or `model_types`.
135 #[serde(default)]
136 pub engine_model_types: Vec<String>,
137 /// 2026-09-30: The object the model's keys live in (`text_config`); the top level's
138 /// other keys are classified by `root`.
139 pub nest: Option<String>,
140 #[serde(default)]
141 pub root: BTreeMap<String, KeyRule>,
142 pub layers: LayersFile,
143 pub dims: BTreeMap<String, DimRule>,
144 #[serde(default)]
145 pub keys: BTreeMap<String, KeyRule>,
146 #[serde(default)]
147 pub variant: BTreeMap<String, VariantFile>,
148}
149
150/// 2026-09-30: Per-`model_type` additions: its dims and key rules replace the base ones of
151/// the same name.
152#[derive(Debug, Clone, Deserialize, Default)]
153#[serde(deny_unknown_fields)]
154pub(crate) struct VariantFile {
155 #[serde(default)]
156 pub dims: BTreeMap<String, DimRule>,
157 #[serde(default)]
158 pub keys: BTreeMap<String, KeyRule>,
159}
160
161/// 2026-09-30: Where the layer kinds come from: exactly one listed source is present.
162#[derive(Debug, Clone, Deserialize)]
163#[serde(deny_unknown_fields)]
164pub(crate) struct LayersFile {
165 /// 2026-09-30: The key holding the layer count.
166 pub count: String,
167 /// 2026-09-30: Every layer is this kind (no per-layer source).
168 pub uniform: Option<String>,
169 #[serde(default)]
170 pub sources: Vec<LayerSource>,
171}
172
173/// 2026-09-30: A per-layer kind source: a list of names, or a pattern string of letters.
174#[derive(Debug, Clone, Deserialize)]
175#[serde(deny_unknown_fields)]
176pub(crate) struct LayerSource {
177 pub key: String,
178 /// 2026-09-30: List entries to layer kinds.
179 #[serde(default)]
180 pub values: BTreeMap<String, String>,
181 /// 2026-09-30: Pattern letters to layer kinds.
182 #[serde(default)]
183 pub chars: BTreeMap<String, String>,
184}
185
186/// 2026-09-30: How one dim is read.
187#[derive(Debug, Clone, Deserialize)]
188#[serde(untagged)]
189pub(crate) enum DimRule {
190 /// 2026-09-30: A required integer key.
191 Key(String),
192 /// 2026-09-30: The general form.
193 Full(DimFull),
194}
195
196#[derive(Debug, Clone, Deserialize)]
197#[serde(deny_unknown_fields)]
198pub(crate) struct DimFull {
199 pub key: Option<String>,
200 /// 2026-09-30: A boolean key read as 0 / 1.
201 #[serde(default)]
202 pub bool: bool,
203 /// 2026-09-30: 1 when the key is present (not null), else 0: a switch on an optional
204 /// feature whose size is another dim (`moe_latent_size`).
205 #[serde(default)]
206 pub bool_present: bool,
207 /// 2026-09-30: A fixed value.
208 #[serde(rename = "const")]
209 pub constant: Option<u64>,
210 /// 2026-09-30: When the key is absent: this mapped dim's value.
211 pub or_dim: Option<String>,
212 /// 2026-09-30: When the key is absent: `a / b` of two mapped dims, exactly.
213 pub or_div: Option<[String; 2]>,
214 /// 2026-09-30: When the key is absent: this value (the HF default).
215 pub default: Option<u64>,
216 /// 2026-09-30: The values the circuit models; any other is refused.
217 pub allowed: Option<Vec<u64>>,
218}
219
220/// 2026-09-30: How one config key is classified.
221#[derive(Debug, Clone, Deserialize)]
222#[serde(untagged)]
223pub(crate) enum KeyRule {
224 /// 2026-09-30: `"ignore"` (not math: ids, dropout, runtime hints) or `"param"` (carried
225 /// into the params as JSON text).
226 Word(String),
227 /// 2026-09-30: The general form.
228 Full(Box<KeyFull>),
229}
230
231#[derive(Debug, Clone, Deserialize)]
232#[serde(deny_unknown_fields)]
233pub(crate) struct KeyFull {
234 /// 2026-09-30: Not math, with the reason.
235 pub ignore: Option<String>,
236 /// 2026-09-30: Refused whenever present (not null), with the reason: math the circuit
237 /// does not model.
238 pub refuse: Option<String>,
239 /// 2026-09-30: Carry the value into the params under this name (`true`: the key's own).
240 pub param: Option<toml::Value>,
241 /// 2026-09-30: The value must be one of these (after `values` renaming).
242 pub require: Option<toml::Value>,
243 /// 2026-09-30: Why a value outside `require` is refused.
244 pub why: Option<String>,
245 /// 2026-09-30: The value when absent (the HF default).
246 pub default: Option<toml::Value>,
247
248 /// 2026-09-30: Rename values (`swish` -> `silu`) before `require` and `param`.
249 #[serde(default)]
250 pub values: BTreeMap<String, String>,
251 /// 2026-09-30: A nested object: rules for its keys, prefixed `<param>.` in the params.
252 pub object: Option<BTreeMap<String, KeyRule>>,
253}
254
255#[cfg(test)]
256#[path = "config_map_tests.rs"]
257mod config_map_tests;