/home/noah/src/realizar/src/speculative.rs
Line | Count | Source |
1 | | //! Speculative Decoding |
2 | | //! |
3 | | //! Per spec ยง8.3: Speculative decoding with draft model for up to 3x speedup. |
4 | | //! Reference: [6] Zheng et al. (2024) "SGLang: Efficient Execution of Structured LM Programs" |
5 | | //! |
6 | | //! ## Algorithm |
7 | | //! |
8 | | //! 1. Generate K speculative tokens with fast draft model |
9 | | //! 2. Verify all K tokens with single forward pass of target model |
10 | | //! 3. Accept tokens matching probability distribution |
11 | | //! 4. Resample at first rejection point |
12 | | //! |
13 | | //! ## Benefits |
14 | | //! |
15 | | //! - Up to 3x speedup for greedy decoding |
16 | | //! - Maintains output quality (mathematically equivalent) |
17 | | //! - Reduces memory-bound bottleneck of autoregressive decoding |
18 | | |
19 | | // Module-level clippy allows |
20 | | #![allow(clippy::must_use_candidate)] |
21 | | #![allow(clippy::return_self_not_must_use)] |
22 | | #![allow(clippy::missing_errors_doc)] |
23 | | |
24 | | use serde::{Deserialize, Serialize}; |
25 | | use std::time::Instant; |
26 | | use thiserror::Error; |
27 | | |
28 | | /// Error type for speculative decoding |
29 | | #[derive(Debug, Error)] |
30 | | pub enum SpeculativeError { |
31 | | /// Draft model error |
32 | | #[error("Draft model error: {0}")] |
33 | | DraftModelError(String), |
34 | | |
35 | | /// Target model error |
36 | | #[error("Target model error: {0}")] |
37 | | TargetModelError(String), |
38 | | |
39 | | /// Invalid speculation length |
40 | | #[error("Invalid speculation length: {0}")] |
41 | | InvalidSpecLength(usize), |
42 | | |
43 | | /// Verification failed |
44 | | #[error("Verification failed at position {position}")] |
45 | | VerificationFailed { |
46 | | /// Position where verification failed |
47 | | position: usize, |
48 | | }, |
49 | | } |
50 | | |
51 | | /// Speculative decoding statistics |
52 | | #[derive(Debug, Clone, Default, Serialize, Deserialize)] |
53 | | pub struct SpeculativeStats { |
54 | | /// Total speculative iterations |
55 | | pub iterations: u64, |
56 | | /// Total tokens speculated |
57 | | pub tokens_speculated: u64, |
58 | | /// Total tokens accepted |
59 | | pub tokens_accepted: u64, |
60 | | /// Average acceptance rate |
61 | | pub acceptance_rate: f32, |
62 | | /// Average speculation length |
63 | | pub avg_spec_length: f32, |
64 | | /// Total time saved (estimated) |
65 | | pub time_saved_ms: f64, |
66 | | /// Draft model time (ms) |
67 | | pub draft_time_ms: f64, |
68 | | /// Target model time (ms) |
69 | | pub target_time_ms: f64, |
70 | | } |
71 | | |
72 | | impl SpeculativeStats { |
73 | | /// Update stats after an iteration |
74 | 5 | pub fn record_iteration( |
75 | 5 | &mut self, |
76 | 5 | speculated: usize, |
77 | 5 | accepted: usize, |
78 | 5 | draft_ms: f64, |
79 | 5 | target_ms: f64, |
80 | 5 | ) { |
81 | 5 | self.iterations += 1; |
82 | 5 | self.tokens_speculated += speculated as u64; |
83 | 5 | self.tokens_accepted += accepted as u64; |
84 | 5 | self.draft_time_ms += draft_ms; |
85 | 5 | self.target_time_ms += target_ms; |
86 | | |
87 | | // Update running averages |
88 | 5 | if self.tokens_speculated > 0 { |
89 | 5 | self.acceptance_rate = self.tokens_accepted as f32 / self.tokens_speculated as f32; |
90 | 5 | }0 |
91 | 5 | if self.iterations > 0 { |
92 | 5 | self.avg_spec_length = self.tokens_speculated as f32 / self.iterations as f32; |
93 | 5 | }0 |
94 | | |
95 | | // Estimate time saved (accepted tokens would have required sequential target calls) |
96 | | // Each accepted token saves ~target_time/batch_size |
97 | 5 | let time_per_token = target_ms / speculated.max(1) as f64; |
98 | 5 | self.time_saved_ms += (accepted.saturating_sub(1)) as f64 * time_per_token; |
99 | 5 | } |
100 | | |
101 | | /// Get speedup ratio |
102 | 1 | pub fn speedup(&self) -> f32 { |
103 | 1 | if self.tokens_accepted == 0 { |
104 | 0 | return 1.0; |
105 | 1 | } |
106 | | // Speedup = (tokens * time_per_token) / actual_time |
107 | | // Assuming draft is ~10x faster than target |
108 | 1 | let draft_tokens_equivalent = self.tokens_speculated as f64 * 0.1; |
109 | 1 | let baseline_time = self.tokens_accepted as f64; |
110 | 1 | let actual_time = draft_tokens_equivalent + self.iterations as f64; |
111 | | |
112 | 1 | if actual_time > 0.0 { |
113 | 1 | (baseline_time / actual_time) as f32 |
114 | | } else { |
115 | 0 | 1.0 |
116 | | } |
117 | 1 | } |
118 | | } |
119 | | |
120 | | /// Token with probability |
121 | | #[derive(Debug, Clone)] |
122 | | pub struct TokenProb { |
123 | | /// Token ID |
124 | | pub token: u32, |
125 | | /// Log probability |
126 | | pub log_prob: f32, |
127 | | } |
128 | | |
129 | | impl TokenProb { |
130 | | /// Create a new token with probability |
131 | 27 | pub fn new(token: u32, log_prob: f32) -> Self { |
132 | 27 | Self { token, log_prob } |
133 | 27 | } |
134 | | |
135 | | /// Get probability (exp of log_prob) |
136 | 2 | pub fn prob(&self) -> f32 { |
137 | 2 | self.log_prob.exp() |
138 | 2 | } |
139 | | } |
140 | | |
141 | | /// Speculative decoding result |
142 | | #[derive(Debug, Clone)] |
143 | | pub struct SpeculativeResult { |
144 | | /// Accepted tokens |
145 | | pub accepted_tokens: Vec<u32>, |
146 | | /// Number of tokens speculated |
147 | | pub num_speculated: usize, |
148 | | /// Number of tokens accepted |
149 | | pub num_accepted: usize, |
150 | | /// Token that was resampled (if any) |
151 | | pub resampled_token: Option<u32>, |
152 | | /// Time taken for draft model (ms) |
153 | | pub draft_time_ms: f64, |
154 | | /// Time taken for target model (ms) |
155 | | pub target_time_ms: f64, |
156 | | } |
157 | | |
158 | | impl SpeculativeResult { |
159 | | /// Get acceptance rate for this iteration |
160 | 1 | pub fn acceptance_rate(&self) -> f32 { |
161 | 1 | if self.num_speculated == 0 { |
162 | 0 | return 0.0; |
163 | 1 | } |
164 | 1 | self.num_accepted as f32 / self.num_speculated as f32 |
165 | 1 | } |
166 | | |
167 | | /// Check if all tokens were accepted |
168 | 2 | pub fn all_accepted(&self) -> bool { |
169 | 2 | self.num_accepted == self.num_speculated |
170 | 2 | } |
171 | | } |
172 | | |
173 | | /// Trait for models that can be used in speculative decoding |
174 | | pub trait SpeculativeModel { |
175 | | /// Generate next token logits |
176 | | fn forward(&self, input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError>; |
177 | | |
178 | | /// Sample from logits |
179 | | fn sample(&self, logits: &[f32]) -> Result<TokenProb, SpeculativeError>; |
180 | | |
181 | | /// Get vocabulary size |
182 | | fn vocab_size(&self) -> usize; |
183 | | |
184 | | /// Get EOS token |
185 | | fn eos_token(&self) -> u32; |
186 | | } |
187 | | |
188 | | /// Speculative decoder with draft and target models |
189 | | pub struct SpeculativeDecoder<D: SpeculativeModel, T: SpeculativeModel> { |
190 | | /// Draft model (smaller, faster) |
191 | | draft: D, |
192 | | /// Target model (larger, slower) |
193 | | target: T, |
194 | | /// Number of tokens to speculate per iteration |
195 | | spec_length: usize, |
196 | | /// Statistics |
197 | | stats: SpeculativeStats, |
198 | | } |
199 | | |
200 | | impl<D: SpeculativeModel, T: SpeculativeModel> SpeculativeDecoder<D, T> { |
201 | | /// Create a new speculative decoder |
202 | 6 | pub fn new(draft: D, target: T, spec_length: usize) -> Result<Self, SpeculativeError> { |
203 | 6 | if spec_length == 0 || spec_length > 325 { |
204 | 1 | return Err(SpeculativeError::InvalidSpecLength(spec_length)); |
205 | 5 | } |
206 | | |
207 | 5 | Ok(Self { |
208 | 5 | draft, |
209 | 5 | target, |
210 | 5 | spec_length, |
211 | 5 | stats: SpeculativeStats::default(), |
212 | 5 | }) |
213 | 6 | } |
214 | | |
215 | | /// Get speculation length |
216 | 2 | pub fn spec_length(&self) -> usize { |
217 | 2 | self.spec_length |
218 | 2 | } |
219 | | |
220 | | /// Set speculation length |
221 | 2 | pub fn set_spec_length(&mut self, spec_length: usize) -> Result<(), SpeculativeError> { |
222 | 2 | if spec_length == 0 || spec_length > 321 { |
223 | 1 | return Err(SpeculativeError::InvalidSpecLength(spec_length)); |
224 | 1 | } |
225 | 1 | self.spec_length = spec_length; |
226 | 1 | Ok(()) |
227 | 2 | } |
228 | | |
229 | | /// Generate one iteration of speculative decoding |
230 | 3 | pub fn decode_iteration( |
231 | 3 | &mut self, |
232 | 3 | context: &[u32], |
233 | 3 | ) -> Result<SpeculativeResult, SpeculativeError> { |
234 | 3 | let mut accepted_tokens = Vec::new(); |
235 | 3 | let mut draft_tokens = Vec::new(); |
236 | 3 | let mut draft_probs = Vec::new(); |
237 | | |
238 | | // 1. Generate speculative tokens with draft model |
239 | 3 | let draft_start = Instant::now(); |
240 | 3 | let mut current_context = context.to_vec(); |
241 | | |
242 | 3 | for _ in 0..self.spec_length { |
243 | 12 | let logits = self |
244 | 12 | .draft |
245 | 12 | .forward(¤t_context) |
246 | 12 | .map_err(|e| SpeculativeError::DraftModelError(e0 .to_string0 ()))?0 ; |
247 | 12 | let token_prob = self |
248 | 12 | .draft |
249 | 12 | .sample(&logits) |
250 | 12 | .map_err(|e| SpeculativeError::DraftModelError(e0 .to_string0 ()))?0 ; |
251 | | |
252 | 12 | let token = token_prob.token; |
253 | 12 | draft_tokens.push(token); |
254 | 12 | draft_probs.push(token_prob); |
255 | 12 | current_context.push(token); |
256 | | |
257 | | // Stop if EOS |
258 | 12 | if token == self.draft.eos_token() { |
259 | 0 | break; |
260 | 12 | } |
261 | | } |
262 | 3 | let draft_time = draft_start.elapsed(); |
263 | | |
264 | | // 2. Verify with target model (single forward pass for all positions) |
265 | 3 | let target_start = Instant::now(); |
266 | 3 | let mut verify_context = context.to_vec(); |
267 | 3 | verify_context.extend(&draft_tokens); |
268 | | |
269 | | // Get target logits for verification |
270 | 3 | let target_logits = self |
271 | 3 | .target |
272 | 3 | .forward(&verify_context) |
273 | 3 | .map_err(|e| SpeculativeError::TargetModelError(e0 .to_string0 ()))?0 ; |
274 | | |
275 | 3 | let target_time = target_start.elapsed(); |
276 | | |
277 | | // 3. Accept/reject based on probability matching |
278 | 3 | let mut resampled_token = None; |
279 | | |
280 | 15 | for draft_prob12 in &draft_probs { |
281 | | // Simple acceptance: compare draft and target probabilities |
282 | | // In practice, use more sophisticated rejection sampling |
283 | 12 | let target_token = self |
284 | 12 | .target |
285 | 12 | .sample(&target_logits) |
286 | 12 | .map_err(|e| SpeculativeError::TargetModelError(e0 .to_string0 ()))?0 ; |
287 | | |
288 | | // Accept if draft token matches target distribution |
289 | | // Simplified: accept if tokens match or random acceptance |
290 | 12 | if self.should_accept(draft_prob, &target_token) { |
291 | 12 | accepted_tokens.push(draft_prob.token); |
292 | 12 | } else { |
293 | | // Resample from target distribution |
294 | 0 | resampled_token = Some(target_token.token); |
295 | 0 | accepted_tokens.push(target_token.token); |
296 | 0 | break; |
297 | | } |
298 | | } |
299 | | |
300 | 3 | let num_speculated = draft_tokens.len(); |
301 | 3 | let num_accepted = accepted_tokens.len(); |
302 | | |
303 | | // Update stats |
304 | 3 | self.stats.record_iteration( |
305 | 3 | num_speculated, |
306 | 3 | num_accepted, |
307 | 3 | draft_time.as_secs_f64() * 1000.0, |
308 | 3 | target_time.as_secs_f64() * 1000.0, |
309 | | ); |
310 | | |
311 | 3 | Ok(SpeculativeResult { |
312 | 3 | accepted_tokens, |
313 | 3 | num_speculated, |
314 | 3 | num_accepted, |
315 | 3 | resampled_token, |
316 | 3 | draft_time_ms: draft_time.as_secs_f64() * 1000.0, |
317 | 3 | target_time_ms: target_time.as_secs_f64() * 1000.0, |
318 | 3 | }) |
319 | 3 | } |
320 | | |
321 | | /// Acceptance criterion for speculative decoding |
322 | | #[allow(clippy::unused_self)] // Will use self for config-based thresholds |
323 | 12 | fn should_accept(&self, draft: &TokenProb, target: &TokenProb) -> bool { |
324 | | // Simple acceptance: tokens must match |
325 | | // More sophisticated: use probability ratio for rejection sampling |
326 | 12 | if draft.token == target.token { |
327 | 12 | return true; |
328 | 0 | } |
329 | | |
330 | | // Probabilistic acceptance based on ratio |
331 | 0 | let ratio = target.prob() / draft.prob().max(1e-10); |
332 | 0 | ratio >= 1.0 || ratio > 0.5 // Simplified threshold |
333 | 12 | } |
334 | | |
335 | | /// Get statistics |
336 | 2 | pub fn stats(&self) -> &SpeculativeStats { |
337 | 2 | &self.stats |
338 | 2 | } |
339 | | |
340 | | /// Reset statistics |
341 | 1 | pub fn reset_stats(&mut self) { |
342 | 1 | self.stats = SpeculativeStats::default(); |
343 | 1 | } |
344 | | } |
345 | | |
346 | | /// Configuration for speculative decoding |
347 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
348 | | pub struct SpeculativeConfig { |
349 | | /// Number of tokens to speculate |
350 | | pub spec_length: usize, |
351 | | /// Minimum acceptance rate before adapting spec_length |
352 | | pub min_acceptance_rate: f32, |
353 | | /// Enable adaptive speculation length |
354 | | pub adaptive: bool, |
355 | | /// Maximum speculation length for adaptive mode |
356 | | pub max_spec_length: usize, |
357 | | } |
358 | | |
359 | | impl Default for SpeculativeConfig { |
360 | 5 | fn default() -> Self { |
361 | 5 | Self { |
362 | 5 | spec_length: 4, |
363 | 5 | min_acceptance_rate: 0.5, |
364 | 5 | adaptive: true, |
365 | 5 | max_spec_length: 8, |
366 | 5 | } |
367 | 5 | } |
368 | | } |
369 | | |
370 | | impl SpeculativeConfig { |
371 | | /// Create a new config with default values |
372 | 4 | pub fn new() -> Self { |
373 | 4 | Self::default() |
374 | 4 | } |
375 | | |
376 | | /// Set speculation length |
377 | 4 | pub fn with_spec_length(mut self, spec_length: usize) -> Self { |
378 | 4 | self.spec_length = spec_length; |
379 | 4 | self |
380 | 4 | } |
381 | | |
382 | | /// Enable/disable adaptive mode |
383 | 2 | pub fn with_adaptive(mut self, adaptive: bool) -> Self { |
384 | 2 | self.adaptive = adaptive; |
385 | 2 | self |
386 | 2 | } |
387 | | |
388 | | /// Adapt speculation length based on acceptance rate |
389 | 3 | pub fn adapt_spec_length(&mut self, acceptance_rate: f32) { |
390 | 3 | if !self.adaptive { |
391 | 1 | return; |
392 | 2 | } |
393 | | |
394 | 2 | if acceptance_rate > 0.8 && self.spec_length < self.max_spec_length1 { |
395 | 1 | // High acceptance: increase speculation |
396 | 1 | self.spec_length = (self.spec_length + 1).min(self.max_spec_length); |
397 | 1 | } else if acceptance_rate < self.min_acceptance_rate && self.spec_length > 1 { |
398 | 1 | // Low acceptance: decrease speculation |
399 | 1 | self.spec_length = (self.spec_length - 1).max(1); |
400 | 1 | }0 |
401 | 3 | } |
402 | | } |
403 | | |
404 | | // ============================================================================ |
405 | | // Tests |
406 | | // ============================================================================ |
407 | | |
408 | | #[cfg(test)] |
409 | | mod tests { |
410 | | use super::*; |
411 | | |
412 | | /// Mock model for testing |
413 | | struct MockModel { |
414 | | vocab_size: usize, |
415 | | eos_token: u32, |
416 | | /// Fixed token to return |
417 | | fixed_token: u32, |
418 | | } |
419 | | |
420 | | impl MockModel { |
421 | 12 | fn new(vocab_size: usize, fixed_token: u32) -> Self { |
422 | 12 | Self { |
423 | 12 | vocab_size, |
424 | 12 | eos_token: 0, |
425 | 12 | fixed_token, |
426 | 12 | } |
427 | 12 | } |
428 | | } |
429 | | |
430 | | impl SpeculativeModel for MockModel { |
431 | 15 | fn forward(&self, _input_ids: &[u32]) -> Result<Vec<f32>, SpeculativeError> { |
432 | | // Return uniform logits |
433 | 15 | Ok(vec![0.0; self.vocab_size]) |
434 | 15 | } |
435 | | |
436 | 24 | fn sample(&self, _logits: &[f32]) -> Result<TokenProb, SpeculativeError> { |
437 | 24 | Ok(TokenProb::new(self.fixed_token, -1.0)) |
438 | 24 | } |
439 | | |
440 | 0 | fn vocab_size(&self) -> usize { |
441 | 0 | self.vocab_size |
442 | 0 | } |
443 | | |
444 | 12 | fn eos_token(&self) -> u32 { |
445 | 12 | self.eos_token |
446 | 12 | } |
447 | | } |
448 | | |
449 | | // === TokenProb Tests === |
450 | | |
451 | | #[test] |
452 | 1 | fn test_token_prob_new() { |
453 | 1 | let tp = TokenProb::new(42, -1.0); |
454 | 1 | assert_eq!(tp.token, 42); |
455 | 1 | assert_eq!(tp.log_prob, -1.0); |
456 | 1 | } |
457 | | |
458 | | #[test] |
459 | 1 | fn test_token_prob_prob() { |
460 | 1 | let tp = TokenProb::new(42, 0.0); |
461 | 1 | assert!((tp.prob() - 1.0).abs() < 0.001); |
462 | | |
463 | 1 | let tp2 = TokenProb::new(42, -1.0); |
464 | 1 | assert!((tp2.prob() - 0.368).abs() < 0.01); |
465 | 1 | } |
466 | | |
467 | | // === SpeculativeStats Tests === |
468 | | |
469 | | #[test] |
470 | 1 | fn test_speculative_stats_default() { |
471 | 1 | let stats = SpeculativeStats::default(); |
472 | 1 | assert_eq!(stats.iterations, 0); |
473 | 1 | assert_eq!(stats.tokens_speculated, 0); |
474 | 1 | assert_eq!(stats.acceptance_rate, 0.0); |
475 | 1 | } |
476 | | |
477 | | #[test] |
478 | 1 | fn test_speculative_stats_record() { |
479 | 1 | let mut stats = SpeculativeStats::default(); |
480 | 1 | stats.record_iteration(4, 3, 1.0, 10.0); |
481 | | |
482 | 1 | assert_eq!(stats.iterations, 1); |
483 | 1 | assert_eq!(stats.tokens_speculated, 4); |
484 | 1 | assert_eq!(stats.tokens_accepted, 3); |
485 | 1 | assert_eq!(stats.acceptance_rate, 0.75); |
486 | 1 | } |
487 | | |
488 | | #[test] |
489 | 1 | fn test_speculative_stats_speedup() { |
490 | 1 | let mut stats = SpeculativeStats::default(); |
491 | 1 | stats.record_iteration(4, 4, 1.0, 10.0); |
492 | | |
493 | 1 | let speedup = stats.speedup(); |
494 | 1 | assert!(speedup > 1.0); |
495 | 1 | } |
496 | | |
497 | | #[test] |
498 | 1 | fn test_speculative_stats_serialization() { |
499 | 1 | let stats = SpeculativeStats { |
500 | 1 | iterations: 10, |
501 | 1 | tokens_speculated: 40, |
502 | 1 | tokens_accepted: 30, |
503 | 1 | acceptance_rate: 0.75, |
504 | 1 | avg_spec_length: 4.0, |
505 | 1 | time_saved_ms: 100.0, |
506 | 1 | draft_time_ms: 10.0, |
507 | 1 | target_time_ms: 100.0, |
508 | 1 | }; |
509 | | |
510 | 1 | let json = serde_json::to_string(&stats).expect("test"); |
511 | 1 | let parsed: SpeculativeStats = serde_json::from_str(&json).expect("test"); |
512 | | |
513 | 1 | assert_eq!(parsed.iterations, stats.iterations); |
514 | 1 | assert_eq!(parsed.acceptance_rate, stats.acceptance_rate); |
515 | 1 | } |
516 | | |
517 | | // === SpeculativeResult Tests === |
518 | | |
519 | | #[test] |
520 | 1 | fn test_speculative_result_acceptance_rate() { |
521 | 1 | let result = SpeculativeResult { |
522 | 1 | accepted_tokens: vec![1, 2, 3], |
523 | 1 | num_speculated: 4, |
524 | 1 | num_accepted: 3, |
525 | 1 | resampled_token: Some(4), |
526 | 1 | draft_time_ms: 1.0, |
527 | 1 | target_time_ms: 10.0, |
528 | 1 | }; |
529 | | |
530 | 1 | assert_eq!(result.acceptance_rate(), 0.75); |
531 | 1 | assert!(!result.all_accepted()); |
532 | 1 | } |
533 | | |
534 | | #[test] |
535 | 1 | fn test_speculative_result_all_accepted() { |
536 | 1 | let result = SpeculativeResult { |
537 | 1 | accepted_tokens: vec![1, 2, 3, 4], |
538 | 1 | num_speculated: 4, |
539 | 1 | num_accepted: 4, |
540 | 1 | resampled_token: None, |
541 | 1 | draft_time_ms: 1.0, |
542 | 1 | target_time_ms: 10.0, |
543 | 1 | }; |
544 | | |
545 | 1 | assert!(result.all_accepted()); |
546 | 1 | } |
547 | | |
548 | | // === SpeculativeConfig Tests === |
549 | | |
550 | | #[test] |
551 | 1 | fn test_speculative_config_default() { |
552 | 1 | let config = SpeculativeConfig::default(); |
553 | 1 | assert_eq!(config.spec_length, 4); |
554 | 1 | assert!(config.adaptive); |
555 | 1 | } |
556 | | |
557 | | #[test] |
558 | 1 | fn test_speculative_config_builder() { |
559 | 1 | let config = SpeculativeConfig::new() |
560 | 1 | .with_spec_length(6) |
561 | 1 | .with_adaptive(false); |
562 | | |
563 | 1 | assert_eq!(config.spec_length, 6); |
564 | 1 | assert!(!config.adaptive); |
565 | 1 | } |
566 | | |
567 | | #[test] |
568 | 1 | fn test_speculative_config_adapt_increase() { |
569 | 1 | let mut config = SpeculativeConfig::new().with_spec_length(4); |
570 | 1 | config.adapt_spec_length(0.9); // High acceptance |
571 | | |
572 | 1 | assert!(config.spec_length >= 4); |
573 | 1 | } |
574 | | |
575 | | #[test] |
576 | 1 | fn test_speculative_config_adapt_decrease() { |
577 | 1 | let mut config = SpeculativeConfig::new().with_spec_length(4); |
578 | 1 | config.adapt_spec_length(0.3); // Low acceptance |
579 | | |
580 | 1 | assert!(config.spec_length <= 4); |
581 | 1 | } |
582 | | |
583 | | #[test] |
584 | 1 | fn test_speculative_config_no_adapt_when_disabled() { |
585 | 1 | let mut config = SpeculativeConfig::new() |
586 | 1 | .with_spec_length(4) |
587 | 1 | .with_adaptive(false); |
588 | | |
589 | 1 | config.adapt_spec_length(0.1); |
590 | 1 | assert_eq!(config.spec_length, 4); // Should not change |
591 | 1 | } |
592 | | |
593 | | // === SpeculativeDecoder Tests === |
594 | | |
595 | | #[test] |
596 | 1 | fn test_speculative_decoder_new() { |
597 | 1 | let draft = MockModel::new(100, 1); |
598 | 1 | let target = MockModel::new(100, 1); |
599 | | |
600 | 1 | let decoder = SpeculativeDecoder::new(draft, target, 4).expect("test"); |
601 | 1 | assert_eq!(decoder.spec_length(), 4); |
602 | 1 | } |
603 | | |
604 | | #[test] |
605 | 1 | fn test_speculative_decoder_invalid_spec_length() { |
606 | 1 | let draft = MockModel::new(100, 1); |
607 | 1 | let target = MockModel::new(100, 1); |
608 | | |
609 | 1 | let result = SpeculativeDecoder::new(draft, target, 0); |
610 | 1 | assert!(matches!0 ( |
611 | 1 | result, |
612 | | Err(SpeculativeError::InvalidSpecLength(0)) |
613 | | )); |
614 | 1 | } |
615 | | |
616 | | #[test] |
617 | 1 | fn test_speculative_decoder_decode_iteration() { |
618 | 1 | let draft = MockModel::new(100, 1); |
619 | 1 | let target = MockModel::new(100, 1); |
620 | | |
621 | 1 | let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test"); |
622 | 1 | let result = decoder.decode_iteration(&[10, 20, 30]).expect("test"); |
623 | | |
624 | 1 | assert!(!result.accepted_tokens.is_empty()); |
625 | 1 | assert!(result.num_speculated > 0); |
626 | 1 | } |
627 | | |
628 | | #[test] |
629 | 1 | fn test_speculative_decoder_set_spec_length() { |
630 | 1 | let draft = MockModel::new(100, 1); |
631 | 1 | let target = MockModel::new(100, 1); |
632 | | |
633 | 1 | let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test"); |
634 | 1 | decoder.set_spec_length(8).expect("test"); |
635 | 1 | assert_eq!(decoder.spec_length(), 8); |
636 | | |
637 | 1 | let err = decoder.set_spec_length(0); |
638 | 1 | assert!(err.is_err()); |
639 | 1 | } |
640 | | |
641 | | #[test] |
642 | 1 | fn test_speculative_decoder_stats() { |
643 | 1 | let draft = MockModel::new(100, 1); |
644 | 1 | let target = MockModel::new(100, 1); |
645 | | |
646 | 1 | let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test"); |
647 | 1 | let _ = decoder.decode_iteration(&[10]).expect("test"); |
648 | | |
649 | 1 | let stats = decoder.stats(); |
650 | 1 | assert_eq!(stats.iterations, 1); |
651 | 1 | } |
652 | | |
653 | | #[test] |
654 | 1 | fn test_speculative_decoder_reset_stats() { |
655 | 1 | let draft = MockModel::new(100, 1); |
656 | 1 | let target = MockModel::new(100, 1); |
657 | | |
658 | 1 | let mut decoder = SpeculativeDecoder::new(draft, target, 4).expect("test"); |
659 | 1 | let _ = decoder.decode_iteration(&[10]).expect("test"); |
660 | | |
661 | 1 | decoder.reset_stats(); |
662 | 1 | assert_eq!(decoder.stats().iterations, 0); |
663 | 1 | } |
664 | | |
665 | | // === Error Tests === |
666 | | |
667 | | #[test] |
668 | 1 | fn test_speculative_error_display() { |
669 | 1 | let err = SpeculativeError::DraftModelError("test".to_string()); |
670 | 1 | assert!(err.to_string().contains("Draft")); |
671 | | |
672 | 1 | let err = SpeculativeError::TargetModelError("test".to_string()); |
673 | 1 | assert!(err.to_string().contains("Target")); |
674 | | |
675 | 1 | let err = SpeculativeError::InvalidSpecLength(0); |
676 | 1 | assert!(err.to_string().contains("0")); |
677 | | |
678 | 1 | let err = SpeculativeError::VerificationFailed { position: 3 }; |
679 | 1 | assert!(err.to_string().contains("3")); |
680 | 1 | } |
681 | | } |