metrale_circuit/
state_kinds.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: What a declared state is ([`StateKind`]) and how long a unit of it lives
4//! ([`Lifetime`]), split from `state.rs`.
5//!
6//! Owner: metrale-circuit.
7//! Invariants:
8//! - Every kind has one spelling ([`StateKind::parse`] of [`StateKind::name`] is the kind) and
9//!   one lifetime ([`Lifetime::of_kind`]); a spelling outside the list is refused by the loader.
10
11#[cfg(doc)]
12use super::{Holding, StatePlan};
13
14/// 2026-09-30: What a state is.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
16pub enum StateKind {
17    /// 2026-09-30: A recurrent state, one unit per sequence slot (GatedDeltaNet and Mamba2
18    /// `h`, their conv windows).
19    Recurrent,
20    /// 2026-09-30: One side of a paged KV cache, one unit per token.
21    PagedKv,
22    /// 2026-10-02: A prefix-cache (Marconi) snapshot of a recurrent state (`of`), or of the
23    /// last-token hidden row, one unit per snapshot slot.
24    PrefixSnapshot,
25    /// 2026-10-02: The decode-rollback ring's copy of a recurrent state (`of`), one unit per ring
26    /// slot per sequence slot.
27    RingSnapshot,
28    /// 2026-10-02: The carried-state verify's per-slot stash of a recurrent layer, one unit per
29    /// verify slot (the dummy included).
30    CarryStash,
31    /// 2026-10-02: The carried-state verify's per-layer pointer and flag tables, one unit per
32    /// carry-table row.
33    CarryTable,
34    /// 2026-10-02: A verify's per-layer WY pointer tables, one unit per verify-table row.
35    VerifyTable,
36    /// 2026-10-02: The hidden rows of accepted drafts a verify stashes for the next propose, one
37    /// unit per verify-table row.
38    AcceptStash,
39    /// 2026-10-02: Hidden rows captured for the draft head (the prompt's, for the drafter
40    /// prefill; DFlash's per-step capture), one unit per captured row.
41    HiddenCapture,
42    /// 2026-10-02: Host: a sequence's prompt-lookup n-gram index, one entry per history token
43    /// (sized by the hash table law in `memory::caches`).
44    PromptLookupIndex,
45    /// 2026-10-02: A token tree's attention mask, one unit per (tree node, tree node) pair of a
46    /// verify slot.
47    TokenTreeMask,
48    /// 2026-10-02: A speculative step's drafted tokens and their scores, one unit per draft
49    /// position of a verify slot (tree nodes under a token tree).
50    DraftTokens,
51}
52
53/// 2026-10-02: Every kind and its spelling in the circuit TOML.
54const KINDS: [(StateKind, &str); 12] = [
55    (StateKind::Recurrent, "recurrent"),
56    (StateKind::PagedKv, "paged_kv"),
57    (StateKind::PrefixSnapshot, "prefix_snapshot"),
58    (StateKind::RingSnapshot, "ring_snapshot"),
59    (StateKind::CarryStash, "carry_stash"),
60    (StateKind::CarryTable, "carry_table"),
61    (StateKind::VerifyTable, "verify_table"),
62    (StateKind::AcceptStash, "accept_stash"),
63    (StateKind::HiddenCapture, "hidden_capture"),
64    (StateKind::PromptLookupIndex, "prompt_lookup_index"),
65    (StateKind::TokenTreeMask, "token_tree_mask"),
66    (StateKind::DraftTokens, "draft_tokens"),
67];
68
69impl StateKind {
70    /// 2026-09-30: The spelling in the circuit TOML.
71    pub fn parse(s: &str) -> Option<Self> {
72        KINDS.iter().find(|(_, n)| *n == s).map(|(k, _)| *k)
73    }
74
75    /// 2026-10-02: The spelling.
76    pub fn name(self) -> &'static str {
77        KINDS
78            .iter()
79            .find(|(k, _)| *k == self)
80            .map(|(_, n)| *n)
81            .unwrap_or("?")
82    }
83
84    /// 2026-10-02: Every spelling, for refusals.
85    pub fn spellings() -> Vec<&'static str> {
86        KINDS.iter().map(|(_, n)| *n).collect()
87    }
88
89    /// 2026-10-02: A cache: a kind no layer node touches, sized by `memory::caches` from the
90    /// cache inputs, never by [`StatePlan`]. It may copy another state of its block (`of`).
91    pub fn is_cache(self) -> bool {
92        !matches!(self, Self::Recurrent | Self::PagedKv)
93    }
94
95    /// 2026-10-02: The kinds whose unit copies a declared state (`of`).
96    pub fn copies_a_state(self) -> bool {
97        matches!(self, Self::PrefixSnapshot | Self::RingSnapshot)
98    }
99
100    /// 2026-10-02: Held in host memory rather than device memory (on a unified-memory device it
101    /// still competes for the same pool, outside the util budget).
102    pub fn is_host(self) -> bool {
103        matches!(self, Self::PromptLookupIndex)
104    }
105}
106
107/// 2026-09-30: How long a unit of state lives (LIFECYCLE-DESIGN.md section 3.3).
108#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
109pub enum Lifetime {
110    /// 2026-09-30: Allocated once with the model (a KV pool; its blocks are per sequence).
111    Model,
112    /// 2026-09-30: Claimed with a sequence's slot and released with it.
113    Sequence,
114    /// 2026-09-30: Valid from a verify's snapshot to its commit or rollback.
115    Verify,
116    /// 2026-10-02: Owned by the prefix cache or the scheduler's ring, in slots of their own.
117    Snapshot,
118}
119
120impl Lifetime {
121    /// 2026-09-30: The spelling in the circuit TOML.
122    pub fn parse(s: &str) -> Option<Self> {
123        match s {
124            "model" => Some(Self::Model),
125            "sequence" => Some(Self::Sequence),
126            "verify" => Some(Self::Verify),
127            "snapshot" => Some(Self::Snapshot),
128            _ => None,
129        }
130    }
131
132    /// 2026-09-30: The spelling.
133    pub fn name(self) -> &'static str {
134        match self {
135            Self::Model => "model",
136            Self::Sequence => "sequence",
137            Self::Verify => "verify",
138            Self::Snapshot => "snapshot",
139        }
140    }
141
142    /// 2026-09-30: The one lifetime a state of `kind` may declare: a recurrent state lives with
143    /// its sequence, a KV side with the model's pool. For a recurrent state `Verify` belongs to
144    /// the verify holdings ([`Holding::lifetime`]), never to its declaration. 2026-10-02: The
145    /// caches: snapshots live in their owner's slots, the verify caches across one verify, the
146    /// verify tables with the model, a prompt-lookup index and a hidden capture with their
147    /// sequence.
148    pub fn of_kind(kind: StateKind) -> Self {
149        match kind {
150            StateKind::Recurrent | StateKind::PromptLookupIndex | StateKind::HiddenCapture => {
151                Self::Sequence
152            }
153            StateKind::PagedKv | StateKind::CarryTable | StateKind::VerifyTable => Self::Model,
154            StateKind::PrefixSnapshot | StateKind::RingSnapshot => Self::Snapshot,
155            StateKind::CarryStash
156            | StateKind::AcceptStash
157            | StateKind::TokenTreeMask
158            | StateKind::DraftTokens => Self::Verify,
159        }
160    }
161}
metrale_circuit/state_kinds.rs

metrale_circuit/
state_kinds.rs
