metrale_circuit/memory/
caches.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The caches a circuit declares ([`StateKind::is_cache`]), sized: how many units
4//! each kind holds is a plan input ([`CacheInputs`]) the engine decides; the bytes of one unit
5//! are the declaration's elements times its element size. A cache whose input is absent holds
6//! nothing; nothing here guesses a count.
7//!
8//! Owner: metrale-circuit (memory).
9//! Invariants:
10//! - One unit rule per kind ([`units_of`]); every byte count but the prompt-lookup index is
11//!   `units x elements x element size`.
12//! - The prompt-lookup index is a host hash table per sequence: its bytes follow the table's
13//!   growth law ([`hash_table_bytes`]), evaluated at the sequence's history length.
14//! - A keyed format must be given ([`crate::state::StateInputs::formats`]); a missing key is
15//!   an error, as in [`crate::state::StatePlan`].
16
17use std::collections::BTreeMap;
18
19use crate::state::{StateDecl, StateDtype, StateError, StateFormat, StateKind};
20
21/// 2026-10-02: The prompt-lookup index of the sequences in flight.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct LookupInputs {
24    /// 2026-10-02: Sequences holding an index (the concurrency).
25    pub sequences: u64,
26    /// 2026-10-02: Tokens of history each one indexes (prompt plus generated).
27    pub history_tokens: u64,
28}
29
30/// 2026-10-02: How many units each cache kind holds; `None` or 0 for a cache the serve does not
31/// keep.
32#[derive(Debug, Clone, Default, PartialEq, Eq)]
33pub struct CacheInputs {
34    /// 2026-10-02: Prefix-cache (Marconi) snapshot slots.
35    pub prefix_snapshot_slots: u64,
36    /// 2026-10-02: Decode-rollback ring slots per sequence slot, and the sequence slots.
37    pub ring: (u64, u64),
38    /// 2026-10-02: Carried-state verify slots (the dummy included) and carry-table rows; 0
39    /// without the carried-state verify.
40    pub carry: (u64, u64),
41    /// 2026-10-02: Verify-table rows (the WY tables, the accepted-hidden stash); 0 without a
42    /// proposer.
43    pub verify_table_rows: u64,
44    /// 2026-10-02: Hidden rows captured for the draft head.
45    pub capture_rows: u64,
46    /// 2026-10-02: The prompt-lookup index; `None` without prompt-lookup decoding.
47    pub lookup: Option<LookupInputs>,
48    /// 2026-10-02: Token-tree verify slots and nodes per tree; `None` without a token tree.
49    pub tree: Option<(u64, u64)>,
50    /// 2026-10-02: Verify slots and draft positions per slot whose drafted tokens are kept.
51    pub drafts: Option<(u64, u64)>,
52}
53
54/// 2026-10-02: One sized cache declaration.
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub struct CacheTerm {
57    /// 2026-10-02: The declaration's id (`l0.gdn.prefix_h`).
58    pub state: String,
59    /// 2026-10-02: Its kind.
60    pub kind: StateKind,
61    /// 2026-10-02: Its element type.
62    pub dtype: StateDtype,
63    /// 2026-10-02: Units held.
64    pub units: u64,
65    /// 2026-10-02: Bytes held.
66    pub bytes: u64,
67    /// 2026-10-02: Held in host memory.
68    pub host: bool,
69}
70
71/// 2026-10-02: The units `inputs` give a cache of `kind`; `None` for a non-cache kind, for the
72/// prompt-lookup index (not linear in its units) and on overflow.
73pub fn units_of(kind: StateKind, inputs: &CacheInputs) -> Option<u64> {
74    match kind {
75        StateKind::Recurrent | StateKind::PagedKv | StateKind::PromptLookupIndex => None,
76        StateKind::PrefixSnapshot => Some(inputs.prefix_snapshot_slots),
77        StateKind::RingSnapshot => inputs.ring.0.checked_mul(inputs.ring.1),
78        StateKind::CarryStash => Some(inputs.carry.0),
79        StateKind::CarryTable => Some(inputs.carry.1),
80        StateKind::VerifyTable | StateKind::AcceptStash => Some(inputs.verify_table_rows),
81        StateKind::HiddenCapture => Some(inputs.capture_rows),
82        StateKind::TokenTreeMask => inputs.tree.map_or(Some(0), |(slots, nodes)| {
83            slots.checked_mul(nodes)?.checked_mul(nodes)
84        }),
85        StateKind::DraftTokens => inputs
86            .drafts
87            .map_or(Some(0), |(slots, pos)| slots.checked_mul(pos)),
88    }
89}
90
91/// 2026-10-02: Bytes of a Rust `HashMap` (hashbrown) holding `entries` entries of
92/// `entry_bytes` each (a multiple of 16), as it grows by inserting them: buckets are the next
93/// power of two of `entries * 8 / 7` (4 or 8 below 8 entries; hashbrown `capacity_to_buckets`),
94/// each bucket one entry plus one control byte, plus one 16-byte control group. 0 entries hold no
95/// allocation.
96pub fn hash_table_bytes(entries: u64, entry_bytes: u64) -> Option<u64> {
97    if entries == 0 {
98        return Some(0);
99    }
100    let buckets = if entries < 4 {
101        4
102    } else if entries < 8 {
103        8
104    } else {
105        (entries.checked_mul(8)? / 7).checked_next_power_of_two()?
106    };
107    buckets
108        .checked_mul(entry_bytes.checked_add(1)?)?
109        .checked_add(16)
110}
111
112/// 2026-10-02: Size every cache declaration of `states` under `inputs`; `formats` resolves keyed
113/// formats.
114pub fn cache_terms(
115    states: &[StateDecl],
116    formats: &BTreeMap<String, StateDtype>,
117    inputs: &CacheInputs,
118) -> Result<Vec<CacheTerm>, StateError> {
119    let mut out = Vec::new();
120    for s in states.iter().filter(|s| s.kind.is_cache()) {
121        let dtype = match &s.format {
122            StateFormat::Fixed(d) => *d,
123            StateFormat::Keyed(k) => *formats.get(k).ok_or_else(|| StateError::MissingFormat {
124                state: s.id.clone(),
125                key: k.clone(),
126            })?,
127        };
128        let over = || StateError::Overflow(s.id.clone());
129        let unit = s.elements.checked_mul(dtype.size()).ok_or_else(over)?;
130        let (units, bytes) = match s.kind {
131            StateKind::PromptLookupIndex => match inputs.lookup {
132                Some(l) => {
133                    let one = hash_table_bytes(l.history_tokens, unit).ok_or_else(over)?;
134                    (
135                        l.sequences.checked_mul(l.history_tokens).ok_or_else(over)?,
136                        l.sequences.checked_mul(one).ok_or_else(over)?,
137                    )
138                }
139                None => (0, 0),
140            },
141            k => {
142                let units = units_of(k, inputs).ok_or_else(over)?;
143                (units, units.checked_mul(unit).ok_or_else(over)?)
144            }
145        };
146        out.push(CacheTerm {
147            state: s.id.clone(),
148            kind: s.kind,
149            dtype,
150            units,
151            bytes,
152            host: s.kind.is_host(),
153        });
154    }
155    Ok(out)
156}
157
158#[cfg(test)]
159#[path = "caches_tests.rs"]
160mod caches_tests;
metrale_circuit/memory/caches.rs

metrale_circuit/memory/
caches.rs
