Skip to main content

qualia_core_db/inference/inference_bench/
counters.rs

1//! Decode/prefill path-selection counters — record which internal path produced
2//! each token (top-k vs argmax, resident vs legacy, speculative steps, sampled
3//! tokens, resident-prefill hits). Each is a process-global atomic with
4//! `record_*` / `*_counts` / `reset_*` helpers. Pure code motion — unchanged.
5
6use std::sync::atomic::{AtomicU64, Ordering};
7
8// ── Output-projection path counters (Codex P0: make the chosen path visible) ───────────────────
9static TOPK_HITS: AtomicU64 = AtomicU64::new(0);
10static ARGMAX_FALLBACKS: AtomicU64 = AtomicU64::new(0);
11
12/// Decode loop: the GPU top-k path produced the next token.
13#[inline]
14pub fn record_topk_hit() {
15    TOPK_HITS.fetch_add(1, Ordering::Relaxed);
16}
17/// Decode loop: fell back to the full-logit-readback argmax path.
18#[inline]
19pub fn record_argmax_fallback() {
20    ARGMAX_FALLBACKS.fetch_add(1, Ordering::Relaxed);
21}
22/// (top-k hits, argmax fallbacks) since the last reset.
23#[inline]
24pub fn output_path_counts() -> (u64, u64) {
25    (
26        TOPK_HITS.load(Ordering::Relaxed),
27        ARGMAX_FALLBACKS.load(Ordering::Relaxed),
28    )
29}
30/// Reset the output-projection path counters.
31#[inline]
32pub fn reset_output_path_counts() {
33    TOPK_HITS.store(0, Ordering::Relaxed);
34    ARGMAX_FALLBACKS.store(0, Ordering::Relaxed);
35}
36
37// ── W1/W9: resident-token decode path counters ─────────────────────────────────
38static RESIDENT_HITS: AtomicU64 = AtomicU64::new(0);
39static RESIDENT_FALLBACKS: AtomicU64 = AtomicU64::new(0);
40
41/// Decode loop: the resident single-fence path produced this token.
42#[inline]
43pub fn record_resident_hit() {
44    RESIDENT_HITS.fetch_add(1, Ordering::Relaxed);
45}
46/// Decode loop: resident path was enabled but ineligible/failed — legacy ran instead.
47#[inline]
48pub fn record_resident_fallback() {
49    RESIDENT_FALLBACKS.fetch_add(1, Ordering::Relaxed);
50}
51/// (resident hits, resident fallbacks) since the last reset.
52#[inline]
53pub fn resident_path_counts() -> (u64, u64) {
54    (
55        RESIDENT_HITS.load(Ordering::Relaxed),
56        RESIDENT_FALLBACKS.load(Ordering::Relaxed),
57    )
58}
59/// Reset the resident-path counters.
60#[inline]
61pub fn reset_resident_path_counts() {
62    RESIDENT_HITS.store(0, Ordering::Relaxed);
63    RESIDENT_FALLBACKS.store(0, Ordering::Relaxed);
64}
65
66// Native CUDA mega-pass execution counters.
67//
68// These are deliberately separate from inference mode and requested backend.
69// A mode/env label is intent; only a successful mega-pass call proves that
70// CUDA produced the token forward.
71static CUDA_MEGA_HITS: AtomicU64 = AtomicU64::new(0);
72static CUDA_MEGA_FALLBACKS: AtomicU64 = AtomicU64::new(0);
73
74/// Decode loop: the CUDA all-layer mega-pass completed this token forward.
75#[inline]
76pub fn record_cuda_mega_hit() {
77    CUDA_MEGA_HITS.fetch_add(1, Ordering::Relaxed);
78}
79
80/// Decode loop: CUDA mega-pass was explicitly requested but failed eligibility
81/// or execution and the ordinary fallback path ran.
82#[inline]
83pub fn record_cuda_mega_fallback() {
84    CUDA_MEGA_FALLBACKS.fetch_add(1, Ordering::Relaxed);
85}
86
87/// (successful CUDA mega-pass forwards, explicit CUDA fallbacks).
88#[inline]
89pub fn cuda_mega_path_counts() -> (u64, u64) {
90    (
91        CUDA_MEGA_HITS.load(Ordering::Relaxed),
92        CUDA_MEGA_FALLBACKS.load(Ordering::Relaxed),
93    )
94}
95
96/// Reset the native CUDA execution counters.
97#[inline]
98pub fn reset_cuda_mega_path_counts() {
99    CUDA_MEGA_HITS.store(0, Ordering::Relaxed);
100    CUDA_MEGA_FALLBACKS.store(0, Ordering::Relaxed);
101}
102
103// ── W2/W9: sampled-token counter ───────────────────────────────────────────────
104static SAMPLED_TOKENS: AtomicU64 = AtomicU64::new(0);
105
106/// Decode loop: the exact CPU sampler produced this token (non-greedy path).
107#[inline]
108pub fn record_sampled_token() {
109    SAMPLED_TOKENS.fetch_add(1, Ordering::Relaxed);
110}
111/// Sampled tokens since the last reset.
112#[inline]
113pub fn sampled_token_count() -> u64 {
114    SAMPLED_TOKENS.load(Ordering::Relaxed)
115}
116/// Reset the sampled-token counter.
117#[inline]
118pub fn reset_sampled_token_count() {
119    SAMPLED_TOKENS.store(0, Ordering::Relaxed);
120}
121
122// ── W6a/W9: speculative-decode counters ────────────────────────────────────────
123static SPEC_STEPS: AtomicU64 = AtomicU64::new(0);
124static SPEC_DRAFTED: AtomicU64 = AtomicU64::new(0);
125static SPEC_ACCEPTED: AtomicU64 = AtomicU64::new(0);
126
127/// One speculative step ran (a draft was proposed + verified).
128#[inline]
129pub fn record_spec_step(drafted: u64, accepted: u64) {
130    SPEC_STEPS.fetch_add(1, Ordering::Relaxed);
131    SPEC_DRAFTED.fetch_add(drafted, Ordering::Relaxed);
132    SPEC_ACCEPTED.fetch_add(accepted, Ordering::Relaxed);
133}
134/// (spec steps, tokens drafted, draft tokens accepted) since the last reset.
135#[inline]
136pub fn spec_decode_counts() -> (u64, u64, u64) {
137    (
138        SPEC_STEPS.load(Ordering::Relaxed),
139        SPEC_DRAFTED.load(Ordering::Relaxed),
140        SPEC_ACCEPTED.load(Ordering::Relaxed),
141    )
142}
143/// Reset the speculative-decode counters.
144#[inline]
145pub fn reset_spec_decode_counts() {
146    SPEC_STEPS.store(0, Ordering::Relaxed);
147    SPEC_DRAFTED.store(0, Ordering::Relaxed);
148    SPEC_ACCEPTED.store(0, Ordering::Relaxed);
149}
150
151// ── W3/W9: resident-prefill path counters ──────────────────────────────────────
152static RESIDENT_PREFILL_HITS: AtomicU64 = AtomicU64::new(0);
153static RESIDENT_PREFILL_FALLBACKS: AtomicU64 = AtomicU64::new(0);
154
155/// Prefill: the resident single-fence-per-chunk arena populated this chunk's KV.
156#[inline]
157pub fn record_resident_prefill_hit() {
158    RESIDENT_PREFILL_HITS.fetch_add(1, Ordering::Relaxed);
159}
160/// Prefill: resident path was enabled but ineligible/failed — legacy chunk ran instead.
161#[inline]
162pub fn record_resident_prefill_fallback() {
163    RESIDENT_PREFILL_FALLBACKS.fetch_add(1, Ordering::Relaxed);
164}
165/// (resident-prefill hits, resident-prefill fallbacks) since the last reset.
166#[inline]
167pub fn resident_prefill_counts() -> (u64, u64) {
168    (
169        RESIDENT_PREFILL_HITS.load(Ordering::Relaxed),
170        RESIDENT_PREFILL_FALLBACKS.load(Ordering::Relaxed),
171    )
172}
173/// Reset the resident-prefill path counters.
174#[inline]
175pub fn reset_resident_prefill_counts() {
176    RESIDENT_PREFILL_HITS.store(0, Ordering::Relaxed);
177    RESIDENT_PREFILL_FALLBACKS.store(0, Ordering::Relaxed);
178}