Skip to main content

ferrum_engine/
speculative.rs

1//! Speculative decoding — draft + verify.
2//!
3//! A small ("draft") model generates `N` candidate tokens autoregressively.
4//! The big ("target") model does ONE forward pass on prompt + drafts and
5//! produces `N+1` logit distributions — one for every position at which a
6//! draft could be accepted, plus one bonus position after full acceptance.
7//! We then apply the DeepMind speculative sampling rule to decide which
8//! drafts survive, yielding 1..=N+1 tokens per target forward pass.
9//!
10//! Original paper: Leviathan et al., 2023 — "Fast Inference from Transformers
11//! via Speculative Decoding" (https://arxiv.org/abs/2211.17192).
12//!
13//! This module is **algorithm-only**: it operates on raw logit vectors and
14//! produces a list of accepted token ids. The engine/scheduler integration
15//! (draft-model loading, KV-cache management, iteration plumbing) is a
16//! separate layer and explicitly out of scope for this file — wiring it is
17//! the follow-up once the algorithm is locked down.
18
19use ferrum_interfaces::{
20    model_executor::{DecodeInput, DecodeOutput},
21    tensor::TensorFactory,
22    KvCacheHandle, ModelExecutor,
23};
24use ferrum_types::{Result, TokenId};
25use rand::RngCore;
26use std::sync::Arc;
27
28/// Softmax a logit vector, in-place modifications avoided — returns new Vec.
29/// `temperature == 0.0` collapses to one-hot at argmax (as elsewhere in the
30/// sampler stack).
31fn softmax(logits: &[f32], temperature: f32) -> Vec<f32> {
32    if temperature == 0.0 {
33        // Greedy: delta at argmax.
34        let (argmax, _) =
35            logits
36                .iter()
37                .enumerate()
38                .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| {
39                    if v > bv {
40                        (i, v)
41                    } else {
42                        (bi, bv)
43                    }
44                });
45        let mut p = vec![0.0f32; logits.len()];
46        p[argmax] = 1.0;
47        return p;
48    }
49    let inv_t = 1.0 / temperature;
50    let mut max = f32::NEG_INFINITY;
51    for &l in logits {
52        if l > max {
53            max = l;
54        }
55    }
56    if !max.is_finite() {
57        max = 0.0;
58    }
59    let mut sum = 0.0f64;
60    let mut out = Vec::with_capacity(logits.len());
61    for &l in logits {
62        let e = ((l - max) * inv_t).exp();
63        out.push(e);
64        sum += e as f64;
65    }
66    let inv_sum = (1.0 / sum) as f32;
67    for x in out.iter_mut() {
68        *x *= inv_sum;
69    }
70    out
71}
72
73/// Sample a token id from a probability distribution using uniform `u` in [0, 1).
74fn sample_categorical(probs: &[f32], u: f32) -> TokenId {
75    let mut acc = 0.0f32;
76    for (i, &p) in probs.iter().enumerate() {
77        acc += p;
78        if u <= acc {
79            return TokenId::new(i as u32);
80        }
81    }
82    TokenId::new((probs.len().saturating_sub(1)) as u32)
83}
84
85/// Draw a fresh `u ~ U[0,1)` from any RngCore.
86fn next_u(rng: &mut dyn RngCore) -> f32 {
87    // Convert a u32 to f32 in [0, 1): divide by 2^32.
88    (rng.next_u32() as f64 / (u32::MAX as f64 + 1.0)) as f32
89}
90
91/// Residual distribution (p_T - p_D) clipped at zero and renormalised.
92/// This is the distribution we sample from when a draft is rejected.
93fn residual(p_target: &[f32], p_draft: &[f32]) -> Vec<f32> {
94    debug_assert_eq!(p_target.len(), p_draft.len());
95    let mut r = Vec::with_capacity(p_target.len());
96    let mut sum = 0.0f64;
97    for (&pt, &pd) in p_target.iter().zip(p_draft.iter()) {
98        let d = (pt - pd).max(0.0);
99        r.push(d);
100        sum += d as f64;
101    }
102    if sum <= 0.0 {
103        // Shouldn't happen unless the target and draft fully agree on the
104        // draft token (probability 1, which the accept rule already catches).
105        // Fall back to the target distribution.
106        return p_target.to_vec();
107    }
108    let inv = (1.0 / sum) as f32;
109    for x in r.iter_mut() {
110        *x *= inv;
111    }
112    r
113}
114
115/// Input to `verify_speculation`: logit vectors at each speculation position.
116///
117/// For a speculation of `N` draft tokens:
118/// - `draft_logits.len() == N` — the distribution the draft model used to
119///   sample each draft token (aligned with `draft_tokens`).
120/// - `target_logits.len() == N + 1` — the target model's distributions at
121///   each position; the extra slot is the bonus-token draw after all drafts
122///   are accepted.
123/// - `draft_tokens.len() == N`.
124pub struct Speculation<'a> {
125    pub draft_tokens: &'a [TokenId],
126    pub draft_logits: &'a [Vec<f32>],
127    pub target_logits: &'a [Vec<f32>],
128    pub temperature: f32,
129}
130
131/// Result of one speculate+verify round.
132#[derive(Debug, Clone, PartialEq)]
133pub struct SpeculationOutcome {
134    /// Tokens that survived verification — always at least 1 (either the
135    /// residual resample on a rejection or the bonus token on full accept).
136    pub tokens: Vec<TokenId>,
137    /// Index of the first rejected draft (0..N for rejection, N for full
138    /// accept). Useful for KV-cache rollback in the engine.
139    pub rejected_at: usize,
140}
141
142/// Execute the DeepMind speculative-sampling accept/reject loop.
143///
144/// Per-draft decision (draft i, i = 0..N):
145///   - Let p_T = target prob of draft[i] at position i.
146///   - Let p_D = draft prob of draft[i] at position i.
147///   - Accept with probability `min(1, p_T / p_D)`.
148///   - If rejected, sample a replacement from the residual `(p_T - p_D)+`.
149///
150/// If every draft is accepted, sample a bonus token from `target_logits[N]`.
151pub fn verify_speculation(
152    spec: Speculation<'_>,
153    rng: &mut dyn RngCore,
154) -> Result<SpeculationOutcome> {
155    let n = spec.draft_tokens.len();
156    assert_eq!(spec.draft_logits.len(), n, "draft_logits count mismatch");
157    assert_eq!(
158        spec.target_logits.len(),
159        n + 1,
160        "target_logits must have N+1 rows (positions 0..N)"
161    );
162
163    let mut accepted: Vec<TokenId> = Vec::with_capacity(n + 1);
164
165    for i in 0..n {
166        let draft_token = spec.draft_tokens[i];
167        let idx = draft_token.get() as usize;
168
169        let p_target = softmax(&spec.target_logits[i], spec.temperature);
170        let p_draft = softmax(&spec.draft_logits[i], spec.temperature);
171
172        if idx >= p_target.len() || idx >= p_draft.len() {
173            // Malformed input — fall back to a greedy target pick and stop.
174            let t = TokenId::new(
175                p_target
176                    .iter()
177                    .enumerate()
178                    .fold((0, f32::NEG_INFINITY), |(bi, bv), (j, &v)| {
179                        if v > bv {
180                            (j, v)
181                        } else {
182                            (bi, bv)
183                        }
184                    })
185                    .0 as u32,
186            );
187            accepted.push(t);
188            return Ok(SpeculationOutcome {
189                tokens: accepted,
190                rejected_at: i,
191            });
192        }
193
194        let pt = p_target[idx];
195        let pd = p_draft[idx].max(1e-20); // avoid division by zero
196        let ratio = (pt / pd).min(1.0);
197        let u = next_u(rng);
198        if u < ratio {
199            // Accepted.
200            accepted.push(draft_token);
201        } else {
202            // Rejected: sample from residual (p_T - p_D)+.
203            let res = residual(&p_target, &p_draft);
204            let replacement = sample_categorical(&res, next_u(rng));
205            accepted.push(replacement);
206            return Ok(SpeculationOutcome {
207                tokens: accepted,
208                rejected_at: i,
209            });
210        }
211    }
212
213    // All drafts accepted — sample a bonus token from the target's trailing
214    // distribution so the round always produces at least one NEW token even
215    // if the draft happened to end exactly where generation should stop.
216    let p_bonus = softmax(&spec.target_logits[n], spec.temperature);
217    let bonus = sample_categorical(&p_bonus, next_u(rng));
218    accepted.push(bonus);
219    Ok(SpeculationOutcome {
220        tokens: accepted,
221        rejected_at: n,
222    })
223}
224
225/// Configuration for speculative decoding.
226#[derive(Debug, Clone)]
227pub struct SpeculativeDecodingConfig {
228    /// Draft model produces this many tokens per target forward pass.
229    /// Typical range: 3-7. Paper used 4-5. Larger N amortises target cost
230    /// more aggressively but raises the probability of early rejection.
231    pub num_speculative_tokens: usize,
232    /// Temperature applied during accept/reject. Matches the sampling
233    /// temperature to keep the target-draft ratio calibrated.
234    pub temperature: f32,
235}
236
237impl Default for SpeculativeDecodingConfig {
238    fn default() -> Self {
239        Self {
240            num_speculative_tokens: 4,
241            temperature: 1.0,
242        }
243    }
244}
245
246/// Drives one round of speculative decoding against a (draft, target) pair
247/// of `ModelExecutor`s. Owns neither executor — the engine keeps them and
248/// hands references per call.
249///
250/// Usage per decode iteration:
251///   - Caller hands the runner the last sampled token for this request,
252///     plus the draft & target KV cache handles.
253///   - `step()` runs N draft decodes, N+1 target decodes (sequentially for
254///     now — performance gain lands once the executors grow multi-position
255///     decode support), then applies `verify_speculation`.
256///   - Returns the list of newly accepted tokens plus the updated KV
257///     handles for draft and target (caller installs them on the sequence
258///     state for the next iteration).
259pub struct SpeculativeRunner<'a> {
260    pub draft: &'a dyn ModelExecutor,
261    pub target: &'a dyn ModelExecutor,
262    pub tensor_factory: Arc<dyn TensorFactory>,
263    pub cfg: SpeculativeDecodingConfig,
264}
265
266/// Result of a single `SpeculativeRunner::step`.
267pub struct SpeculativeStepOutcome {
268    pub tokens: Vec<TokenId>,
269    pub draft_kv: Arc<dyn KvCacheHandle>,
270    pub target_kv: Arc<dyn KvCacheHandle>,
271    /// True when a draft was rejected (caller may want to roll target KV
272    /// back to `rejected_at` to stay token-aligned with the draft model).
273    pub rejected: bool,
274    pub rejected_at: usize,
275    /// Draft KV is `N` writes; target KV is `N+1` writes. In the full-accept
276    /// case the draft is exactly one position behind target. Caller
277    /// catches up by feeding this token (the final draft input to target,
278    /// i.e. `draft_tokens[N-1]`) into the draft executor once. `None` on
279    /// rejection (caller handles rollback separately).
280    pub draft_catchup_token: Option<TokenId>,
281}
282
283impl<'a> SpeculativeRunner<'a> {
284    /// Run one draft+verify cycle. `last_token` is the token that was sampled
285    /// from the previous iteration (or the end of prefill) — it's the
286    /// starting point both executors advance from.
287    pub async fn step(
288        &self,
289        last_token: TokenId,
290        draft_kv: Arc<dyn KvCacheHandle>,
291        target_kv: Arc<dyn KvCacheHandle>,
292        rng: &mut (dyn RngCore + Send),
293    ) -> Result<SpeculativeStepOutcome> {
294        let n = self.cfg.num_speculative_tokens.max(1);
295
296        // ── Draft: N sequential decodes, one token at a time ─────────────
297        let mut draft_tokens: Vec<TokenId> = Vec::with_capacity(n);
298        let mut draft_logits: Vec<Vec<f32>> = Vec::with_capacity(n);
299        let mut draft_kv_cur = draft_kv;
300        let mut draft_prev_token = last_token;
301        for _ in 0..n {
302            let input_tensor = tokens_to_tensor(&self.tensor_factory, &[draft_prev_token.get()])?;
303            let input = DecodeInput::new(input_tensor, draft_kv_cur.clone());
304            let output = self.draft.decode(&input).await?;
305            let logits = output.logits.to_vec_f32()?;
306            let next_token = argmax_token(&logits);
307            draft_tokens.push(next_token);
308            draft_logits.push(logits);
309            draft_kv_cur = output.kv_cache.clone();
310            draft_prev_token = next_token;
311        }
312
313        // ── Target: ONE multi-position forward over N+1 tokens ──────────
314        // Feeds [last_token, draft_0, ..., draft_{N-1}] into a single
315        // forward pass and gets N+1 logit rows back — one per position.
316        // Dramatically cheaper than N+1 sequential decodes because the
317        // weight matrices are read from HBM exactly once instead of N+1
318        // times (the decode hot path is memory-bound).
319        let mut verify_tokens = Vec::with_capacity(n + 1);
320        verify_tokens.push(last_token);
321        for i in 0..n {
322            verify_tokens.push(draft_tokens[i]);
323        }
324        let mut verify_inputs = Vec::with_capacity(verify_tokens.len());
325        let mut kv_for_verify = target_kv.clone();
326        for tok in &verify_tokens {
327            let input_tensor = tokens_to_tensor(&self.tensor_factory, &[tok.get()])?;
328            verify_inputs.push(DecodeInput::new(input_tensor, kv_for_verify.clone()));
329            // Subsequent inputs share the same handle; `forward_verify` only
330            // reads cache_id + starting seq from the first one.
331            kv_for_verify = kv_for_verify.clone();
332        }
333        let verify_outputs: Vec<DecodeOutput> = self.target.forward_verify(&verify_inputs).await?;
334        assert_eq!(verify_outputs.len(), n + 1);
335        let mut target_logits: Vec<Vec<f32>> = Vec::with_capacity(n + 1);
336        for out in &verify_outputs {
337            target_logits.push(out.logits.to_vec_f32()?);
338        }
339        let target_kv_cur = verify_outputs
340            .last()
341            .map(|o| o.kv_cache.clone())
342            .unwrap_or(target_kv);
343
344        let spec = Speculation {
345            draft_tokens: &draft_tokens,
346            draft_logits: &draft_logits,
347            target_logits: &target_logits,
348            temperature: self.cfg.temperature,
349        };
350        let outcome = verify_speculation(spec, rng)?;
351        let rejected = outcome.rejected_at < n;
352        // For full-accept, the token needed to realign draft with target is
353        // the last draft input consumed by target: draft_tokens[N-1].
354        let draft_catchup_token = if !rejected {
355            draft_tokens.last().copied()
356        } else {
357            None
358        };
359        Ok(SpeculativeStepOutcome {
360            tokens: outcome.tokens,
361            draft_kv: draft_kv_cur,
362            target_kv: target_kv_cur,
363            rejected,
364            rejected_at: outcome.rejected_at,
365            draft_catchup_token,
366        })
367    }
368}
369
370fn tokens_to_tensor(
371    factory: &Arc<dyn TensorFactory>,
372    token_ids: &[u32],
373) -> Result<ferrum_interfaces::tensor::TensorRef> {
374    use ferrum_types::{DataType, Device};
375    let f32_data: Vec<f32> = token_ids.iter().map(|&v| v as f32).collect();
376    let len = f32_data.len();
377    factory.from_slice(&f32_data, &[1, len], DataType::FP32, Device::CPU)
378}
379
380fn argmax_token(logits: &[f32]) -> TokenId {
381    let (idx, _) =
382        logits
383            .iter()
384            .enumerate()
385            .fold((0usize, f32::NEG_INFINITY), |(bi, bv), (i, &v)| {
386                if v > bv {
387                    (i, v)
388                } else {
389                    (bi, bv)
390                }
391            });
392    TokenId::new(idx as u32)
393}
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398    use ferrum_interfaces::KvCacheHandle;
399    use ferrum_testkit::{ConfigurableModelExecutor, MockKvCacheHandle, MockTensorFactory};
400    use ferrum_types::RequestId;
401    use rand::{rngs::StdRng, SeedableRng};
402
403    /// Build a logit vector biased toward `favored` with `strength` over a
404    /// `vocab_size`-entry uniform baseline.
405    fn biased_logits(vocab_size: usize, favored: u32, strength: f32) -> Vec<f32> {
406        let mut v = vec![0.0f32; vocab_size];
407        if (favored as usize) < vocab_size {
408            v[favored as usize] = strength;
409        }
410        v
411    }
412
413    /// Greedy agreement: draft picks the same token as target → full accept
414    /// plus bonus, regardless of rng.
415    #[test]
416    fn full_accept_when_draft_matches_target() {
417        let vocab = 32;
418        let drafts = vec![TokenId::new(3), TokenId::new(7), TokenId::new(11)];
419        let dl: Vec<Vec<f32>> = drafts
420            .iter()
421            .map(|t| biased_logits(vocab, t.get(), 20.0))
422            .collect();
423        // Target agrees perfectly on positions 0..2 and at the bonus slot 3.
424        let tl: Vec<Vec<f32>> = drafts
425            .iter()
426            .map(|t| biased_logits(vocab, t.get(), 20.0))
427            .chain(std::iter::once(biased_logits(vocab, 19, 20.0)))
428            .collect();
429
430        let mut rng = StdRng::seed_from_u64(1);
431        let out = verify_speculation(
432            Speculation {
433                draft_tokens: &drafts,
434                draft_logits: &dl,
435                target_logits: &tl,
436                temperature: 1.0,
437            },
438            &mut rng,
439        )
440        .unwrap();
441        assert_eq!(out.rejected_at, 3, "no rejections → rejected_at == N");
442        assert_eq!(
443            out.tokens,
444            vec![
445                TokenId::new(3),
446                TokenId::new(7),
447                TokenId::new(11),
448                TokenId::new(19)
449            ],
450            "should accept all three drafts + sample the bonus (token 19)"
451        );
452    }
453
454    /// First-draft rejection: draft puts mass on token A, target puts ~0 on
455    /// A and heavy on B → ratio ≈ 0 → reject → residual resamples toward B.
456    #[test]
457    fn first_draft_rejected_residual_prefers_target() {
458        let vocab = 16;
459        let drafts = vec![TokenId::new(2)];
460        let dl = vec![biased_logits(vocab, 2, 20.0)];
461        // Target says token 5 is the right answer, not token 2.
462        let tl = vec![
463            biased_logits(vocab, 5, 20.0),
464            biased_logits(vocab, 0, 0.0), // unused — rejection stops the loop
465        ];
466
467        let mut rng = StdRng::seed_from_u64(7);
468        let out = verify_speculation(
469            Speculation {
470                draft_tokens: &drafts,
471                draft_logits: &dl,
472                target_logits: &tl,
473                temperature: 1.0,
474            },
475            &mut rng,
476        )
477        .unwrap();
478
479        assert_eq!(out.rejected_at, 0);
480        assert_eq!(out.tokens.len(), 1);
481        assert_eq!(
482            out.tokens[0],
483            TokenId::new(5),
484            "residual should pick target's preferred token"
485        );
486    }
487
488    /// Partial acceptance: first draft matches target, second doesn't.
489    /// Expect exactly 2 tokens out: accepted[0] + residual replacement.
490    #[test]
491    fn partial_acceptance_second_draft_rejected() {
492        let vocab = 16;
493        let drafts = vec![TokenId::new(4), TokenId::new(9)];
494        let dl = vec![
495            biased_logits(vocab, 4, 20.0), // draft 0: prefers token 4
496            biased_logits(vocab, 9, 20.0), // draft 1: prefers token 9
497        ];
498        let tl = vec![
499            biased_logits(vocab, 4, 20.0), // target agrees at pos 0
500            biased_logits(vocab, 1, 20.0), // target disagrees at pos 1 (wants 1)
501            biased_logits(vocab, 0, 0.0),  // unused
502        ];
503
504        let mut rng = StdRng::seed_from_u64(42);
505        let out = verify_speculation(
506            Speculation {
507                draft_tokens: &drafts,
508                draft_logits: &dl,
509                target_logits: &tl,
510                temperature: 1.0,
511            },
512            &mut rng,
513        )
514        .unwrap();
515
516        assert_eq!(out.rejected_at, 1);
517        assert_eq!(out.tokens.len(), 2);
518        assert_eq!(out.tokens[0], TokenId::new(4));
519        assert_eq!(
520            out.tokens[1],
521            TokenId::new(1),
522            "replacement should be the target's preferred token at position 1"
523        );
524    }
525
526    /// Empty speculation (N=0): just draws the target's bonus token.
527    #[test]
528    fn zero_drafts_returns_bonus_only() {
529        let vocab = 8;
530        let tl = vec![biased_logits(vocab, 7, 20.0)];
531        let mut rng = StdRng::seed_from_u64(0);
532        let out = verify_speculation(
533            Speculation {
534                draft_tokens: &[],
535                draft_logits: &[],
536                target_logits: &tl,
537                temperature: 1.0,
538            },
539            &mut rng,
540        )
541        .unwrap();
542        assert_eq!(out.rejected_at, 0);
543        assert_eq!(out.tokens, vec![TokenId::new(7)]);
544    }
545
546    /// Temperature 0 (greedy) with matching argmaxes should behave like the
547    /// agreement case — always full accept + bonus argmax.
548    #[test]
549    fn greedy_temperature_full_accept_deterministic() {
550        let vocab = 16;
551        let drafts = vec![TokenId::new(2), TokenId::new(5)];
552        let dl: Vec<Vec<f32>> = drafts
553            .iter()
554            .map(|t| biased_logits(vocab, t.get(), 10.0))
555            .collect();
556        let tl: Vec<Vec<f32>> = drafts
557            .iter()
558            .map(|t| biased_logits(vocab, t.get(), 10.0))
559            .chain(std::iter::once(biased_logits(vocab, 13, 10.0)))
560            .collect();
561
562        let mut rng = StdRng::seed_from_u64(999);
563        let out = verify_speculation(
564            Speculation {
565                draft_tokens: &drafts,
566                draft_logits: &dl,
567                target_logits: &tl,
568                temperature: 0.0,
569            },
570            &mut rng,
571        )
572        .unwrap();
573
574        assert_eq!(out.rejected_at, 2);
575        assert_eq!(
576            out.tokens,
577            vec![TokenId::new(2), TokenId::new(5), TokenId::new(13)]
578        );
579    }
580
581    // ── SpeculativeRunner integration tests (mock executors) ─────────
582
583    fn mock_kv(num_layers: usize) -> Arc<dyn KvCacheHandle> {
584        Arc::new(MockKvCacheHandle::new(RequestId::new(), num_layers, 0))
585    }
586
587    /// Draft and target both use the same executor logic (matching logits)
588    /// → every draft should be accepted and a bonus token should follow.
589    #[tokio::test]
590    async fn runner_full_accept_when_models_agree() {
591        let vocab = 64;
592        let draft: Arc<ConfigurableModelExecutor> = Arc::new(
593            ConfigurableModelExecutor::with_token_sequence(vocab, vec![13, 13, 13, 13, 13]),
594        );
595        let target: Arc<ConfigurableModelExecutor> = Arc::new(
596            ConfigurableModelExecutor::with_token_sequence(vocab, vec![13, 13, 13, 13, 13]),
597        );
598        let tf: Arc<dyn TensorFactory> = Arc::new(MockTensorFactory);
599        let runner = SpeculativeRunner {
600            draft: draft.as_ref(),
601            target: target.as_ref(),
602            tensor_factory: tf,
603            cfg: SpeculativeDecodingConfig {
604                num_speculative_tokens: 3,
605                temperature: 1.0,
606            },
607        };
608        let mut rng = StdRng::seed_from_u64(0);
609        let out = runner
610            .step(TokenId::new(5), mock_kv(12), mock_kv(12), &mut rng)
611            .await
612            .expect("step");
613
614        assert!(!out.rejected, "agreeing models should not reject");
615        assert_eq!(out.rejected_at, 3);
616        assert_eq!(out.tokens.len(), 4, "3 drafts + 1 bonus");
617        for &t in &out.tokens {
618            assert_eq!(
619                t.get(),
620                13,
621                "agreeing models should all emit the biased token 13"
622            );
623        }
624    }
625
626    /// Draft biases token A, target biases token B → first draft rejected
627    /// → step returns 1 replacement token (target's preferred). Remaining
628    /// draft positions don't contribute.
629    #[tokio::test]
630    async fn runner_rejects_when_models_disagree() {
631        let vocab = 64;
632        let draft = Arc::new(ConfigurableModelExecutor::with_token_sequence(
633            vocab,
634            vec![7, 7, 7],
635        ));
636        let target = Arc::new(ConfigurableModelExecutor::with_token_sequence(
637            vocab,
638            vec![21, 21, 21, 21],
639        ));
640        let tf: Arc<dyn TensorFactory> = Arc::new(MockTensorFactory);
641        let runner = SpeculativeRunner {
642            draft: draft.as_ref(),
643            target: target.as_ref(),
644            tensor_factory: tf,
645            cfg: SpeculativeDecodingConfig {
646                num_speculative_tokens: 3,
647                temperature: 1.0,
648            },
649        };
650        let mut rng = StdRng::seed_from_u64(1);
651        let out = runner
652            .step(TokenId::new(0), mock_kv(12), mock_kv(12), &mut rng)
653            .await
654            .expect("step");
655
656        assert!(out.rejected);
657        assert_eq!(out.rejected_at, 0, "first draft should be rejected");
658        assert_eq!(out.tokens.len(), 1);
659        assert_eq!(
660            out.tokens[0].get(),
661            21,
662            "residual should sample target's preferred token"
663        );
664    }
665
666    /// Sanity: unbiased-vs-unbiased (both distributions uniform) — the
667    /// algorithm always accepts because ratio = 1. Ensures no accidental
668    /// rejection due to floating-point wobble on equal distributions.
669    #[test]
670    fn equal_distributions_always_accept() {
671        let vocab = 8;
672        let drafts = vec![TokenId::new(3)];
673        let dl = vec![vec![0.0f32; vocab]];
674        let tl = vec![vec![0.0f32; vocab], vec![0.0f32; vocab]];
675        for seed in 0..20u64 {
676            let mut rng = StdRng::seed_from_u64(seed);
677            let out = verify_speculation(
678                Speculation {
679                    draft_tokens: &drafts,
680                    draft_logits: &dl,
681                    target_logits: &tl,
682                    temperature: 1.0,
683                },
684                &mut rng,
685            )
686            .unwrap();
687            assert_eq!(out.rejected_at, 1, "seed {seed}: should accept");
688            assert_eq!(out.tokens.len(), 2);
689            assert_eq!(out.tokens[0], TokenId::new(3));
690        }
691    }
692}