Home › Neural network training & inference › Inference
algorithm low

Speculative decoding

Problem Generating text from a large language model one token at a time is slow for a single user, because each step costs a full pass of the model. Operators want lower latency without changing the output.

A cheap draft proposes γ tokens, the big model verifies them in one pass, and rejection sampling keeps the target's exact distribution at 2-3x speed.

expected tokens per target pass (1 − α^(γ+1))/(1 − α) for acceptance rate α; speedup bounded by that divided by (1 + γ·c), with c = draft/target cost
llm-inferencelatencydraft-modelrejection-samplingmedusaeagle
Overview

At small batch, one decode step of a big model costs about the same for 5 tokens as for 1. A cheap draft guesses γ tokens, and the target scores all γ+1 positions in one pass. Each guess is accepted with probability min(1, p/q). At the first rejection, a token is drawn from the normalized residual max(0, p − q). The output distribution is exactly the target's.

Key facts

Use when

  • Latency-sensitive, low-batch serving of large models with predictable outputs (code, RAG).
  • Translation and summarization with encoder-decoder models (Leviathan et al.: T5-XXL), and code completion (Chen et al.: HumanEval with Chinchilla).

Avoid when

  • The server already runs compute-bound at high batch.
  • No same-tokenizer draft agrees well with the target: use Medusa/EAGLE heads.

Deep dive

How it works

In: a prefix, a target model p and a cheap drafter q, draft length γ · Out: 1 to γ+1 new tokens per target pass, distributed exactly as p · State: KV caches of both models, rolled back on rejection; running acceptance rate if γ is adaptive · Cost: one target pass over γ+1 positions plus γ draft steps; expected tokens per pass (1 − α^(γ+1))/(1 − α)

  1. Draft: autoregressively sample x₁..x_γ from q (draft model), keeping each q_i(·).
  2. Verify: run the target on prefix + x₁..x_γ in one forward pass. It returns p₁..p_{γ+1}, the target distributions at every position, including the one after x_γ.
  3. Accept/reject for i = 1..γ: with u ~ U(0,1), accept x_i if u < p_i(x_i)/q_i(x_i). At the first rejection, sample from norm(max(0, p_i − q_i)) and stop.
  4. Bonus: if all γ are accepted, sample one extra token from p_{γ+1}. Every target pass therefore yields at least 1 and at most γ+1 tokens.
  5. Roll back the KV-cache entries of rejected draft tokens (truncate the length) in both models.

Why it's exact: P(output = x) = q(x)·min(1, p(x)/q(x)) + P(reject)·residual(x) = min(q, p)(x) + (p(x) − min(q, p)(x)) = p(x).

Expected tokens per target pass, with i.i.d. acceptance α: E = (1 − α^(γ+1)) / (1 − α). At α = 0.8: γ = 4 → 3.36, γ = 8 → 4.33. Wall-clock speedup ≈ E / (1 + γ·c), where c is the draft cost relative to one target step. Pick γ about 3-8 where that ratio peaks.

Greedy decoding reduces to "accept while the draft's token equals the target's argmax".

Code

The verification step, plus a statistical check that the output distribution equals the target's.

struct Rng(u64);
impl Rng {
    fn f(&mut self) -> f32 {                                  // xorshift64*, uniform [0,1)
        self.0 ^= self.0 >> 12; self.0 ^= self.0 << 25; self.0 ^= self.0 >> 27;
        ((self.0.wrapping_mul(0x2545F4914F6CDD1D) >> 40) as f32) / (1u64 << 24) as f32
    }
    fn sample(&mut self, probs: &[f32]) -> usize {
        let mut u = self.f() * probs.iter().sum::<f32>();       // tolerates unnormalized input
        for (i, &p) in probs.iter().enumerate() { u -= p; if u < 0.0 { return i; } }
        probs.len() - 1
    }
}

/// draft[i] was sampled from q[i]; p has γ+1 rows (target dists, last = bonus position).
/// Returns the accepted tokens plus exactly one corrected or bonus token.
fn verify(draft: &[usize], q: &[Vec<f32>], p: &[Vec<f32>], rng: &mut Rng) -> Vec<usize> {
    let mut out = Vec::with_capacity(draft.len() + 1);
    for (i, &x) in draft.iter().enumerate() {
        if rng.f() * q[i][x] < p[i][x] {                       // u < p/q without dividing
            out.push(x);
        } else {
            let residual: Vec<f32> = p[i].iter().zip(&q[i]).map(|(a, b)| (a - b).max(0.0)).collect();
            out.push(rng.sample(&residual));                   // sample from norm(max(0, p − q))
            return out;
        }
    }
    out.push(rng.sample(&p[draft.len()]));                     // all accepted: free bonus token
    out
}

fn main() {
    let p = vec![0.5f32, 0.3, 0.15, 0.05];                     // target next-token distribution
    let q = vec![0.25f32, 0.25, 0.25, 0.25];                   // a poor draft
    let mut rng = Rng(0x1234_5678_9abc_def1);
    let (n, mut counts, mut accepted) = (400_000, [0usize; 4], 0usize);
    for _ in 0..n {
        let x = rng.sample(&q);
        let toks = verify(&[x], &[q.clone()], &[p.clone(), p.clone()], &mut rng);
        counts[toks[0]] += 1;                                  // first emitted token must follow p
        accepted += (toks.len() == 2) as usize;
    }
    for (c, &pi) in counts.iter().zip(&p) { assert!((*c as f32 / n as f32 - pi).abs() < 0.005); }
    let alpha = accepted as f32 / n as f32;                    // = Σ min(p, q) = 0.25+0.25+0.15+0.05
    println!("empirical {:?}, acceptance {alpha:.3} (expected 0.700)", counts.map(|c| c as f32 / n as f32));
}

Complexity and constants

  • Acceptance rate in expectation: α = Σ_x min(p(x), q(x)) = 1 − TV(p, q). The draft's quality is exactly its total-variation distance to the target.
  • Verification costs 1 target pass over γ+1 tokens. At small batch, this costs about the same as 1 decode step (weight-read bound). Attention over the KV grows slightly.
  • Typical drafts are 10-50x smaller than the target (e.g. 1B for 70B), so c ≈ 0.02-0.1. EAGLE-style heads are one extra transformer layer, so c ≈ 0.02-0.05.
  • Reported speedups: Leviathan et al. 2-3x (T5-XXL); Chen et al. 2-2.5x (Chinchilla 70B); Medusa-2 2.3-3.6x; EAGLE about 3x (Vicuna/LLaMA2-Chat at greedy).

Efficiency

  • Cost: a verify pass over γ+1 tokens costs about one decode step at small batch (weight-read bound). Expected tokens per pass at α = 0.8 is 3.36 for γ = 4 and 4.33 for γ = 8 (formula above). Wall-clock speedup is E / (1 + γ·c), with c about 0.02-0.1 for a 10-50x smaller draft (order of). Bound: memory bandwidth at small batch; compute once the batch is large.
  • Levers:
    • tree speculation with a tree attention mask — more accepted tokens per pass for extra verify FLOPs (SpecInfer, Medusa, EAGLE-2) — small batch
    • draft heads on the target's hidden features (EAGLE, Medusa) — c about 0.02-0.05, higher acceptance than a separate model; EAGLE about 3x, Medusa-2 2.3-3.6x (papers) — when no small same-tokenizer model exists
    • n-gram or prompt-lookup drafting — zero draft cost — edit, summarization and RAG outputs that copy the input
    • adaptive γ, off at large batch — avoids slowdown once the target is compute-bound — continuous-batching servers
    • exact residual sampling — output distribution stays p, so no quality trade (Leviathan et al.) — any sampling temperature
  • Naive → standard → state of the art:
    • one token per target pass · 1.0 tokens per pass
    • separate draft model, chain of γ tokens · 2-3x (T5-XXL; Chinchilla 70B 2-2.5x)
    • tree-verified feature-level drafters (EAGLE family) · about 3x on Vicuna and LLaMA2-Chat, greedy

Implementation notes

  • The draft must use the target's tokenizer and vocabulary. Logits from both go through the same temperature, top-k and top-p processing before computing p and q, or exactness breaks.
  • Tree speculation (SpecInfer, Medusa, EAGLE-2): verify several candidate continuations in one pass with a tree attention mask. More accepted tokens per pass at the cost of verification FLOPs.
  • Prompt lookup / n-gram drafting: copy the continuation of the last matching n-gram from the prompt. Zero cost, very effective for edit-style tasks.
  • Under continuous batching, make γ adaptive per request (from running acceptance) and switch speculation off when the batch is large enough to be compute-bound.
  • KV management: allocate γ slots ahead, then free the rejected ones (paged KV makes this a length truncate).

Pitfalls

  • Using greedy acceptance (match argmax) while sampling with temperature: the output is no longer the target distribution.
  • Bugs in the residual distribution (forgetting the max(0,·) or the normalization) silently bias outputs. Test empirically, as the code does.
  • Floating-point mismatch between a batched verify pass and single-token decode can flip greedy argmax ties. Expect tiny divergences from non-speculative greedy output.
  • A draft that is too large, or too long a γ: the speedup becomes a slowdown when c·γ approaches E − 1.
  • Ignoring the draft's own KV-cache prefill cost on long prompts.

Variants and alternatives

Draft model (Leviathan et al.; Chen et al.)the classic setup. Needs a same-tokenizer small model.
Medusaextra decoding heads on the target predict tokens t+2..t+k, with tree verification.
EAGLE / EAGLE-2 / EAGLE-3a lightweight autoregressive head on the target's hidden features. Higher acceptance than Medusa.
Lookahead decodingJacobi iteration with an n-gram pool, no draft model.
Self-speculative / layer skippingdraft with a subset of the target's layers.
Multi-token prediction heads trained during pretraining (DeepSeek-V3 MTP) reused as the drafter.

In production

vLLM (draft models, n-gram, Medusa, EAGLE, MTP), SGLang (EAGLE), TensorRT-LLM (draft-target, Medusa, EAGLE, lookahead), llama.cpp (--model-draft, lookup decoding), Hugging Face Transformers assisted generation.

Implementations

namelangrepoinstallverdictlicencehealthas_ofnotes
llama.cppC / C++ggml-org/llama.cppbrew install llama.cpp————draft model (--model-draft) and lookup decoding
vLLMPython / C++ / CUDAvllm-project/vllmpip install vllm————draft model, n-gram, Medusa, EAGLE, MTP
SGLangPython / CUDAsgl-project/sglangpip install sglang————EAGLE speculative decoding
TensorRT-LLMC++ / PythonNVIDIA/TensorRT-LLMpip install tensorrt_llm————draft-target, Medusa, EAGLE, lookahead
EAGLEPythonSafeAILab/EAGLEn/a (clone the repo)————reference code for EAGLE, EAGLE-2, EAGLE-3

Sources

  • Leviathan, Kalman, Matias, Fast Inference from Transformers via Speculative Decoding, ICML 2023. paper
  • Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling, 2023. paper
  • Cai et al., Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads, ICML 2024. paper · code
  • Li, Wei, Zhang, Zhang, EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty, ICML 2024. paper · code
  • Miao et al., SpecInfer: Accelerating Large Language Model Serving with Tree-based Speculative Inference and Verification, ASPLOS 2024. paper · code
  • Fu, Bailis, Stoica, Zhang, Break the Sequential Dependency of LLM Inference Using Lookahead Decoding, ICML 2024. paper · code

Builds on

rejection sampling · plannedBest-of-n and filtered sampling: generate many candidates, keep those passing a verifier or reward threshold
Transformer blockAnyone building or sizing a language model needs the standard structure of its layers, and the exact sizes, compute and memory cost that follow from that structure.Token samplingA model scores every candidate next item and a program must pick one: sometimes the best, sometimes a varied but sensible one. Always picking the best repeats itself, picking at random produces nonsense.

Used by

Systems (5):

llm inferenceon device inferencellm api platformchat assistantcode completion

Alternatives and related

Continuous batchingA server answering language-model requests of very different lengths leaves expensive GPUs partly idle because short requests wait for long ones in the same batch. Operators want higher throughput and lower waiting time.Paged KV cacheA language-model server wastes much of its GPU memory on cached attention data reserved for sequences of unknown final length, limiting how many requests it can serve at once.Weight quantizationRunning a large language model for one user at a time is slow and needs a lot of GPU memory, because every weight must be read for each generated token. Engineers want faster, smaller deployments.FlashAttentionTraining or running a transformer on long sequences is slow and runs out of GPU memory, because the attention step builds and moves a huge intermediate table. Engineers want exact attention that is faster and leaner.