metrale_circuit/pipeline/
vocab.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The step vocabulary: which steps each op's pipeline has, in order, and the shape
4//! and spelling of each step's value. The one table every requirement, declaration and
5//! rendering reads.
6//!
7//! Owner: metrale-circuit (pipeline).
8//! Invariants:
9//! - [`steps_of`] matches every [`OpKind`] without a wildcard arm, so a new op does not compile
10//!   until its pipeline shape is defined.
11//! - Projections (a linear of any role, the LM head, the router and both expert projections)
12//!   share one shape; the routed gate/up projection adds the row gather by expert id.
13
14use super::{Num, StepKind, Value};
15use crate::format::Format;
16use crate::ir::{LinearRole, OpKind};
17use crate::state::StateDtype;
18
19const PROJECTION: &[StepKind] = &[
20    StepKind::Act,
21    StepKind::Weight,
22    StepKind::Mma,
23    StepKind::Accumulate,
24    StepKind::Scale,
25];
26const GATHERED_PROJECTION: &[StepKind] = &[
27    StepKind::Gather,
28    StepKind::Act,
29    StepKind::Weight,
30    StepKind::Mma,
31    StepKind::Accumulate,
32    StepKind::Scale,
33];
34const COMPUTE: &[StepKind] = &[StepKind::Compute];
35const RECURRENT: &[StepKind] = &[StepKind::State, StepKind::Compute];
36const ATTENTION: &[StepKind] = &[
37    StepKind::Cache,
38    StepKind::Scores,
39    StepKind::Softmax,
40    StepKind::Accumulate,
41];
42
43/// 2026-10-02: The steps of `op`'s pipeline, in order.
44pub fn steps_of(op: &OpKind) -> &'static [StepKind] {
45    match op {
46        OpKind::Linear(_) | OpKind::LmHead | OpKind::Router | OpKind::ExpertDown => PROJECTION,
47        OpKind::ExpertGateUp => GATHERED_PROJECTION,
48        OpKind::Embed => &[StepKind::Gather],
49        OpKind::RmsNorm
50        | OpKind::FinalNorm
51        | OpKind::QkNorm
52        | OpKind::GatedRmsNorm
53        | OpKind::L2Norm
54        | OpKind::ResidualAdd
55        | OpKind::SiluMul
56        | OpKind::Relu2
57        | OpKind::SigmoidGateMul
58        | OpKind::GdnGates
59        | OpKind::Rope
60        | OpKind::ActQuant(_) => COMPUTE,
61        OpKind::Copy | OpKind::Concat | OpKind::Split => &[StepKind::Move],
62        OpKind::KvWrite => &[StepKind::Cache],
63        OpKind::PagedAttention => ATTENTION,
64        OpKind::Conv1dUpdate | OpKind::GdnRecurrence | OpKind::SsmUpdate => RECURRENT,
65        OpKind::StateSnapshot => &[StepKind::State],
66        OpKind::TopK => &[StepKind::Score],
67        OpKind::Blend => &[StepKind::Scatter, StepKind::Combine],
68        OpKind::EpReduce => &[StepKind::Reduce],
69        OpKind::Argmax => &[StepKind::Compare],
70    }
71}
72
73/// 2026-10-02: The steps of the op a manifest names by base name (`linear`, `act_quant`,
74/// `rms_norm`, ...); `None` for a name outside the vocabulary.
75pub fn steps_for_base(base: &str) -> Option<&'static [StepKind]> {
76    let op = match base {
77        "linear" => OpKind::Linear(LinearRole::Q),
78        "act_quant" => OpKind::ActQuant(Format::Bf16),
79        other => OpKind::parse(other, None, None).ok()?,
80    };
81    Some(steps_of(&op))
82}
83
84/// 2026-10-02: The shape of a step's value.
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum Shape {
87    /// 2026-10-02: [`Value::Format`].
88    Format,
89    /// 2026-10-02: [`Value::Weight`].
90    Weight,
91    /// 2026-10-02: [`Value::Mma`].
92    Mma,
93    /// 2026-10-02: [`Value::Num`].
94    Num,
95    /// 2026-10-02: [`Value::OptNum`].
96    OptNum,
97    /// 2026-10-02: [`Value::Kv`].
98    Kv,
99    /// 2026-10-02: [`Value::State`].
100    State,
101}
102
103const STEPS: [(StepKind, &str, Shape); 17] = [
104    (StepKind::Gather, "gather", Shape::Format),
105    (StepKind::Act, "act", Shape::Format),
106    (StepKind::Weight, "weight", Shape::Weight),
107    (StepKind::Mma, "mma", Shape::Mma),
108    (StepKind::Accumulate, "accumulate", Shape::Num),
109    (StepKind::Scale, "scale", Shape::OptNum),
110    (StepKind::Compute, "compute", Shape::Num),
111    (StepKind::Move, "move", Shape::Format),
112    (StepKind::Cache, "cache", Shape::Kv),
113    (StepKind::Scores, "scores", Shape::Num),
114    (StepKind::Softmax, "softmax", Shape::Num),
115    (StepKind::State, "state", Shape::State),
116    (StepKind::Score, "score", Shape::Num),
117    (StepKind::Scatter, "scatter", Shape::Num),
118    (StepKind::Combine, "combine", Shape::OptNum),
119    (StepKind::Reduce, "reduce", Shape::Num),
120    (StepKind::Compare, "compare", Shape::Format),
121];
122
123/// 2026-10-02: The step's spelling.
124pub fn step_name(k: StepKind) -> &'static str {
125    STEPS
126        .iter()
127        .find(|(s, _, _)| *s == k)
128        .map_or("?", |(_, n, _)| n)
129}
130
131/// 2026-10-02: The step a spelling names.
132pub fn parse_step(s: &str) -> Option<StepKind> {
133    STEPS.iter().find(|(_, n, _)| *n == s).map(|(k, _, _)| *k)
134}
135
136/// 2026-10-02: The shape of `k`'s value.
137pub fn shape_of(k: StepKind) -> Shape {
138    STEPS
139        .iter()
140        .find(|(s, _, _)| *s == k)
141        .map_or(Shape::Num, |(_, _, sh)| *sh)
142}
143
144fn num(s: &str) -> Result<Num, String> {
145    Num::parse(s).ok_or_else(|| format!("`{s}` is no precision (bf16, f16, f32, e4m3, e2m1)"))
146}
147
148/// 2026-10-02: Parse `text` as the value of step `k`.
149pub fn parse_value(k: StepKind, text: &str) -> Result<Value, String> {
150    let fmt = |s: &str| Format::parse(s).map_err(|e| e.to_string());
151    let v = match shape_of(k) {
152        Shape::Format => Value::Format(fmt(text)?),
153        Shape::Weight => {
154            let (stored, operand) = text
155                .split_once("->")
156                .ok_or_else(|| format!("`{text}` is not `<stored format>-><precision>`"))?;
157            Value::Weight {
158                stored: fmt(stored)?,
159                operand: num(operand)?,
160            }
161        }
162        Shape::Mma => {
163            let (a, b) = text
164                .split_once('*')
165                .ok_or_else(|| format!("`{text}` is not `<precision>*<precision>`"))?;
166            Value::Mma {
167                a: num(a)?,
168                b: num(b)?,
169            }
170        }
171        Shape::Num => Value::Num(num(text)?),
172        Shape::OptNum => match text {
173            "none" => Value::OptNum(None),
174            other => Value::OptNum(Some(num(other)?)),
175        },
176        Shape::Kv => {
177            if text.is_empty() || !text.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
178                return Err(format!("`{text}` is no KV-cache dtype spelling"));
179            }
180            Value::Kv(text.to_string())
181        }
182        Shape::State => Value::State(
183            StateDtype::parse(text)
184                .ok_or_else(|| format!("`{text}` is no state dtype (f32, f16, bf16, fp8)"))?,
185        ),
186    };
187    Ok(v)
188}
metrale_circuit/pipeline/vocab.rs

metrale_circuit/pipeline/
vocab.rs
