metrale_circuit/
ir.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The circuit IR: a closed op vocabulary, edges with a format and a shape, nodes
4//! bound to checkpoint modules, and the instantiated graph of one model.
5//!
6//! Owner: metrale-circuit.
7//! Invariants:
8//! - [`OpKind`] is closed. A template naming anything else is a load error, never an
9//!   opaque pass-through.
10//! - Heavy ops ([`OpKind::is_heavy`]) are opaque: the fuser fuses at their edges or through a
11//!   pre-written kernel, and never generates them.
12//! - Nodes are stored in execution order, which is a topological order: every input of a
13//!   node is produced by an earlier node or is a circuit input.
14//! - An edge's fused or materialised state belongs to a plan (`FusionPlan::edge_states`),
15//!   not to the circuit, because one circuit has one plan per mode and row count.
16
17use std::collections::BTreeMap;
18
19use crate::dims::DimExpr;
20use crate::format::Format;
21
22/// 2026-09-28: What a linear (GEMV/GEMM) node projects. Closed, like [`OpKind`].
23#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
24pub enum LinearRole {
25    /// 2026-09-28: Attention Q with its output gate, interleaved per head.
26    Q,
27    /// 2026-09-28: Attention K.
28    K,
29    /// 2026-09-28: Attention V.
30    V,
31    /// 2026-09-28: Attention output projection.
32    O,
33    /// 2026-09-28: GatedDeltaNet `in_proj_qkv` and `in_proj_z` in one projection.
34    Qkvz,
35    /// 2026-09-28: GatedDeltaNet `in_proj_b` and `in_proj_a` (the beta and decay inputs).
36    Ba,
37    /// 2026-09-28: GatedDeltaNet `out_proj`.
38    GdnOut,
39    /// 2026-09-28: MLP gate and up projections.
40    GateUp,
41    /// 2026-09-28: MLP down projection.
42    Down,
43    /// 2026-09-28: MoE shared-expert gate and up projections.
44    SharedGateUp,
45    /// 2026-09-28: MoE shared-expert down projection.
46    SharedDown,
47    /// 2026-09-28: MoE shared-expert scalar gate (`shared_expert_gate`).
48    SharedGate,
49    /// 2026-09-28: MTP head input projection (`fc`, embedding and hidden concatenated).
50    MtpFc,
51    /// 2026-09-29: Mamba2 `in_proj`: z, x, B, C and dt in one projection.
52    MambaIn,
53    /// 2026-09-29: Mamba2 `out_proj`.
54    MambaOut,
55    /// 2026-09-29: An ungated MoE shared expert's up projection.
56    SharedUp,
57    /// 2026-09-30: A latent MoE's projection from the hidden width into the experts' latent
58    /// width (`fc1_latent_proj`, Nemotron-3 Super).
59    MoeLatentIn,
60    /// 2026-09-30: A latent MoE's projection of the routed sum back to the hidden width
61    /// (`fc2_latent_proj`).
62    MoeLatentOut,
63}
64
65const ROLES: [(LinearRole, &str); 18] = [
66    (LinearRole::Q, "q"),
67    (LinearRole::K, "k"),
68    (LinearRole::V, "v"),
69    (LinearRole::O, "o"),
70    (LinearRole::Qkvz, "qkvz"),
71    (LinearRole::Ba, "ba"),
72    (LinearRole::GdnOut, "gdn_out"),
73    (LinearRole::GateUp, "gate_up"),
74    (LinearRole::Down, "down"),
75    (LinearRole::SharedGateUp, "shared_gate_up"),
76    (LinearRole::SharedDown, "shared_down"),
77    (LinearRole::SharedGate, "shared_gate"),
78    (LinearRole::MtpFc, "mtp_fc"),
79    (LinearRole::MambaIn, "mamba_in"),
80    (LinearRole::MambaOut, "mamba_out"),
81    (LinearRole::SharedUp, "shared_up"),
82    (LinearRole::MoeLatentIn, "moe_latent_in"),
83    (LinearRole::MoeLatentOut, "moe_latent_out"),
84];
85
86impl LinearRole {
87    /// 2026-09-28: The role spelled in templates and rules.
88    pub fn parse(s: &str) -> Option<Self> {
89        ROLES.iter().find(|(_, n)| *n == s).map(|(r, _)| *r)
90    }
91
92    /// 2026-09-28: The canonical spelling.
93    pub fn name(self) -> &'static str {
94        ROLES
95            .iter()
96            .find(|(r, _)| *r == self)
97            .map(|(_, n)| *n)
98            .unwrap_or("?")
99    }
100}
101
102/// 2026-09-28: The closed op vocabulary.
103#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
104pub enum OpKind {
105    /// 2026-09-28: Token embedding lookup; a source node.
106    Embed,
107    /// 2026-09-28: RMSNorm with a learned weight.
108    RmsNorm,
109    /// 2026-09-28: RMSNorm of one input gated by SiLU of another (GDN output norm).
110    GatedRmsNorm,
111    /// 2026-09-28: Elementwise sum of two streams.
112    ResidualAdd,
113    /// 2026-09-28: A copy between buffers.
114    Copy,
115    /// 2026-09-28: Concatenate inputs along the feature dim.
116    Concat,
117    /// 2026-09-28: Split or deinterleave a packed projection output.
118    Split,
119    /// 2026-09-28: A weight-bearing projection.
120    Linear(LinearRole),
121    /// 2026-09-28: `silu(gate) * up` over a packed gate|up input.
122    SiluMul,
123    /// 2026-09-29: `relu(x)^2` of an ungated up projection.
124    Relu2,
125    /// 2026-09-28: Quantize activations to the given format.
126    ActQuant(Format),
127    /// 2026-09-28: Per-head RMSNorm of Q or K.
128    QkNorm,
129    /// 2026-09-28: Rotary position embedding.
130    Rope,
131    /// 2026-09-28: Append K and V to the paged cache.
132    KvWrite,
133    /// 2026-09-28: Paged softmax attention over the cache.
134    PagedAttention,
135    /// 2026-09-28: `x * sigmoid(gate)` (attention output gate).
136    SigmoidGateMul,
137    /// 2026-09-28: GDN beta and decay from the `ba` projection.
138    GdnGates,
139    /// 2026-09-28: Causal depthwise conv1d step with its rolling state.
140    Conv1dUpdate,
141    /// 2026-09-28: L2 normalisation of Q and K heads.
142    L2Norm,
143    /// 2026-09-28: The gated delta rule recurrence over the recurrent state.
144    GdnRecurrence,
145    /// 2026-09-29: The Mamba2 selective state-space update: one step of the SSM state per row
146    /// (`dt` softplus, `A`, the grouped `B`/`C` projections and the `D` skip).
147    SsmUpdate,
148    /// 2026-09-28: MoE router logits.
149    Router,
150    /// 2026-09-28: Top-k expert selection and weights.
151    TopK,
152    /// 2026-09-28: Routed experts' gate and up projections.
153    ExpertGateUp,
154    /// 2026-09-28: Routed experts' down projections.
155    ExpertDown,
156    /// 2026-09-28: Weighted sum of the routed experts and the gated shared expert.
157    Blend,
158    /// 2026-09-28: Cross-rank reduction of an expert-parallel FFN output.
159    EpReduce,
160    /// 2026-09-28: The model's final RMSNorm.
161    FinalNorm,
162    /// 2026-09-28: Vocabulary projection.
163    LmHead,
164    /// 2026-09-28: Greedy token selection over the logits.
165    Argmax,
166    /// 2026-09-29: A snapshot of the recurrent state its input's producer updates, taken after
167    /// each row but the last into the step's checkpoint slots (the verify's rollback points).
168    /// It has no output.
169    StateSnapshot,
170}
171
172const PLAIN_OPS: [(OpKind, &str); 29] = [
173    (OpKind::Embed, "embed"),
174    (OpKind::RmsNorm, "rms_norm"),
175    (OpKind::GatedRmsNorm, "gated_rms_norm"),
176    (OpKind::ResidualAdd, "residual_add"),
177    (OpKind::Copy, "copy"),
178    (OpKind::Concat, "concat"),
179    (OpKind::Split, "split"),
180    (OpKind::SiluMul, "silu_mul"),
181    (OpKind::Relu2, "relu2"),
182    (OpKind::QkNorm, "qk_norm"),
183    (OpKind::Rope, "rope"),
184    (OpKind::KvWrite, "kv_write"),
185    (OpKind::PagedAttention, "paged_attention"),
186    (OpKind::SigmoidGateMul, "sigmoid_gate_mul"),
187    (OpKind::GdnGates, "gdn_gates"),
188    (OpKind::Conv1dUpdate, "conv1d_update"),
189    (OpKind::L2Norm, "l2_norm"),
190    (OpKind::GdnRecurrence, "gdn_recurrence"),
191    (OpKind::SsmUpdate, "ssm_update"),
192    (OpKind::Router, "router"),
193    (OpKind::TopK, "top_k"),
194    (OpKind::ExpertGateUp, "expert_gate_up"),
195    (OpKind::ExpertDown, "expert_down"),
196    (OpKind::Blend, "blend"),
197    (OpKind::EpReduce, "ep_reduce"),
198    (OpKind::FinalNorm, "final_norm"),
199    (OpKind::LmHead, "lm_head"),
200    (OpKind::Argmax, "argmax"),
201    (OpKind::StateSnapshot, "state_snapshot"),
202];
203
204impl OpKind {
205    /// 2026-09-28: Resolve an op name. `linear` needs `role`, `act_quant` needs `format`;
206    /// every other op refuses both.
207    pub fn parse(
208        name: &str,
209        role: Option<&str>,
210        format: Option<Format>,
211    ) -> Result<Self, OpParseError> {
212        match (name, role, format) {
213            ("linear", Some(r), None) => LinearRole::parse(r)
214                .map(OpKind::Linear)
215                .ok_or_else(|| OpParseError::UnknownRole(r.to_string())),
216            ("linear", None, _) => Err(OpParseError::MissingRole),
217            ("act_quant", None, Some(f)) => Ok(OpKind::ActQuant(f)),
218            ("act_quant", _, None) => Err(OpParseError::MissingFormat),
219            _ => {
220                let op = PLAIN_OPS
221                    .iter()
222                    .find(|(_, n)| *n == name)
223                    .map(|(k, _)| *k)
224                    .ok_or_else(|| OpParseError::UnknownOp(name.to_string()))?;
225                if role.is_some() || format.is_some() {
226                    return Err(OpParseError::StrayQualifier(name.to_string()));
227                }
228                Ok(op)
229            }
230        }
231    }
232
233    /// 2026-09-28: The op's base name (`linear`, `act_quant`, `rms_norm`, ...).
234    pub fn base_name(&self) -> &'static str {
235        match self {
236            OpKind::Linear(_) => "linear",
237            OpKind::ActQuant(_) => "act_quant",
238            other => PLAIN_OPS
239                .iter()
240                .find(|(k, _)| k == other)
241                .map(|(_, n)| *n)
242                .unwrap_or("?"),
243        }
244    }
245
246    /// 2026-09-28: The canonical spelling with its qualifier, e.g. `linear:down`,
247    /// `act_quant:fp8/token`.
248    pub fn name(&self) -> String {
249        match self {
250            OpKind::Linear(r) => format!("linear:{}", r.name()),
251            OpKind::ActQuant(f) => format!("act_quant:{}", f.name()),
252            _ => self.base_name().to_string(),
253        }
254    }
255
256    /// 2026-09-28: An op that is never generated: projections, attention, the recurrence,
257    /// the experts and the vocabulary projection. Fusion happens at its edges.
258    pub fn is_heavy(&self) -> bool {
259        matches!(
260            self,
261            OpKind::Linear(_)
262                | OpKind::PagedAttention
263                | OpKind::GdnRecurrence
264                | OpKind::SsmUpdate
265                | OpKind::Router
266                | OpKind::ExpertGateUp
267                | OpKind::ExpertDown
268                | OpKind::LmHead
269        )
270    }
271
272    /// 2026-09-28: Ops that read a weight whose format the precision resolver answers.
273    pub fn reads_linear_weight(&self) -> bool {
274        matches!(
275            self,
276            OpKind::Linear(_)
277                | OpKind::ExpertGateUp
278                | OpKind::ExpertDown
279                | OpKind::LmHead
280                | OpKind::Router
281        )
282    }
283}
284
285/// 2026-09-28: Why an op spelling was refused.
286#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
287pub enum OpParseError {
288    /// 2026-09-28: Not in the vocabulary.
289    #[error("unknown op `{0}`")]
290    UnknownOp(String),
291    /// 2026-09-28: `linear` with a role outside [`LinearRole`].
292    #[error("unknown linear role `{0}`")]
293    UnknownRole(String),
294    /// 2026-09-28: `linear` without `role`.
295    #[error("`linear` needs a `role`")]
296    MissingRole,
297    /// 2026-09-28: `act_quant` without `format`.
298    #[error("`act_quant` needs a `format`")]
299    MissingFormat,
300    /// 2026-09-28: `role` or `format` on an op that takes neither.
301    #[error("op `{0}` takes no `role` or `format`")]
302    StrayQualifier(String),
303}
304
305/// 2026-09-28: Index of an edge in [`Circuit::edges`].
306pub type EdgeIdx = usize;
307/// 2026-09-28: Index of a node in [`Circuit::nodes`].
308pub type NodeIdx = usize;
309
310/// 2026-09-28: One tensor between nodes: `rows x dim` in `format`. `rows` is an expression
311/// over `n`, the step's row count, and the arch dims (e.g. `n*top_k`).
312#[derive(Debug, Clone, PartialEq, Eq)]
313pub struct Edge {
314    /// 2026-09-28: Unique id, e.g. `l3.attn.q` or `stream.4`.
315    pub id: String,
316    /// 2026-09-28: Element format.
317    pub format: Format,
318    /// 2026-09-28: Row expression; reads `n`.
319    pub rows: DimExpr,
320    /// 2026-09-28: Feature dimension expression.
321    pub dim: DimExpr,
322    /// 2026-09-28: `dim` evaluated under the arch dims.
323    pub dim_value: u64,
324    /// 2026-09-28: The node that writes it; `None` for none (refused by the loader).
325    pub producer: Option<NodeIdx>,
326    /// 2026-09-28: The nodes that read it, in node order.
327    pub consumers: Vec<NodeIdx>,
328    /// 2026-09-28: Read outside the circuit (the logits); never fused away.
329    pub is_output: bool,
330    /// 2026-09-30: The model buffer a declared output lands in; `None` for an output declared
331    /// without one (the executor refuses such a program) and for every other edge.
332    pub binds: Option<crate::model_buffer::ModelBuffer>,
333}
334
335/// 2026-09-28: One op instance.
336#[derive(Debug, Clone, PartialEq, Eq)]
337pub struct Node {
338    /// 2026-09-28: Unique id, e.g. `l3.attn.o_proj`.
339    pub id: String,
340    /// 2026-09-28: The id inside its block template, e.g. `o_proj`.
341    pub local: String,
342    /// 2026-09-28: The op.
343    pub op: OpKind,
344    /// 2026-09-28: Input edges, in the op's argument order.
345    pub inputs: Vec<EdgeIdx>,
346    /// 2026-09-28: Output edges.
347    pub outputs: Vec<EdgeIdx>,
348    /// 2026-09-28: Weight format, for ops that read a linear weight.
349    pub weight: Option<Format>,
350    /// 2026-09-28: Checkpoint modules this node reads (`layers.3.mlp.down_proj`), in order.
351    pub binding: Vec<String>,
352    /// 2026-09-28: Op parameters from the template, e.g. `top_k = "8"`.
353    pub params: BTreeMap<String, String>,
354    /// 2026-09-28: The layer this node belongs to; `None` in the prologue and epilogue.
355    pub layer: Option<usize>,
356    /// 2026-09-28: The block template it was instantiated from.
357    pub block: String,
358    /// 2026-09-30: The states it touches: indices into [`Circuit::states`], and how.
359    pub state: Vec<(usize, crate::state::StateAccess)>,
360}
361
362/// 2026-09-28: A decoder layer's kind, as `ModelConfig::layer_type` reports it.
363#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
364pub enum LayerKind {
365    /// 2026-09-28: A GatedDeltaNet layer.
366    LinearAttention,
367    /// 2026-09-28: A softmax attention layer.
368    FullAttention,
369    /// 2026-09-29: A Mamba2 layer (Nemotron-H `mamba`).
370    Mamba,
371    /// 2026-09-29: A layer whose only mixer is a MoE FFN (Nemotron-H `moe`).
372    Moe,
373}
374
375impl LayerKind {
376    /// 2026-09-28: The spelling in HF configs and circuit layouts.
377    pub fn name(self) -> &'static str {
378        match self {
379            LayerKind::LinearAttention => "linear_attention",
380            LayerKind::FullAttention => "full_attention",
381            LayerKind::Mamba => "mamba",
382            LayerKind::Moe => "moe",
383        }
384    }
385
386    /// 2026-09-28: Parse the HF spelling.
387    pub fn parse(s: &str) -> Option<Self> {
388        match s {
389            "linear_attention" => Some(LayerKind::LinearAttention),
390            "full_attention" => Some(LayerKind::FullAttention),
391            "mamba" => Some(LayerKind::Mamba),
392            "moe" => Some(LayerKind::Moe),
393            _ => None,
394        }
395    }
396}
397
398/// 2026-09-28: What instantiation needs from a model config: its layer kinds in order and
399/// its named dims. The caller fills it from `ModelConfig`, so this crate does not depend on
400/// the config tree.
401#[derive(Debug, Clone, PartialEq, Eq)]
402pub struct ArchShape {
403    /// 2026-09-28: Kind of each text-stack layer, index = layer.
404    pub layer_kinds: Vec<LayerKind>,
405    /// 2026-09-28: Named dims (`hidden`, `head_dim`, ...), the names the templates read.
406    pub dims: BTreeMap<String, u64>,
407}
408
409/// 2026-09-28: Which forward a block belongs to.
410#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
411pub enum Section {
412    /// 2026-09-28: The target model's forward.
413    Main,
414    /// 2026-09-28: The speculative draft head (MTP), which reads the main stream.
415    Draft,
416}
417
418/// 2026-09-28: One instantiated block: its template and the node range it produced.
419#[derive(Debug, Clone, PartialEq, Eq)]
420pub struct BlockInstance {
421    /// 2026-09-28: Template name.
422    pub template: String,
423    /// 2026-09-28: Layer index; `None` for the prologue, epilogue and draft.
424    pub layer: Option<usize>,
425    /// 2026-09-28: Main or draft.
426    pub section: Section,
427    /// 2026-09-28: First node.
428    pub first: NodeIdx,
429    /// 2026-09-28: One past the last node.
430    pub end: NodeIdx,
431    /// 2026-09-28: The stream edge entering the block, if it reads one.
432    pub stream_in: Option<EdgeIdx>,
433    /// 2026-09-28: The stream edge leaving it, if it writes one.
434    pub stream_out: Option<EdgeIdx>,
435}
436
437/// 2026-09-28: An instantiated circuit.
438#[derive(Debug, Clone, PartialEq, Eq)]
439pub struct Circuit {
440    /// 2026-09-28: The arch id of the TOML it came from (`qwen3_5`).
441    pub arch: String,
442    /// 2026-09-28: The TOML's one-line description.
443    pub description: String,
444    /// 2026-09-28: Nodes in execution order.
445    pub nodes: Vec<Node>,
446    /// 2026-09-28: All edges.
447    pub edges: Vec<Edge>,
448    /// 2026-09-28: Blocks in execution order.
449    pub blocks: Vec<BlockInstance>,
450    /// 2026-09-28: Layer kinds, from the arch shape.
451    pub layer_kinds: Vec<LayerKind>,
452    /// 2026-09-28: The dims every expression was evaluated under.
453    pub dims: BTreeMap<String, u64>,
454    /// 2026-09-30: The state every block keeps between steps, in block order.
455    pub states: Vec<crate::state::StateDecl>,
456}
457
458impl Circuit {
459    /// 2026-09-28: Find an edge by id.
460    pub fn edge(&self, id: &str) -> Option<EdgeIdx> {
461        self.edges.iter().position(|e| e.id == id)
462    }
463
464    /// 2026-09-28: Find a node by id.
465    pub fn node(&self, id: &str) -> Option<NodeIdx> {
466        self.nodes.iter().position(|n| n.id == id)
467    }
468}
metrale_circuit/ir.rs

metrale_circuit/
ir.rs
