metrale_circuit/memory/
budget.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: The memory budget of a device class and the queries over it: the driver's share
4//! (`kernels/<class>/HARDWARE.toml [memory]`, calibrated per class), the util budget, and the
5//! inverse questions "the most sequences at this length" and "the longest length at this many
6//! sequences", answered by bisection over the one evaluator the forward report uses.
7//!
8//! Owner: metrale-circuit (memory).
9//! Invariants:
10//! - Every driver term is read from the class's `[memory]` table; a missing key is an error,
11//!   never a default (#72's `runtime_headroom.rs` constants are this table's first source).
12//! - The inverse queries assume the footprint never shrinks as sequences or tokens grow, and
13//!   they check the answer: the bound they return fits, the next value does not (or is the cap).
14
15use serde::Deserialize;
16
17/// 2026-10-02: What the driver holds beyond the allocation ledger, for one device class.
18#[derive(Debug, Clone, Copy, PartialEq, Deserialize)]
19#[serde(deny_unknown_fields)]
20pub struct DriverTerms {
21    /// 2026-10-02: The CUDA context, loaded modules, the local-memory reservation and the
22    /// graphs, bytes.
23    pub driver_fixed_bytes: u64,
24    /// 2026-10-02: The driver's per-allocation bookkeeping, per mille of the util budget.
25    pub driver_budget_per_mille: u64,
26    /// 2026-10-02: Host memory is the device's memory (a unified-memory SoC): host caches compete
27    /// for the same pool, outside the util budget.
28    pub unified: bool,
29    /// 2026-10-02: The largest `--gpu-memory-utilization` a serve of this class may run at.
30    pub util_ceiling: f64,
31}
32
33/// 2026-10-02: Why the class's memory terms could not be read.
34#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
35pub enum BudgetError {
36    /// 2026-10-02: The HARDWARE.toml text, or its `[memory]` table.
37    #[error("kernels/{class}/HARDWARE.toml [memory]: {detail}")]
38    Terms {
39        /// 2026-10-02: The class.
40        class: String,
41        /// 2026-10-02: What was wrong.
42        detail: String,
43    },
44    /// 2026-10-02: An inverse query whose lower bound does not fit, or whose evaluator failed.
45    #[error("{0}")]
46    Query(String),
47}
48
49impl DriverTerms {
50    /// 2026-10-02: The `[memory]` table of `class`'s HARDWARE.toml text.
51    pub fn parse(class: &str, hardware_toml: &str) -> Result<Self, BudgetError> {
52        let err = |detail: String| BudgetError::Terms {
53            class: class.to_string(),
54            detail,
55        };
56        let mut t: toml::Table = toml::from_str(hardware_toml).map_err(|e| err(e.to_string()))?;
57        let mem = t
58            .remove("memory")
59            .ok_or_else(|| err("the class declares no [memory] table".into()))?;
60        let terms: Self = mem
61            .try_into()
62            .map_err(|e: toml::de::Error| err(e.to_string()))?;
63        if terms.driver_budget_per_mille > 1000
64            || !(terms.util_ceiling > 0.0 && terms.util_ceiling <= 1.0)
65        {
66            return Err(err(
67                "driver_budget_per_mille is at most 1000 and util_ceiling in (0, 1]".into(),
68            ));
69        }
70        Ok(terms)
71    }
72
73    /// 2026-10-02: The driver's bytes under a util budget of `budget_bytes`, with
74    /// `chunk_slack` bytes of small-allocation chunks (measured from a ledger, or 0 when no
75    /// ledger is known).
76    pub fn bytes(&self, budget_bytes: u64, chunk_slack: u64) -> u64 {
77        self.driver_fixed_bytes
78            .saturating_add(budget_bytes / 1000 * self.driver_budget_per_mille)
79            .saturating_add(chunk_slack)
80    }
81}
82
83/// 2026-10-02: The util budget: `device_bytes x util`, as the engine computes it (truncated).
84pub fn util_budget(device_bytes: u64, util: f64) -> u64 {
85    (device_bytes as f64 * util) as u64
86}
87
88/// 2026-10-02: The largest `x` in `lo..=hi` with `fits(x)`, by bisection; `None` when `lo` does
89/// not fit. `fits` must be monotone (true up to some `x`, false after).
90pub fn largest_fitting(
91    lo: u64,
92    hi: u64,
93    fits: &mut dyn FnMut(u64) -> Result<bool, BudgetError>,
94) -> Result<Option<u64>, BudgetError> {
95    if lo > hi {
96        return Err(BudgetError::Query(format!("empty range {lo}..={hi}")));
97    }
98    if !fits(lo)? {
99        return Ok(None);
100    }
101    if fits(hi)? {
102        return Ok(Some(hi));
103    }
104    // 2026-10-02: `good` fits, `bad` does not; halve the gap.
105    let (mut good, mut bad) = (lo, hi);
106    while bad - good > 1 {
107        let mid = good + (bad - good) / 2;
108        if fits(mid)? {
109            good = mid;
110        } else {
111            bad = mid;
112        }
113    }
114    Ok(Some(good))
115}
116
117#[cfg(test)]
118#[path = "budget_tests.rs"]
119mod budget_tests;
metrale_circuit/memory/budget.rs

metrale_circuit/memory/
budget.rs
