Initial commit

crates/production/rtx-inference/src/speculative/advanced.rs

@@ -0,0 +1,103 @@
+//! 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()
+        }
+    }
+}