metrale_circuit/
checkpoint.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: A circuit from a checkpoint alone: its `config.json` and quantization metadata
4//! pick the architecture, map to an arch shape through the architecture's declarative config
5//! map (`kernels/circuits/<arch>.config.toml`), and give every projection its declared
6//! formats. The model axis of `met circuit plan --checkpoint --hardware`.
7//!
8//! Owner: metrale-circuit.
9//! Invariants:
10//! - Pure: the circuit, block and config-map TOMLs are embedded at build time; the caller
11//!   passes the checkpoint's file text. No I/O, no environment.
12//! - The architecture is chosen by the config's top-level `model_type` alone; a model type no
13//!   map serves is refused, never approximated by a neighbour.
14//! - Precision comes from config.json's own `quantization_config` when it has one, else from
15//!   the sidecar `hf_quant_config.json` (the serve rule, `merge_sidecar_quant_config`).
16
17use std::collections::BTreeMap;
18
19use metrale_config::DeclaredPrecisionPlan;
20use serde_json::Value;
21
22use crate::circuit_toml::CircuitError;
23use crate::config_map::{ConfigMap, ConfigMapError, MappedConfig, map_config};
24use crate::declared_precision::DeclaredPrecision;
25use crate::format::{Format, Scale};
26use crate::instantiate::instantiate;
27use crate::ir::{ArchShape, Circuit};
28use crate::precision::{LinearFormats, PrecisionError};
29use crate::precision_policy::{PolicyPrecision, caps_named, tier_named};
30
31/// 2026-09-30: The quantization metadata a checkpoint ships beside its `config.json`.
32#[derive(Debug, Clone, Copy, Default)]
33pub struct QuantMetadata<'a> {
34    /// 2026-09-30: `hf_quant_config.json` (ModelOpt), verbatim; `None` when absent. Read only
35    /// when `config.json` has no `quantization_config`.
36    pub hf_quant_config: Option<&'a str>,
37}
38
39/// 2026-09-30: Which formats the circuit's edges take.
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub enum ServePrecision {
42    /// 2026-09-30: The checkpoint's declared formats ([`DeclaredPrecision`]).
43    Declared,
44    /// 2026-09-30: The engine's serving policy over the declared plan: the
45    /// `--weight-quantization` tier, the kernel capabilities present, and the formats the
46    /// engine chooses itself (as `INSTANCES.toml` states them for a golden instance).
47    Policy {
48        /// 2026-09-30: `declared` or `nvfp4`.
49        tier: String,
50        /// 2026-09-30: `KernelCaps` field names.
51        caps: Vec<String>,
52        /// 2026-09-30: Module glob to formats, first match wins.
53        engine: Vec<(String, LinearFormats)>,
54    },
55}
56
57/// 2026-09-30: A checkpoint resolved to its circuit.
58#[derive(Debug, Clone)]
59pub struct ResolvedCheckpoint {
60    /// 2026-09-30: The circuit arch (`qwen3_5`, `qwen3_6_moe`, `nemotron_h`, `dense_gqa`).
61    pub arch: String,
62    /// 2026-09-30: The config's top-level `model_type`.
63    pub model_type: String,
64    /// 2026-09-30: Layer kinds and dims, from config.json.
65    pub shape: ArchShape,
66    /// 2026-09-30: The math parameters config.json states (RoPE, norm eps, routing, ...), as
67    /// JSON text by name.
68    pub params: BTreeMap<String, String>,
69    /// 2026-09-30: The declared KV-cache format; `None` is 16-bit.
70    pub kv_cache: Option<Format>,
71    /// 2026-09-30: The instantiated circuit.
72    pub circuit: Circuit,
73}
74
75/// 2026-09-30: Why a checkpoint has no circuit.
76#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
77pub enum CheckpointError {
78    /// 2026-09-30: A file is not the JSON it should be.
79    #[error("{file}: {detail}")]
80    Json {
81        /// 2026-09-30: `config.json` or `hf_quant_config.json`.
82        file: &'static str,
83        /// 2026-09-30: What was wrong.
84        detail: String,
85    },
86    /// 2026-09-30: No config map serves this `model_type`.
87    #[error("no circuit serves model_type `{model_type}` (served: {served})")]
88    UnknownModelType {
89        /// 2026-09-30: The config's `model_type`.
90        model_type: String,
91        /// 2026-09-30: The model types the embedded maps serve.
92        served: String,
93    },
94    /// 2026-09-30: The config does not map.
95    #[error(transparent)]
96    Map(#[from] ConfigMapError),
97    /// 2026-09-30: The quantization metadata is malformed, or declares a format the circuit
98    /// has no edge or weight format for.
99    #[error("quantization: {0}")]
100    Quant(String),
101    /// 2026-09-30: A serving-policy name is unknown.
102    #[error(transparent)]
103    Precision(#[from] PrecisionError),
104    /// 2026-09-30: The circuit does not instantiate under the mapped shape.
105    #[error(transparent)]
106    Circuit(#[from] CircuitError),
107}
108
109mod states;
110pub use states::recurrent_states;
111
112/// 2026-09-30: One architecture's embedded files.
113struct Arch {
114    circuit: &'static str,
115    config_map: &'static str,
116}
117
118macro_rules! circuits_file {
119    ($f:literal) => {
120        include_str!(concat!("../../../kernels/circuits/", $f))
121    };
122}
123
124const ARCHES: [Arch; 4] = [
125    Arch {
126        circuit: circuits_file!("qwen3_5.toml"),
127        config_map: circuits_file!("qwen3_5.config.toml"),
128    },
129    Arch {
130        circuit: circuits_file!("qwen3_6_moe.toml"),
131        config_map: circuits_file!("qwen3_6_moe.config.toml"),
132    },
133    Arch {
134        circuit: circuits_file!("nemotron_h.toml"),
135        config_map: circuits_file!("nemotron_h.config.toml"),
136    },
137    Arch {
138        circuit: circuits_file!("dense_gqa.toml"),
139        config_map: circuits_file!("dense_gqa.config.toml"),
140    },
141];
142
143/// 2026-09-30: The block libraries the embedded circuits include.
144const BLOCKS: [(&str, &str); 1] = [("qwen3_hybrid", circuits_file!("blocks/qwen3_hybrid.toml"))];
145
146/// 2026-09-30: The embedded config maps, parsed.
147pub fn config_maps() -> Result<Vec<ConfigMap>, ConfigMapError> {
148    ARCHES
149        .iter()
150        .map(|a| ConfigMap::parse(a.config_map))
151        .collect()
152}
153
154/// 2026-09-30: The circuit of the checkpoint `config_json` describes, every edge at the
155/// checkpoint's declared formats.
156pub fn instantiate_from_checkpoint(
157    config_json: &str,
158    quant: QuantMetadata<'_>,
159) -> Result<Circuit, CheckpointError> {
160    Ok(resolve_checkpoint(config_json, quant, &ServePrecision::Declared)?.circuit)
161}
162
163/// 2026-09-30: The checkpoint's arch, shape, params, KV-cache format and circuit, with the
164/// edge formats `serve` gives.
165pub fn resolve_checkpoint(
166    config_json: &str,
167    quant: QuantMetadata<'_>,
168    serve: &ServePrecision,
169) -> Result<ResolvedCheckpoint, CheckpointError> {
170    let config = parse_config(config_json)?;
171    let (i, mapped) = map_checkpoint_value(&config)?;
172    let qc = quantization_config(&config, quant)?;
173    let plan = match &qc {
174        Some(qc) => DeclaredPrecisionPlan::from_quantization_config(qc)
175            .map_err(|e| CheckpointError::Quant(format!("{e:#}")))?,
176        None => DeclaredPrecisionPlan::default(),
177    };
178    if let Some(qc) = &qc
179        && plan.rules.is_empty()
180    {
181        return Err(CheckpointError::Quant(format!(
182            "the quantization block declares no scheme for any layer (a missing quant group): {}",
183            serde_json::to_string(qc)
184                .unwrap_or_default()
185                .chars()
186                .take(200)
187                .collect::<String>()
188        )));
189    }
190    let kv_cache = kv_cache_format(qc.as_ref())?;
191    let text = ARCHES[i].circuit;
192    let circuit = match serve {
193        ServePrecision::Declared => {
194            let precision = DeclaredPrecision::new(&plan);
195            let circuit = instantiate(text, &BLOCKS, &mapped.shape, &precision)?;
196            let refused = precision.refusals();
197            if !refused.is_empty() {
198                return Err(CheckpointError::Quant(refused.join("; ")));
199            }
200            circuit
201        }
202        ServePrecision::Policy { tier, caps, engine } => {
203            let policy =
204                metrale_config::WeightQuantPolicy::new(tier_named(tier)?, &plan, caps_named(caps)?);
205            instantiate(
206                text,
207                &BLOCKS,
208                &mapped.shape,
209                &PolicyPrecision::new(policy, engine),
210            )?
211        }
212    };
213    Ok(ResolvedCheckpoint {
214        arch: mapped.arch,
215        model_type: mapped.model_type,
216        shape: mapped.shape,
217        params: mapped.params,
218        kv_cache,
219        circuit,
220    })
221}
222
223/// 2026-09-30: `config.json` mapped to its circuit arch, arch shape and params, without
224/// instantiating: the shape `INSTANCES.toml` must restate for a golden instance.
225pub fn map_checkpoint(config_json: &str) -> Result<MappedConfig, CheckpointError> {
226    Ok(map_checkpoint_value(&parse_config(config_json)?)?.1)
227}
228
229fn parse_config(config_json: &str) -> Result<Value, CheckpointError> {
230    serde_json::from_str(config_json).map_err(|e| CheckpointError::Json {
231        file: "config.json",
232        detail: e.to_string(),
233    })
234}
235
236/// 2026-09-30: The index of the arch whose map serves `config`'s model type, and the mapping.
237fn map_checkpoint_value(config: &Value) -> Result<(usize, MappedConfig), CheckpointError> {
238    let model_type = config
239        .get("model_type")
240        .and_then(Value::as_str)
241        .unwrap_or_default()
242        .to_string();
243    let maps = config_maps()?;
244    let Some((i, map)) = maps
245        .iter()
246        .enumerate()
247        .find(|(_, m)| m.model_types().contains(&model_type))
248    else {
249        let mut served: Vec<&str> = maps
250            .iter()
251            .flat_map(|m| m.model_types().iter().map(String::as_str))
252            .collect();
253        served.sort_unstable();
254        return Err(CheckpointError::UnknownModelType {
255            model_type,
256            served: served.join(", "),
257        });
258    };
259    Ok((i, map_config(map, config)?))
260}
261
262/// 2026-09-30: The checkpoint's `quantization_config`: config.json's own, else the sidecar's.
263fn quantization_config(
264    config: &Value,
265    quant: QuantMetadata<'_>,
266) -> Result<Option<Value>, CheckpointError> {
267    if let Some(qc) = config.get("quantization_config").filter(|v| !v.is_null()) {
268        return Ok(Some(qc.clone()));
269    }
270    quant
271        .hf_quant_config
272        .map(|text| {
273            serde_json::from_str(text).map_err(|e| CheckpointError::Json {
274                file: "hf_quant_config.json",
275                detail: e.to_string(),
276            })
277        })
278        .transpose()
279}
280
281/// 2026-09-30: The KV-cache format a quantization block declares: compressed-tensors
282/// `kv_cache_scheme`, or ModelOpt `kv_cache_quant_algo` (flat, or nested under `quantization`
283/// in a sidecar).
284fn kv_cache_format(qc: Option<&Value>) -> Result<Option<Format>, CheckpointError> {
285    let Some(qc) = qc else {
286        return Ok(None);
287    };
288    let fp8 = Format::Fp8E4m3 {
289        scale: Scale::PerTensor,
290    };
291    if let Some(s) = qc.get("kv_cache_scheme").filter(|v| !v.is_null()) {
292        let bits = s.get("num_bits").and_then(Value::as_u64);
293        let ty = s.get("type").and_then(Value::as_str);
294        return match (bits, ty) {
295            (Some(8), Some("float")) => Ok(Some(fp8)),
296            _ => Err(CheckpointError::Quant(format!(
297                "kv_cache_scheme {s} has no KV-cache format in the circuit"
298            ))),
299        };
300    }
301    let algo = qc
302        .get("kv_cache_quant_algo")
303        .or_else(|| qc.get("quantization")?.get("kv_cache_quant_algo"))
304        .filter(|v| !v.is_null());
305    match algo.map(|a| a.as_str()) {
306        None => Ok(None),
307        Some(Some("FP8")) => Ok(Some(fp8)),
308        Some(other) => Err(CheckpointError::Quant(format!(
309            "kv_cache_quant_algo {other:?} has no KV-cache format in the circuit"
310        ))),
311    }
312}
metrale_circuit/checkpoint.rs

metrale_circuit/
checkpoint.rs
