metrale_speculative/
spec_stats.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Per-run speculation counters: draft acceptance and per-position
4//! draft matches for the K=2, K=3 and K=4 verify arms, the B1 low-margin
5//! gauge, decode timing, and one-shot log latches. Nothing here changes
6//! generation.
7//!
8//! Owner: speculative.
9//! Invariants: none beyond the types.
10
11use std::sync::atomic::{AtomicU64, Ordering};
12
13pub const SUMMARY_PERIOD: u64 = 512;
14
15/// 2026-09-25: The counters of one run, held by the scheduler's telemetry
16/// (`SysTelemetry::stats`). K=2 counts accepts and rejects; K=3 and K=4 also
17/// bucket steps by drafts accepted and count per-position draft matches.
18#[derive(Debug, Default)]
19pub struct SpecStats {
20    /// 2026-09-25: One-shot log latches, keyed by a `&'static str`.
21    fired: std::sync::Mutex<std::collections::BTreeSet<&'static str>>,
22    pub k2_accepts: AtomicU64,
23    pub k2_rejects: AtomicU64,
24
25    pub k3_accept: [AtomicU64; 3],
26    pub k3_steps: AtomicU64,
27    pub k3_d1_match: AtomicU64,
28    pub k3_d2_match_uncond: AtomicU64,
29    pub k3_d2_match_cond: AtomicU64,
30
31    pub k4_accept: [AtomicU64; 4],
32    pub k4_steps: AtomicU64,
33    pub k4_d1: AtomicU64,
34    pub k4_d2_uncond: AtomicU64,
35    pub k4_d3_uncond: AtomicU64,
36    pub k4_d2_cond: AtomicU64,
37    pub k4_d3_cond: AtomicU64,
38
39    /// 2026-09-25: B1 gauge: decode positions inside a tool parameter body
40    /// whose top-1/top-2 logit gap is below the low-margin threshold
41    /// (`logit_processors::b1_margin`).
42    pub b1_low_margin: AtomicU64,
43
44    // 2026-09-25: Decode timing, recorded under `METRALE_DECODE_TIMING`.
45    pub decode_copy_us: AtomicU64,
46    pub decode_sample_us: AtomicU64,
47    pub decode_count: AtomicU64,
48}
49
50/// 2026-09-25: Adds one to `counter` (relaxed) and returns the previous
51/// value.
52#[inline]
53pub fn bump(counter: &AtomicU64) -> u64 {
54    counter.fetch_add(1, Ordering::Relaxed)
55}
56
57#[inline]
58pub fn get(counter: &AtomicU64) -> u64 {
59    counter.load(Ordering::Relaxed)
60}
61
62#[inline]
63pub fn reset(counter: &AtomicU64) {
64    counter.store(0, Ordering::Relaxed);
65}
66
67impl SpecStats {
68    pub fn new() -> Self {
69        Self::default()
70    }
71
72    /// 2026-09-25: `true` the first time this `SpecStats` sees `key`, `false`
73    /// after. The scheduler's counterpart to `ModelStats::once`.
74    pub fn once(&self, key: &'static str) -> bool {
75        self.fired.lock().expect("run latches poisoned").insert(key)
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn a_fresh_run_starts_at_zero() {
85        let s = SpecStats::new();
86        assert_eq!(get(&s.k4_steps), 0);
87        assert_eq!(get(&s.k2_accepts), 0);
88        assert!(s.k3_accept.iter().all(|c| get(c) == 0));
89    }
90
91    #[test]
92    fn two_runs_count_independently() {
93        let a = SpecStats::new();
94        let b = SpecStats::new();
95        for _ in 0..7 {
96            bump(&a.k4_steps);
97        }
98        assert_eq!(get(&a.k4_steps), 7);
99        assert_eq!(get(&b.k4_steps), 0, "a second run starts clean");
100    }
101
102    #[test]
103    fn reset_clears_a_counter_for_the_next_summary_window() {
104        let s = SpecStats::new();
105        bump(&s.k3_steps);
106        bump(&s.k3_steps);
107        assert_eq!(get(&s.k3_steps), 2);
108        reset(&s.k3_steps);
109        assert_eq!(get(&s.k3_steps), 0);
110    }
111}