metrale_telemetry/
sched.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: L2 scheduler and L3 cache instruments (queue depths, batch
4//! shape, loop phase timings, blocking syncs and D2H copies, CUDA-graph
5//! replays and captures, KV-block and SSM-slot occupancy) and the L4
6//! [`SpecMatrix`].
7//!
8//! Fed through the scheduler's `TelemetryIo` router and the GPU backend's
9//! copy, sync and graph entry points. The prefix-cache hit and miss counters
10//! live in [`crate::run_metrics`], not here.
11//!
12//! Owner: telemetry.
13//! Invariants: none beyond the types.
14
15use crate::instrument::{CountHistogram, Counter, Gauge, NanosHistogram};
16
17/// 2026-09-26: The most loop phases `Telemetry::set_phase_names` keeps.
18pub const MAX_PHASES: usize = 32;
19
20/// 2026-09-26: One tick's scheduler shape, as the scheduler publishes it.
21#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize)]
22pub struct SchedShape {
23    pub pending: u64,
24    pub active: u64,
25    pub prefilling: u64,
26    pub swapped: u64,
27    pub kv_blocks_free: u64,
28    pub kv_blocks_total: u64,
29    pub ssm_slots_used: u64,
30    pub ssm_slots_total: u64,
31    /// 2026-09-26: Ticks the publishing scheduler has counted since it was built.
32    pub ticks: u64,
33}
34
35#[derive(Debug)]
36pub struct SchedInstruments {
37    pub phase: [NanosHistogram; MAX_PHASES],
38    pub pending: Gauge,
39    pub active: Gauge,
40    pub prefilling: Gauge,
41    pub swapped: Gauge,
42    pub kv_blocks_free: Gauge,
43    pub kv_blocks_total: Gauge,
44    pub ssm_slots_used: Gauge,
45    pub ssm_slots_total: Gauge,
46    pub ticks: Gauge,
47    /// 2026-09-26: Active sequences, on ticks that have any.
48    pub decode_rows: CountHistogram,
49    /// 2026-09-26: Prefilling sequences, on ticks that have any.
50    pub prefill_seqs: CountHistogram,
51    pub stream_syncs: Counter,
52    pub blocking_d2h: Counter,
53    pub graph_replays: Counter,
54    pub graph_captures: Counter,
55}
56
57impl Default for SchedInstruments {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63impl SchedInstruments {
64    pub const fn new() -> Self {
65        Self {
66            phase: [const { NanosHistogram::new() }; MAX_PHASES],
67            pending: Gauge::new(),
68            active: Gauge::new(),
69            prefilling: Gauge::new(),
70            swapped: Gauge::new(),
71            kv_blocks_free: Gauge::new(),
72            kv_blocks_total: Gauge::new(),
73            ssm_slots_used: Gauge::new(),
74            ssm_slots_total: Gauge::new(),
75            ticks: Gauge::new(),
76            decode_rows: CountHistogram::new(),
77            prefill_seqs: CountHistogram::new(),
78            stream_syncs: Counter::new(),
79            blocking_d2h: Counter::new(),
80            graph_replays: Counter::new(),
81            graph_captures: Counter::new(),
82        }
83    }
84
85    pub(crate) fn record(&self, s: &SchedShape) {
86        self.pending.set(s.pending);
87        self.active.set(s.active);
88        self.prefilling.set(s.prefilling);
89        self.swapped.set(s.swapped);
90        self.kv_blocks_free.set(s.kv_blocks_free);
91        self.kv_blocks_total.set(s.kv_blocks_total);
92        self.ssm_slots_used.set(s.ssm_slots_used);
93        self.ssm_slots_total.set(s.ssm_slots_total);
94        self.ticks.set(s.ticks);
95        if s.active > 0 {
96            self.decode_rows.observe(s.active);
97        }
98        if s.prefilling > 0 {
99            self.prefill_seqs.observe(s.prefilling);
100        }
101    }
102
103    pub(crate) fn shape(&self) -> SchedShape {
104        SchedShape {
105            pending: self.pending.get(),
106            active: self.active.get(),
107            prefilling: self.prefilling.get(),
108            swapped: self.swapped.get(),
109            kv_blocks_free: self.kv_blocks_free.get(),
110            kv_blocks_total: self.kv_blocks_total.get(),
111            ssm_slots_used: self.ssm_slots_used.get(),
112            ssm_slots_total: self.ssm_slots_total.get(),
113            ticks: self.ticks.get(),
114        }
115    }
116}
117
118/// 2026-09-26: The widest verify step [`SpecMatrix`] places.
119pub const MAX_DRAFTS: usize = 16;
120
121#[derive(Debug)]
122pub struct SpecMatrix {
123    /// 2026-09-26: `cells[d][a]`: verify steps that checked `d` drafts and accepted `a`.
124    cells: [[Counter; MAX_DRAFTS + 1]; MAX_DRAFTS + 1],
125    /// 2026-09-26: Steps not placed: zero drafts, more than [`MAX_DRAFTS`], or
126    /// more accepted than checked.
127    pub overflow: Counter,
128}
129
130impl Default for SpecMatrix {
131    fn default() -> Self {
132        Self::new()
133    }
134}
135
136impl SpecMatrix {
137    pub const fn new() -> Self {
138        Self {
139            cells: [const { [const { Counter::new() }; MAX_DRAFTS + 1] }; MAX_DRAFTS + 1],
140            overflow: Counter::new(),
141        }
142    }
143
144    pub(crate) fn record(&self, drafts: usize, accepted: usize) {
145        if drafts == 0 || drafts > MAX_DRAFTS || accepted > drafts {
146            self.overflow.add(1);
147            return;
148        }
149        self.cells[drafts][accepted].add(1);
150    }
151
152    pub fn steps(&self, drafts: usize, accepted: usize) -> u64 {
153        self.cells
154            .get(drafts)
155            .and_then(|r| r.get(accepted))
156            .map_or(0, Counter::get)
157    }
158
159    /// 2026-09-26: For each depth `1..=drafts`, the fraction of `drafts`-wide
160    /// verify steps that accepted at least that deep. Empty when no such step
161    /// ran.
162    pub fn acceptance_by_depth(&self, drafts: usize) -> Vec<f64> {
163        if drafts == 0 || drafts > MAX_DRAFTS {
164            return Vec::new();
165        }
166        let row: Vec<u64> = (0..=drafts).map(|a| self.steps(drafts, a)).collect();
167        let total: u64 = row.iter().sum();
168        if total == 0 {
169            return Vec::new();
170        }
171        (1..=drafts)
172            .map(|depth| row[depth..].iter().sum::<u64>() as f64 / total as f64)
173            .collect()
174    }
175
176    /// 2026-09-26: Draft widths that have recorded any step.
177    pub fn widths(&self) -> Vec<usize> {
178        (1..=MAX_DRAFTS)
179            .filter(|d| (0..=*d).any(|a| self.steps(*d, a) > 0))
180            .collect()
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn acceptance_by_depth_is_the_survival_curve_of_accepted_depth() {
190        let m = SpecMatrix::new();
191        // 2026-09-26: Width 3: 4 steps accepting 0, 1, 3, 3.
192        for a in [0, 1, 3, 3] {
193            m.record(3, a);
194        }
195        assert_eq!(m.acceptance_by_depth(3), vec![0.75, 0.5, 0.5]);
196        assert_eq!(m.widths(), vec![3]);
197        assert!(m.acceptance_by_depth(2).is_empty(), "no width-2 step ran");
198    }
199
200    #[test]
201    fn impossible_outcomes_are_counted_as_overflow_not_placed() {
202        let m = SpecMatrix::new();
203        m.record(2, 3);
204        m.record(0, 0);
205        m.record(MAX_DRAFTS + 1, 1);
206        assert_eq!(m.overflow.get(), 3);
207        assert!(m.widths().is_empty());
208    }
209
210    #[test]
211    fn a_tick_with_no_decode_rows_is_not_a_batch_shape_sample() {
212        let s = SchedInstruments::new();
213        s.record(&SchedShape {
214            pending: 2,
215            ..Default::default()
216        });
217        assert_eq!(s.decode_rows.count(), 0);
218        assert_eq!(s.pending.get(), 2);
219        s.record(&SchedShape {
220            active: 5,
221            ..Default::default()
222        });
223        assert_eq!(s.decode_rows.count(), 1);
224    }
225}
metrale_telemetry/sched.rs

metrale_telemetry/
sched.rs
