metrale_telemetry/
kernel.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: L1: GPU time from CUDA event pairs: every serve lane on every
4//! step, and every eager kernel launch on one step in N.
5//!
6//! [`SpanRing`] is the bookkeeping, generic over [`EventOps`] so it is tested
7//! here with a scripted event source; the CUDA implementation is
8//! `metrale-gpu-runtime`'s `timing` module. A span lands when a later
9//! [`SpanRing::harvest`] polls it complete; the ring never waits on the GPU.
10//! Kernels replayed inside a CUDA graph get no per-kernel span; only the lane
11//! span covers a replay.
12//!
13//! Owner: telemetry.
14//! Invariants: a ring slot is handed out again only after its span landed,
15//! failed, or could not start.
16
17use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering};
18
19use crate::instrument::{Counter, NanosHistogram};
20
21/// 2026-09-26: The serve lanes timed on the GPU each step at level `Kernel`,
22/// indexed as the scheduler's `Lane::telemetry_index` numbers them.
23pub const LANES: [&str; 3] = ["admit", "prefill", "decode"];
24/// 2026-09-26: [`LANES`] as NVTX range names (NUL-terminated), in the same order.
25#[cfg(feature = "nvtx")]
26const LANES_NVTX: [&std::ffi::CStr; LANES.len()] = [c"admit", c"prefill", c"decode"];
27
28/// 2026-09-26: Open lane `lane`'s NVTX range on the calling thread; an index
29/// outside [`LANES`] is ignored.
30#[cfg(feature = "nvtx")]
31pub fn nvtx_push_lane(lane: usize) {
32    if let Some(name) = LANES_NVTX.get(lane) {
33        metrale_gpu_sys::nvtx::push(name);
34    }
35}
36
37/// 2026-09-26: Close the innermost NVTX range on the calling thread.
38#[cfg(feature = "nvtx")]
39pub fn nvtx_pop() {
40    metrale_gpu_sys::nvtx::pop();
41}
42
43/// 2026-09-26: Distinct kernels the per-kernel table can hold.
44pub const KERNEL_TABLE_SLOTS: usize = 512;
45
46/// 2026-09-26: What a span measured.
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48pub enum SpanKey {
49    Lane(usize),
50    /// 2026-09-26: A kernel, by its function handle.
51    Kernel(u64),
52}
53
54const LANE_BIT: u64 = 1 << 63;
55
56impl SpanKey {
57    fn encode(self) -> u64 {
58        match self {
59            Self::Lane(i) => LANE_BIT | i as u64,
60            Self::Kernel(f) => f & !LANE_BIT,
61        }
62    }
63    fn decode(v: u64) -> Self {
64        if v & LANE_BIT != 0 {
65            Self::Lane((v & !LANE_BIT) as usize)
66        } else {
67            Self::Kernel(v)
68        }
69    }
70}
71
72/// 2026-09-26: One edge of a span.
73#[derive(Clone, Copy, Debug, PartialEq, Eq)]
74pub enum Edge {
75    Start,
76    End,
77}
78
79/// 2026-09-26: A span's completion state.
80#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum SpanPoll {
82    /// 2026-09-26: Both events completed; GPU nanoseconds between them.
83    Ready(u64),
84    /// 2026-09-26: The end event has not completed yet.
85    Pending,
86    /// 2026-09-26: The events cannot be read.
87    Failed,
88}
89
90/// 2026-09-26: A pool of event pairs addressed by slot.
91pub trait EventOps: Send + Sync {
92    /// 2026-09-26: Record `slot`'s `edge` event on `stream`. `false` = the
93    /// record failed.
94    fn record(&self, slot: usize, edge: Edge, stream: u64) -> bool;
95    /// 2026-09-26: Non-blocking completion query for `slot`'s pair.
96    fn poll(&self, slot: usize) -> SpanPoll;
97}
98
99/// 2026-09-26: Where span durations land.
100#[derive(Debug)]
101pub struct KernelInstruments {
102    pub lane: [NanosHistogram; LANES.len()],
103    keys: [AtomicU64; KERNEL_TABLE_SLOTS],
104    total_ns: [AtomicU64; KERNEL_TABLE_SLOTS],
105    samples: [AtomicU64; KERNEL_TABLE_SLOTS],
106    /// 2026-09-26: Kernel spans not placed: the table had no free slot, or
107    /// the handle was 0.
108    pub table_full: Counter,
109    /// 2026-09-26: Spans not started because the next ring slot was still in
110    /// use.
111    pub spans_dropped: Counter,
112    /// 2026-09-26: Spans whose events could not be recorded or read.
113    pub span_failures: Counter,
114    pub(crate) armed: AtomicBool,
115    pub(crate) step: AtomicU64,
116}
117
118impl Default for KernelInstruments {
119    fn default() -> Self {
120        Self::new()
121    }
122}
123
124/// 2026-09-26: One row of the per-kernel table.
125#[derive(Clone, Copy, Debug, PartialEq, Eq)]
126pub struct KernelTime {
127    pub func: u64,
128    pub total_ns: u64,
129    pub samples: u64,
130}
131
132impl KernelInstruments {
133    pub const fn new() -> Self {
134        Self {
135            lane: [const { NanosHistogram::new() }; LANES.len()],
136            keys: [const { AtomicU64::new(0) }; KERNEL_TABLE_SLOTS],
137            total_ns: [const { AtomicU64::new(0) }; KERNEL_TABLE_SLOTS],
138            samples: [const { AtomicU64::new(0) }; KERNEL_TABLE_SLOTS],
139            table_full: Counter::new(),
140            spans_dropped: Counter::new(),
141            span_failures: Counter::new(),
142            armed: AtomicBool::new(false),
143            step: AtomicU64::new(0),
144        }
145    }
146
147    /// 2026-09-26: Land one completed span; a lane index outside [`LANES`] is
148    /// ignored.
149    pub fn record(&self, key: SpanKey, ns: u64) {
150        match key {
151            SpanKey::Lane(i) => {
152                if let Some(h) = self.lane.get(i) {
153                    h.observe(ns);
154                }
155            }
156            SpanKey::Kernel(func) => self.record_kernel(func, ns),
157        }
158    }
159
160    fn record_kernel(&self, func: u64, ns: u64) {
161        // 2026-09-26: Key 0 marks an empty slot, so handle 0 cannot be placed.
162        if func == 0 {
163            self.table_full.add(1);
164            return;
165        }
166        let start = (func.wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 55) as usize;
167        for probe in 0..KERNEL_TABLE_SLOTS {
168            let i = (start + probe) % KERNEL_TABLE_SLOTS;
169            let k = self.keys[i].load(Ordering::Acquire);
170            let owned = k == func
171                || (k == 0
172                    && match self.keys[i].compare_exchange(
173                        0,
174                        func,
175                        Ordering::AcqRel,
176                        Ordering::Acquire,
177                    ) {
178                        Ok(_) => true,
179                        Err(now) => now == func,
180                    });
181            if owned {
182                self.total_ns[i].fetch_add(ns, Ordering::Relaxed);
183                self.samples[i].fetch_add(1, Ordering::Relaxed);
184                return;
185            }
186        }
187        self.table_full.add(1);
188    }
189
190    /// 2026-09-26: Every kernel with at least one sample.
191    pub fn kernels(&self) -> Vec<KernelTime> {
192        (0..KERNEL_TABLE_SLOTS)
193            .filter_map(|i| {
194                let func = self.keys[i].load(Ordering::Acquire);
195                (func != 0).then(|| KernelTime {
196                    func,
197                    total_ns: self.total_ns[i].load(Ordering::Relaxed),
198                    samples: self.samples[i].load(Ordering::Relaxed),
199                })
200            })
201            .filter(|k| k.samples > 0)
202            .collect()
203    }
204}
205
206const FREE: u8 = 0;
207const OPEN: u8 = 1;
208const CLOSED: u8 = 2;
209
210/// 2026-09-26: A span handed out by [`SpanRing::begin`].
211#[derive(Clone, Copy, Debug, PartialEq, Eq)]
212pub struct SpanToken(usize);
213
214impl SpanToken {
215    /// 2026-09-26: The ring slot, for a caller that parks an open span in an
216    /// atomic.
217    pub fn slot(self) -> usize {
218        self.0
219    }
220    /// 2026-09-26: The token for a slot this caller got from [`SpanToken::slot`].
221    pub fn from_slot(slot: usize) -> Self {
222        Self(slot)
223    }
224}
225
226/// 2026-09-26: The in-flight spans over a fixed pool of event pairs.
227pub struct SpanRing<E: EventOps> {
228    ops: E,
229    state: Box<[AtomicU8]>,
230    key: Box<[AtomicU64]>,
231    head: AtomicUsize,
232}
233
234impl<E: EventOps> SpanRing<E> {
235    /// 2026-09-26: `slots` event pairs, which `ops` must already hold.
236    pub fn new(ops: E, slots: usize) -> Self {
237        assert!(slots > 0, "a span ring needs at least one slot");
238        Self {
239            ops,
240            state: (0..slots).map(|_| AtomicU8::new(FREE)).collect(),
241            key: (0..slots).map(|_| AtomicU64::new(0)).collect(),
242            head: AtomicUsize::new(0),
243        }
244    }
245
246    pub fn slots(&self) -> usize {
247        self.state.len()
248    }
249
250    /// 2026-09-26: Start a span on `stream`. `None` (and counted) when the next
251    /// slot is still in use or its start event could not be recorded.
252    pub fn begin(&self, key: SpanKey, stream: u64, sink: &KernelInstruments) -> Option<SpanToken> {
253        let i = self.head.fetch_add(1, Ordering::Relaxed) % self.state.len();
254        if self.state[i]
255            .compare_exchange(FREE, OPEN, Ordering::Acquire, Ordering::Relaxed)
256            .is_err()
257        {
258            sink.spans_dropped.add(1);
259            return None;
260        }
261        self.key[i].store(key.encode(), Ordering::Relaxed);
262        if !self.ops.record(i, Edge::Start, stream) {
263            self.state[i].store(FREE, Ordering::Release);
264            sink.span_failures.add(1);
265            return None;
266        }
267        Some(SpanToken(i))
268    }
269
270    /// 2026-09-26: Close a span on `stream`. If the end event cannot be
271    /// recorded, the slot is freed and the failure counted.
272    pub fn end(&self, tok: SpanToken, stream: u64, sink: &KernelInstruments) {
273        let i = tok.0;
274        if self.ops.record(i, Edge::End, stream) {
275            self.state[i].store(CLOSED, Ordering::Release);
276        } else {
277            self.state[i].store(FREE, Ordering::Release);
278            sink.span_failures.add(1);
279        }
280    }
281
282    /// 2026-09-26: Land every completed span into `sink`; returns how many
283    /// landed. A pending span is left for a later harvest.
284    pub fn harvest(&self, sink: &KernelInstruments) -> usize {
285        let mut landed = 0;
286        for i in 0..self.state.len() {
287            if self.state[i].load(Ordering::Acquire) != CLOSED {
288                continue;
289            }
290            match self.ops.poll(i) {
291                SpanPoll::Pending => continue,
292                SpanPoll::Ready(ns) => {
293                    sink.record(SpanKey::decode(self.key[i].load(Ordering::Relaxed)), ns);
294                    landed += 1;
295                }
296                SpanPoll::Failed => sink.span_failures.add(1),
297            }
298            self.state[i].store(FREE, Ordering::Release);
299        }
300        landed
301    }
302
303    pub fn ops(&self) -> &E {
304        &self.ops
305    }
306}
307
308#[cfg(test)]
309#[path = "kernel_tests.rs"]
310mod tests;
metrale_telemetry/kernel.rs

metrale_telemetry/
kernel.rs
