metrale_model_engine/traits/model/
circuit.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: `ModelCircuit`, the supertrait of `Model` that selects which forward decode runs:
4//! the hand-written layer loops (`legacy`) or a program compiled from the model's circuit plan
5//! (`circuit`), and discloses the choice for records.
6//!
7//! Owner: model-engine.
8//! Invariants:
9//! - The default is the legacy forward, and a model without a circuit executor refuses
10//!   `circuit` rather than running legacy under that name.
11
12use anyhow::{Result, bail};
13use metrale_circuit::Instance;
14use metrale_model_layers::circuit_exec::{Fusions, TargetModules};
15
16/// 2026-09-28: The forward a model's decode runs.
17#[derive(Debug, Clone)]
18pub enum ForwardSelect {
19    /// 2026-09-28: The hand-written layer loops.
20    Legacy,
21    /// 2026-09-28: The program compiled from `instance`'s circuit.
22    Circuit {
23        /// 2026-09-28: The circuit instance serving this checkpoint on this target.
24        instance: Box<Instance>,
25        /// 2026-09-28: Which rules the plan may select.
26        fusions: Fusions,
27        /// 2026-09-28: The served target's compiled modules: which kernels exist.
28        modules: TargetModules,
29        /// 2026-09-30: The checkpoint's `config.json` as the server read it; the executor maps it
30        /// to the served shape through the circuit's config map.
31        config_json: String,
32    },
33}
34
35/// 2026-09-28: What a record states about the forward.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct ForwardDisclosure {
38    /// 2026-09-28: `legacy`, `circuit`, or `circuit-reference` (no `bit_identical` fusions).
39    pub forward: &'static str,
40    /// 2026-09-28: The digest over every compiled plan (decode, then each multi-sequence
41    /// width; `CircuitExec::plans_digest`); `None` under legacy.
42    pub plan_digest: Option<String>,
43    /// 2026-09-28: Launches per decode step of the plan; `None` under legacy.
44    pub launches_per_step: Option<usize>,
45}
46
47impl ForwardDisclosure {
48    /// 2026-09-28: The legacy forward.
49    pub fn legacy() -> Self {
50        ForwardDisclosure {
51            forward: "legacy",
52            plan_digest: None,
53            launches_per_step: None,
54        }
55    }
56}
57
58/// 2026-09-28: Forward selection; see the module header.
59pub trait ModelCircuit {
60    /// 2026-09-28: Make decode run `sel` from the next step. Captured decode graphs are dropped.
61    /// Default: only `Legacy` is accepted.
62    fn set_forward(&self, sel: &ForwardSelect) -> Result<()> {
63        match sel {
64            ForwardSelect::Legacy => Ok(()),
65            ForwardSelect::Circuit { .. } => bail!("this model has no circuit executor"),
66        }
67    }
68
69    /// 2026-09-28: The forward decode runs now. Default: legacy.
70    fn forward_disclosure(&self) -> ForwardDisclosure {
71        ForwardDisclosure::legacy()
72    }
73}