Initial commit

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

@@ -0,0 +1,113 @@
+//! Medusa-style prediction head configuration
+
+use super::tree::CandidateTree;
+use crate::InferenceResult;
+
+/// Medusa-style prediction head configuration
+/// Multiple heads predict tokens at different positions simultaneously
+#[derive(Debug, Clone)]
+pub struct MedusaConfig {
+    /// Number of prediction heads (each predicts one future position)
+    pub num_heads: usize,
+    /// Top-k candidates per head
+    pub top_k_per_head: usize,
+    /// Temperature for sampling from heads
+    pub temperature: f32,
+    /// Whether to use tree attention for verification
+    pub use_tree_attention: bool,
+    /// Typical acceptance threshold for Medusa heads
+    pub typical_acceptance: f32,
+}
+
+impl Default for MedusaConfig {
+    fn default() -> Self {
+        Self {
+            num_heads: 4,
+            top_k_per_head: 10,
+            temperature: 1.0,
+            use_tree_attention: true,
+            typical_acceptance: 0.6,
+        }
+    }
+}
+
+impl MedusaConfig {
+    /// Create config for fast inference (fewer heads, lower k)
+    #[must_use]
+    pub fn fast() -> Self {
+        Self {
+            num_heads: 2,
+            top_k_per_head: 5,
+            temperature: 0.8,
+            use_tree_attention: true,
+            typical_acceptance: 0.5,
+        }
+    }
+
+    /// Create config for high quality (more heads, higher k)
+    #[must_use]
+    pub fn high_quality() -> Self {
+        Self {
+            num_heads: 5,
+            top_k_per_head: 15,
+            temperature: 1.0,
+            use_tree_attention: true,
+            typical_acceptance: 0.7,
+        }
+    }
+}
+
+/// Trait for Medusa-style multi-head draft prediction
+#[async_trait::async_trait]
+pub trait MedusaDraftModel: Send + Sync {
+    /// Generate predictions from all heads simultaneously
+    /// Returns Vec<Vec<(`token_id`, probability)>> - one vec per head
+    async fn predict_all_heads(
+        &self,
+        context: &[u32],
+        config: &MedusaConfig,
+    ) -> InferenceResult<Vec<Vec<(u32, f32)>>>;
+
+    /// Build a candidate tree from head predictions
+    async fn build_candidate_tree(
+        &self,
+        context: &[u32],
+        config: &MedusaConfig,
+    ) -> InferenceResult<CandidateTree> {
+        let head_predictions = self.predict_all_heads(context, config).await?;
+
+        let mut tree = CandidateTree::new();
+
+        if head_predictions.is_empty() {
+            return Ok(tree);
+        }
+
+        // Add root candidates from first head
+        let first_head = &head_predictions[0];
+        let root_indices: Vec<usize> = first_head
+            .iter()
+            .take(config.top_k_per_head)
+            .map(|(token_id, prob)| tree.add_root(*token_id, *prob))
+            .collect();
+
+        // For each subsequent head, add children to existing nodes
+        let mut current_level = root_indices;
+
+        for head_preds in head_predictions.iter().skip(1) {
+            let mut next_level = Vec::new();
+
+            for &parent_idx in &current_level {
+                for (token_id, prob) in head_preds.iter().take(config.top_k_per_head) {
+                    let child_idx = tree.add_child(parent_idx, *token_id, *prob);
+                    next_level.push(child_idx);
+                }
+            }
+
+            current_level = next_level;
+        }
+
+        tree.total_candidates = tree.get_all_paths().len();
+
+        Ok(tree)
+    }
+}