Decode is memory-bound. Every single token forces you to load the entire model's weights from HBM, and the GPU's compute units sit mostly idle while they wait. Speculative decoding is one of the cleverest ideas in inference: instead of wasting those idle cycles, use them to run a small draft model that guesses ahead, then verify all the guesses in parallel with the big model. If the draft model guesses right, you just produced multiple tokens for the cost of one big-model forward pass.
The key insight, and the reason this is not just "approximate decoding," is that the acceptance sampling scheme guarantees the final token distribution is exactly the same as if you had decoded autoregressively with the target model alone. Zero quality loss. It is a free lunch, paid for by the spare compute that memory-bound decode was leaving on the table anyway.
How it works at a high level
The algorithm has three phases that repeat in a loop:
- Draft phase: the small model (e.g., a 68M parameter model for a 7B target) autoregressively generates K candidate tokens. This is fast because the draft model is tiny.
- Verify phase: feed the original context plus all K draft tokens into the target model in a single forward pass. This gives you the target model's probability distribution at each of the K+1 positions (K draft positions plus one bonus position).
- Accept/reject phase: walk through the K draft tokens left to right. For each one, compare the draft model's probability q(x) to the target model's probability p(x). Accept the token with probability min(1, p(x)/q(x)). If you reject a token, resample from an adjusted distribution and stop.
The beautiful part is that even in the worst case, where every draft token is rejected, you still get one token from the target model (the resampled token at the first position). So speculative decoding never makes things worse than standard decoding; it can only help.
The acceptance sampling math
Let p(x) be the target model's probability of token x, and q(x) be the draft model's probability. The draft model proposes token x. We accept it with probability:
accept_prob = min(1, p(x) / q(x))If the draft model already assigned the token a probability higher than the target model would have (q(x) > p(x)), we accept with probability p(x)/q(x) < 1. If the target model likes the token at least as much as the draft model does (p(x) >= q(x)), we always accept.
When we reject, we resample from the residual distribution:
p_adjusted(x) = max(0, p(x) - q(x)) / sum_over_vocab(max(0, p(x) - q(x)))This is the normalized difference between the target and draft distributions, zeroing out any token where the draft already "overshot." The proof that this produces samples from exactly p(x) follows from the rejection sampling theorem: the mixture of "accepted draft tokens" and "resampled tokens from the residual" equals the target distribution.
Implementation in PyTorch
Here is my implementation. The core loop is surprisingly compact:
import torch
def speculative_decode(target_model, draft_model, input_ids, K=5, temperature=1.0):
"""Generate tokens using speculative decoding.
Returns the same distribution as target_model autoregressive decoding,
but potentially produces multiple tokens per target-model forward pass.
"""
device = input_ids.device
generated = input_ids.clone()
while True: # outer generation loop
# Phase 1: Draft K tokens autoregressively with the small model
draft_tokens = []
draft_probs = []
draft_input = generated.clone()
for _ in range(K):
with torch.no_grad():
logits = draft_model(draft_input).logits[:, -1, :] / temperature
q = torch.softmax(logits, dim=-1)
token = torch.multinomial(q, num_samples=1)
draft_tokens.append(token)
draft_probs.append(q)
draft_input = torch.cat([draft_input, token], dim=-1)
draft_tokens = torch.cat(draft_tokens, dim=-1) # shape: (1, K)
# Phase 2: Verify all K tokens with the target model in one pass
verify_input = torch.cat([generated, draft_tokens], dim=-1)
with torch.no_grad():
target_logits = target_model(verify_input).logits / temperature
# Extract target probs at each draft position
n = generated.shape[1]
# Phase 3: Accept/reject with sampling
accepted = 0
for i in range(K):
target_p = torch.softmax(target_logits[:, n + i - 1, :], dim=-1)
draft_q = draft_probs[i]
token_id = draft_tokens[:, i]
p_token = target_p[0, token_id[0]]
q_token = draft_q[0, token_id[0]]
# Accept with probability min(1, p/q)
ratio = p_token / q_token
if torch.rand(1, device=device) < ratio:
accepted += 1
else:
# Reject: sample from adjusted distribution
residual = torch.clamp(target_p - draft_q, min=0)
residual = residual / residual.sum()
new_token = torch.multinomial(residual, num_samples=1)
generated = torch.cat([generated, draft_tokens[:, :i], new_token], dim=-1)
break
else:
# All K tokens accepted; bonus: sample one more from target
bonus_p = torch.softmax(target_logits[:, n + K - 1, :], dim=-1)
bonus_token = torch.multinomial(bonus_p, num_samples=1)
generated = torch.cat([generated, draft_tokens, bonus_token], dim=-1)
yield generated, accepted # for the caller to check stopping criteriaWhy the verification is cheap
The verify step feeds K+1 tokens (the original context plus K draft tokens) through the target model. But wait, does that not cost K+1 times more than a single decode step? No, and this is the crucial point. The verification pass is a prefill-style pass: you process all K+1 tokens in parallel with a single matrix multiplication per layer. As we discussed in the roofline post, prefill is compute-bound, and during decode you have idle compute. So the verification pass uses compute that would have been wasted anyway.
The breakeven point depends on the draft model's acceptance rate. If the draft model's tokens are accepted with probability alpha on average, you produce roughly K * alpha + 1 tokens per verify pass instead of 1. The speedup is approximately:
speedup = (K * alpha + 1) / (1 + cost_ratio)
# where cost_ratio = draft_model_time / target_model_timeIn practice, a well-chosen draft model (same tokenizer, similar training data, 10-20x smaller) achieves acceptance rates of 0.7-0.85 on natural text. With K=5 and alpha=0.8, you get about 5 tokens per iteration instead of 1, for a 2-3x wall-clock speedup.
If the draft model is too different from the target (low acceptance rate), or if the target model is already compute-bound (large batch sizes), or if the draft model is not cheap enough relative to the target, the overhead of running both models eats the gains. Speculative decoding shines at batch size 1 with a well-matched draft model.
Choosing K: the speculation length
K is a tuning parameter. Too small and you do not amortize the draft overhead. Too large and later draft tokens have lower acceptance rates (errors compound), wasting draft computation. The original papers (Leviathan et al. 2023 and Chen et al. 2023) found K=4 to K=8 to be the sweet spot for most model pairs. I found K=5 works well as a starting default.
There is a nice intuition here: if each token has independent acceptance probability alpha, the expected number of accepted tokens is alpha / (1 - alpha) capped at K. For alpha = 0.8, that is 4 tokens. Going beyond K=5 or 6 barely helps because the probability of accepting all of them drops exponentially: 0.8^6 = 0.26.
The guarantee that matters
Speculative decoding produces samples from exactly the target model's distribution. It is not an approximation. The acceptance-rejection scheme is a mathematically rigorous way to convert cheap draft samples into exact target samples, using the target model as a verifier rather than a generator.
This guarantee is what makes speculative decoding practical for production. You can deploy it without worrying about quality regressions, A/B testing output quality, or explaining to stakeholders why the outputs are "slightly different." They are not different at all.
Tomorrow I am switching gears to memory management: building a KV cache manager with block allocation and eviction policies, the infrastructure that makes serving thousands of concurrent requests possible.