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;