Initial commit

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

@@ -0,0 +1,138 @@
+//! Configuration types for speculative decoding
+
+use super::types::DraftModelType;
+use crate::{InferenceError, InferenceResult};
+
+/// Adaptive configuration for draft length optimization
+#[derive(Debug, Clone)]
+pub struct AdaptiveConfig {
+    /// Minimum draft length to try
+    pub min_length: usize,
+    /// Maximum draft length to try
+    pub max_length: usize,
+    /// Target acceptance rate for adaptation
+    pub target_acceptance_rate: f32,
+    /// How aggressively to adjust draft length
+    pub adjustment_factor: f32,
+}
+
+impl Default for AdaptiveConfig {
+    fn default() -> Self {
+        Self {
+            min_length: 1,
+            max_length: 8,
+            target_acceptance_rate: 0.8,
+            adjustment_factor: 0.1,
+        }
+    }
+}
+
+/// Configuration for speculative decoding
+#[derive(Debug, Clone)]
+pub struct SpeculativeConfig {
+    /// Type of draft model to use
+    pub draft_model_type: DraftModelType,
+    /// Maximum number of tokens to generate in draft
+    pub max_draft_tokens: usize,
+    /// Minimum confidence threshold for acceptance
+    pub acceptance_threshold: f32,
+    /// Number of additional lookahead tokens
+    pub lookahead_tokens: usize,
+    /// Whether to track performance metrics
+    pub performance_tracking: bool,
+    /// Adaptive draft length configuration
+    pub adaptive_config: Option<AdaptiveConfig>,
+}
+
+impl SpeculativeConfig {
+    /// Create a new configuration with sensible defaults
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            draft_model_type: DraftModelType::SmallModel("default".to_string()),
+            max_draft_tokens: 4,
+            acceptance_threshold: 0.7,
+            lookahead_tokens: 1,
+            performance_tracking: false,
+            adaptive_config: None,
+        }
+    }
+
+    /// Set the draft model type
+    #[must_use]
+    pub fn with_draft_model(mut self, model_type: DraftModelType) -> Self {
+        self.draft_model_type = model_type;
+        self
+    }
+
+    /// Set maximum draft tokens
+    #[must_use]
+    pub fn with_max_draft_tokens(mut self, max_tokens: usize) -> Self {
+        self.max_draft_tokens = max_tokens;
+        self
+    }
+
+    /// Set acceptance threshold
+    #[must_use]
+    pub fn with_acceptance_threshold(mut self, threshold: f32) -> Self {
+        self.acceptance_threshold = threshold;
+        self
+    }
+
+    /// Set lookahead tokens
+    #[must_use]
+    pub fn with_lookahead_tokens(mut self, lookahead: usize) -> Self {
+        self.lookahead_tokens = lookahead;
+        self
+    }
+
+    /// Enable performance tracking
+    #[must_use]
+    pub fn with_performance_tracking(mut self, enabled: bool) -> Self {
+        self.performance_tracking = enabled;
+        self
+    }
+
+    /// Enable adaptive draft length
+    #[must_use]
+    pub fn with_adaptive_draft_length(mut self, config: AdaptiveConfig) -> Self {
+        self.adaptive_config = Some(config);
+        self
+    }
+
+    /// Validate configuration parameters
+    pub fn validate(&self) -> InferenceResult<()> {
+        if self.max_draft_tokens == 0 {
+            return Err(InferenceError::invalid_request(
+                "max_draft_tokens must be greater than 0",
+            ));
+        }
+
+        if !(0.0..=1.0).contains(&self.acceptance_threshold) {
+            return Err(InferenceError::invalid_request(
+                "acceptance_threshold must be between 0.0 and 1.0",
+            ));
+        }
+
+        if let Some(ref adaptive) = self.adaptive_config {
+            if adaptive.min_length == 0 || adaptive.min_length > adaptive.max_length {
+                return Err(InferenceError::invalid_request(
+                    "adaptive config: min_length must be > 0 and <= max_length",
+                ));
+            }
+            if !(0.0..=1.0).contains(&adaptive.target_acceptance_rate) {
+                return Err(InferenceError::invalid_request(
+                    "adaptive config: target_acceptance_rate must be between 0.0 and 1.0",
+                ));
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl Default for SpeculativeConfig {
+    fn default() -> Self {
+        Self::new()
+    }
+}