104 lines
3.3 KiB
Rust
104 lines
3.3 KiB
Rust
//! Advanced speculative decoder supporting multiple strategies
|
|
|
|
use super::config::SpeculativeConfig;
|
|
use super::lookahead::{LookaheadConfig, NgramPool};
|
|
use super::medusa::MedusaConfig;
|
|
use super::metrics::PerformanceMetrics;
|
|
use super::self_spec::SelfSpeculativeConfig;
|
|
use std::sync::Arc;
|
|
use tokio::sync::{Mutex, RwLock};
|
|
|
|
/// Extended metrics for advanced speculative decoding
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct AdvancedMetrics {
|
|
/// Base speculative metrics
|
|
pub base_metrics: PerformanceMetrics,
|
|
/// Tree attention speedup
|
|
pub tree_attention_speedup: f32,
|
|
/// Average accepted path length in tree
|
|
pub avg_accepted_path_length: f32,
|
|
/// Medusa head acceptance rates
|
|
pub medusa_head_acceptance: Vec<f32>,
|
|
/// Self-speculative early exit rates per layer
|
|
pub early_exit_rates: Vec<f32>,
|
|
/// N-gram cache hit rate
|
|
pub ngram_hit_rate: f32,
|
|
}
|
|
|
|
/// Combined speculative decoding engine supporting multiple strategies
|
|
pub struct AdvancedSpeculativeDecoder {
|
|
/// Basic speculative config
|
|
pub base_config: SpeculativeConfig,
|
|
/// Medusa config (if enabled)
|
|
pub medusa_config: Option<MedusaConfig>,
|
|
/// Self-speculative config (if enabled)
|
|
pub self_spec_config: Option<SelfSpeculativeConfig>,
|
|
/// Lookahead config (if enabled)
|
|
pub lookahead_config: Option<LookaheadConfig>,
|
|
/// N-gram pool for lookahead decoding
|
|
pub ngram_pool: Option<Arc<RwLock<NgramPool>>>,
|
|
/// Performance metrics
|
|
pub metrics: Arc<Mutex<AdvancedMetrics>>,
|
|
}
|
|
|
|
impl AdvancedSpeculativeDecoder {
|
|
/// Create a new advanced speculative decoder with basic config
|
|
#[must_use]
|
|
pub fn new(base_config: SpeculativeConfig) -> Self {
|
|
Self {
|
|
base_config,
|
|
medusa_config: None,
|
|
self_spec_config: None,
|
|
lookahead_config: None,
|
|
ngram_pool: None,
|
|
metrics: Arc::new(Mutex::new(AdvancedMetrics::default())),
|
|
}
|
|
}
|
|
|
|
/// Enable Medusa-style prediction
|
|
#[must_use]
|
|
pub fn with_medusa(mut self, config: MedusaConfig) -> Self {
|
|
self.medusa_config = Some(config);
|
|
self
|
|
}
|
|
|
|
/// Enable self-speculative decoding
|
|
#[must_use]
|
|
pub fn with_self_speculative(mut self, config: SelfSpeculativeConfig) -> Self {
|
|
self.self_spec_config = Some(config);
|
|
self
|
|
}
|
|
|
|
/// Enable lookahead decoding
|
|
#[must_use]
|
|
pub fn with_lookahead(mut self, config: LookaheadConfig) -> Self {
|
|
let pool = NgramPool::new(config.ngram_size, config.max_pool_size);
|
|
self.ngram_pool = Some(Arc::new(RwLock::new(pool)));
|
|
self.lookahead_config = Some(config);
|
|
self
|
|
}
|
|
|
|
/// Get current metrics
|
|
pub async fn get_metrics(&self) -> AdvancedMetrics {
|
|
self.metrics.lock().await.clone()
|
|
}
|
|
|
|
/// Add tokens to n-gram pool for lookahead decoding
|
|
pub async fn observe_tokens(&self, tokens: &[u32]) {
|
|
if let Some(ref pool) = self.ngram_pool {
|
|
let mut pool = pool.write().await;
|
|
pool.add_observation(tokens);
|
|
}
|
|
}
|
|
|
|
/// Get lookahead predictions from n-gram pool
|
|
pub async fn get_lookahead_predictions(&self, context: &[u32]) -> Vec<u32> {
|
|
if let Some(ref pool) = self.ngram_pool {
|
|
let pool = pool.read().await;
|
|
pool.predict(context)
|
|
} else {
|
|
Vec::new()
|
|
}
|
|
}
|
|
}
|