Speculative Decoding
A cheap draft proposes k tokens; the big model verifies all k in one pass. Rejection sampling makes the output distribution provably identical.
- draft model
- n-gram lookup
- verification
- rejection sampling
- acceptance rate
- EAGLE/Medusa
Motivation
The problem
Decode is memory-bound: you read 16 GB of weights to produce one token, and the arithmetic units are idle almost the whole time. That has a striking implication.
Checking five candidate tokens costs almost exactly the same as generating one. The weight read is the expensive part, and it is shared. The extra arithmetic for four more positions comes free: it fits in the idle capacity you already paid for.
So if something cheap could guess the next few tokens, the big model could check all the guesses in a single pass. The only question is whether the checking preserves the output distribution, or whether you are quietly serving a worse model.
Core idea
The solution
A cheap proposer generates k candidate tokens. The target model runs once over all k+1 positions. A rejection-sampling test then accepts the longest correct prefix, and the test is constructed so that the accepted tokens are distributed exactly as the target model would have produced them alone.
def verify(draft_tokens, q_probs, p_probs, rng):
"""q = draft's distribution, p = target's. Output is distributed as p."""
accepted = []
for i, tok in enumerate(draft_tokens):
p, q = p_probs[i][tok], q_probs[i][tok]
if rng.random() < min(1.0, p / q): # accept
accepted.append(tok)
continue
# rejected: resample from the RESIDUAL distribution
residual = np.maximum(0.0, p_probs[i] - q_probs[i])
return accepted + [sample(residual / residual.sum(), rng)]
# every draft token accepted -> take the target's own next token free
return accepted + [sample(p_probs[len(draft_tokens)], rng)]Lossless by construction
min(1, p/q) combined with the residual max(0, p−q) is a modified rejection sampler: for every token, the probability of accepting it plus the probability of resampling it from the residual sum to exactly p. That is a direct identity, proved in a few lines in Leviathan et al. and Chen et al. (2023). Speculative decoding is not a quality/speed trade-off. If your implementation changes outputs, it has a bug.Mechanics
How it works
The economics
With a per-token acceptance rate α, the expected number of tokens emitted per target pass is geometric:
E[tokens] = (1 - α^(k+1)) / (1 - α)
α = 0.9, k = 6 -> 5.2 tokens per target pass
α = 0.8, k = 4 -> 3.4
α = 0.5, k = 4 -> 1.9
α = 0.3, k = 4 -> 1.4 (barely worth the draft cost)One rejection discards everything after it, because those proposals were conditioned on a token that no longer exists. Returns therefore diminish sharply in k. Past an optimal k, extra proposals cost draft time without ever being reached, and speculating too far leaves you slower than plain decoding.
Three ways to propose
- Draft model — a small model of the same family (Llama-1B drafting for Llama-70B). Best acceptance, but you pay real compute and manage a second model with its own KV cache.
- N-gram / prompt lookup — search the prompt and recent output for the current suffix and copy what followed it last time. No model at all. Acceptance is poor for open-ended prose and remarkable for summarisation, code editing and RAG, where the output legitimately repeats the input.
- EAGLE / Medusa — small extra heads trained on the target model's own hidden states to predict several future positions. Very cheap and high-acceptance, at the cost of needing training per model. EAGLE-style methods are the flagship speculative option in vLLM.
Why it fights with batching
Speculative decoding spends spare compute capacity. Batching spends the same capacity. At batch 1 there is plenty to spare and speculation is a large win; at batch 128 the GPU is already compute-saturated and speculation mostly adds work.
Production engines therefore make the speculation length dynamic: speculate aggressively when the batch is small, back off as concurrency rises, and track measured acceptance to adjust k on the fly.
Explore
Try it
- tokens / target pass
- 2.79
- theory
- 3.05
- speedup
- 1.89×
- best k here
- 4 (2.06×)
- accepted draft token
- first rejection — round ends
- bonus token, free from the same pass
Rejection is prefix-terminating: once one token is rejected, every proposal after it is discarded unchecked, because it was conditioned on a token that is no longer there. This is why expected tokens per round is geometric in α, not linear in k.
The curve has an interior maximum. Past it, extra proposals are almost never reached and still cost draft time, so speculating too far is a slowdown. With a weak proposer (try n-gram) the optimum can be k = 1 or 2.
- draft model — a small model of the same family. Highest acceptance, but costs real compute and needs its own weights and KV cache.
- n-gram lookup — copy continuations found in the prompt. Nearly free and needs no model; acceptance is poor for open-ended text and excellent for summarisation, code editing and RAG, where output repeats input.
- EAGLE head — extra heads on the target model that predict several positions from its own hidden state. Very cheap, high acceptance, but must be trained per model.
Switch the proposer and watch both cost and acceptance move together. The best choice depends on whether your workload's output resembles its input.
Theoretical speedup at these settings: 2.06×. Empirical from 24 rounds: 1.89×. They converge as you run more rounds. The variance is real, and it is why speculative decoding improves mean latency while slightly worsening latency variance.
Set acceptance to 0.35 and k to 10. The speedup drops below 1×: you are now slower than plain decoding, having paid for ten draft passes to accept about one token. Switch the proposer to n-gram and the cost falls enough that even poor acceptance stays profitable. Prompt-lookup decoding is worth shipping for that reason, bad acceptance and all.
Build it
Implementation
The bookkeeping is where implementations go wrong. After a partial acceptance you must roll back the KV cache for every rejected position, in both models.
def step(self, req):
draft_tokens, q_probs = self.proposer.propose(req, k=self.k)
# one target pass over [context + all k drafts]
p_probs = self.target.forward(req.tokens + draft_tokens)
accepted = verify(draft_tokens, q_probs, p_probs, self.rng)
# CRITICAL: undo the KV written for tokens we did not accept
n_rejected = len(draft_tokens) - (len(accepted) - 1)
if n_rejected > 0:
self.target.kv.truncate(req, n_rejected)
self.proposer.kv.truncate(req, n_rejected)
req.tokens.extend(accepted)
self.stats.record(proposed=len(draft_tokens), accepted=len(accepted) - 1)
return acceptedForgetting to roll back the KV cache
temperature=0 and a fixed seed, with and without speculation, and assert the outputs are token-identical.$ python code/s14_speculative_decoding.pyOnly NumPy is required — setup instructions.
Production notes
In production
- vLLM — supports n-gram, draft-model and EAGLE proposers, with the speculation length adapted to batch size.
- Medusa — multiple heads producing a tree of candidates rather than a single chain, verified in one pass with a tree attention mask. Higher acceptance per pass.
- EAGLE-2/3 — the current state of the art; dynamically shaped draft trees with acceptance rates high enough for 3–4× end-to-end speedups.
- The self-speculation trick — skip layers of the target model to form a draft. No extra weights at all, at the cost of lower acceptance.
Exercises
- 1Implement prompt-lookup decoding and measure acceptance on two workloads: open-ended creative writing versus "summarise this document". The gap should be dramatic.
- 2Implement tree-based speculation: propose a branching tree instead of a chain, build the tree attention mask, and verify the whole tree in one pass. Compare expected tokens per pass at equal draft cost.
- 3Add adaptive k: track a running acceptance estimate and adjust speculation length to maximise measured throughput. Show it converging to the optimum the static sweep found.
Continue
What's next
Speculation constrains when tokens are produced. S15 constrains which tokens are allowed at all: a JSON schema compiles into a state machine, every logit the grammar forbids is masked, and invalid output becomes impossible rather than unlikely.
Check yourself
Questions
Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.
Why does verifying k draft tokens cost about the same as generating one token?