metrale_speculative/
spec_cost.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-04: The measured speculative-cost table (`--spec-cost-model measured`): what a
4//! verify step and a draft step cost, in wall time and GPU-rail energy, at each batch width
5//! and draft depth, on one box class for one recipe; and the drafter's confidence→acceptance
6//! calibration. `met bench spec-cost` measures both and writes them; the serve loads them and
7//! plans draft depth from them. Pure: parsing and arithmetic only, no I/O.
8//!
9//! Owner: speculative.
10//! Invariants:
11//! - A table is used only under the key it was measured under ([`TableKey::check`]): the box
12//!   class, the recipe, and the plan digest of every mode it measured. A calibration is used
13//!   only for the drafter it was fitted on ([`DrafterKey`]). A mismatch is refused, never
14//!   approximated.
15//! - A parsed table has, for every width it lists, a cell at every draft depth 0..=its maximum,
16//!   all with finite positive costs ([`CostTable::parse`]).
17//! - Costs between measured widths are interpolated linearly in width; a width outside the
18//!   measured range takes the nearest measured width ([`CostTable::cell`]).
19
20use std::collections::BTreeMap;
21
22use serde::Deserialize;
23
24/// 2026-10-04: Schema version of the table files; bumped when a field changes meaning.
25pub const SCHEMA: u32 = 1;
26
27/// 2026-10-04: What a table was measured under. Equal keys are the only licence to use it.
28#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
29pub struct TableKey {
30    pub schema: u32,
31    /// 2026-10-04: `HARDWARE.toml` box class, e.g. `gb10`.
32    pub box_class: String,
33    /// 2026-10-04: Recipe id, e.g. `qwen3.6/qwen3.6-35b-a3b-fp8-nvfp4head`.
34    pub recipe: String,
35    /// 2026-10-04: Circuit plan digest per measured mode (`verify`, `verify_batch`, `draft`),
36    /// the plan the serve compiles for that mode; a kernel change outside these plans does
37    /// not stale the table.
38    pub plan_digests: BTreeMap<String, String>,
39}
40
41/// 2026-10-04: One reason a table key does not match the serve's.
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub enum KeyMismatch {
44    Schema {
45        table: u32,
46        serve: u32,
47    },
48    BoxClass {
49        table: String,
50        serve: String,
51    },
52    Recipe {
53        table: String,
54        serve: String,
55    },
56    /// 2026-10-04: A mode whose digest differs, or is present on one side only.
57    PlanDigest {
58        mode: String,
59        table: Option<String>,
60        serve: Option<String>,
61    },
62}
63
64impl TableKey {
65    /// 2026-10-04: Every way `self` (the table's key) differs from `serve`'s; empty when the
66    /// table may be used.
67    pub fn check(&self, serve: &TableKey) -> Vec<KeyMismatch> {
68        let mut out = Vec::new();
69        if self.schema != serve.schema {
70            out.push(KeyMismatch::Schema {
71                table: self.schema,
72                serve: serve.schema,
73            });
74        }
75        if self.box_class != serve.box_class {
76            out.push(KeyMismatch::BoxClass {
77                table: self.box_class.clone(),
78                serve: serve.box_class.clone(),
79            });
80        }
81        if self.recipe != serve.recipe {
82            out.push(KeyMismatch::Recipe {
83                table: self.recipe.clone(),
84                serve: serve.recipe.clone(),
85            });
86        }
87        let modes: std::collections::BTreeSet<&String> = self
88            .plan_digests
89            .keys()
90            .chain(serve.plan_digests.keys())
91            .collect();
92        for mode in modes {
93            let (t, s) = (self.plan_digests.get(mode), serve.plan_digests.get(mode));
94            if t != s {
95                out.push(KeyMismatch::PlanDigest {
96                    mode: mode.clone(),
97                    table: t.cloned(),
98                    serve: s.cloned(),
99                });
100            }
101        }
102        out
103    }
104}
105
106/// 2026-10-04: The drafter a confidence calibration was fitted on.
107#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
108pub struct DrafterKey {
109    /// 2026-10-04: sha256 over the MTP head's tensors as loaded.
110    pub weights_sha256: String,
111    /// 2026-10-04: `--mtp-vocab`.
112    pub vocab: usize,
113    /// 2026-10-04: `--mtp-quantization`.
114    pub quantization: String,
115    /// 2026-10-04: Whether the drafter's prompt context was active.
116    pub context: bool,
117}
118
119/// 2026-10-04: One measured cell: `n` sequences, `k` drafts each (`k = 0` is a plain decode
120/// step, no draft and no verify rows beyond the one token).
121#[derive(Debug, Clone, Copy, PartialEq, Deserialize)]
122pub struct Cell {
123    pub n: usize,
124    pub k: usize,
125    /// 2026-10-04: One verify (or, at `k = 0`, decode) step: median wall ms and GPU-rail J.
126    pub verify_ms: f64,
127    pub verify_j: f64,
128    /// 2026-10-04: The `k` chained draft steps that feed the next verify; 0 at `k = 0`.
129    pub draft_ms: f64,
130    pub draft_j: f64,
131}
132
133impl Cell {
134    /// 2026-10-04: Wall ms of a whole speculative step at this depth.
135    pub fn step_ms(&self) -> f64 {
136        self.verify_ms + self.draft_ms
137    }
138
139    /// 2026-10-04: GPU-rail joules of a whole speculative step at this depth.
140    pub fn step_j(&self) -> f64 {
141        self.verify_j + self.draft_j
142    }
143}
144
145#[derive(Debug, Deserialize)]
146struct TableFile {
147    key: TableKey,
148    cells: Vec<Cell>,
149}
150
151/// 2026-10-04: A parsed, validated cost table.
152#[derive(Debug, Clone, PartialEq)]
153pub struct CostTable {
154    pub key: TableKey,
155    /// 2026-10-04: Measured widths, ascending.
156    widths: Vec<usize>,
157    /// 2026-10-04: Deepest measured draft count (the same at every width).
158    max_k: usize,
159    /// 2026-10-04: `cells[(n, k)]`.
160    cells: BTreeMap<(usize, usize), Cell>,
161}
162
163impl CostTable {
164    /// 2026-10-04: Parses a table file's text and validates it: schema [`SCHEMA`], at least
165    /// one width, every listed width measured at every depth `0..=max_k`, no duplicate cell,
166    /// every cost finite and positive (a draft cost of exactly 0 is required at `k = 0`).
167    pub fn parse(text: &str) -> Result<Self, String> {
168        let file: TableFile = toml::from_str(text).map_err(|e| format!("spec-cost table: {e}"))?;
169        if file.key.schema != SCHEMA {
170            return Err(format!(
171                "spec-cost table schema {} (this serve reads {SCHEMA})",
172                file.key.schema
173            ));
174        }
175        let mut cells = BTreeMap::new();
176        for c in &file.cells {
177            let finite = [c.verify_ms, c.verify_j, c.draft_ms, c.draft_j]
178                .iter()
179                .all(|v| v.is_finite());
180            let positive = c.verify_ms > 0.0 && c.verify_j > 0.0;
181            let draft_ok = if c.k == 0 {
182                c.draft_ms == 0.0 && c.draft_j == 0.0
183            } else {
184                c.draft_ms > 0.0 && c.draft_j > 0.0
185            };
186            if c.n == 0 || !finite || !positive || !draft_ok {
187                return Err(format!("spec-cost table: bad cell {c:?}"));
188            }
189            if cells.insert((c.n, c.k), *c).is_some() {
190                return Err(format!(
191                    "spec-cost table: duplicate cell n={} k={}",
192                    c.n, c.k
193                ));
194            }
195        }
196        let widths: Vec<usize> = {
197            let mut w: Vec<usize> = cells.keys().map(|&(n, _)| n).collect();
198            w.dedup();
199            w
200        };
201        let max_k = cells
202            .keys()
203            .map(|&(_, k)| k)
204            .max()
205            .ok_or("spec-cost table: no cells")?;
206        for &n in &widths {
207            for k in 0..=max_k {
208                if !cells.contains_key(&(n, k)) {
209                    return Err(format!("spec-cost table: width {n} lacks depth {k}"));
210                }
211            }
212        }
213        Ok(Self {
214            key: file.key,
215            widths,
216            max_k,
217            cells,
218        })
219    }
220
221    /// 2026-10-04: The table file text for `key` and `cells`, in the format [`Self::parse`]
222    /// reads; validated by parsing it back, so a grid that would be refused at serve time is
223    /// refused when written.
224    pub fn render(key: &TableKey, cells: &[Cell]) -> Result<String, String> {
225        use std::fmt::Write as _;
226        let mut s = String::new();
227        let _ = writeln!(s, "[key]\nschema = {}", key.schema);
228        let _ = writeln!(
229            s,
230            "box_class = {:?}\nrecipe = {:?}",
231            key.box_class, key.recipe
232        );
233        let _ = writeln!(s, "\n[key.plan_digests]");
234        for (mode, digest) in &key.plan_digests {
235            let _ = writeln!(s, "{mode} = {digest:?}");
236        }
237        let mut sorted = cells.to_vec();
238        sorted.sort_by_key(|c| (c.n, c.k));
239        for c in &sorted {
240            let _ = write!(
241                s,
242                "\n[[cells]]\nn = {}\nk = {}\nverify_ms = {:?}\nverify_j = {:?}\ndraft_ms = {:?}\ndraft_j = {:?}\n",
243                c.n, c.k, c.verify_ms, c.verify_j, c.draft_ms, c.draft_j
244            );
245        }
246        let parsed = Self::parse(&s)?;
247        if parsed.key != *key {
248            return Err("spec-cost table: the rendered key does not read back".into());
249        }
250        Ok(s)
251    }
252
253    /// 2026-10-04: Deepest draft count the table measured.
254    pub fn max_k(&self) -> usize {
255        self.max_k
256    }
257
258    /// 2026-10-04: The cost of a step at width `n` and depth `k` (`k <= max_k`), linearly
259    /// interpolated between the two measured widths around `n`; the nearest measured width
260    /// outside the measured range. `None` above `max_k`.
261    pub fn cell(&self, n: usize, k: usize) -> Option<Cell> {
262        if k > self.max_k {
263            return None;
264        }
265        let lo = self.widths.iter().copied().filter(|&w| w <= n).max();
266        let hi = self.widths.iter().copied().filter(|&w| w >= n).min();
267        let at = |w| self.cells[&(w, k)];
268        Some(match (lo, hi) {
269            (Some(a), Some(b)) if a == b => at(a),
270            (Some(a), Some(b)) => {
271                let t = (n - a) as f64 / (b - a) as f64;
272                let (x, y) = (at(a), at(b));
273                let mix = |p: f64, q: f64| p + (q - p) * t;
274                Cell {
275                    n,
276                    k,
277                    verify_ms: mix(x.verify_ms, y.verify_ms),
278                    verify_j: mix(x.verify_j, y.verify_j),
279                    draft_ms: mix(x.draft_ms, y.draft_ms),
280                    draft_j: mix(x.draft_j, y.draft_j),
281                }
282            }
283            (Some(a), None) => Cell { n, ..at(a) },
284            (None, Some(b)) => Cell { n, ..at(b) },
285            (None, None) => unreachable!("parse guarantees at least one width"),
286        })
287    }
288}
289
290#[derive(Debug, Deserialize)]
291struct CalibrationFile {
292    drafter: DrafterKey,
293    acceptance: AcceptanceFields,
294}
295
296#[derive(Debug, Deserialize)]
297struct AcceptanceFields {
298    edges: Vec<f32>,
299    p_accept: Vec<f64>,
300    prior_by_position: Vec<f64>,
301}
302
303/// 2026-10-04: How likely a draft is to be accepted, from the drafter's top-1 log-probability
304/// for it, fitted on one drafter ([`DrafterKey`]); and per-position priors for drafts whose
305/// confidence was not measured. A draft is accepted only if every draft before it was, so a
306/// chain's expected accepted length is the sum of the running products.
307#[derive(Debug, Clone, PartialEq)]
308pub struct AcceptanceCalibration {
309    pub drafter: DrafterKey,
310    /// 2026-10-04: Ascending log-probability upper bounds; bucket `i` holds `lp <= edges[i]`
311    /// (and above `edges[i - 1]`). The last edge is at least 0, so every log-probability lands.
312    edges: Vec<f32>,
313    p_accept: Vec<f64>,
314    /// 2026-10-04: P(draft j accepted | drafts before it accepted), for j = 1, 2, ..
315    prior_by_position: Vec<f64>,
316}
317
318impl AcceptanceCalibration {
319    /// 2026-10-04: Parses and validates a calibration file: as many probabilities as edges,
320    /// edges strictly ascending with the last at least 0, every probability in 0..=1, at least
321    /// one position prior.
322    pub fn parse(text: &str) -> Result<Self, String> {
323        let f: CalibrationFile =
324            toml::from_str(text).map_err(|e| format!("acceptance calibration: {e}"))?;
325        let a = f.acceptance;
326        let in01 = |p: &f64| (0.0..=1.0).contains(p);
327        if a.edges.is_empty()
328            || a.edges.len() != a.p_accept.len()
329            || !a.edges.windows(2).all(|w| w[0] < w[1])
330            || a.edges.last().is_some_and(|&e| e < 0.0)
331            || !a.p_accept.iter().all(in01)
332            || a.prior_by_position.is_empty()
333            || !a.prior_by_position.iter().all(in01)
334        {
335            return Err("acceptance calibration: malformed buckets or priors".into());
336        }
337        Ok(Self {
338            drafter: f.drafter,
339            edges: a.edges,
340            p_accept: a.p_accept,
341            prior_by_position: a.prior_by_position,
342        })
343    }
344
345    /// 2026-10-04: P(accept) of a draft whose top-1 log-probability is `lp`.
346    pub fn p_given_lp(&self, lp: f32) -> f64 {
347        let i = self
348            .edges
349            .iter()
350            .position(|&e| lp <= e)
351            .unwrap_or(self.edges.len() - 1);
352        self.p_accept[i]
353    }
354
355    /// 2026-10-04: The prior for draft position `j` (1-based); positions past the last
356    /// fitted one take the last prior.
357    pub fn prior(&self, j: usize) -> f64 {
358        let i = j.saturating_sub(1).min(self.prior_by_position.len() - 1);
359        self.prior_by_position[i]
360    }
361
362    /// 2026-10-04: Expected accepted drafts of a chain whose drafts have confidences `lps`
363    /// (missing entries take the position prior), counting only the first `k`.
364    pub fn expected_accepted(&self, lps: &[f32], k: usize) -> f64 {
365        let mut run = 1.0;
366        let mut total = 0.0;
367        for j in 1..=k {
368            let p = lps
369                .get(j - 1)
370                .map_or_else(|| self.prior(j), |&lp| self.p_given_lp(lp));
371            run *= p;
372            total += run;
373        }
374        total
375    }
376}
377
378/// 2026-10-04: `--spec-cost-model measured`'s resolved state, as the serve loaded and checked it
379/// at boot (`TableKey::check` and a `DrafterKey` comparison both empty/equal): the table, the
380/// drafter's acceptance calibration, and `--spec-cost-slack`.
381pub struct SpecCostState {
382    pub table: CostTable,
383    pub calibration: AcceptanceCalibration,
384    pub slack: f64,
385}
386
387#[path = "spec_cost_fit.rs"]
388mod fit;
389pub use fit::MIN_OUTCOMES;
390#[path = "spec_cost_plan.rs"]
391pub mod plan;
392
393#[cfg(test)]
394#[path = "spec_cost_tests.rs"]
395mod tests;
metrale_speculative/spec_cost.rs

metrale_speculative/
spec_cost.rs
