osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
463 lines
15 KiB
Rust
463 lines
15 KiB
Rust
//! Performance metrics for speculative decoding
|
||
|
||
use std::collections::VecDeque;
|
||
use std::time::Duration;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// WindowedAcceptanceTracker
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Sliding-window acceptance rate over the last `window_size` decode steps.
|
||
///
|
||
/// Each element in the window is the per-step acceptance rate
|
||
/// (`accepted_tokens / draft_tokens`) for that step, clamped to `[0.0, 1.0]`.
|
||
/// Oldest entries are evicted automatically once the window is full.
|
||
///
|
||
/// # Example
|
||
/// ```
|
||
/// use rtx_inference::speculative::WindowedAcceptanceTracker;
|
||
///
|
||
/// let mut tracker = WindowedAcceptanceTracker::new(3);
|
||
/// tracker.push(0.5);
|
||
/// tracker.push(0.8);
|
||
/// tracker.push(0.7);
|
||
/// let rate = tracker.windowed_rate();
|
||
/// assert!((rate - (0.5_f32 + 0.8 + 0.7) / 3.0).abs() < 1e-5);
|
||
/// ```
|
||
#[derive(Debug, Clone)]
|
||
pub struct WindowedAcceptanceTracker {
|
||
window_size: usize,
|
||
buffer: VecDeque<f32>,
|
||
sum: f32,
|
||
}
|
||
|
||
impl WindowedAcceptanceTracker {
|
||
/// Create a new tracker with the given sliding-window capacity.
|
||
///
|
||
/// # Panics
|
||
/// Panics if `window_size` is 0.
|
||
pub fn new(window_size: usize) -> Self {
|
||
assert!(window_size > 0, "window_size must be > 0");
|
||
Self {
|
||
window_size,
|
||
buffer: VecDeque::with_capacity(window_size),
|
||
sum: 0.0,
|
||
}
|
||
}
|
||
|
||
/// Push the per-step acceptance rate for one decode step.
|
||
///
|
||
/// `step_rate` is clamped to `[0.0, 1.0]` before storage.
|
||
/// If the window is full the oldest value is evicted and its contribution
|
||
/// subtracted from the running sum before the new value is inserted.
|
||
pub fn push(&mut self, step_rate: f32) {
|
||
let clamped = step_rate.clamp(0.0, 1.0);
|
||
if self.buffer.len() == self.window_size {
|
||
// Evict the oldest entry.
|
||
if let Some(oldest) = self.buffer.pop_front() {
|
||
self.sum -= oldest;
|
||
}
|
||
}
|
||
self.buffer.push_back(clamped);
|
||
self.sum += clamped;
|
||
}
|
||
|
||
/// Mean acceptance rate over all values currently in the window.
|
||
///
|
||
/// Returns `0.0` when the window is empty.
|
||
pub fn windowed_rate(&self) -> f32 {
|
||
if self.buffer.is_empty() {
|
||
0.0
|
||
} else {
|
||
// Guard against accumulated floating-point drift going negative.
|
||
(self.sum / self.buffer.len() as f32).max(0.0)
|
||
}
|
||
}
|
||
|
||
/// The `p`-th percentile of values currently in the window (0.0 – 1.0).
|
||
///
|
||
/// Uses nearest-rank method: index = `(p * (len - 1)).round()`.
|
||
/// Returns `0.0` when the window is empty.
|
||
pub fn percentile(&self, p: f32) -> f32 {
|
||
if self.buffer.is_empty() {
|
||
return 0.0;
|
||
}
|
||
let mut sorted: Vec<f32> = self.buffer.iter().copied().collect();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||
let idx = (p * (sorted.len() as f32 - 1.0)).round() as usize;
|
||
sorted[idx.min(sorted.len() - 1)]
|
||
}
|
||
|
||
/// 50th-percentile (median) of the current window.
|
||
pub fn p50(&self) -> f32 {
|
||
self.percentile(0.5)
|
||
}
|
||
|
||
/// 95th-percentile of the current window.
|
||
pub fn p95(&self) -> f32 {
|
||
self.percentile(0.95)
|
||
}
|
||
|
||
/// Minimum value in the current window, or `0.0` if empty.
|
||
pub fn min(&self) -> f32 {
|
||
if self.buffer.is_empty() {
|
||
return 0.0;
|
||
}
|
||
self.buffer.iter().copied().fold(f32::INFINITY, f32::min)
|
||
}
|
||
|
||
/// Maximum value in the current window, or `0.0` if empty.
|
||
pub fn max(&self) -> f32 {
|
||
if self.buffer.is_empty() {
|
||
return 0.0;
|
||
}
|
||
self.buffer
|
||
.iter()
|
||
.copied()
|
||
.fold(f32::NEG_INFINITY, f32::max)
|
||
}
|
||
|
||
/// Number of samples currently in the window.
|
||
pub fn count(&self) -> usize {
|
||
self.buffer.len()
|
||
}
|
||
|
||
/// `true` once `window_size` samples have been pushed (or more).
|
||
pub fn is_full(&self) -> bool {
|
||
self.buffer.len() == self.window_size
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// AcceptanceTrend
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Direction of acceptance rate movement relative to the global average.
|
||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||
pub enum AcceptanceTrend {
|
||
/// Windowed rate is more than 5 percentage points above the global rate.
|
||
Rising,
|
||
/// Windowed rate is more than 5 percentage points below the global rate.
|
||
Falling,
|
||
/// Windowed rate is within ±5 percentage points of the global rate.
|
||
Stable,
|
||
}
|
||
|
||
impl AcceptanceTrend {
|
||
fn from_rates(windowed: f32, global: f32) -> Self {
|
||
const THRESHOLD: f32 = 0.05;
|
||
let delta = windowed - global;
|
||
if delta > THRESHOLD {
|
||
Self::Rising
|
||
} else if delta < -THRESHOLD {
|
||
Self::Falling
|
||
} else {
|
||
Self::Stable
|
||
}
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// AcceptanceDashboard
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Snapshot of acceptance rate statistics for real-time monitoring.
|
||
///
|
||
/// Obtained via [`PerformanceMetrics::dashboard`].
|
||
#[derive(Debug, Clone)]
|
||
pub struct AcceptanceDashboard {
|
||
/// Global running acceptance rate since the decoder started.
|
||
pub global_rate: f32,
|
||
/// Mean acceptance rate over the last `window_size` steps.
|
||
pub windowed_rate: f32,
|
||
/// Median (p50) of the windowed samples.
|
||
pub p50: f32,
|
||
/// 95th-percentile of the windowed samples.
|
||
pub p95: f32,
|
||
/// Minimum acceptance rate observed in the current window.
|
||
pub min: f32,
|
||
/// Maximum acceptance rate observed in the current window.
|
||
pub max: f32,
|
||
/// Total decode steps recorded (lifetime of the decoder).
|
||
pub total_steps: u64,
|
||
/// Configured window size.
|
||
pub window_size: usize,
|
||
/// `true` once `window_size` samples have accumulated.
|
||
pub window_full: bool,
|
||
/// Trend of the windowed rate relative to the global rate.
|
||
pub trend: AcceptanceTrend,
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// PerformanceMetrics (extended)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Performance metrics for speculative decoding
|
||
#[derive(Debug, Clone, Default)]
|
||
pub struct PerformanceMetrics {
|
||
/// Total number of decoding steps
|
||
pub total_steps: u64,
|
||
/// Total tokens generated by draft model
|
||
pub tokens_generated: u64,
|
||
/// Total tokens accepted by target model
|
||
pub tokens_accepted: u64,
|
||
/// Current acceptance rate
|
||
pub acceptance_rate: f32,
|
||
/// Speedup ratio compared to sequential generation
|
||
pub speedup_ratio: f32,
|
||
/// Average time for draft generation
|
||
pub draft_generation_time: Duration,
|
||
/// Average time for verification
|
||
pub verification_time: Duration,
|
||
/// Total time spent in speculative decoding
|
||
pub total_time: Duration,
|
||
/// Sliding-window acceptance tracker (default window = 100 steps).
|
||
pub(crate) windowed: WindowedAcceptanceTracker,
|
||
}
|
||
|
||
impl Default for WindowedAcceptanceTracker {
|
||
fn default() -> Self {
|
||
Self::new(100)
|
||
}
|
||
}
|
||
|
||
impl PerformanceMetrics {
|
||
pub(crate) fn new() -> Self {
|
||
Self {
|
||
total_steps: 0,
|
||
tokens_generated: 0,
|
||
tokens_accepted: 0,
|
||
acceptance_rate: 0.0,
|
||
speedup_ratio: 0.0,
|
||
draft_generation_time: Duration::ZERO,
|
||
verification_time: Duration::ZERO,
|
||
total_time: Duration::ZERO,
|
||
windowed: WindowedAcceptanceTracker::new(100),
|
||
}
|
||
}
|
||
|
||
pub(crate) fn update(&mut self, step_result: &StepMetrics) {
|
||
self.total_steps += 1;
|
||
self.tokens_generated += step_result.draft_tokens as u64;
|
||
self.tokens_accepted += step_result.accepted_tokens as u64;
|
||
|
||
// Update running averages
|
||
let weighted_draft = self.draft_generation_time.as_nanos()
|
||
* u128::from(self.total_steps - 1)
|
||
+ step_result.draft_time.as_nanos();
|
||
self.draft_generation_time =
|
||
Duration::from_nanos((weighted_draft / u128::from(self.total_steps)) as u64);
|
||
|
||
let weighted_verification = self.verification_time.as_nanos()
|
||
* u128::from(self.total_steps - 1)
|
||
+ step_result.verification_time.as_nanos();
|
||
self.verification_time =
|
||
Duration::from_nanos((weighted_verification / u128::from(self.total_steps)) as u64);
|
||
self.total_time += step_result.total_time;
|
||
|
||
// Update acceptance rate
|
||
self.acceptance_rate = if self.tokens_generated > 0 {
|
||
self.tokens_accepted as f32 / self.tokens_generated as f32
|
||
} else {
|
||
0.0
|
||
};
|
||
|
||
// Push per-step acceptance rate into the sliding window.
|
||
let step_rate = step_result.accepted_tokens as f32 / step_result.draft_tokens.max(1) as f32;
|
||
self.windowed.push(step_rate);
|
||
|
||
// Estimate speedup ratio
|
||
// This is simplified - in reality would compare against baseline sequential generation
|
||
let draft_efficiency =
|
||
step_result.draft_tokens as f32 / step_result.draft_time.as_secs_f32();
|
||
let acceptance_benefit = self.acceptance_rate * step_result.draft_tokens as f32;
|
||
self.speedup_ratio = 1.0 + (acceptance_benefit / draft_efficiency).min(10.0);
|
||
}
|
||
|
||
/// Build a real-time [`AcceptanceDashboard`] from the current state.
|
||
pub fn dashboard(&self) -> AcceptanceDashboard {
|
||
let windowed_rate = self.windowed.windowed_rate();
|
||
let trend = AcceptanceTrend::from_rates(windowed_rate, self.acceptance_rate);
|
||
AcceptanceDashboard {
|
||
global_rate: self.acceptance_rate,
|
||
windowed_rate,
|
||
p50: self.windowed.p50(),
|
||
p95: self.windowed.p95(),
|
||
min: self.windowed.min(),
|
||
max: self.windowed.max(),
|
||
total_steps: self.total_steps,
|
||
window_size: self.windowed.window_size,
|
||
window_full: self.windowed.is_full(),
|
||
trend,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Metrics for a single decoding step
|
||
pub(crate) struct StepMetrics {
|
||
pub draft_tokens: usize,
|
||
pub accepted_tokens: usize,
|
||
pub draft_time: Duration,
|
||
pub verification_time: Duration,
|
||
pub total_time: Duration,
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Tests
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::time::Duration;
|
||
|
||
// ------------------------------------------------------------------
|
||
// WindowedAcceptanceTracker tests
|
||
// ------------------------------------------------------------------
|
||
|
||
#[test]
|
||
fn test_windowed_tracker_empty() {
|
||
let tracker = WindowedAcceptanceTracker::new(10);
|
||
assert_eq!(tracker.windowed_rate(), 0.0);
|
||
assert_eq!(tracker.count(), 0);
|
||
assert!(!tracker.is_full());
|
||
}
|
||
|
||
#[test]
|
||
fn test_windowed_tracker_push_and_evict() {
|
||
let mut tracker = WindowedAcceptanceTracker::new(3);
|
||
tracker.push(0.1); // [0.1]
|
||
tracker.push(0.2); // [0.1, 0.2]
|
||
tracker.push(0.3); // [0.1, 0.2, 0.3] — full
|
||
tracker.push(0.4); // [0.2, 0.3, 0.4] — 0.1 evicted
|
||
assert_eq!(tracker.count(), 3);
|
||
assert!(tracker.is_full());
|
||
// Sum should be 0.2 + 0.3 + 0.4 = 0.9 → mean ≈ 0.3
|
||
let rate = tracker.windowed_rate();
|
||
assert!((rate - 0.3).abs() < 1e-5, "expected ~0.3, got {rate}");
|
||
}
|
||
|
||
#[test]
|
||
fn test_windowed_rate_accuracy() {
|
||
let mut tracker = WindowedAcceptanceTracker::new(10);
|
||
tracker.push(0.5);
|
||
tracker.push(0.8);
|
||
tracker.push(0.7);
|
||
let expected = (0.5_f32 + 0.8 + 0.7) / 3.0;
|
||
let got = tracker.windowed_rate();
|
||
assert!(
|
||
(got - expected).abs() < 1e-5,
|
||
"expected {expected}, got {got}"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_percentile_p50() {
|
||
let mut tracker = WindowedAcceptanceTracker::new(10);
|
||
tracker.push(0.4);
|
||
tracker.push(0.6);
|
||
tracker.push(0.8);
|
||
// sorted: [0.4, 0.6, 0.8]; idx = round(0.5 * 2) = 1 → 0.6
|
||
assert!((tracker.p50() - 0.6).abs() < 1e-5);
|
||
}
|
||
|
||
#[test]
|
||
fn test_percentile_p95() {
|
||
let mut tracker = WindowedAcceptanceTracker::new(100);
|
||
// Push 20 evenly spaced values: 0.0, 1/19, 2/19 ... 1.0
|
||
for i in 0..20_u32 {
|
||
tracker.push(i as f32 / 19.0);
|
||
}
|
||
// p95: idx = round(0.95 * 19) = round(18.05) = 18 → value = 18/19 ≈ 0.947
|
||
let p95 = tracker.p95();
|
||
assert!(p95 > 0.9, "p95 should be close to 1.0, got {p95}");
|
||
}
|
||
|
||
#[test]
|
||
fn test_trend_rising() {
|
||
let trend = AcceptanceTrend::from_rates(0.70, 0.50);
|
||
assert_eq!(trend, AcceptanceTrend::Rising);
|
||
}
|
||
|
||
#[test]
|
||
fn test_trend_falling() {
|
||
let trend = AcceptanceTrend::from_rates(0.50, 0.70);
|
||
assert_eq!(trend, AcceptanceTrend::Falling);
|
||
}
|
||
|
||
#[test]
|
||
fn test_trend_stable() {
|
||
let trend = AcceptanceTrend::from_rates(0.52, 0.50);
|
||
assert_eq!(trend, AcceptanceTrend::Stable);
|
||
}
|
||
|
||
#[test]
|
||
fn test_dashboard_builds_correctly() {
|
||
let mut metrics = PerformanceMetrics::new();
|
||
let step = StepMetrics {
|
||
draft_tokens: 4,
|
||
accepted_tokens: 3,
|
||
draft_time: Duration::from_millis(5),
|
||
verification_time: Duration::from_millis(2),
|
||
total_time: Duration::from_millis(7),
|
||
};
|
||
metrics.update(&step);
|
||
metrics.update(&step);
|
||
metrics.update(&step);
|
||
|
||
let dash = metrics.dashboard();
|
||
assert_eq!(dash.total_steps, 3);
|
||
assert_eq!(dash.window_size, 100);
|
||
assert!(!dash.window_full); // need 100 steps for full
|
||
// Each step: accepted=3, drafted=4 → step_rate = 0.75
|
||
assert!((dash.windowed_rate - 0.75).abs() < 1e-4);
|
||
assert!((dash.global_rate - 0.75).abs() < 1e-4);
|
||
assert_eq!(dash.trend, AcceptanceTrend::Stable);
|
||
}
|
||
|
||
#[test]
|
||
fn test_window_full_flag() {
|
||
let mut tracker = WindowedAcceptanceTracker::new(5);
|
||
for i in 0..5 {
|
||
assert!(!tracker.is_full(), "should not be full after {i} pushes");
|
||
tracker.push(0.5);
|
||
}
|
||
assert!(tracker.is_full());
|
||
}
|
||
|
||
#[test]
|
||
fn test_min_max() {
|
||
let mut tracker = WindowedAcceptanceTracker::new(10);
|
||
tracker.push(0.3);
|
||
tracker.push(0.7);
|
||
tracker.push(0.9);
|
||
assert!(
|
||
(tracker.min() - 0.3).abs() < 1e-5,
|
||
"min got {}",
|
||
tracker.min()
|
||
);
|
||
assert!(
|
||
(tracker.max() - 0.9).abs() < 1e-5,
|
||
"max got {}",
|
||
tracker.max()
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_min_empty_returns_zero() {
|
||
let tracker = WindowedAcceptanceTracker::new(5);
|
||
assert_eq!(tracker.min(), 0.0, "min() of empty window must be 0.0");
|
||
assert_eq!(tracker.max(), 0.0, "max() of empty window must be 0.0");
|
||
assert_eq!(tracker.windowed_rate(), 0.0);
|
||
assert_eq!(tracker.count(), 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_clamp_out_of_range() {
|
||
let mut tracker = WindowedAcceptanceTracker::new(5);
|
||
tracker.push(-0.5); // clamped to 0.0
|
||
tracker.push(1.5); // clamped to 1.0
|
||
assert!((tracker.windowed_rate() - 0.5).abs() < 1e-5);
|
||
}
|
||
}
|