mirror of
https://github.com/ruvnet/RuView
synced 2026-08-07 20:01:43 +00:00
Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'
This commit is contained in:
ruv
@@ -0,0 +1,821 @@
|
||||
//! Inference Metrics for Optimization Decisions
|
||||
//!
|
||||
//! This module provides comprehensive metrics collection for LLM inference,
|
||||
//! enabling data-driven optimization decisions.
|
||||
//!
|
||||
//! ## Tracked Metrics
|
||||
//!
|
||||
//! - **TTFT (Time to First Token)**: Latency until first token generation
|
||||
//! - **TPS (Tokens Per Second)**: Generation throughput
|
||||
//! - **KV Cache Hit Rate**: Cache efficiency metric
|
||||
//! - **Memory Usage**: Current memory consumption
|
||||
//! - **Request Statistics**: Active requests, queue depth
|
||||
|
||||
use parking_lot::RwLock;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Moving average calculator with configurable window
|
||||
#[derive(Debug)]
|
||||
pub struct MovingAverage {
|
||||
/// Circular buffer of values
|
||||
values: RwLock<VecDeque<f32>>,
|
||||
/// Window size
|
||||
window_size: usize,
|
||||
/// Running sum for O(1) average calculation
|
||||
running_sum: RwLock<f32>,
|
||||
}
|
||||
|
||||
impl MovingAverage {
|
||||
/// Create a new moving average calculator
|
||||
pub fn new(window_size: usize) -> Self {
|
||||
Self {
|
||||
values: RwLock::new(VecDeque::with_capacity(window_size)),
|
||||
window_size,
|
||||
running_sum: RwLock::new(0.0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a value to the moving average
|
||||
pub fn add(&self, value: f32) {
|
||||
let mut values = self.values.write();
|
||||
let mut sum = self.running_sum.write();
|
||||
|
||||
// Remove oldest if at capacity
|
||||
if values.len() >= self.window_size {
|
||||
if let Some(old) = values.pop_front() {
|
||||
*sum -= old;
|
||||
}
|
||||
}
|
||||
|
||||
values.push_back(value);
|
||||
*sum += value;
|
||||
}
|
||||
|
||||
/// Get the current average
|
||||
pub fn average(&self) -> f32 {
|
||||
let values = self.values.read();
|
||||
let sum = self.running_sum.read();
|
||||
|
||||
if values.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
*sum / values.len() as f32
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the minimum value in the window
|
||||
pub fn min(&self) -> f32 {
|
||||
let values = self.values.read();
|
||||
values.iter().cloned().fold(f32::INFINITY, f32::min)
|
||||
}
|
||||
|
||||
/// Get the maximum value in the window
|
||||
pub fn max(&self) -> f32 {
|
||||
let values = self.values.read();
|
||||
values.iter().cloned().fold(f32::NEG_INFINITY, f32::max)
|
||||
}
|
||||
|
||||
/// Get the standard deviation
|
||||
pub fn std_dev(&self) -> f32 {
|
||||
let values = self.values.read();
|
||||
if values.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mean = self.average();
|
||||
let variance: f32 =
|
||||
values.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / (values.len() - 1) as f32;
|
||||
|
||||
variance.sqrt()
|
||||
}
|
||||
|
||||
/// Get the current window size (number of samples)
|
||||
pub fn count(&self) -> usize {
|
||||
self.values.read().len()
|
||||
}
|
||||
|
||||
/// Clear all values
|
||||
pub fn clear(&self) {
|
||||
let mut values = self.values.write();
|
||||
let mut sum = self.running_sum.write();
|
||||
values.clear();
|
||||
*sum = 0.0;
|
||||
}
|
||||
|
||||
/// Get percentile value (0-100)
|
||||
pub fn percentile(&self, p: f32) -> f32 {
|
||||
let values = self.values.read();
|
||||
if values.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let mut sorted: Vec<f32> = values.iter().copied().collect();
|
||||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
let idx = ((p / 100.0) * (sorted.len() - 1) as f32).round() as usize;
|
||||
sorted[idx.min(sorted.len() - 1)]
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for MovingAverage {
|
||||
fn default() -> Self {
|
||||
Self::new(100)
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for MovingAverage {
|
||||
fn clone(&self) -> Self {
|
||||
let values = self.values.read();
|
||||
let sum = self.running_sum.read();
|
||||
|
||||
Self {
|
||||
values: RwLock::new(values.clone()),
|
||||
window_size: self.window_size,
|
||||
running_sum: RwLock::new(*sum),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Latency histogram for distribution analysis
|
||||
#[derive(Debug)]
|
||||
pub struct LatencyHistogram {
|
||||
/// Bucket boundaries in milliseconds
|
||||
buckets: Vec<f32>,
|
||||
/// Counts per bucket
|
||||
counts: Vec<AtomicU64>,
|
||||
/// Total count
|
||||
total: AtomicU64,
|
||||
/// Sum for mean calculation
|
||||
sum: RwLock<f64>,
|
||||
}
|
||||
|
||||
impl LatencyHistogram {
|
||||
/// Create a new histogram with default buckets
|
||||
pub fn new() -> Self {
|
||||
Self::with_buckets(vec![
|
||||
1.0, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0, 200.0, 500.0, 1000.0, 2000.0, 5000.0,
|
||||
])
|
||||
}
|
||||
|
||||
/// Create a histogram with custom bucket boundaries
|
||||
pub fn with_buckets(buckets: Vec<f32>) -> Self {
|
||||
let counts = buckets.iter().map(|_| AtomicU64::new(0)).collect();
|
||||
Self {
|
||||
buckets,
|
||||
counts,
|
||||
total: AtomicU64::new(0),
|
||||
sum: RwLock::new(0.0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record a latency value in milliseconds
|
||||
pub fn record(&self, latency_ms: f32) {
|
||||
// Find the appropriate bucket
|
||||
let bucket_idx = self
|
||||
.buckets
|
||||
.iter()
|
||||
.position(|&b| latency_ms <= b)
|
||||
.unwrap_or(self.buckets.len() - 1);
|
||||
|
||||
self.counts[bucket_idx].fetch_add(1, Ordering::Relaxed);
|
||||
self.total.fetch_add(1, Ordering::Relaxed);
|
||||
|
||||
let mut sum = self.sum.write();
|
||||
*sum += latency_ms as f64;
|
||||
}
|
||||
|
||||
/// Get the mean latency
|
||||
pub fn mean(&self) -> f32 {
|
||||
let total = self.total.load(Ordering::Relaxed);
|
||||
if total == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
let sum = self.sum.read();
|
||||
(*sum / total as f64) as f32
|
||||
}
|
||||
|
||||
/// Get approximate percentile (linear interpolation between buckets)
|
||||
pub fn percentile(&self, p: f32) -> f32 {
|
||||
let total = self.total.load(Ordering::Relaxed);
|
||||
if total == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let target = (p / 100.0 * total as f32) as u64;
|
||||
let mut cumulative = 0u64;
|
||||
|
||||
for (i, count) in self.counts.iter().enumerate() {
|
||||
let bucket_count = count.load(Ordering::Relaxed);
|
||||
cumulative += bucket_count;
|
||||
|
||||
if cumulative >= target {
|
||||
// Found the bucket containing the percentile
|
||||
if i == 0 {
|
||||
return self.buckets[0];
|
||||
}
|
||||
// Linear interpolation
|
||||
let prev_cumulative = cumulative - bucket_count;
|
||||
let fraction = (target - prev_cumulative) as f32 / bucket_count.max(1) as f32;
|
||||
let prev_bucket = if i > 0 { self.buckets[i - 1] } else { 0.0 };
|
||||
return prev_bucket + fraction * (self.buckets[i] - prev_bucket);
|
||||
}
|
||||
}
|
||||
|
||||
*self.buckets.last().unwrap_or(&0.0)
|
||||
}
|
||||
|
||||
/// Get bucket counts for visualization
|
||||
pub fn bucket_counts(&self) -> Vec<(f32, u64)> {
|
||||
self.buckets
|
||||
.iter()
|
||||
.zip(self.counts.iter())
|
||||
.map(|(b, c)| (*b, c.load(Ordering::Relaxed)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Reset all counts
|
||||
pub fn reset(&self) {
|
||||
for count in &self.counts {
|
||||
count.store(0, Ordering::Relaxed);
|
||||
}
|
||||
self.total.store(0, Ordering::Relaxed);
|
||||
*self.sum.write() = 0.0;
|
||||
}
|
||||
|
||||
/// Get total count
|
||||
pub fn count(&self) -> u64 {
|
||||
self.total.load(Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LatencyHistogram {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for LatencyHistogram {
|
||||
fn clone(&self) -> Self {
|
||||
let counts: Vec<AtomicU64> = self
|
||||
.counts
|
||||
.iter()
|
||||
.map(|c| AtomicU64::new(c.load(Ordering::Relaxed)))
|
||||
.collect();
|
||||
let sum = *self.sum.read();
|
||||
|
||||
Self {
|
||||
buckets: self.buckets.clone(),
|
||||
counts,
|
||||
total: AtomicU64::new(self.total.load(Ordering::Relaxed)),
|
||||
sum: RwLock::new(sum),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Comprehensive inference metrics
|
||||
#[derive(Debug)]
|
||||
pub struct InferenceMetrics {
|
||||
/// Time to first token (milliseconds)
|
||||
pub ttft_ms: MovingAverage,
|
||||
/// Tokens per second throughput
|
||||
pub tps: MovingAverage,
|
||||
/// KV cache hit rate (0.0 - 1.0)
|
||||
kv_cache_hits: AtomicU64,
|
||||
kv_cache_misses: AtomicU64,
|
||||
/// Memory usage in bytes
|
||||
memory_usage_bytes: AtomicUsize,
|
||||
/// Peak memory usage
|
||||
peak_memory_bytes: AtomicUsize,
|
||||
/// Active request count
|
||||
active_requests: AtomicUsize,
|
||||
/// Total requests processed
|
||||
total_requests: AtomicU64,
|
||||
/// Total tokens generated
|
||||
total_tokens: AtomicU64,
|
||||
/// Request latency histogram
|
||||
pub latency_histogram: LatencyHistogram,
|
||||
/// Queue depth for pending requests
|
||||
queue_depth: AtomicUsize,
|
||||
/// Start time for uptime calculation
|
||||
start_time: Instant,
|
||||
/// Last update time
|
||||
last_update: RwLock<Instant>,
|
||||
/// Inter-token latency
|
||||
pub inter_token_latency_ms: MovingAverage,
|
||||
/// Batch size history
|
||||
pub batch_sizes: MovingAverage,
|
||||
}
|
||||
|
||||
impl InferenceMetrics {
|
||||
/// Create new inference metrics
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
ttft_ms: MovingAverage::new(100),
|
||||
tps: MovingAverage::new(100),
|
||||
kv_cache_hits: AtomicU64::new(0),
|
||||
kv_cache_misses: AtomicU64::new(0),
|
||||
memory_usage_bytes: AtomicUsize::new(0),
|
||||
peak_memory_bytes: AtomicUsize::new(0),
|
||||
active_requests: AtomicUsize::new(0),
|
||||
total_requests: AtomicU64::new(0),
|
||||
total_tokens: AtomicU64::new(0),
|
||||
latency_histogram: LatencyHistogram::new(),
|
||||
queue_depth: AtomicUsize::new(0),
|
||||
start_time: Instant::now(),
|
||||
last_update: RwLock::new(Instant::now()),
|
||||
inter_token_latency_ms: MovingAverage::new(100),
|
||||
batch_sizes: MovingAverage::new(50),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record time to first token
|
||||
pub fn record_ttft(&self, ttft_ms: f32) {
|
||||
self.ttft_ms.add(ttft_ms);
|
||||
self.latency_histogram.record(ttft_ms);
|
||||
*self.last_update.write() = Instant::now();
|
||||
}
|
||||
|
||||
/// Record tokens per second for a generation
|
||||
pub fn record_tps(&self, tokens: usize, duration: Duration) {
|
||||
if duration.as_secs_f32() > 0.0 {
|
||||
let tps = tokens as f32 / duration.as_secs_f32();
|
||||
self.tps.add(tps);
|
||||
}
|
||||
self.total_tokens
|
||||
.fetch_add(tokens as u64, Ordering::Relaxed);
|
||||
*self.last_update.write() = Instant::now();
|
||||
}
|
||||
|
||||
/// Record inter-token latency
|
||||
pub fn record_inter_token_latency(&self, latency_ms: f32) {
|
||||
self.inter_token_latency_ms.add(latency_ms);
|
||||
}
|
||||
|
||||
/// Record batch size
|
||||
pub fn record_batch_size(&self, size: usize) {
|
||||
self.batch_sizes.add(size as f32);
|
||||
}
|
||||
|
||||
/// Record KV cache hit
|
||||
pub fn record_kv_cache_hit(&self) {
|
||||
self.kv_cache_hits.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Record KV cache miss
|
||||
pub fn record_kv_cache_miss(&self) {
|
||||
self.kv_cache_misses.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Get current KV cache hit rate
|
||||
pub fn kv_cache_hit_rate(&self) -> f32 {
|
||||
let hits = self.kv_cache_hits.load(Ordering::Relaxed);
|
||||
let misses = self.kv_cache_misses.load(Ordering::Relaxed);
|
||||
let total = hits + misses;
|
||||
|
||||
if total == 0 {
|
||||
1.0 // No accesses yet, assume perfect
|
||||
} else {
|
||||
hits as f32 / total as f32
|
||||
}
|
||||
}
|
||||
|
||||
/// Update memory usage
|
||||
pub fn update_memory_usage(&self, bytes: usize) {
|
||||
self.memory_usage_bytes.store(bytes, Ordering::Relaxed);
|
||||
|
||||
// Update peak if necessary
|
||||
let current_peak = self.peak_memory_bytes.load(Ordering::Relaxed);
|
||||
if bytes > current_peak {
|
||||
self.peak_memory_bytes.store(bytes, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current memory usage
|
||||
pub fn memory_usage_bytes(&self) -> usize {
|
||||
self.memory_usage_bytes.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Get peak memory usage
|
||||
pub fn peak_memory_bytes(&self) -> usize {
|
||||
self.peak_memory_bytes.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Increment active requests
|
||||
pub fn request_started(&self) {
|
||||
self.active_requests.fetch_add(1, Ordering::Relaxed);
|
||||
self.total_requests.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Decrement active requests
|
||||
pub fn request_completed(&self) {
|
||||
self.active_requests.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Get active request count
|
||||
pub fn active_requests(&self) -> usize {
|
||||
self.active_requests.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Get total requests
|
||||
pub fn total_requests(&self) -> u64 {
|
||||
self.total_requests.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Get total tokens generated
|
||||
pub fn total_tokens(&self) -> u64 {
|
||||
self.total_tokens.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Update queue depth
|
||||
pub fn set_queue_depth(&self, depth: usize) {
|
||||
self.queue_depth.store(depth, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Get queue depth
|
||||
pub fn queue_depth(&self) -> usize {
|
||||
self.queue_depth.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Get uptime duration
|
||||
pub fn uptime(&self) -> Duration {
|
||||
self.start_time.elapsed()
|
||||
}
|
||||
|
||||
/// Get time since last update
|
||||
pub fn time_since_update(&self) -> Duration {
|
||||
self.last_update.read().elapsed()
|
||||
}
|
||||
|
||||
/// Take a snapshot of current metrics
|
||||
pub fn snapshot(&self) -> MetricsSnapshot {
|
||||
MetricsSnapshot {
|
||||
ttft_avg_ms: self.ttft_ms.average(),
|
||||
ttft_p50_ms: self.ttft_ms.percentile(50.0),
|
||||
ttft_p95_ms: self.ttft_ms.percentile(95.0),
|
||||
ttft_p99_ms: self.ttft_ms.percentile(99.0),
|
||||
tps_avg: self.tps.average(),
|
||||
tps_min: self.tps.min(),
|
||||
tps_max: self.tps.max(),
|
||||
kv_cache_hit_rate: self.kv_cache_hit_rate(),
|
||||
memory_usage_bytes: self.memory_usage_bytes(),
|
||||
peak_memory_bytes: self.peak_memory_bytes(),
|
||||
active_requests: self.active_requests(),
|
||||
total_requests: self.total_requests(),
|
||||
total_tokens: self.total_tokens(),
|
||||
queue_depth: self.queue_depth(),
|
||||
uptime_secs: self.uptime().as_secs_f32(),
|
||||
inter_token_latency_avg_ms: self.inter_token_latency_ms.average(),
|
||||
avg_batch_size: self.batch_sizes.average(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset all metrics
|
||||
pub fn reset(&self) {
|
||||
self.ttft_ms.clear();
|
||||
self.tps.clear();
|
||||
self.kv_cache_hits.store(0, Ordering::Relaxed);
|
||||
self.kv_cache_misses.store(0, Ordering::Relaxed);
|
||||
self.peak_memory_bytes.store(
|
||||
self.memory_usage_bytes.load(Ordering::Relaxed),
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
self.total_requests.store(0, Ordering::Relaxed);
|
||||
self.total_tokens.store(0, Ordering::Relaxed);
|
||||
self.latency_histogram.reset();
|
||||
self.inter_token_latency_ms.clear();
|
||||
self.batch_sizes.clear();
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for InferenceMetrics {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Snapshot of metrics at a point in time
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct MetricsSnapshot {
|
||||
/// Average time to first token (ms)
|
||||
pub ttft_avg_ms: f32,
|
||||
/// P50 TTFT
|
||||
pub ttft_p50_ms: f32,
|
||||
/// P95 TTFT
|
||||
pub ttft_p95_ms: f32,
|
||||
/// P99 TTFT
|
||||
pub ttft_p99_ms: f32,
|
||||
/// Average tokens per second
|
||||
pub tps_avg: f32,
|
||||
/// Minimum TPS observed
|
||||
pub tps_min: f32,
|
||||
/// Maximum TPS observed
|
||||
pub tps_max: f32,
|
||||
/// KV cache hit rate (0.0 - 1.0)
|
||||
pub kv_cache_hit_rate: f32,
|
||||
/// Current memory usage (bytes)
|
||||
pub memory_usage_bytes: usize,
|
||||
/// Peak memory usage (bytes)
|
||||
pub peak_memory_bytes: usize,
|
||||
/// Active requests
|
||||
pub active_requests: usize,
|
||||
/// Total requests processed
|
||||
pub total_requests: u64,
|
||||
/// Total tokens generated
|
||||
pub total_tokens: u64,
|
||||
/// Queue depth
|
||||
pub queue_depth: usize,
|
||||
/// Uptime in seconds
|
||||
pub uptime_secs: f32,
|
||||
/// Average inter-token latency
|
||||
pub inter_token_latency_avg_ms: f32,
|
||||
/// Average batch size
|
||||
pub avg_batch_size: f32,
|
||||
}
|
||||
|
||||
impl MetricsSnapshot {
|
||||
/// Check if metrics indicate healthy performance
|
||||
pub fn is_healthy(&self, max_ttft_ms: f32, min_tps: f32) -> bool {
|
||||
self.ttft_avg_ms <= max_ttft_ms && self.tps_avg >= min_tps
|
||||
}
|
||||
|
||||
/// Calculate throughput efficiency
|
||||
pub fn throughput_efficiency(&self, target_tps: f32) -> f32 {
|
||||
if target_tps <= 0.0 {
|
||||
return 1.0;
|
||||
}
|
||||
(self.tps_avg / target_tps).min(1.0)
|
||||
}
|
||||
|
||||
/// Calculate latency score (0-1, higher is better)
|
||||
pub fn latency_score(&self, target_ttft_ms: f32) -> f32 {
|
||||
if self.ttft_avg_ms <= 0.0 {
|
||||
return 1.0;
|
||||
}
|
||||
(target_ttft_ms / self.ttft_avg_ms).min(1.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Metrics collector with periodic aggregation
|
||||
pub struct MetricsCollector {
|
||||
/// Current metrics
|
||||
metrics: InferenceMetrics,
|
||||
/// Historical snapshots
|
||||
history: RwLock<VecDeque<(Instant, MetricsSnapshot)>>,
|
||||
/// Maximum history size
|
||||
max_history: usize,
|
||||
/// Snapshot interval
|
||||
snapshot_interval: Duration,
|
||||
/// Last snapshot time
|
||||
last_snapshot: RwLock<Instant>,
|
||||
}
|
||||
|
||||
impl MetricsCollector {
|
||||
/// Create a new metrics collector
|
||||
pub fn new(max_history: usize, snapshot_interval: Duration) -> Self {
|
||||
Self {
|
||||
metrics: InferenceMetrics::new(),
|
||||
history: RwLock::new(VecDeque::with_capacity(max_history)),
|
||||
max_history,
|
||||
snapshot_interval,
|
||||
last_snapshot: RwLock::new(Instant::now()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get reference to current metrics
|
||||
pub fn metrics(&self) -> &InferenceMetrics {
|
||||
&self.metrics
|
||||
}
|
||||
|
||||
/// Record TTFT and auto-snapshot if needed
|
||||
pub fn record_ttft(&self, ttft_ms: f32) {
|
||||
self.metrics.record_ttft(ttft_ms);
|
||||
self.maybe_snapshot();
|
||||
}
|
||||
|
||||
/// Record TPS and auto-snapshot if needed
|
||||
pub fn record_tps(&self, tokens: usize, duration: Duration) {
|
||||
self.metrics.record_tps(tokens, duration);
|
||||
self.maybe_snapshot();
|
||||
}
|
||||
|
||||
/// Check if snapshot is needed and take it
|
||||
fn maybe_snapshot(&self) {
|
||||
let last = *self.last_snapshot.read();
|
||||
if last.elapsed() >= self.snapshot_interval {
|
||||
self.take_snapshot();
|
||||
}
|
||||
}
|
||||
|
||||
/// Force a snapshot
|
||||
pub fn take_snapshot(&self) {
|
||||
let snapshot = self.metrics.snapshot();
|
||||
let now = Instant::now();
|
||||
|
||||
let mut history = self.history.write();
|
||||
if history.len() >= self.max_history {
|
||||
history.pop_front();
|
||||
}
|
||||
history.push_back((now, snapshot));
|
||||
|
||||
*self.last_snapshot.write() = now;
|
||||
}
|
||||
|
||||
/// Get recent snapshots
|
||||
pub fn get_history(&self, count: usize) -> Vec<MetricsSnapshot> {
|
||||
let history = self.history.read();
|
||||
history
|
||||
.iter()
|
||||
.rev()
|
||||
.take(count)
|
||||
.map(|(_, s)| s.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get trend analysis (positive = improving, negative = degrading)
|
||||
pub fn ttft_trend(&self) -> f32 {
|
||||
let history = self.history.read();
|
||||
if history.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let recent: Vec<f32> = history
|
||||
.iter()
|
||||
.rev()
|
||||
.take(10)
|
||||
.map(|(_, s)| s.ttft_avg_ms)
|
||||
.collect();
|
||||
|
||||
if recent.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
// Simple linear regression slope
|
||||
let n = recent.len() as f32;
|
||||
let sum_x: f32 = (0..recent.len()).map(|i| i as f32).sum();
|
||||
let sum_y: f32 = recent.iter().sum();
|
||||
let sum_xy: f32 = recent.iter().enumerate().map(|(i, y)| i as f32 * y).sum();
|
||||
let sum_xx: f32 = (0..recent.len()).map(|i| (i * i) as f32).sum();
|
||||
|
||||
let slope = (n * sum_xy - sum_x * sum_y) / (n * sum_xx - sum_x * sum_x);
|
||||
|
||||
// Negative slope means TTFT is decreasing (improving)
|
||||
-slope
|
||||
}
|
||||
|
||||
/// Get TPS trend
|
||||
pub fn tps_trend(&self) -> f32 {
|
||||
let history = self.history.read();
|
||||
if history.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let recent: Vec<f32> = history
|
||||
.iter()
|
||||
.rev()
|
||||
.take(10)
|
||||
.map(|(_, s)| s.tps_avg)
|
||||
.collect();
|
||||
|
||||
if recent.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
let n = recent.len() as f32;
|
||||
let sum_x: f32 = (0..recent.len()).map(|i| i as f32).sum();
|
||||
let sum_y: f32 = recent.iter().sum();
|
||||
let sum_xy: f32 = recent.iter().enumerate().map(|(i, y)| i as f32 * y).sum();
|
||||
let sum_xx: f32 = (0..recent.len()).map(|i| (i * i) as f32).sum();
|
||||
|
||||
(n * sum_xy - sum_x * sum_y) / (n * sum_xx - sum_x * sum_x)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for MetricsCollector {
|
||||
fn default() -> Self {
|
||||
Self::new(1000, Duration::from_secs(60))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_moving_average() {
|
||||
let ma = MovingAverage::new(3);
|
||||
|
||||
ma.add(1.0);
|
||||
ma.add(2.0);
|
||||
ma.add(3.0);
|
||||
|
||||
assert!((ma.average() - 2.0).abs() < 0.01);
|
||||
|
||||
// Adding 4th value should evict 1.0
|
||||
ma.add(4.0);
|
||||
assert!((ma.average() - 3.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_moving_average_percentile() {
|
||||
let ma = MovingAverage::new(10);
|
||||
|
||||
for i in 1..=10 {
|
||||
ma.add(i as f32);
|
||||
}
|
||||
|
||||
let p50 = ma.percentile(50.0);
|
||||
assert!(p50 >= 5.0 && p50 <= 6.0);
|
||||
|
||||
let p90 = ma.percentile(90.0);
|
||||
assert!(p90 >= 9.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_latency_histogram() {
|
||||
let hist = LatencyHistogram::new();
|
||||
|
||||
hist.record(5.0);
|
||||
hist.record(15.0);
|
||||
hist.record(50.0);
|
||||
|
||||
assert_eq!(hist.count(), 3);
|
||||
assert!((hist.mean() - 23.33).abs() < 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_inference_metrics() {
|
||||
let metrics = InferenceMetrics::new();
|
||||
|
||||
metrics.record_ttft(10.0);
|
||||
metrics.record_ttft(20.0);
|
||||
|
||||
assert!((metrics.ttft_ms.average() - 15.0).abs() < 0.01);
|
||||
|
||||
metrics.record_kv_cache_hit();
|
||||
metrics.record_kv_cache_hit();
|
||||
metrics.record_kv_cache_miss();
|
||||
|
||||
assert!((metrics.kv_cache_hit_rate() - 0.667).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_metrics_snapshot() {
|
||||
let metrics = InferenceMetrics::new();
|
||||
|
||||
metrics.record_ttft(10.0);
|
||||
metrics.record_tps(100, Duration::from_secs(1));
|
||||
metrics.update_memory_usage(1024 * 1024);
|
||||
metrics.request_started();
|
||||
|
||||
let snapshot = metrics.snapshot();
|
||||
|
||||
assert!((snapshot.ttft_avg_ms - 10.0).abs() < 0.01);
|
||||
assert!((snapshot.tps_avg - 100.0).abs() < 0.01);
|
||||
assert_eq!(snapshot.memory_usage_bytes, 1024 * 1024);
|
||||
assert_eq!(snapshot.active_requests, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_metrics_collector() {
|
||||
let collector = MetricsCollector::new(100, Duration::from_millis(10));
|
||||
|
||||
for i in 1..=5 {
|
||||
collector.record_ttft(i as f32 * 10.0);
|
||||
}
|
||||
|
||||
collector.take_snapshot();
|
||||
|
||||
let history = collector.get_history(1);
|
||||
assert_eq!(history.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_snapshot_health_check() {
|
||||
let snapshot = MetricsSnapshot {
|
||||
ttft_avg_ms: 50.0,
|
||||
ttft_p50_ms: 45.0,
|
||||
ttft_p95_ms: 80.0,
|
||||
ttft_p99_ms: 100.0,
|
||||
tps_avg: 150.0,
|
||||
tps_min: 100.0,
|
||||
tps_max: 200.0,
|
||||
kv_cache_hit_rate: 0.95,
|
||||
memory_usage_bytes: 1024 * 1024,
|
||||
peak_memory_bytes: 2 * 1024 * 1024,
|
||||
active_requests: 5,
|
||||
total_requests: 1000,
|
||||
total_tokens: 100000,
|
||||
queue_depth: 2,
|
||||
uptime_secs: 3600.0,
|
||||
inter_token_latency_avg_ms: 5.0,
|
||||
avg_batch_size: 8.0,
|
||||
};
|
||||
|
||||
assert!(snapshot.is_healthy(100.0, 100.0));
|
||||
assert!(!snapshot.is_healthy(30.0, 100.0)); // TTFT too high
|
||||
assert!(!snapshot.is_healthy(100.0, 200.0)); // TPS too low
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
//! Real-time Optimization System for RuvLLM
|
||||
//!
|
||||
//! This module provides the optimization infrastructure for LLM inference,
|
||||
//! integrating SONA learning with MicroLoRA and custom kernels.
|
||||
//!
|
||||
//! ## Quick Start
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! use ruvllm::optimization::{
|
||||
//! SonaLlm, SonaLlmConfig, RealtimeOptimizer, RealtimeConfig,
|
||||
//! MetricsCollector, ConsolidationStrategy,
|
||||
//! };
|
||||
//!
|
||||
//! // Create SONA integration for three-tier learning
|
||||
//! let sona_config = SonaLlmConfig {
|
||||
//! instant_lr: 0.01,
|
||||
//! background_interval_ms: 100,
|
||||
//! background_min_samples: 10,
|
||||
//! consolidation_strategy: ConsolidationStrategy::EwcMerge,
|
||||
//! ..Default::default()
|
||||
//! };
|
||||
//! let sona = SonaLlm::new(sona_config);
|
||||
//!
|
||||
//! // During inference: instant adaptation
|
||||
//! let result = sona.instant_adapt(&query_embedding, &response_embedding, 0.85);
|
||||
//! println!("Adapt latency: {}us", result.latency_us);
|
||||
//!
|
||||
//! // Periodic: background consolidation
|
||||
//! if let Some(bg_result) = sona.maybe_background() {
|
||||
//! println!("Consolidated {} samples", bg_result.samples_used);
|
||||
//! }
|
||||
//!
|
||||
//! // Triggered: deep optimization
|
||||
//! if sona.should_trigger_deep() {
|
||||
//! let deep_result = sona.deep_optimize(&samples);
|
||||
//! println!("Quality delta: {:.3}", deep_result.quality_delta);
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! ## Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! +-------------------+ +-------------------+
|
||||
//! | Inference Request |---->| RealtimeOptimizer |
|
||||
//! | (tokens, params) | | - Batch sizing |
|
||||
//! +-------------------+ | - KV management |
|
||||
//! | - Token budgets |
|
||||
//! +--------+----------+
|
||||
//! |
|
||||
//! v (metrics)
|
||||
//! +--------+----------+
|
||||
//! | InferenceMetrics |
|
||||
//! | - TTFT tracking |
|
||||
//! | - TPS monitoring |
|
||||
//! | - Memory usage |
|
||||
//! +--------+----------+
|
||||
//! |
|
||||
//! v (feedback)
|
||||
//! +--------+----------+
|
||||
//! | SonaLlm |
|
||||
//! | - Instant adapt |
|
||||
//! | - Background loop |
|
||||
//! | - Deep optimize |
|
||||
//! +-------------------+
|
||||
//! ```
|
||||
//!
|
||||
//! ## SONA Learning Tiers
|
||||
//!
|
||||
//! | Tier | Latency | Trigger | Action |
|
||||
//! |------|---------|---------|--------|
|
||||
//! | Instant | <1ms | Every request | MicroLoRA gradient update |
|
||||
//! | Background | ~100ms | Timer/threshold | Pattern consolidation |
|
||||
//! | Deep | Minutes | Manual/scheduled | Full training pipeline |
|
||||
//!
|
||||
//! ## Features
|
||||
//!
|
||||
//! - **Real-time Optimization**: Dynamic batch sizing and KV cache management
|
||||
//! - **SONA Integration**: Three-tier learning loops for continuous improvement
|
||||
//! - **Metrics Collection**: Comprehensive inference telemetry
|
||||
//! - **Speculative Decoding**: Draft model integration for faster generation
|
||||
//!
|
||||
//! ## Consolidation Strategies
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! use ruvllm::optimization::ConsolidationStrategy;
|
||||
//!
|
||||
//! // EWC++ merge (default) - preserves important weights
|
||||
//! let strategy = ConsolidationStrategy::EwcMerge;
|
||||
//!
|
||||
//! // Quality-weighted - higher quality samples have more influence
|
||||
//! let strategy = ConsolidationStrategy::QualityWeighted;
|
||||
//!
|
||||
//! // Best only - keep top 20% by quality
|
||||
//! let strategy = ConsolidationStrategy::BestOnly;
|
||||
//! ```
|
||||
|
||||
pub mod metrics;
|
||||
pub mod realtime;
|
||||
pub mod sona_llm;
|
||||
|
||||
// Re-exports
|
||||
pub use metrics::{
|
||||
InferenceMetrics, LatencyHistogram, MetricsCollector, MetricsSnapshot, MovingAverage,
|
||||
};
|
||||
pub use realtime::{
|
||||
BatchSizeStrategy, KvCachePressurePolicy, OptimizationDecision, RealtimeConfig,
|
||||
RealtimeOptimizer, SpeculativeConfig, TokenBudgetAllocation,
|
||||
};
|
||||
pub use sona_llm::{
|
||||
AdaptationResult, ConsolidationStrategy, LearningLoopStats, OptimizationTrigger, SonaLlm,
|
||||
SonaLlmConfig, TrainingSample,
|
||||
};
|
||||
@@ -0,0 +1,956 @@
|
||||
//! Real-time Optimization for LLM Inference
|
||||
//!
|
||||
//! Features:
|
||||
//! - Dynamic batch sizing based on latency targets
|
||||
//! - KV cache pressure management
|
||||
//! - Token budget allocation
|
||||
//! - Speculative decoding integration
|
||||
|
||||
use crate::error::{Result, RuvLLMError};
|
||||
use crate::optimization::metrics::{InferenceMetrics, MetricsSnapshot};
|
||||
use parking_lot::RwLock;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Configuration for the realtime optimizer
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RealtimeConfig {
|
||||
/// Target latency for TTFT (milliseconds)
|
||||
pub latency_target_ms: f32,
|
||||
/// Target throughput (tokens per second)
|
||||
pub throughput_target_tps: f32,
|
||||
/// Minimum batch size
|
||||
pub min_batch_size: usize,
|
||||
/// Maximum batch size
|
||||
pub max_batch_size: usize,
|
||||
/// KV cache pressure threshold (0.0 - 1.0)
|
||||
pub kv_cache_pressure_threshold: f32,
|
||||
/// Enable speculative decoding
|
||||
pub enable_speculative: bool,
|
||||
/// Speculative decoding configuration
|
||||
pub speculative: SpeculativeConfig,
|
||||
/// Batch sizing strategy
|
||||
pub batch_strategy: BatchSizeStrategy,
|
||||
/// KV cache pressure policy
|
||||
pub kv_policy: KvCachePressurePolicy,
|
||||
/// Maximum memory budget (bytes)
|
||||
pub max_memory_bytes: usize,
|
||||
/// Optimization interval (how often to recompute decisions)
|
||||
pub optimization_interval_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for RealtimeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
latency_target_ms: 100.0,
|
||||
throughput_target_tps: 50.0,
|
||||
min_batch_size: 1,
|
||||
max_batch_size: 64,
|
||||
kv_cache_pressure_threshold: 0.8,
|
||||
enable_speculative: true, // Enabled by default for 2-3x decode speedup
|
||||
speculative: SpeculativeConfig::default(),
|
||||
batch_strategy: BatchSizeStrategy::Adaptive,
|
||||
kv_policy: KvCachePressurePolicy::Evict,
|
||||
max_memory_bytes: 8 * 1024 * 1024 * 1024, // 8GB default
|
||||
optimization_interval_ms: 100,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch size selection strategy
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum BatchSizeStrategy {
|
||||
/// Fixed batch size
|
||||
Fixed,
|
||||
/// Adaptive based on latency
|
||||
Adaptive,
|
||||
/// Aggressive (maximize throughput)
|
||||
Aggressive,
|
||||
/// Conservative (minimize latency)
|
||||
Conservative,
|
||||
/// Hybrid (balance throughput and latency)
|
||||
Hybrid,
|
||||
}
|
||||
|
||||
impl Default for BatchSizeStrategy {
|
||||
fn default() -> Self {
|
||||
Self::Adaptive
|
||||
}
|
||||
}
|
||||
|
||||
/// Policy for handling KV cache pressure
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum KvCachePressurePolicy {
|
||||
/// Evict oldest entries
|
||||
Evict,
|
||||
/// Quantize more aggressively
|
||||
Quantize,
|
||||
/// Reject new requests
|
||||
Reject,
|
||||
/// Spill to disk
|
||||
Spill,
|
||||
/// Hybrid approach
|
||||
Hybrid,
|
||||
}
|
||||
|
||||
impl Default for KvCachePressurePolicy {
|
||||
fn default() -> Self {
|
||||
Self::Evict
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration for speculative decoding
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpeculativeConfig {
|
||||
/// Draft model identifier
|
||||
pub draft_model: Option<String>,
|
||||
/// Number of speculative tokens
|
||||
pub num_speculative_tokens: usize,
|
||||
/// Acceptance threshold
|
||||
pub acceptance_threshold: f32,
|
||||
/// Enable tree-based speculation
|
||||
pub tree_speculation: bool,
|
||||
/// Maximum tree depth
|
||||
pub max_tree_depth: usize,
|
||||
}
|
||||
|
||||
impl Default for SpeculativeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
draft_model: None,
|
||||
num_speculative_tokens: 4,
|
||||
acceptance_threshold: 0.8,
|
||||
tree_speculation: false,
|
||||
max_tree_depth: 3,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Token budget allocation for a request
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TokenBudgetAllocation {
|
||||
/// Request identifier
|
||||
pub request_id: String,
|
||||
/// Maximum tokens to generate
|
||||
pub max_tokens: usize,
|
||||
/// Priority level (higher = more resources)
|
||||
pub priority: f32,
|
||||
/// Deadline (optional)
|
||||
pub deadline: Option<Duration>,
|
||||
/// Allocated batch slot
|
||||
pub batch_slot: Option<usize>,
|
||||
/// Estimated completion time
|
||||
pub estimated_completion_ms: f32,
|
||||
}
|
||||
|
||||
/// Request representation for optimization
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Request {
|
||||
/// Request identifier
|
||||
pub id: String,
|
||||
/// Input token count
|
||||
pub input_tokens: usize,
|
||||
/// Maximum output tokens
|
||||
pub max_output_tokens: usize,
|
||||
/// Priority (0.0 - 1.0)
|
||||
pub priority: f32,
|
||||
/// Arrival time
|
||||
pub arrival_time: Instant,
|
||||
/// Deadline (optional)
|
||||
pub deadline: Option<Duration>,
|
||||
}
|
||||
|
||||
/// Optimization decision output
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OptimizationDecision {
|
||||
/// Recommended batch size
|
||||
pub batch_size: usize,
|
||||
/// Whether to evict KV cache
|
||||
pub should_evict_kv_cache: bool,
|
||||
/// Number of entries to evict
|
||||
pub evict_count: usize,
|
||||
/// Whether to enable speculative decoding
|
||||
pub enable_speculative: bool,
|
||||
/// Token budgets for requests
|
||||
pub token_budgets: Vec<TokenBudgetAllocation>,
|
||||
/// Quantization recommendation
|
||||
pub quantization_recommendation: Option<String>,
|
||||
/// Estimated latency for current batch
|
||||
pub estimated_latency_ms: f32,
|
||||
/// Estimated throughput for current batch
|
||||
pub estimated_tps: f32,
|
||||
/// Confidence in this decision (0.0 - 1.0)
|
||||
pub confidence: f32,
|
||||
/// Reason for this decision
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
impl Default for OptimizationDecision {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
batch_size: 1,
|
||||
should_evict_kv_cache: false,
|
||||
evict_count: 0,
|
||||
enable_speculative: false,
|
||||
token_budgets: Vec::new(),
|
||||
quantization_recommendation: None,
|
||||
estimated_latency_ms: 0.0,
|
||||
estimated_tps: 0.0,
|
||||
confidence: 0.5,
|
||||
reason: "Default decision".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Real-time optimizer for LLM inference
|
||||
pub struct RealtimeOptimizer {
|
||||
/// Configuration
|
||||
config: RwLock<RealtimeConfig>,
|
||||
/// Current batch size
|
||||
current_batch_size: AtomicUsize,
|
||||
/// Current KV cache pressure (0.0 - 1.0)
|
||||
kv_cache_pressure: RwLock<f32>,
|
||||
/// Recent latency measurements
|
||||
recent_latencies: RwLock<VecDeque<f32>>,
|
||||
/// Recent throughput measurements
|
||||
recent_throughputs: RwLock<VecDeque<f32>>,
|
||||
/// Whether speculative decoding is active
|
||||
speculative_active: AtomicBool,
|
||||
/// Draft model identifier (if loaded)
|
||||
draft_model: RwLock<Option<String>>,
|
||||
/// Last optimization time
|
||||
last_optimization: RwLock<Instant>,
|
||||
/// Pending requests
|
||||
pending_requests: RwLock<Vec<Request>>,
|
||||
/// Current memory usage
|
||||
current_memory_bytes: AtomicUsize,
|
||||
}
|
||||
|
||||
impl RealtimeOptimizer {
|
||||
/// Create a new realtime optimizer
|
||||
pub fn new(config: RealtimeConfig) -> Self {
|
||||
let initial_batch_size = match config.batch_strategy {
|
||||
BatchSizeStrategy::Fixed => config.max_batch_size,
|
||||
BatchSizeStrategy::Aggressive => config.max_batch_size,
|
||||
BatchSizeStrategy::Conservative => config.min_batch_size,
|
||||
_ => (config.min_batch_size + config.max_batch_size) / 2,
|
||||
};
|
||||
|
||||
Self {
|
||||
config: RwLock::new(config),
|
||||
current_batch_size: AtomicUsize::new(initial_batch_size),
|
||||
kv_cache_pressure: RwLock::new(0.0),
|
||||
recent_latencies: RwLock::new(VecDeque::with_capacity(100)),
|
||||
recent_throughputs: RwLock::new(VecDeque::with_capacity(100)),
|
||||
speculative_active: AtomicBool::new(false),
|
||||
draft_model: RwLock::new(None),
|
||||
last_optimization: RwLock::new(Instant::now()),
|
||||
pending_requests: RwLock::new(Vec::new()),
|
||||
current_memory_bytes: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Optimize batch size based on recent latency measurements
|
||||
pub fn optimize_batch_size(&self, recent_latencies: &[f32]) -> usize {
|
||||
let config = self.config.read();
|
||||
|
||||
// Update internal latency tracking
|
||||
{
|
||||
let mut latencies = self.recent_latencies.write();
|
||||
for &l in recent_latencies {
|
||||
if latencies.len() >= 100 {
|
||||
latencies.pop_front();
|
||||
}
|
||||
latencies.push_back(l);
|
||||
}
|
||||
}
|
||||
|
||||
let current_batch = self.current_batch_size.load(Ordering::Relaxed);
|
||||
|
||||
let new_batch_size = match config.batch_strategy {
|
||||
BatchSizeStrategy::Fixed => current_batch,
|
||||
|
||||
BatchSizeStrategy::Adaptive => self.adaptive_batch_size(&config, recent_latencies),
|
||||
|
||||
BatchSizeStrategy::Aggressive => {
|
||||
// Maximize batch size while staying under latency target
|
||||
let avg_latency = self.average_latency();
|
||||
if avg_latency < config.latency_target_ms * 0.7 {
|
||||
(current_batch + 4).min(config.max_batch_size)
|
||||
} else if avg_latency > config.latency_target_ms {
|
||||
(current_batch.saturating_sub(2)).max(config.min_batch_size)
|
||||
} else {
|
||||
current_batch
|
||||
}
|
||||
}
|
||||
|
||||
BatchSizeStrategy::Conservative => {
|
||||
// Minimize latency, slowly increase batch size
|
||||
let avg_latency = self.average_latency();
|
||||
if avg_latency < config.latency_target_ms * 0.5 {
|
||||
(current_batch + 1).min(config.max_batch_size)
|
||||
} else if avg_latency > config.latency_target_ms * 0.8 {
|
||||
(current_batch.saturating_sub(1)).max(config.min_batch_size)
|
||||
} else {
|
||||
current_batch
|
||||
}
|
||||
}
|
||||
|
||||
BatchSizeStrategy::Hybrid => {
|
||||
// Balance throughput and latency using a utility function
|
||||
self.hybrid_batch_size(&config)
|
||||
}
|
||||
};
|
||||
|
||||
self.current_batch_size
|
||||
.store(new_batch_size, Ordering::Relaxed);
|
||||
new_batch_size
|
||||
}
|
||||
|
||||
/// Adaptive batch sizing based on PID-like control
|
||||
fn adaptive_batch_size(&self, config: &RealtimeConfig, recent_latencies: &[f32]) -> usize {
|
||||
let current_batch = self.current_batch_size.load(Ordering::Relaxed);
|
||||
|
||||
if recent_latencies.is_empty() {
|
||||
return current_batch;
|
||||
}
|
||||
|
||||
let avg_latency: f32 = recent_latencies.iter().sum::<f32>() / recent_latencies.len() as f32;
|
||||
let target = config.latency_target_ms;
|
||||
|
||||
// Error term (positive = too slow, negative = too fast)
|
||||
let error = avg_latency - target;
|
||||
let error_ratio = error / target;
|
||||
|
||||
// PID-like adjustment
|
||||
let adjustment = if error_ratio.abs() < 0.1 {
|
||||
// Within 10% of target, no change
|
||||
0
|
||||
} else if error_ratio > 0.0 {
|
||||
// Too slow, reduce batch size
|
||||
let reduction = (error_ratio * 4.0).ceil() as i32;
|
||||
-reduction.min(4)
|
||||
} else {
|
||||
// Too fast, increase batch size
|
||||
let increase = (-error_ratio * 2.0).ceil() as i32;
|
||||
increase.min(2)
|
||||
};
|
||||
|
||||
let new_batch = (current_batch as i32 + adjustment)
|
||||
.max(config.min_batch_size as i32)
|
||||
.min(config.max_batch_size as i32) as usize;
|
||||
|
||||
new_batch
|
||||
}
|
||||
|
||||
/// Hybrid batch sizing using utility maximization
|
||||
fn hybrid_batch_size(&self, config: &RealtimeConfig) -> usize {
|
||||
let current_batch = self.current_batch_size.load(Ordering::Relaxed);
|
||||
let avg_latency = self.average_latency();
|
||||
let avg_throughput = self.average_throughput();
|
||||
|
||||
// Utility = alpha * throughput_normalized - beta * latency_normalized
|
||||
let alpha = 0.6; // Weight for throughput
|
||||
let beta = 0.4; // Weight for latency
|
||||
|
||||
let latency_normalized = (avg_latency / config.latency_target_ms).min(2.0);
|
||||
let throughput_normalized = (avg_throughput / config.throughput_target_tps).min(2.0);
|
||||
|
||||
let current_utility = alpha * throughput_normalized - beta * latency_normalized;
|
||||
|
||||
// Try neighboring batch sizes and pick the one with best predicted utility
|
||||
let candidates = [
|
||||
current_batch.saturating_sub(2),
|
||||
current_batch.saturating_sub(1),
|
||||
current_batch,
|
||||
current_batch + 1,
|
||||
current_batch + 2,
|
||||
];
|
||||
|
||||
let mut best_batch = current_batch;
|
||||
let mut best_utility = current_utility;
|
||||
|
||||
for &candidate in &candidates {
|
||||
if candidate < config.min_batch_size || candidate > config.max_batch_size {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Predict utility for this batch size
|
||||
let batch_ratio = candidate as f32 / current_batch as f32;
|
||||
let predicted_latency = avg_latency * batch_ratio.sqrt(); // Latency grows sub-linearly
|
||||
let predicted_throughput = avg_throughput * batch_ratio; // Throughput grows linearly
|
||||
|
||||
let pred_latency_norm = (predicted_latency / config.latency_target_ms).min(2.0);
|
||||
let pred_throughput_norm =
|
||||
(predicted_throughput / config.throughput_target_tps).min(2.0);
|
||||
|
||||
let predicted_utility = alpha * pred_throughput_norm - beta * pred_latency_norm;
|
||||
|
||||
if predicted_utility > best_utility {
|
||||
best_utility = predicted_utility;
|
||||
best_batch = candidate;
|
||||
}
|
||||
}
|
||||
|
||||
best_batch
|
||||
}
|
||||
|
||||
/// Check if KV cache eviction is needed
|
||||
pub fn should_evict_kv_cache(&self) -> bool {
|
||||
let config = self.config.read();
|
||||
let pressure = *self.kv_cache_pressure.read();
|
||||
pressure >= config.kv_cache_pressure_threshold
|
||||
}
|
||||
|
||||
/// Update KV cache pressure
|
||||
pub fn update_kv_cache_pressure(&self, pressure: f32) {
|
||||
*self.kv_cache_pressure.write() = pressure.clamp(0.0, 1.0);
|
||||
}
|
||||
|
||||
/// Get KV cache pressure
|
||||
pub fn kv_cache_pressure(&self) -> f32 {
|
||||
*self.kv_cache_pressure.read()
|
||||
}
|
||||
|
||||
/// Allocate token budgets for a set of requests
|
||||
pub fn allocate_token_budget(&self, requests: &[Request]) -> Vec<TokenBudgetAllocation> {
|
||||
let config = self.config.read();
|
||||
let batch_size = self.current_batch_size.load(Ordering::Relaxed);
|
||||
let memory_budget = config.max_memory_bytes;
|
||||
|
||||
// Sort requests by priority and deadline
|
||||
let mut sorted_requests: Vec<(usize, &Request)> = requests.iter().enumerate().collect();
|
||||
sorted_requests.sort_by(|(_, a), (_, b)| {
|
||||
// Higher priority first
|
||||
let priority_cmp = b
|
||||
.priority
|
||||
.partial_cmp(&a.priority)
|
||||
.unwrap_or(std::cmp::Ordering::Equal);
|
||||
if priority_cmp != std::cmp::Ordering::Equal {
|
||||
return priority_cmp;
|
||||
}
|
||||
// Earlier deadline first
|
||||
match (&a.deadline, &b.deadline) {
|
||||
(Some(da), Some(db)) => da.cmp(db),
|
||||
(Some(_), None) => std::cmp::Ordering::Less,
|
||||
(None, Some(_)) => std::cmp::Ordering::Greater,
|
||||
(None, None) => std::cmp::Ordering::Equal,
|
||||
}
|
||||
});
|
||||
|
||||
let mut allocations = Vec::with_capacity(requests.len());
|
||||
let mut total_memory = 0usize;
|
||||
let mut assigned_slots = 0usize;
|
||||
|
||||
for (original_idx, request) in sorted_requests {
|
||||
// Estimate memory for this request
|
||||
let estimated_memory = self.estimate_request_memory(request);
|
||||
|
||||
let (max_tokens, batch_slot) = if assigned_slots < batch_size
|
||||
&& total_memory + estimated_memory <= memory_budget
|
||||
{
|
||||
total_memory += estimated_memory;
|
||||
let slot = assigned_slots;
|
||||
assigned_slots += 1;
|
||||
(request.max_output_tokens, Some(slot))
|
||||
} else {
|
||||
// Request is queued, reduced token budget
|
||||
let reduced = (request.max_output_tokens / 2).max(1);
|
||||
(reduced, None)
|
||||
};
|
||||
|
||||
let estimated_completion = self.estimate_completion_time(request, batch_slot);
|
||||
|
||||
allocations.push((
|
||||
original_idx,
|
||||
TokenBudgetAllocation {
|
||||
request_id: request.id.clone(),
|
||||
max_tokens,
|
||||
priority: request.priority,
|
||||
deadline: request.deadline,
|
||||
batch_slot,
|
||||
estimated_completion_ms: estimated_completion,
|
||||
},
|
||||
));
|
||||
}
|
||||
|
||||
// Sort back to original order
|
||||
allocations.sort_by_key(|(idx, _)| *idx);
|
||||
allocations.into_iter().map(|(_, alloc)| alloc).collect()
|
||||
}
|
||||
|
||||
/// Estimate memory requirement for a request
|
||||
fn estimate_request_memory(&self, request: &Request) -> usize {
|
||||
// Rough estimate: 2 bytes per token for KV cache (FP16)
|
||||
// Plus overhead for attention computation
|
||||
let kv_memory = (request.input_tokens + request.max_output_tokens) * 2 * 128; // head_dim
|
||||
let attention_overhead = request.input_tokens * 32; // Attention scores
|
||||
kv_memory + attention_overhead
|
||||
}
|
||||
|
||||
/// Estimate completion time for a request
|
||||
fn estimate_completion_time(&self, request: &Request, batch_slot: Option<usize>) -> f32 {
|
||||
let avg_tps = self.average_throughput().max(1.0);
|
||||
let base_time = request.max_output_tokens as f32 / avg_tps * 1000.0;
|
||||
|
||||
// Add queue time if not in current batch
|
||||
if batch_slot.is_none() {
|
||||
let queue_size = self.pending_requests.read().len();
|
||||
base_time + (queue_size as f32 * self.average_latency())
|
||||
} else {
|
||||
base_time
|
||||
}
|
||||
}
|
||||
|
||||
/// Enable speculative decoding
|
||||
pub fn enable_speculative_decoding(&self, draft_model: &str) {
|
||||
*self.draft_model.write() = Some(draft_model.to_string());
|
||||
self.speculative_active.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Disable speculative decoding
|
||||
pub fn disable_speculative_decoding(&self) {
|
||||
self.speculative_active.store(false, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Update speculation statistics for learning/monitoring
|
||||
///
|
||||
/// This records the acceptance rate of speculative decoding rounds
|
||||
/// to help tune the lookahead parameter adaptively.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `accepted_count` - Number of draft tokens that were accepted
|
||||
/// * `total_drafted` - Total number of draft tokens generated
|
||||
pub fn update_speculation_stats(&self, accepted_count: usize, total_drafted: usize) {
|
||||
if total_drafted == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
// Calculate acceptance rate
|
||||
let acceptance_rate = accepted_count as f32 / total_drafted as f32;
|
||||
|
||||
// Use acceptance rate to adjust future speculative decoding behavior
|
||||
// High acceptance (>0.8) suggests we can increase lookahead
|
||||
// Low acceptance (<0.5) suggests we should reduce lookahead or disable
|
||||
let mut config = self.config.write();
|
||||
|
||||
if acceptance_rate > 0.9 && config.speculative.num_speculative_tokens < 8 {
|
||||
// Excellent acceptance, try more tokens
|
||||
config.speculative.num_speculative_tokens += 1;
|
||||
} else if acceptance_rate < 0.3 && config.speculative.num_speculative_tokens > 2 {
|
||||
// Poor acceptance, reduce speculation
|
||||
config.speculative.num_speculative_tokens -= 1;
|
||||
}
|
||||
|
||||
// Update acceptance threshold based on recent performance
|
||||
// This implements a simple exponential moving average
|
||||
let alpha = 0.1; // Learning rate
|
||||
config.speculative.acceptance_threshold =
|
||||
config.speculative.acceptance_threshold * (1.0 - alpha) + acceptance_rate * alpha;
|
||||
}
|
||||
|
||||
/// Check if speculative decoding is active
|
||||
pub fn is_speculative_active(&self) -> bool {
|
||||
self.speculative_active.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Get the draft model identifier
|
||||
pub fn draft_model(&self) -> Option<String> {
|
||||
self.draft_model.read().clone()
|
||||
}
|
||||
|
||||
/// Record a latency measurement
|
||||
pub fn record_latency(&self, latency_ms: f32) {
|
||||
let mut latencies = self.recent_latencies.write();
|
||||
if latencies.len() >= 100 {
|
||||
latencies.pop_front();
|
||||
}
|
||||
latencies.push_back(latency_ms);
|
||||
}
|
||||
|
||||
/// Record a throughput measurement
|
||||
pub fn record_throughput(&self, tps: f32) {
|
||||
let mut throughputs = self.recent_throughputs.write();
|
||||
if throughputs.len() >= 100 {
|
||||
throughputs.pop_front();
|
||||
}
|
||||
throughputs.push_back(tps);
|
||||
}
|
||||
|
||||
/// Get average latency
|
||||
pub fn average_latency(&self) -> f32 {
|
||||
let latencies = self.recent_latencies.read();
|
||||
if latencies.is_empty() {
|
||||
return 50.0; // Default estimate
|
||||
}
|
||||
latencies.iter().sum::<f32>() / latencies.len() as f32
|
||||
}
|
||||
|
||||
/// Get average throughput
|
||||
pub fn average_throughput(&self) -> f32 {
|
||||
let throughputs = self.recent_throughputs.read();
|
||||
if throughputs.is_empty() {
|
||||
return 50.0; // Default estimate
|
||||
}
|
||||
throughputs.iter().sum::<f32>() / throughputs.len() as f32
|
||||
}
|
||||
|
||||
/// Update memory usage
|
||||
pub fn update_memory_usage(&self, bytes: usize) {
|
||||
self.current_memory_bytes.store(bytes, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Get memory pressure (0.0 - 1.0)
|
||||
pub fn memory_pressure(&self) -> f32 {
|
||||
let config = self.config.read();
|
||||
let current = self.current_memory_bytes.load(Ordering::Relaxed);
|
||||
current as f32 / config.max_memory_bytes as f32
|
||||
}
|
||||
|
||||
/// Make a comprehensive optimization decision
|
||||
pub fn optimize(&self, metrics: &InferenceMetrics) -> OptimizationDecision {
|
||||
let config = self.config.read();
|
||||
let snapshot = metrics.snapshot();
|
||||
|
||||
// Check if we need to optimize
|
||||
let last_opt = *self.last_optimization.read();
|
||||
if last_opt.elapsed().as_millis() < config.optimization_interval_ms as u128 {
|
||||
return OptimizationDecision {
|
||||
batch_size: self.current_batch_size.load(Ordering::Relaxed),
|
||||
confidence: 0.3,
|
||||
reason: "Skipping optimization (too recent)".to_string(),
|
||||
..Default::default()
|
||||
};
|
||||
}
|
||||
*self.last_optimization.write() = Instant::now();
|
||||
|
||||
// Determine batch size
|
||||
let latencies: Vec<f32> = self.recent_latencies.read().iter().copied().collect();
|
||||
let batch_size = self.optimize_batch_size(&latencies);
|
||||
|
||||
// Determine KV cache action
|
||||
let kv_pressure = *self.kv_cache_pressure.read();
|
||||
let (should_evict, evict_count) = if kv_pressure >= config.kv_cache_pressure_threshold {
|
||||
let excess_pressure = kv_pressure - config.kv_cache_pressure_threshold;
|
||||
let evict_ratio =
|
||||
(excess_pressure / (1.0 - config.kv_cache_pressure_threshold)).min(0.5);
|
||||
(true, (evict_ratio * 1000.0) as usize) // Evict proportionally
|
||||
} else {
|
||||
(false, 0)
|
||||
};
|
||||
|
||||
// Determine speculative decoding
|
||||
let enable_speculative = config.enable_speculative
|
||||
&& snapshot.ttft_avg_ms < config.latency_target_ms * 0.5
|
||||
&& self.draft_model.read().is_some();
|
||||
|
||||
// Token budget allocation for pending requests
|
||||
let pending = self.pending_requests.read().clone();
|
||||
let token_budgets = self.allocate_token_budget(&pending);
|
||||
|
||||
// Quantization recommendation
|
||||
let quantization_recommendation = if self.memory_pressure() > 0.8 {
|
||||
Some("Q4".to_string())
|
||||
} else if self.memory_pressure() > 0.6 {
|
||||
Some("Q8".to_string())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Estimate outcomes
|
||||
let batch_ratio =
|
||||
batch_size as f32 / self.current_batch_size.load(Ordering::Relaxed).max(1) as f32;
|
||||
let estimated_latency = snapshot.ttft_avg_ms * batch_ratio.sqrt();
|
||||
let estimated_tps = snapshot.tps_avg * batch_ratio;
|
||||
|
||||
// Calculate confidence based on data quality
|
||||
let sample_count = latencies.len();
|
||||
let confidence = if sample_count < 10 {
|
||||
0.3
|
||||
} else if sample_count < 50 {
|
||||
0.6
|
||||
} else {
|
||||
0.9
|
||||
};
|
||||
|
||||
// Generate reason
|
||||
let reason = format!(
|
||||
"Batch: {} (latency={:.1}ms, target={:.1}ms), KV pressure: {:.1}%, Memory: {:.1}%",
|
||||
batch_size,
|
||||
snapshot.ttft_avg_ms,
|
||||
config.latency_target_ms,
|
||||
kv_pressure * 100.0,
|
||||
self.memory_pressure() * 100.0
|
||||
);
|
||||
|
||||
OptimizationDecision {
|
||||
batch_size,
|
||||
should_evict_kv_cache: should_evict,
|
||||
evict_count,
|
||||
enable_speculative,
|
||||
token_budgets,
|
||||
quantization_recommendation,
|
||||
estimated_latency_ms: estimated_latency,
|
||||
estimated_tps,
|
||||
confidence,
|
||||
reason,
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a pending request
|
||||
pub fn add_request(&self, request: Request) {
|
||||
self.pending_requests.write().push(request);
|
||||
}
|
||||
|
||||
/// Remove a completed request
|
||||
pub fn remove_request(&self, request_id: &str) {
|
||||
self.pending_requests.write().retain(|r| r.id != request_id);
|
||||
}
|
||||
|
||||
/// Get pending request count
|
||||
pub fn pending_request_count(&self) -> usize {
|
||||
self.pending_requests.read().len()
|
||||
}
|
||||
|
||||
/// Get current batch size
|
||||
pub fn current_batch_size(&self) -> usize {
|
||||
self.current_batch_size.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Update configuration
|
||||
pub fn update_config(&self, config: RealtimeConfig) {
|
||||
*self.config.write() = config;
|
||||
}
|
||||
|
||||
/// Get current configuration
|
||||
pub fn config(&self) -> RealtimeConfig {
|
||||
self.config.read().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl RealtimeOptimizer {
|
||||
/// Check if speculative decoding should be used for these generation parameters
|
||||
///
|
||||
/// Returns true when:
|
||||
/// - Temperature is low (< 0.5) - deterministic generation benefits most
|
||||
/// - Greedy decoding (top_k = 1)
|
||||
/// - Speculative decoding is enabled in config
|
||||
pub fn should_use_speculative(&self, params: &crate::backends::GenerateParams) -> bool {
|
||||
let config = self.config.read();
|
||||
if !config.enable_speculative {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Speculative decoding is most effective for:
|
||||
// 1. Low temperature (more deterministic)
|
||||
// 2. Greedy decoding
|
||||
// 3. When not using high top-p sampling
|
||||
params.temperature < 0.5 || params.top_k == 1
|
||||
}
|
||||
|
||||
/// Get recommended speculative decoding configuration based on current metrics
|
||||
pub fn get_speculative_config(&self) -> SpeculativeConfig {
|
||||
let config = self.config.read();
|
||||
let avg_latency = self.average_latency();
|
||||
let memory_pressure = self.memory_pressure();
|
||||
|
||||
// Adjust speculative config based on system state
|
||||
let mut spec_config = config.speculative.clone();
|
||||
|
||||
// Reduce lookahead under memory pressure
|
||||
if memory_pressure > 0.8 {
|
||||
spec_config.num_speculative_tokens = (spec_config.num_speculative_tokens / 2).max(2);
|
||||
}
|
||||
|
||||
// Increase acceptance threshold when latency is high
|
||||
if avg_latency > config.latency_target_ms {
|
||||
spec_config.acceptance_threshold = (spec_config.acceptance_threshold + 0.1).min(0.95);
|
||||
}
|
||||
|
||||
spec_config
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for RealtimeOptimizer {
|
||||
fn default() -> Self {
|
||||
Self::new(RealtimeConfig::default())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_realtime_config_default() {
|
||||
let config = RealtimeConfig::default();
|
||||
assert!((config.latency_target_ms - 100.0).abs() < 0.01);
|
||||
assert!((config.throughput_target_tps - 50.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimizer_creation() {
|
||||
let config = RealtimeConfig {
|
||||
min_batch_size: 1,
|
||||
max_batch_size: 32,
|
||||
batch_strategy: BatchSizeStrategy::Adaptive,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let optimizer = RealtimeOptimizer::new(config);
|
||||
assert!(optimizer.current_batch_size() >= 1);
|
||||
assert!(optimizer.current_batch_size() <= 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_batch_size_optimization() {
|
||||
let config = RealtimeConfig {
|
||||
latency_target_ms: 100.0,
|
||||
min_batch_size: 1,
|
||||
max_batch_size: 16,
|
||||
batch_strategy: BatchSizeStrategy::Adaptive,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let optimizer = RealtimeOptimizer::new(config);
|
||||
|
||||
// High latency should reduce batch size
|
||||
let high_latencies = vec![150.0, 160.0, 140.0];
|
||||
let batch = optimizer.optimize_batch_size(&high_latencies);
|
||||
assert!(batch <= 8, "High latency should reduce batch size");
|
||||
|
||||
// Low latency should increase batch size
|
||||
let low_latencies = vec![30.0, 35.0, 25.0];
|
||||
let batch = optimizer.optimize_batch_size(&low_latencies);
|
||||
assert!(batch >= 4, "Low latency should allow larger batch size");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_kv_cache_pressure() {
|
||||
let config = RealtimeConfig {
|
||||
kv_cache_pressure_threshold: 0.8,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let optimizer = RealtimeOptimizer::new(config);
|
||||
|
||||
optimizer.update_kv_cache_pressure(0.5);
|
||||
assert!(!optimizer.should_evict_kv_cache());
|
||||
|
||||
optimizer.update_kv_cache_pressure(0.9);
|
||||
assert!(optimizer.should_evict_kv_cache());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_token_budget_allocation() {
|
||||
let optimizer = RealtimeOptimizer::new(RealtimeConfig::default());
|
||||
|
||||
let requests = vec![
|
||||
Request {
|
||||
id: "req1".to_string(),
|
||||
input_tokens: 100,
|
||||
max_output_tokens: 200,
|
||||
priority: 0.9,
|
||||
arrival_time: Instant::now(),
|
||||
deadline: None,
|
||||
},
|
||||
Request {
|
||||
id: "req2".to_string(),
|
||||
input_tokens: 50,
|
||||
max_output_tokens: 100,
|
||||
priority: 0.5,
|
||||
arrival_time: Instant::now(),
|
||||
deadline: Some(Duration::from_secs(1)),
|
||||
},
|
||||
];
|
||||
|
||||
let allocations = optimizer.allocate_token_budget(&requests);
|
||||
assert_eq!(allocations.len(), 2);
|
||||
|
||||
// Higher priority request should get more resources
|
||||
let high_priority = allocations.iter().find(|a| a.request_id == "req1").unwrap();
|
||||
assert!(high_priority.batch_slot.is_some() || high_priority.max_tokens >= 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_speculative_decoding() {
|
||||
let optimizer = RealtimeOptimizer::new(RealtimeConfig {
|
||||
enable_speculative: true,
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
assert!(!optimizer.is_speculative_active());
|
||||
|
||||
optimizer.enable_speculative_decoding("draft-model-1");
|
||||
assert!(optimizer.is_speculative_active());
|
||||
assert_eq!(optimizer.draft_model(), Some("draft-model-1".to_string()));
|
||||
|
||||
optimizer.disable_speculative_decoding();
|
||||
assert!(!optimizer.is_speculative_active());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_optimization_decision() {
|
||||
let optimizer = RealtimeOptimizer::new(RealtimeConfig::default());
|
||||
let metrics = InferenceMetrics::new();
|
||||
|
||||
// Record some metrics
|
||||
for i in 1..=10 {
|
||||
metrics.record_ttft(i as f32 * 10.0);
|
||||
optimizer.record_latency(i as f32 * 10.0);
|
||||
optimizer.record_throughput(50.0 + i as f32);
|
||||
}
|
||||
|
||||
let decision = optimizer.optimize(&metrics);
|
||||
assert!(decision.batch_size >= 1);
|
||||
assert!(decision.confidence > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_memory_pressure() {
|
||||
let config = RealtimeConfig {
|
||||
max_memory_bytes: 1024 * 1024 * 1024, // 1GB
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let optimizer = RealtimeOptimizer::new(config);
|
||||
|
||||
optimizer.update_memory_usage(512 * 1024 * 1024); // 512MB
|
||||
assert!((optimizer.memory_pressure() - 0.5).abs() < 0.01);
|
||||
|
||||
optimizer.update_memory_usage(800 * 1024 * 1024); // 800MB
|
||||
assert!((optimizer.memory_pressure() - 0.78).abs() < 0.02);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_batch_strategies() {
|
||||
let strategies = vec![
|
||||
BatchSizeStrategy::Fixed,
|
||||
BatchSizeStrategy::Adaptive,
|
||||
BatchSizeStrategy::Aggressive,
|
||||
BatchSizeStrategy::Conservative,
|
||||
BatchSizeStrategy::Hybrid,
|
||||
];
|
||||
|
||||
for strategy in strategies {
|
||||
let config = RealtimeConfig {
|
||||
batch_strategy: strategy,
|
||||
min_batch_size: 1,
|
||||
max_batch_size: 16,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let optimizer = RealtimeOptimizer::new(config);
|
||||
let latencies = vec![50.0, 55.0, 45.0];
|
||||
let batch = optimizer.optimize_batch_size(&latencies);
|
||||
|
||||
assert!(
|
||||
batch >= 1 && batch <= 16,
|
||||
"Strategy {:?} produced invalid batch size",
|
||||
strategy
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user