metrale_circuit/
state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The state a circuit keeps between steps, declared per block
4//! (`[[block.<name>.state]]`), and the one function that sizes it (M5, LIFECYCLE-DESIGN.md
5//! section 3.3): the recurrent slots of GatedDeltaNet and Mamba2 layers with their verify
6//! intermediates and checkpoints, and the paged KV blocks of attention layers.
7//!
8//! A declaration says what one unit of a state holds: one sequence slot of a recurrent state,
9//! one token of one side of a KV cache. How many units exist (slots, the padding dummy, verify
10//! intermediates per slot, KV blocks) is a plan input ([`StateInputs`]), decided by the engine
11//! and passed in explicitly; this module multiplies and never reads a process global.
12//!
13//! Owner: metrale-circuit.
14//! Invariants:
15//! - Every byte count is `units x elements x element size`, the element count evaluated from
16//!   the circuit's dims when it is instantiated; one function ([`StatePlan::new`]) produces
17//!   every term, so the reserve and the allocation cannot size a state two ways.
18//! - A format named by a key (`{kv_cache_dtype}`) must be given in [`StateInputs::formats`]; a
19//!   missing key is an error, never a default.
20
21use std::collections::BTreeMap;
22
23#[cfg(doc)]
24use crate::ir::Circuit;
25use crate::ir::Section;
26
27#[path = "state_kinds.rs"]
28mod kinds;
29pub use kinds::{Lifetime, StateKind};
30
31/// 2026-09-30: Which verify intermediate count a recurrent state keeps per verify slot.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
33pub enum VerifySteps {
34    /// 2026-09-30: [`VerifyInputs::h_steps`] of the slot (`K - 1`: the state after the last
35    /// row stays live).
36    H,
37    /// 2026-09-30: [`VerifyInputs::conv_steps`] (all `K` windows).
38    Conv,
39}
40
41impl VerifySteps {
42    /// 2026-09-30: The spelling in the circuit TOML (`verify = "h_steps"`).
43    pub fn parse(s: &str) -> Option<Self> {
44        match s {
45            "h_steps" => Some(Self::H),
46            "conv_steps" => Some(Self::Conv),
47            _ => None,
48        }
49    }
50}
51
52/// 2026-09-30: How a node touches a state of its block (`state = { h = "update" }`).
53#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
54pub enum StateAccess {
55    /// 2026-09-30: Reads it (attention over the KV cache).
56    Read,
57    /// 2026-09-30: Writes new units without reading them (the KV write of this step's rows).
58    Write,
59    /// 2026-09-30: Reads and replaces it in place (a recurrence, a conv window).
60    Update,
61    /// 2026-09-30: Copies it after each row but the last into the verify intermediates.
62    Snapshot,
63}
64
65impl StateAccess {
66    /// 2026-09-30: The spelling in the circuit TOML.
67    pub fn parse(s: &str) -> Option<Self> {
68        match s {
69            "read" => Some(Self::Read),
70            "write" => Some(Self::Write),
71            "update" => Some(Self::Update),
72            "snapshot" => Some(Self::Snapshot),
73            _ => None,
74        }
75    }
76}
77
78/// 2026-09-30: The storage element of a state.
79#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
80pub enum StateDtype {
81    F32,
82    F16,
83    Bf16,
84    /// 2026-09-30: One E4M3 byte per element; its scales live outside the state (a static
85    /// per-layer KV scale).
86    Fp8,
87    /// 2026-10-02: Integer elements of the caches: masks and flags (`u8`), token ids and
88    /// counts (`i32`, `u32`), keys and device pointers (`u64`).
89    U8,
90    I32,
91    U32,
92    U64,
93}
94
95impl StateDtype {
96    /// 2026-09-30: The spelling in the circuit TOML and in [`StateInputs::formats`].
97    pub fn parse(s: &str) -> Option<Self> {
98        match s {
99            "f32" => Some(Self::F32),
100            "f16" => Some(Self::F16),
101            "bf16" => Some(Self::Bf16),
102            "fp8" => Some(Self::Fp8),
103            "u8" => Some(Self::U8),
104            "i32" => Some(Self::I32),
105            "u32" => Some(Self::U32),
106            "u64" => Some(Self::U64),
107            _ => None,
108        }
109    }
110
111    /// 2026-10-02: The spelling.
112    pub fn name(self) -> &'static str {
113        match self {
114            Self::F32 => "f32",
115            Self::F16 => "f16",
116            Self::Bf16 => "bf16",
117            Self::Fp8 => "fp8",
118            Self::U8 => "u8",
119            Self::I32 => "i32",
120            Self::U32 => "u32",
121            Self::U64 => "u64",
122        }
123    }
124
125    /// 2026-09-30: Bytes per element.
126    pub fn size(self) -> u64 {
127        match self {
128            Self::F32 | Self::I32 | Self::U32 => 4,
129            Self::F16 | Self::Bf16 => 2,
130            Self::Fp8 | Self::U8 => 1,
131            Self::U64 => 8,
132        }
133    }
134}
135
136/// 2026-09-30: A state's storage format: fixed by the circuit, or named by a key the plan
137/// inputs give.
138#[derive(Debug, Clone, PartialEq, Eq)]
139pub enum StateFormat {
140    Fixed(StateDtype),
141    Keyed(String),
142}
143
144impl StateFormat {
145    /// 2026-09-30: `{key}` or a dtype spelling.
146    pub fn parse(s: &str) -> Option<Self> {
147        match s.strip_prefix('{').and_then(|r| r.strip_suffix('}')) {
148            Some(key) if !key.is_empty() => Some(Self::Keyed(key.to_string())),
149            Some(_) => None,
150            None => StateDtype::parse(s).map(Self::Fixed),
151        }
152    }
153}
154
155/// 2026-09-30: One declared state of one block instance.
156#[derive(Debug, Clone, PartialEq, Eq)]
157pub struct StateDecl {
158    /// 2026-09-30: `<block prefix>.<local>` (`l0.gdn.h`, `draft.attn.k`).
159    pub id: String,
160    /// 2026-09-30: The template-local id (`h`, `conv`, `k`, `v`).
161    pub local: String,
162    /// 2026-09-30: The block template.
163    pub block: String,
164    /// 2026-09-30: The layer, when the block is a layer's.
165    pub layer: Option<usize>,
166    /// 2026-09-30: Main or draft.
167    pub section: Section,
168    /// 2026-09-30: What it is.
169    pub kind: StateKind,
170    /// 2026-09-30: Storage format.
171    pub format: StateFormat,
172    /// 2026-09-30: Elements per unit (slot or token).
173    pub elements: u64,
174    /// 2026-09-30: The verify intermediates a recurrent state keeps; `None` for a KV side.
175    pub verify: Option<VerifySteps>,
176    /// 2026-09-30: How long a live unit lives ([`Lifetime::of_kind`]).
177    pub lifetime: Lifetime,
178}
179
180/// 2026-09-30: Why a state plan could not be made.
181#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
182pub enum StateError {
183    /// 2026-09-30: A keyed format the inputs do not give.
184    #[error("state `{state}` reads format key `{key}`, which the plan inputs do not give")]
185    MissingFormat {
186        /// 2026-09-30: The state.
187        state: String,
188        /// 2026-09-30: The key.
189        key: String,
190    },
191    /// 2026-09-30: Inputs that do not fit the circuit.
192    #[error("state plan: {0}")]
193    Inputs(String),
194    /// 2026-09-30: A size that does not fit u64.
195    #[error("state `{0}` overflows")]
196    Overflow(String),
197}
198
199/// 2026-09-30: The verify pools of a speculative serve over the recurrent states.
200#[derive(Debug, Clone, PartialEq, Eq)]
201pub struct VerifyInputs {
202    /// 2026-09-30: Per verify slot (the dummy, when there is one, last): the h intermediates it
203    /// holds (`K - 1` for a K-row verify, tiered by the engine's ladder).
204    pub h_steps: Vec<u64>,
205    /// 2026-09-30: Conv intermediates every verify slot holds (all K).
206    pub conv_steps: u64,
207}
208
209/// 2026-09-30: One paged KV cache: blocks of `block_size` tokens.
210#[derive(Debug, Clone, Copy, PartialEq, Eq)]
211pub struct KvInputs {
212    pub blocks: u64,
213    pub block_size: u64,
214}
215
216/// 2026-09-30: The plan inputs the engine decides.
217#[derive(Debug, Clone, PartialEq, Eq)]
218pub struct StateInputs {
219    /// 2026-09-30: Storage format per key the circuit names (`kv_cache_dtype`, `ssm_h_storage`).
220    pub formats: BTreeMap<String, StateDtype>,
221    /// 2026-09-30: Recurrent slots, the padding dummy included when the plan has one.
222    pub slots: u64,
223    /// 2026-09-30: The speculative verify pools; `None` without speculation.
224    pub verify: Option<VerifyInputs>,
225    /// 2026-09-30: The target's KV cache; `None` leaves the attention states out.
226    pub kv: Option<KvInputs>,
227    /// 2026-09-30: The draft head's own KV cache; `None` leaves the draft's states out.
228    pub draft_kv: Option<KvInputs>,
229}
230
231/// 2026-09-30: What a term of the plan holds.
232#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
233pub enum Holding {
234    /// 2026-09-30: The live state of every slot.
235    Live,
236    /// 2026-09-30: Verify intermediates (the per-row states of a verify).
237    Steps,
238    /// 2026-09-30: The pre-verify checkpoint, one per verify slot.
239    Checkpoint,
240    /// 2026-09-30: KV blocks.
241    Blocks,
242}
243
244impl Holding {
245    /// 2026-09-30: How long this holding of a state declared `declared` lives: the verify
246    /// intermediates and checkpoints only across one verify, the rest as declared.
247    pub fn lifetime(self, declared: Lifetime) -> Lifetime {
248        match self {
249            Holding::Steps | Holding::Checkpoint => Lifetime::Verify,
250            Holding::Live | Holding::Blocks => declared,
251        }
252    }
253}
254
255/// 2026-09-30: One sized term: a state, what it holds, its units and bytes.
256#[derive(Debug, Clone, PartialEq, Eq)]
257pub struct StateTerm {
258    pub state: String,
259    pub holding: Holding,
260    /// 2026-09-30: `holding.lifetime(declared)`.
261    pub lifetime: Lifetime,
262    pub dtype: StateDtype,
263    pub units: u64,
264    pub bytes: u64,
265}
266
267/// 2026-09-30: Every state of a circuit, sized.
268#[derive(Debug, Clone, PartialEq, Eq)]
269pub struct StatePlan {
270    pub terms: Vec<StateTerm>,
271}
272
273impl StatePlan {
274    /// 2026-09-30: Size every state of `states` (a circuit's [`Circuit::states`], or the
275    /// declarations of one block) under `inputs`. 2026-10-02: Cache kinds
276    /// ([`StateKind::is_cache`]) are left out: `memory::caches` sizes them.
277    pub fn new(states: &[StateDecl], inputs: &StateInputs) -> Result<Self, StateError> {
278        let mut terms = Vec::new();
279        for s in states {
280            let dtype = match &s.format {
281                StateFormat::Fixed(d) => *d,
282                StateFormat::Keyed(k) => {
283                    *inputs
284                        .formats
285                        .get(k)
286                        .ok_or_else(|| StateError::MissingFormat {
287                            state: s.id.clone(),
288                            key: k.clone(),
289                        })?
290                }
291            };
292            let unit = s
293                .elements
294                .checked_mul(dtype.size())
295                .ok_or_else(|| StateError::Overflow(s.id.clone()))?;
296            let mut push = |holding, units: u64| -> Result<(), StateError> {
297                let bytes = unit
298                    .checked_mul(units)
299                    .ok_or_else(|| StateError::Overflow(s.id.clone()))?;
300                terms.push(StateTerm {
301                    state: s.id.clone(),
302                    holding,
303                    lifetime: holding.lifetime(s.lifetime),
304                    dtype,
305                    units,
306                    bytes,
307                });
308                Ok(())
309            };
310            match s.kind {
311                StateKind::Recurrent => {
312                    push(Holding::Live, inputs.slots)?;
313                    if let Some(v) = &inputs.verify {
314                        let steps = match s.verify {
315                            Some(VerifySteps::H) => v.h_steps.iter().sum(),
316                            Some(VerifySteps::Conv) => v.conv_steps * v.h_steps.len() as u64,
317                            None => {
318                                return Err(StateError::Inputs(format!(
319                                    "recurrent state `{}` declares no verify intermediates",
320                                    s.id
321                                )));
322                            }
323                        };
324                        push(Holding::Steps, steps)?;
325                        push(Holding::Checkpoint, v.h_steps.len() as u64)?;
326                    }
327                }
328                StateKind::PagedKv => {
329                    let kv = match s.section {
330                        Section::Main => inputs.kv,
331                        Section::Draft => inputs.draft_kv,
332                    };
333                    if let Some(kv) = kv {
334                        let tokens = kv
335                            .blocks
336                            .checked_mul(kv.block_size)
337                            .ok_or_else(|| StateError::Overflow(s.id.clone()))?;
338                        push(Holding::Blocks, tokens)?;
339                    }
340                }
341                // 2026-10-02: Caches are sized by `memory::caches` from their own inputs.
342                StateKind::PrefixSnapshot
343                | StateKind::RingSnapshot
344                | StateKind::CarryStash
345                | StateKind::CarryTable
346                | StateKind::VerifyTable
347                | StateKind::AcceptStash
348                | StateKind::HiddenCapture
349                | StateKind::PromptLookupIndex
350                | StateKind::TokenTreeMask
351                | StateKind::DraftTokens => {}
352            }
353        }
354        Ok(Self { terms })
355    }
356
357    /// 2026-09-30: Total bytes of the terms `keep` selects.
358    pub fn bytes_where(&self, keep: impl Fn(&StateTerm) -> bool) -> u64 {
359        self.terms.iter().filter(|t| keep(t)).map(|t| t.bytes).sum()
360    }
361
362    /// 2026-09-30: Total bytes.
363    pub fn bytes(&self) -> u64 {
364        self.bytes_where(|_| true)
365    }
366}
367
368#[cfg(test)]
369#[path = "state_tests.rs"]
370mod state_tests;
metrale_circuit/state.rs

metrale_circuit/
state.rs
