metrale_circuit/hardware/
model_checkpoint.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The model source over the checkpoint itself: `config.json` (and the sidecar
4//! `hf_quant_config.json`) resolved by [`crate::resolve_checkpoint`] at the declared formats,
5//! with the serving policy derived from the checkpoint and the engine's stated defaults. A
6//! `recipe` request goes to the recipe's golden instance ([`InstancesSource`]), whose formats
7//! are the recipe's.
8//!
9//! Owner: metrale-circuit (hardware).
10//! Invariants:
11//! - Every derived setting names where its value comes from ([`POLICY_SOURCES`]); the class's
12//!   `[defaults]` settings are placeholders here and are always re-read from the device's
13//!   class ([`super::model::policy_on_class`]).
14//! - A checkpoint the model axis refuses is refused here with its reason, never planned from a
15//!   guess.
16
17use std::collections::{BTreeMap, BTreeSet};
18
19use super::HwError;
20use super::model::{CircuitSource, InstancesSource, ModelSpec, ModelUnderPlan, PrecisionChoice};
21use super::sources::KernelTree;
22use crate::format::Format;
23use crate::fuser::Policy;
24use crate::ir::{Circuit, OpKind};
25use crate::{QuantMetadata, ServePrecision, resolve_checkpoint};
26
27/// 2026-09-30: The kernel quant axis every target is built for: the build compiles
28/// `METRALE_TARGET_QUANT=nvfp4` unless told otherwise, and the FP8 35B-A3B golden instance
29/// targets `gb10/qwen3.6-35b-a3b/nvfp4` too.
30pub const KERNEL_QUANT: &str = "nvfp4";
31
32/// 2026-09-30: Where each derived setting comes from, printed in the report.
33pub const POLICY_SOURCES: [(&str, &str); 11] = [
34    (
35        "row_tiers",
36        "canonical for FP8 routed experts or a MoE with FP8 projections, else by_rows \
37         (ml/row_tiers.rs:64-78)",
38    ),
39    (
40        "kv_cache_dtype",
41        "the checkpoint's declared KV-cache format (bf16 when none)",
42    ),
43    ("lm_head_dtype", "the lm_head's declared weight format"),
44    (
45        "ssm_h_dtype",
46        "f32, the --ssm-h-dtype default (crates/server/src/cli/serve_args.rs)",
47    ),
48    (
49        "gemv_sw",
50        "on unless METRALE_NO_GEMV_SW=1 (ml/ops/model_levers_resolve.rs:53)",
51    ),
52    (
53        "w4a16_tc",
54        "on unless METRALE_NO_W4A16_TC (ml/ops/gemv_tc.rs:79-82)",
55    ),
56    (
57        "ssm_batched_recurrent",
58        "the device class's HARDWARE.toml [defaults]",
59    ),
60    (
61        "ssm_ba_gates_hopper",
62        "the device class's HARDWARE.toml [defaults]",
63    ),
64    (
65        "decode_split_silu",
66        "the device class's HARDWARE.toml [defaults]",
67    ),
68    (
69        "rms_norm_act_quant",
70        "off: the executor has no fused norm-quantize launch (2026-09-30)",
71    ),
72    (
73        "activation_quantization",
74        "adaptive: the per-row-count routing FUSIONS.toml encodes, which the rules plan; a fixed \
75         --activation-quantization is not modelled by the rules yet (2026-10-02)",
76    ),
77];
78
79/// 2026-09-30: The source over checkpoints, falling back to the recipes for `recipe` formats.
80pub struct CheckpointSource<'a> {
81    /// 2026-09-30: The kernel tree (MODEL.toml target resolution).
82    pub tree: &'a dyn KernelTree,
83}
84
85impl CircuitSource for CheckpointSource<'_> {
86    fn model(&self, spec: &ModelSpec<'_>) -> Result<ModelUnderPlan, HwError> {
87        if spec.precision == PrecisionChoice::Recipe {
88            return InstancesSource {
89                repo: self.tree.as_repo(),
90            }
91            .model(spec);
92        }
93        let config = spec.config_json.ok_or_else(|| {
94            HwError::Model(format!(
95                "{}: declared formats need the checkpoint's config.json (a checkpoint directory, \
96                 the local Hugging Face cache, or --allow-network)",
97                spec.checkpoint
98            ))
99        })?;
100        let rc = resolve_checkpoint(
101            config,
102            QuantMetadata {
103                hf_quant_config: spec.hf_quant,
104            },
105            &ServePrecision::Declared,
106        )
107        .map_err(|e| HwError::Model(format!("{}: {e}", spec.checkpoint)))?;
108        // 2026-09-30: The engine's own resolution. A config the engine cannot parse has no
109        // kernel target: the plan says why and reads the class's common layer only; it never
110        // falls back to the raw config.json model_type (which the engine may rewrite).
111        let (target, why_none) = match self.tree.kernel_target(config, &[spec.checkpoint]) {
112            Ok(Some(t)) => (Some(t), String::new()),
113            Ok(None) => (None, "none: no MODEL.toml claims it".to_string()),
114            Err(e) => (
115                None,
116                format!("unresolved: the engine's config parse refuses this checkpoint ({e})"),
117            ),
118        };
119        let policy = derive_policy(&rc.circuit, rc.kv_cache)?;
120        let kernel_model = target.clone().unwrap_or_else(|| "(none)".into());
121        let settings: Vec<String> = policy
122            .settings
123            .iter()
124            .map(|(k, v)| format!("{k}={v}"))
125            .collect();
126        Ok(ModelUnderPlan {
127            label: spec.checkpoint.to_string(),
128            checkpoint: spec.checkpoint.to_string(),
129            header: vec![
130                ("recipe".into(), "none (planned from the checkpoint)".into()),
131                ("checkpoint".into(), spec.checkpoint.to_string()),
132                (
133                    "target".into(),
134                    format!("gb10/{kernel_model}/{KERNEL_QUANT}"),
135                ),
136                ("settings".into(), settings.join(" ")),
137                ("opt-in levers".into(), "none".into()),
138            ],
139            kernel_model,
140            kernel_quant: KERNEL_QUANT.into(),
141            circuit: rc.circuit,
142            policy,
143            precision: format!(
144                "declared by the checkpoint (arch `{}`, model_type `{}`; kernel target {})",
145                rc.arch,
146                rc.model_type,
147                target.map_or(why_none, |t| format!("`{t}`"))
148            ),
149            precision_choice: PrecisionChoice::Declared,
150            policy_sources: POLICY_SOURCES
151                .iter()
152                .map(|(k, v)| (k.to_string(), v.to_string()))
153                .collect(),
154            settings_class: None,
155        })
156    }
157}
158
159fn dtype_name(f: Format) -> Result<&'static str, HwError> {
160    match f {
161        Format::Bf16 => Ok("bf16"),
162        Format::Fp8E4m3 { .. } => Ok("fp8"),
163        Format::Nvfp4 { .. } => Ok("nvfp4"),
164        other => Err(HwError::Model(format!(
165            "{} is no serving dtype",
166            other.name()
167        ))),
168    }
169}
170
171/// 2026-09-30: The serving policy of a checkpoint-derived circuit ([`POLICY_SOURCES`]).
172pub fn derive_policy(c: &Circuit, kv_cache: Option<Format>) -> Result<Policy, HwError> {
173    let fp8_experts = c
174        .nodes
175        .iter()
176        .any(|n| n.op == OpKind::ExpertGateUp && matches!(n.weight, Some(Format::Fp8E4m3 { .. })));
177    // 2026-10-02: A MoE whose other projections are FP8 (nvidia/Qwen3.6-35B-A3B-NVFP4) is
178    // canonical too (crates/server/src/main_modules/serve_load/model_setup.rs `publish_row_tiers`).
179    let moe = c.nodes.iter().any(|n| n.op == OpKind::ExpertGateUp);
180    let fp8_projections = c.nodes.iter().any(|n| {
181        matches!(n.op, OpKind::Linear(_)) && matches!(n.weight, Some(Format::Fp8E4m3 { .. }))
182    });
183    let canonical = fp8_experts || (moe && fp8_projections);
184    let head = c
185        .nodes
186        .iter()
187        .find(|n| n.op == OpKind::LmHead)
188        .and_then(|n| n.weight)
189        .ok_or_else(|| HwError::Model("the circuit has no lm_head weight".into()))?;
190    let kv = match kv_cache {
191        None => "bf16",
192        Some(f) => dtype_name(f)?,
193    };
194    // 2026-09-30: The two class settings are re-read from the device's class before any plan;
195    // "class" never reaches a rule (policy_on_class refuses a class that does not state them).
196    let settings: BTreeMap<String, String> = [
197        ("row_tiers", if canonical { "canonical" } else { "by_rows" }),
198        ("kv_cache_dtype", kv),
199        ("lm_head_dtype", dtype_name(head)?),
200        ("ssm_h_dtype", "f32"),
201        ("gemv_sw", "on"),
202        ("w4a16_tc", "on"),
203        ("ssm_batched_recurrent", "class"),
204        ("ssm_ba_gates_hopper", "class"),
205        ("decode_split_silu", "class"),
206        ("rms_norm_act_quant", "off"),
207        ("activation_quantization", "adaptive"),
208    ]
209    .into_iter()
210    .map(|(k, v)| (k.to_string(), v.to_string()))
211    .collect();
212    Ok(Policy {
213        opt_in_levers: BTreeSet::new(),
214        settings,
215    })
216}
metrale_circuit/hardware/model_checkpoint.rs

metrale_circuit/hardware/
model_checkpoint.rs
