metrale_circuit/memory/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The circuit memory model (owner directive 2026-10-02): for each node, the memory it
4//! holds (weights and their load-time copies, the activations it materializes at the peak run,
5//! its kernel family's workspace, the state it touches), for each declared state and cache its
6//! bytes as a function of the runtime parameters, and their sum against a device's budget.
7//!
8//! - [`weights`] / [`copies`]: stored weights at the declared formats and the derived copies.
9//! - [`caches`]: the cache kinds of [`crate::state::StateKind`], sized from [`CacheInputs`].
10//! - [`budget`]: the class's driver terms (`HARDWARE.toml [memory]`), the util budget and the
11//!   inverse queries.
12//! - [`render`]: the tables and the JSON of `met circuit memory`.
13//!
14//! Owner: metrale-circuit (memory).
15//! Invariants:
16//! - Pure (SBIO): every text, plan and count arrives as an argument. Counts the engine decides
17//!   (slots, verify slots, snapshot slots, KV blocks, the legacy arena) are inputs, so the
18//!   engine's own functions decide them and this module only multiplies.
19//! - Absent inputs size nothing; no count or format is defaulted here.
20//! - Activations are the buffer planner's arena for each run given ([`ActivationRun`]); the
21//!   peak run sets the term. A legacy arena, when given, is what the budget charges instead
22//!   (it is what the default forward allocates), and the workspaces it already holds
23//!   (`arena = true`) are not charged twice. A workspace is sized at the widest run.
24//! - A state is attributed to the node that owns it (updates or writes it); caches, which no
25//!   node touches, appear only in the state table.
26
27pub mod budget;
28pub mod caches;
29pub mod copies;
30pub mod render;
31pub mod weights;
32mod workspace;
33
34use std::collections::BTreeMap;
35
36pub use budget::{BudgetError, DriverTerms};
37pub use caches::{CacheInputs, CacheTerm, LookupInputs};
38pub use copies::{CopyRule, parse_copies};
39pub use weights::NodeWeights;
40pub use workspace::WorkspaceTerm;
41
42use crate::fuser::FusionPlan;
43use crate::ir::{Circuit, NodeIdx};
44use crate::planner;
45use crate::state::{StateAccess, StateError, StateInputs, StatePlan};
46use crate::venn::families::Families;
47
48/// 2026-10-02: Why a memory model could not be evaluated.
49#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
50pub enum MemoryError {
51    /// 2026-10-02: A dim a term needs.
52    #[error("node `{node}`: the memory model needs dim `{dim}`, which the circuit does not define")]
53    MissingDim {
54        /// 2026-10-02: Node id.
55        node: String,
56        /// 2026-10-02: The dim.
57        dim: String,
58    },
59    /// 2026-10-02: A size that does not evaluate (overflow, or a shape off its format's groups).
60    #[error("node `{node}`: {what} has no byte size")]
61    Size {
62        /// 2026-10-02: Node id.
63        node: String,
64        /// 2026-10-02: What was sized.
65        what: String,
66    },
67    /// 2026-10-02: The state or cache plan.
68    #[error(transparent)]
69    State(#[from] StateError),
70    /// 2026-10-02: The buffer planner.
71    #[error(transparent)]
72    Plan(#[from] planner::PlanError),
73    /// 2026-10-02: A workspace expression.
74    #[error("family `{family}` workspace `{name}`: {detail}")]
75    Workspace {
76        /// 2026-10-02: Family id.
77        family: String,
78        /// 2026-10-02: Workspace name.
79        name: String,
80        /// 2026-10-02: What was wrong.
81        detail: String,
82    },
83    /// 2026-10-02: The budget or a query over it.
84    #[error(transparent)]
85    Budget(#[from] BudgetError),
86}
87
88/// 2026-10-02: One fused plan whose activations the model sizes, at its row count.
89#[derive(Debug, Clone, Copy)]
90pub struct ActivationRun<'a> {
91    /// 2026-10-02: What the run is (`decode C=16`, `verify 16x4`, `prefill T=8192`).
92    pub label: &'a str,
93    /// 2026-10-02: Rows the plan is sized for.
94    pub rows: u64,
95    /// 2026-10-02: The plan.
96    pub plan: &'a FusionPlan,
97}
98
99/// 2026-10-02: Everything a memory model is evaluated from.
100#[derive(Clone, Copy)]
101pub struct MemoryInputs<'a> {
102    /// 2026-10-02: The circuit at the served formats.
103    pub served: &'a Circuit,
104    /// 2026-10-02: The same model at its checkpoint's declared formats.
105    pub declared: &'a Circuit,
106    /// 2026-10-02: Derived-copy rules.
107    pub copies: &'a [CopyRule],
108    /// 2026-10-02: Serve settings the copy rules read.
109    pub settings: &'a BTreeMap<String, String>,
110    /// 2026-10-02: Recurrent slots and KV blocks (the engine's counts).
111    pub states: &'a StateInputs,
112    /// 2026-10-02: The draft head's KV element type when it differs from the target's (an
113    /// MTP head keeps its own pool, `kv_cache_dtype` of the draft section's states).
114    pub draft_kv_dtype: Option<crate::state::StateDtype>,
115    /// 2026-10-02: Cache units (the engine's counts).
116    pub caches: &'a CacheInputs,
117    /// 2026-10-02: Activation runs; the largest arena is the term.
118    pub runs: &'a [ActivationRun<'a>],
119    /// 2026-10-02: Kernel families for the workspaces, and the device's SM count they read.
120    pub families: Option<(&'a Families, u64)>,
121    /// 2026-10-02: The legacy buffer arena the default forward allocates, bytes.
122    pub legacy_arena: Option<u64>,
123    /// 2026-10-02: The class's driver terms and the util budget, bytes.
124    pub driver: DriverTerms,
125    /// 2026-10-02: `device memory x util`.
126    pub budget_bytes: u64,
127    /// 2026-10-02: Small-allocation chunk slack measured from a ledger; 0 when unknown.
128    pub chunk_slack: u64,
129}
130
131/// 2026-10-02: One node's memory.
132#[derive(Debug, Clone, PartialEq, Eq)]
133pub struct NodeMemory {
134    /// 2026-10-02: The node.
135    pub node: NodeIdx,
136    /// 2026-10-02: Stored weight bytes.
137    pub stored: u64,
138    /// 2026-10-02: Derived copy bytes.
139    pub derived: u64,
140    /// 2026-10-02: Bytes of the edges it materializes in the peak run.
141    pub activations: u64,
142    /// 2026-10-02: Its family's workspace (shared by every node of the family).
143    pub workspace: u64,
144    /// 2026-10-02: Bytes of the states it owns (updates or writes).
145    pub state: u64,
146}
147
148/// 2026-10-02: One activation run's arena.
149#[derive(Debug, Clone, PartialEq, Eq)]
150pub struct RunArena {
151    /// 2026-10-02: The run's label.
152    pub label: String,
153    /// 2026-10-02: Rows.
154    pub rows: u64,
155    /// 2026-10-02: Planned arena (liveness-shared) bytes.
156    pub arena: u64,
157    /// 2026-10-02: Materialized bytes before reuse.
158    pub materialized: u64,
159}
160
161/// 2026-10-02: The sums.
162#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
163pub struct Totals {
164    pub weights_stored: u64,
165    /// 2026-10-02: Of the stored bytes, the checkpoint tensors no weight node binds (norms, conv
166    /// and gate parameters, a vision tower), measured by the caller ([`MemoryReport::with_outside`]).
167    pub weights_outside: u64,
168    pub weights_derived: u64,
169    /// 2026-10-02: Of the derived bytes, the leaked ones.
170    pub weights_leaked: u64,
171    /// 2026-10-02: The planner's peak arena.
172    pub activations_planned: u64,
173    /// 2026-10-02: The activations the budget charges (the legacy arena when given).
174    pub activations: u64,
175    /// 2026-10-02: Workspaces the budget charges.
176    pub workspace: u64,
177    /// 2026-10-02: Recurrent pool and KV.
178    pub states: u64,
179    /// 2026-10-02: Device caches.
180    pub caches: u64,
181    /// 2026-10-02: Host caches.
182    pub host: u64,
183    pub driver: u64,
184    /// 2026-10-02: Everything on the device.
185    pub device: u64,
186}
187
188/// 2026-10-02: An evaluated memory model.
189#[derive(Debug, Clone, PartialEq)]
190pub struct MemoryReport {
191    pub nodes: Vec<NodeMemory>,
192    pub weights: Vec<NodeWeights>,
193    pub states: StatePlan,
194    pub caches: Vec<CacheTerm>,
195    pub runs: Vec<RunArena>,
196    pub workspaces: Vec<WorkspaceTerm>,
197    pub totals: Totals,
198    pub budget_bytes: u64,
199}
200
201impl MemoryReport {
202    /// 2026-10-02: Charge `bytes` of checkpoint tensors no weight node binds to the stored
203    /// weights (the caller reads them from the checkpoint's headers).
204    pub fn with_outside(mut self, bytes: u64) -> Self {
205        self.totals.weights_outside += bytes;
206        self.totals.weights_stored += bytes;
207        self.totals.device += bytes;
208        self
209    }
210
211    /// 2026-10-02: `budget - device` (negative: over budget).
212    pub fn headroom(&self) -> i128 {
213        i128::from(self.budget_bytes) - i128::from(self.totals.device)
214    }
215}
216
217/// 2026-10-02: KV blocks `sequences` sequences of `tokens` tokens each hold, as the scheduler's
218/// admission reserves them (`tokens / block_size + 1` per sequence: the block the next token
219/// writes into is reserved with the prompt), plus the pool's one dummy block. `None` for a zero
220/// block size or on overflow.
221pub fn kv_blocks_for(sequences: u64, tokens: u64, block_size: u64) -> Option<u64> {
222    let per_seq = tokens.checked_div(block_size)?.checked_add(1)?;
223    sequences.checked_mul(per_seq)?.checked_add(1)
224}
225
226/// 2026-10-02: Every non-cache state of `c` sized under `inputs`, the draft section's at
227/// `draft_kv` when given.
228fn state_plan(
229    c: &Circuit,
230    inputs: &StateInputs,
231    draft_kv: Option<crate::state::StateDtype>,
232) -> Result<StatePlan, MemoryError> {
233    let Some(dt) = draft_kv else {
234        return Ok(StatePlan::new(&c.states, inputs)?);
235    };
236    let (draft, main): (Vec<_>, Vec<_>) = c
237        .states
238        .iter()
239        .cloned()
240        .partition(|s| s.section == crate::ir::Section::Draft);
241    let mut di = inputs.clone();
242    di.formats.insert("kv_cache_dtype".into(), dt);
243    let mut plan = StatePlan::new(&main, inputs)?;
244    plan.terms.extend(StatePlan::new(&draft, &di)?.terms);
245    Ok(plan)
246}
247
248/// 2026-10-02: Evaluate the model.
249pub fn evaluate(inp: &MemoryInputs<'_>) -> Result<MemoryReport, MemoryError> {
250    let c = inp.served;
251    let weights = weights::node_weights(c, inp.declared, inp.copies, inp.settings)?;
252    let states = state_plan(c, inp.states, inp.draft_kv_dtype)?;
253    let caches = caches::cache_terms(&c.states, &inp.states.formats, inp.caches)?;
254    let mut runs = Vec::new();
255    let mut peak: Option<(u64, planner::BufferPlan)> = None;
256    for r in inp.runs {
257        let b = planner::plan_buffers(c, r.plan, r.rows)?;
258        runs.push(RunArena {
259            label: r.label.to_string(),
260            rows: r.rows,
261            arena: b.arena_bytes,
262            materialized: b.materialized_bytes,
263        });
264        if peak.as_ref().is_none_or(|p| b.arena_bytes > p.0) {
265            peak = Some((b.arena_bytes, b));
266        }
267    }
268    let workspaces = match inp.families {
269        Some((fams, sm)) => {
270            let mut all = Vec::new();
271            for r in inp.runs {
272                all.push(workspace::workspace_terms(c, r.plan, fams, r.rows, sm)?);
273            }
274            workspace::widest(all)
275        }
276        None => Vec::new(),
277    };
278    let mut nodes: Vec<NodeMemory> = (0..c.nodes.len())
279        .map(|node| NodeMemory {
280            node,
281            stored: 0,
282            derived: 0,
283            activations: 0,
284            workspace: 0,
285            state: 0,
286        })
287        .collect();
288    for w in &weights {
289        nodes[w.node].stored = w.stored;
290        nodes[w.node].derived = w.derived.iter().map(|d| d.bytes).sum();
291    }
292    if let Some((_, b)) = &peak {
293        for s in &b.slots {
294            if let Some(p) = c.edges[s.edge].producer {
295                nodes[p].activations += s.bytes;
296            }
297        }
298    }
299    for w in &workspaces {
300        for &n in &w.nodes {
301            nodes[n].workspace = nodes[n].workspace.max(w.bytes);
302        }
303    }
304    // 2026-10-02: A state belongs to the node that updates or writes it; readers and
305    // snapshotters touch it without owning it, so the per-node column sums to the states.
306    for (i, n) in c.nodes.iter().enumerate() {
307        nodes[i].state = n
308            .state
309            .iter()
310            .filter(|(_, a)| matches!(a, StateAccess::Update | StateAccess::Write))
311            .map(|&(si, _)| states.bytes_where(|t| t.state == c.states[si].id))
312            .sum();
313    }
314    let mut t = Totals::default();
315    for w in &weights {
316        t.weights_stored += w.stored;
317        for d in &w.derived {
318            t.weights_derived += d.bytes;
319            if d.leaked {
320                t.weights_leaked += d.bytes;
321            }
322        }
323    }
324    t.activations_planned = peak.as_ref().map_or(0, |p| p.0);
325    t.activations = inp.legacy_arena.unwrap_or(t.activations_planned);
326    t.workspace = workspaces
327        .iter()
328        .filter(|w| inp.legacy_arena.is_none() || !w.arena)
329        .map(|w| w.bytes)
330        .sum();
331    t.states = states.bytes();
332    t.caches = caches.iter().filter(|x| !x.host).map(|x| x.bytes).sum();
333    t.host = caches.iter().filter(|x| x.host).map(|x| x.bytes).sum();
334    t.driver = inp.driver.bytes(inp.budget_bytes, inp.chunk_slack);
335    t.device = t.weights_stored
336        + t.weights_derived
337        + t.activations
338        + t.workspace
339        + t.states
340        + t.caches
341        + t.driver;
342    Ok(MemoryReport {
343        nodes,
344        weights,
345        states,
346        caches,
347        runs,
348        workspaces,
349        totals: t,
350        budget_bytes: inp.budget_bytes,
351    })
352}
353
354#[cfg(test)]
355#[path = "memory_tests.rs"]
356mod memory_tests;
metrale_circuit/memory/mod.rs

metrale_circuit/memory/
mod.rs
