qualia_core_db/inference/inference_bench/
metrics.rs

1//! Shared per-phase timing metrics — written exactly once per phase by
2//! `infer_local_model_inner` (NOT per token, so zero hot-path cost) and read by
3//! the harness via [`phase_snapshot`]. Process-global because the decode loop
4//! runs on a spawned engine thread. Also the small unit-conversion helpers used
5//! by the runner and probes. Pure code motion — behaviour unchanged.
6
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::Duration;
9
10use serde::Serialize;
11
12// ── Shared phase metrics ──────────────────────────────────────────────────────
13// Written exactly once per phase by `infer_local_model_inner` (NOT per token, so
14// there is zero hot-path cost), read by the harness via `phase_snapshot`. Process
15// global because the decode loop runs on a spawned engine thread.
16
17static LOAD_NS: AtomicU64 = AtomicU64::new(0);
18static PREFILL_NS: AtomicU64 = AtomicU64::new(0);
19static PREFILL_TOKENS: AtomicU64 = AtomicU64::new(0);
20static DECODE_NS: AtomicU64 = AtomicU64::new(0);
21static DECODE_TOKENS: AtomicU64 = AtomicU64::new(0);
22// Decode-profiler accumulators (summed ACROSS the decode loop; one atomic add per
23// token-phase — nanosecond cost, off the GPU critical path). Localize where the
24// per-token wall-clock goes: transformer forward (32 layers) vs output projection.
25static DECODE_FORWARD_NS: AtomicU64 = AtomicU64::new(0);
26static DECODE_OUTPUT_NS: AtomicU64 = AtomicU64::new(0);
27// One-shot empty submit→poll(Wait) baseline (total ns over N round-trips), measured once per
28// profiled decode so the bench can compare per-round-trip fence latency to real forward time.
29static DECODE_EMPTY_RT_NS: AtomicU64 = AtomicU64::new(0);
30static DECODE_EMPTY_RT_N: AtomicU64 = AtomicU64::new(0);
31
32/// Record the empty-round-trip baseline: `total_ns` measured over `n` empty submit→wait cycles.
33#[inline]
34pub fn record_empty_rt(total_ns: u64, n: u64) {
35    DECODE_EMPTY_RT_NS.store(total_ns, Ordering::Relaxed);
36    DECODE_EMPTY_RT_N.store(n, Ordering::Relaxed);
37}
38/// Read the empty-round-trip baseline as `(total_ns, n)`.
39#[inline]
40pub fn empty_rt() -> (u64, u64) {
41    (
42        DECODE_EMPTY_RT_NS.load(Ordering::Relaxed),
43        DECODE_EMPTY_RT_N.load(Ordering::Relaxed),
44    )
45}
46
47/// Re-export: GPU `submit → poll(Wait)` round-trips counted during the last run (see `gguf_bridge`).
48#[inline]
49pub fn gpu_wait_count() -> u64 {
50    crate::gguf_bridge::gpu_wait_count()
51}
52
53/// Accumulate one token's transformer-forward (32-layer) wall-clock.
54#[inline]
55pub fn add_decode_forward_ns(ns: u64) {
56    DECODE_FORWARD_NS.fetch_add(ns, Ordering::Relaxed);
57}
58/// Accumulate one token's output-projection (argmax/top-k) wall-clock.
59#[inline]
60pub fn add_decode_output_ns(ns: u64) {
61    DECODE_OUTPUT_NS.fetch_add(ns, Ordering::Relaxed);
62}
63/// Intra-layer split (summed over all 32 layers of one token): attention vs FFN — localizes which
64/// shader (fused_attention vs the GEMM kernel) bleeds the compute-bound forward time.
65static DECODE_ATTN_NS: AtomicU64 = AtomicU64::new(0);
66static DECODE_FFN_NS: AtomicU64 = AtomicU64::new(0);
67/// Accumulate one layer's attention (QKV-proj + SDPA + O-proj) wall-clock.
68#[inline]
69pub fn add_decode_attn_ns(ns: u64) {
70    DECODE_ATTN_NS.fetch_add(ns, Ordering::Relaxed);
71}
72/// Accumulate one layer's FFN (pre-norm SwiGLU) wall-clock.
73#[inline]
74pub fn add_decode_ffn_ns(ns: u64) {
75    DECODE_FFN_NS.fetch_add(ns, Ordering::Relaxed);
76}
77/// Read the intra-layer attention/FFN accumulators as `(attn_ns, ffn_ns)` (summed over the run).
78#[inline]
79pub fn decode_attn_ffn() -> (u64, u64) {
80    (
81        DECODE_ATTN_NS.load(Ordering::Relaxed),
82        DECODE_FFN_NS.load(Ordering::Relaxed),
83    )
84}
85
86/// Clear the phase counters before a measured run.
87#[inline]
88pub fn reset_phase_metrics() {
89    LOAD_NS.store(0, Ordering::Relaxed);
90    PREFILL_NS.store(0, Ordering::Relaxed);
91    PREFILL_TOKENS.store(0, Ordering::Relaxed);
92    DECODE_NS.store(0, Ordering::Relaxed);
93    DECODE_TOKENS.store(0, Ordering::Relaxed);
94    DECODE_FORWARD_NS.store(0, Ordering::Relaxed);
95    DECODE_OUTPUT_NS.store(0, Ordering::Relaxed);
96    DECODE_ATTN_NS.store(0, Ordering::Relaxed);
97    DECODE_FFN_NS.store(0, Ordering::Relaxed);
98    DECODE_EMPTY_RT_NS.store(0, Ordering::Relaxed);
99    DECODE_EMPTY_RT_N.store(0, Ordering::Relaxed);
100    crate::gguf_bridge::reset_gpu_wait_count();
101}
102
103/// Record model load/mmap-adopt + pipeline-build time (engine ready → before prefill).
104#[inline]
105pub fn record_load_ns(ns: u64) {
106    LOAD_NS.store(ns, Ordering::Relaxed);
107}
108
109/// Record prefill (prompt KV population) time + tokens prefilled.
110#[inline]
111pub fn record_prefill(ns: u64, tokens: u64) {
112    PREFILL_NS.store(ns, Ordering::Relaxed);
113    PREFILL_TOKENS.store(tokens, Ordering::Relaxed);
114}
115
116/// Record the autoregressive decode loop time + tokens generated.
117#[inline]
118pub fn record_decode(ns: u64, tokens: u64) {
119    DECODE_NS.store(ns, Ordering::Relaxed);
120    DECODE_TOKENS.store(tokens, Ordering::Relaxed);
121}
122
123/// Snapshot of the phase counters (last completed inference).
124#[derive(Debug, Clone, Copy, Serialize)]
125pub struct LlmPhaseSnapshot {
126    pub load_ns: u64,
127    pub prefill_ns: u64,
128    pub prefill_tokens: u64,
129    pub decode_ns: u64,
130    pub decode_tokens: u64,
131    pub decode_forward_ns: u64,
132    pub decode_output_ns: u64,
133}
134
135/// Read the phase counters recorded by the last inference call.
136#[inline]
137pub fn phase_snapshot() -> LlmPhaseSnapshot {
138    LlmPhaseSnapshot {
139        load_ns: LOAD_NS.load(Ordering::Relaxed),
140        prefill_ns: PREFILL_NS.load(Ordering::Relaxed),
141        prefill_tokens: PREFILL_TOKENS.load(Ordering::Relaxed),
142        decode_ns: DECODE_NS.load(Ordering::Relaxed),
143        decode_tokens: DECODE_TOKENS.load(Ordering::Relaxed),
144        decode_forward_ns: DECODE_FORWARD_NS.load(Ordering::Relaxed),
145        decode_output_ns: DECODE_OUTPUT_NS.load(Ordering::Relaxed),
146    }
147}
148
149// ── Unit conversion helpers ───────────────────────────────────────────────────
150// Shared by the runner (`run_bench`) and the probes (`decode_with_metrics`).
151// Widened from module-private `fn` to `pub(super)` so sibling submodules reach
152// them; not part of the crate-public surface.
153
154#[inline]
155pub(super) fn ms(d: Duration) -> f64 {
156    d.as_secs_f64() * 1000.0
157}
158
159#[inline]
160pub(super) fn ns_to_ms(ns: u64) -> f64 {
161    ns as f64 / 1_000_000.0
162}
163
164#[inline]
165pub(super) fn tok_per_s(tokens: u64, ns: u64) -> f64 {
166    if ns == 0 || tokens == 0 {
167        0.0
168    } else {
169        tokens as f64 / (ns as f64 / 1_000_000_000.0)
170    }
171}
qualia_core_db/inference/inference_bench/metrics.rs

qualia_core_db/inference/inference_bench/
metrics.rs
