metrale_circuit/hardware/
model.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The model side of a hardware plan: a circuit at its checkpoint's formats, the
4//! serving policy its rules read, and the kernel target (`<model>/<quant>`) whose sources a
5//! class compiles for it. [`CircuitSource`] is the one seam between the hardware axis and how
6//! a circuit is instantiated (today `kernels/circuits/INSTANCES.toml`; the checkpoint
7//! instantiation plugs in behind the same trait).
8//!
9//! Owner: metrale-circuit (hardware).
10//! Invariants:
11//! - Settings a class's `HARDWARE.toml [defaults]` decides are re-read from the device's class
12//!   ([`policy_on_class`]), except on the class a recipe was written for
13//!   ([`ModelUnderPlan::settings_class`]), where the recipe's value stands: it may pin one away
14//!   from the class default.
15//! - `recipe` precision is the instance's own (the golden plans' formats); `declared` is the
16//!   checkpoint's declared formats. Neither is chosen silently: the report prints which.
17
18use std::collections::BTreeMap;
19
20use super::HwError;
21use super::class::ClassInfo;
22use crate::fuser::Policy;
23use crate::instances::{Instance, PrecisionSpec, parse_instances};
24use crate::ir::Circuit;
25use crate::render::Header;
26use crate::venn::repo::{Repo, load_instance, resolve};
27
28/// 2026-09-30: Which formats a plan serves.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum PrecisionChoice {
31    /// 2026-09-30: The recipe's pinned formats (a golden instance's tier and engine choices).
32    Recipe,
33    /// 2026-09-30: The checkpoint's declared formats.
34    Declared,
35}
36
37impl PrecisionChoice {
38    /// 2026-09-30: The CLI and report spelling.
39    pub fn name(self) -> &'static str {
40        match self {
41            PrecisionChoice::Recipe => "recipe",
42            PrecisionChoice::Declared => "declared",
43        }
44    }
45}
46
47/// 2026-09-30: What names the model.
48#[derive(Debug, Clone, Copy)]
49pub struct ModelSpec<'a> {
50    /// 2026-09-30: Checkpoint id (`org/name`) or recipe id.
51    pub checkpoint: &'a str,
52    /// 2026-09-30: The checkpoint's `config.json`, when read.
53    pub config_json: Option<&'a str>,
54    /// 2026-09-30: Its `hf_quant_config.json`, when it ships one.
55    pub hf_quant: Option<&'a str>,
56    /// 2026-09-30: Formats served.
57    pub precision: PrecisionChoice,
58}
59
60/// 2026-09-30: A model ready to plan.
61#[derive(Debug, Clone)]
62pub struct ModelUnderPlan {
63    /// 2026-09-30: Recipe id, or the checkpoint id when no recipe serves it.
64    pub label: String,
65    /// 2026-09-30: Checkpoint id.
66    pub checkpoint: String,
67    /// 2026-09-30: Kernel model directory (`qwen3.8-27b`).
68    pub kernel_model: String,
69    /// 2026-09-30: Kernel quant directory (`nvfp4`).
70    pub kernel_quant: String,
71    /// 2026-09-30: The circuit.
72    pub circuit: Circuit,
73    /// 2026-09-30: Policy as the model's source states it.
74    pub policy: Policy,
75    /// 2026-09-30: Plan header lines (`target` is re-spelled per class).
76    pub header: Header,
77    /// 2026-09-30: Where the formats come from, one line.
78    pub precision: String,
79    /// 2026-09-30: Which formats were asked for.
80    pub precision_choice: PrecisionChoice,
81    /// 2026-09-30: Where each policy setting comes from, when derived rather than stated by a
82    /// recipe (empty: the recipe states them in INSTANCES.toml).
83    pub policy_sources: Vec<(String, String)>,
84    /// 2026-09-30: The class a recipe's settings were stated for (its INSTANCES.toml target's
85    /// hardware); `None` for a model planned from its checkpoint, which takes every class's
86    /// defaults.
87    pub settings_class: Option<String>,
88}
89
90/// 2026-09-30: Instantiates the model a spec names.
91pub trait CircuitSource {
92    /// 2026-09-30: The model.
93    fn model(&self, spec: &ModelSpec<'_>) -> Result<ModelUnderPlan, HwError>;
94}
95
96/// 2026-09-30: The source over `kernels/circuits/INSTANCES.toml`.
97pub struct InstancesSource<'a> {
98    /// 2026-09-30: The repository.
99    pub repo: &'a dyn Repo,
100}
101
102impl CircuitSource for InstancesSource<'_> {
103    fn model(&self, spec: &ModelSpec<'_>) -> Result<ModelUnderPlan, HwError> {
104        let text = self
105            .repo
106            .read("kernels/circuits/INSTANCES.toml")
107            .map_err(HwError::Model)?;
108        let all = parse_instances(&text).map_err(|e| HwError::Model(e.to_string()))?;
109        let inst = resolve(&all, spec.checkpoint).map_err(|e| HwError::Model(e.to_string()))?;
110        let (inst, precision) = match (spec.precision, &inst.precision) {
111            (PrecisionChoice::Recipe, p) => {
112                let line = format!("recipe {} ({})", inst.recipe, describe(p));
113                (inst, line)
114            }
115            (PrecisionChoice::Declared, PrecisionSpec::Policy { .. }) => {
116                let mut i = inst.clone();
117                if let PrecisionSpec::Policy { tier, .. } = &mut i.precision {
118                    *tier = "declared".into();
119                }
120                let line = format!("declared ({})", describe(&i.precision));
121                (i, line)
122            }
123            (PrecisionChoice::Declared, PrecisionSpec::Table(t)) => {
124                let line = format!("declared, as stated by the precision table `{t}`");
125                (inst, line)
126            }
127        };
128        model_of(self.repo, &inst, precision, spec.precision)
129    }
130}
131
132fn describe(p: &PrecisionSpec) -> String {
133    match p {
134        PrecisionSpec::Table(t) => format!("precision table `{t}`"),
135        PrecisionSpec::Policy {
136            checkpoint_plan,
137            tier,
138            ..
139        } => format!("checkpoint plan `{checkpoint_plan}` at tier `{tier}`"),
140    }
141}
142
143/// 2026-09-30: `inst` loaded as a model to plan.
144pub fn model_of(
145    repo: &dyn Repo,
146    inst: &Instance,
147    precision: String,
148    precision_choice: PrecisionChoice,
149) -> Result<ModelUnderPlan, HwError> {
150    let loaded = load_instance(repo, inst).map_err(|e| HwError::Model(e.to_string()))?;
151    let mut parts = inst.target.split('/');
152    let (hw, model, quant) = (parts.next(), parts.next(), parts.next());
153    let (Some(hw), Some(model), Some(quant)) = (hw, model, quant) else {
154        return Err(HwError::Model(format!(
155            "{}: target `{}` is not hw/model/quant",
156            inst.recipe, inst.target
157        )));
158    };
159    Ok(ModelUnderPlan {
160        label: inst.recipe.clone(),
161        checkpoint: inst.checkpoint.clone(),
162        kernel_model: model.to_string(),
163        kernel_quant: quant.to_string(),
164        circuit: loaded.circuit,
165        policy: inst.policy.clone(),
166        header: crate::header(inst),
167        precision,
168        precision_choice,
169        policy_sources: Vec::new(),
170        settings_class: Some(hw.to_string()),
171    })
172}
173
174/// 2026-09-30: Policy settings a class's `HARDWARE.toml [defaults]` decides: setting to the
175/// defaults key (booleans map to `on` / `off`). INSTANCES.toml cites these rows for the values
176/// its recipes state.
177pub const CLASS_DEFAULT_SETTINGS: [(&str, &str); 3] = [
178    ("ssm_batched_recurrent", "ssm_batched_recurrent"),
179    ("ssm_ba_gates_hopper", "ssm_ba_gates_hopper"),
180    ("decode_split_silu", "decode_split_silu"),
181];
182
183/// 2026-09-30: `policy` with the class-decided settings re-read from `class`, unless `class` is
184/// the `settings_class` the policy was stated for; the changed keys are returned for the report.
185/// A class that does not state one is an error either way.
186pub fn policy_on_class(
187    policy: &Policy,
188    settings_class: Option<&str>,
189    class: &ClassInfo,
190) -> Result<(Policy, BTreeMap<String, String>), HwError> {
191    let own = settings_class == Some(class.name.as_str());
192    let mut out = policy.clone();
193    let mut changed = BTreeMap::new();
194    for (setting, key) in CLASS_DEFAULT_SETTINGS {
195        if !policy.settings.contains_key(setting) {
196            continue;
197        }
198        let raw = class.defaults.get(key).ok_or_else(|| {
199            HwError::Class(format!(
200                "kernels/{}/HARDWARE.toml [defaults] does not state `{key}`",
201                class.name
202            ))
203        })?;
204        let v = match raw.as_str() {
205            "true" => "on",
206            "false" => "off",
207            other => other,
208        }
209        .to_string();
210        if own {
211            continue;
212        }
213        if policy.settings.get(setting) != Some(&v) {
214            changed.insert(setting.to_string(), v.clone());
215        }
216        out.settings.insert(setting.to_string(), v);
217    }
218    Ok((out, changed))
219}
220
221/// 2026-09-30: `header` with the target re-spelled for `class`.
222pub fn header_on(model: &ModelUnderPlan, class: &str) -> Header {
223    model
224        .header
225        .iter()
226        .map(|(k, v)| {
227            if k == "target" {
228                (
229                    k.clone(),
230                    format!("{class}/{}/{}", model.kernel_model, model.kernel_quant),
231                )
232            } else {
233                (k.clone(), v.clone())
234            }
235        })
236        .collect()
237}