Speculative Decoding.
The Concept
Picture a senior engineer reviewing pull requests. The senior is slow, careful, and expensive. A junior engineer drafts the changes (fast, cheap, occasionally wrong). Most of the time the junior's diff is fine and the senior approves it with a glance. Sometimes the junior is wrong and the senior rewrites the section. Either way, the senior reviews many diffs in the time it would have taken to write any one of them from scratch. This is the entire shape of speculative decoding.
The senior is your main model, say a 7-billion parameter transformer. The junior is a draft model, which is just a smaller, cheaper version of the same kind of model, usually trained on the same data with maybe one tenth the parameters. The draft model proposes a short run of tokens; the main model reviews them all at once. When the main model agrees, those tokens are free. Well, almost free. They cost the main model exactly one parallel forward pass, regardless of how many tokens the draft proposed. When the main model disagrees, you fall back to a single token from the main model, throw away the rest of the draft, and start over.
The first time I saw this scheme described, my honest reaction was suspicion. It felt like a sleight of hand. Two models giving you the answer of one, faster, for free. Once I worked through the rejection-rule proof in Step 4 it went from suspicious to inevitable, but I want to flag that confusion as legitimate: this trick is harder to believe than it is to implement.
Two pieces of jargon to settle now. The number of tokens the draft proposes in each round is K, usually somewhere between 4 and 8. The fraction of proposed tokens the main model accepts on average is the acceptance rate, written as α. A run with K=5 and α=0.7 means that on average 3.5 of the 5 draft tokens get accepted per round, which means the main model produces roughly 4.5 tokens per forward pass on average (3.5 accepted draft tokens plus 1 bonus token the main model emits after rejection). Compare that to standard autoregressive decoding, which produces exactly one token per forward pass, and you can see where the speedup comes from.
The cheap part is intuitive. The careful part, the part this chapter exists to teach, is the rejection rule. The naive instinct is to say: if the draft picked token t and the main model would have also picked t, accept it. That sounds right and it is dead wrong. The main model is not picking one token. It is producing a distribution over the entire vocabulary at every position, and a sampler turns that distribution into a single token. To say "the main model picked t" you would have had to roll the dice, and rolling the dice already commits you to one of many possible answers. To preserve the main model's distribution exactly, you cannot just check what the main model picked. You have to check the probability the main model would have assigned to t and apply a precise rule that accepts or rejects accordingly.
Here is the rule, stated cleanly and then unpacked. For each draft token t at position i, let p be the main model's probability of t at that position and q be the draft model's probability of t at that position. Sample a uniform random number u between 0 and 1. Accept t if u < min(1, p/q). If you reject, do not just take the main model's most likely token; instead, sample from a particular residual distribution that we will define carefully below.
Why this rule? Walk through it for a moment. If p is greater than or equal to q, the main model likes the token at least as much as the draft did, and you accept with probability 1. That is the easy case. If p is less than q, the draft was over-eager about t, and you accept with probability p/q. That ratio is exactly the correction needed to bring the joint distribution of "what the draft proposed and you accepted" back into line with what the main model would have done on its own. The math here is short and worth seeing later. For now, the intuition: the rule throws away exactly enough of the over-eager draft picks to make the surviving picks look like they were sampled from the main model directly.
The residual distribution on rejection is the second half of the magic. When you reject a token, you have leaked information. The main model has revealed that it disagreed with the draft on this particular token. To preserve the main distribution, you cannot just sample fresh from the main model. You have to sample from a distribution that compensates for what you have already learned. Specifically, you sample from a distribution proportional to max(0, p(x) − q(x)) over the vocabulary. This is the part of the main model's probability mass that the draft did not already cover. We will work through this carefully in Step 4.
The diagram in your head should be something like this. The draft runs forward K times, producing K candidate tokens autoregressively. Then the main model runs forward exactly once, with all K candidate tokens packed into its input sequence. Attention is parallel along the sequence dimension, so the main model produces a full output distribution at each of the K positions in that single forward pass, at the cost of one forward, not K. Then you walk left-to-right through the K positions and apply the rejection rule. The first position where you reject is where this round ends; everything to the right of that position is discarded. After rejection at position i, you emit one token from the residual distribution at position i. You also get to emit the token from the main model's distribution at position K+1 if all K draft tokens were accepted. That is the bonus token that gives you the "K+1 tokens per forward" payoff in the best case.
A couple of named variants are worth flagging up front so you recognize them later. The version above, with a separate small model as the drafter, is the Leviathan-Chen style of speculative decoding from the two 2023 papers that introduced it. A later variant, called Medusa, drops the separate draft model entirely. Instead, it bolts a small number of extra heads onto the main model itself, trained to predict positions K+1, K+2, K+3 in parallel. The drafter is now built into the main model. You get a similar speedup without the cost of training and serving a second network. We will look at Medusa in Step 7.
If you have spent any time around multi-model routing systems, this story will feel familiar. In past multi-model routing work spanning frontier API models, the practical question was always the same: when does a cheap draft answer beat a careful expensive answer, and how do you decide without running both? Speculative decoding is a tighter version of that exact question, with provably no quality loss. The routing work did not give us the same guarantee, which is part of why this trick excites me more than it should.
Why It Matters
Without speculative decoding, every output token of a large language model costs one forward pass through the entire stack. That cost is the dominant term in inference latency at production scale, and it is what determines how many tokens per second any given GPU can produce. Project 13: Fast Inference: The KV Cache brought the per-step cost down by reusing intermediate state, but it could not get you below one forward pass per token. Speculative decoding breaks that floor. With a well-matched draft model, you get two to three output tokens per main-model forward pass on average. The same hardware now serves two to three times the throughput.
A second reason this trick matters is that it costs you nothing in output quality. The acceptance and resample rules are mathematically constructed so that the distribution of generated tokens is identical to what the main model would have produced if you had run it the slow way. This is not "almost the same" or "close enough for most purposes." The proof is short and exact, and we walk through it in Step 4. If your main model produces a particular probability distribution at each position, the speculative procedure produces samples from that exact distribution. No drift. No subtle degradation. The user gets the main model's tokens, just faster.
I want to underline this because most engineering tradeoffs do not look like this. Usually faster means worse, in some dimension you have to defend in a design review. Speculative decoding really does give you a free lunch in the output-quality dimension. The catch is somewhere else, and the BREAK IT section is where it lives.
The trick is everywhere in production serving stacks now. Most large model APIs you call through OpenAI, Anthropic, or open-source servers like vLLM use some flavor of it. Project 15: Grouped Query Attention is about making the per-forward cost smaller; speculative decoding is about needing fewer forwards. The two stack cleanly. Project 17: Production Serving: Continuous Batching, PagedAttention wires them together with batched scheduling so that a single GPU can serve many concurrent conversations, each using speculative decoding, with KV caches paged in and out of memory as needed. You cannot understand modern inference serving without this chapter.
There is also a clean intellectual reason to learn it. Speculative decoding is one of the cleanest examples in deep learning of a free-lunch technique that works only when a non-obvious condition holds. The condition is that the draft model has to agree with the main model often enough that the work of running the draft is paid back by the work saved on the main. When agreement collapses (when the draft is poorly matched to the main) the trick becomes worse than useless. The BREAK IT section forces you to see exactly when and how the free lunch turns into a tax.
The Build
build.py$python projects/14_speculative-decoding/build.py --tiny
"""
Project 14: Speculative Decoding
Canonical runnable build for this project. Generated by tools/extract_code.py
from the book's chapter markdown — re-run that tool to regenerate.
"""
# === Step 1: Train a tiny draft model on the same data ===
def measure_tokens_per_second(model, prompt_ids, num_tokens=128):
model.eval()
start = time.time()
ids = prompt_ids.clone()
with torch.no_grad():
for _ in range(num_tokens):
logits = model(ids)[:, -1, :]
next_id = torch.argmax(logits, dim=-1, keepdim=True)
ids = torch.cat([ids, next_id], dim=1)
elapsed = time.time() - start
return num_tokens / elapsed
# === Step 2: Naive draft generation of K=4 tokens ===
def draft_propose(draft_model, ids, K=4):
candidate_ids = []
candidate_qs = []
cur = ids.clone()
for _ in range(K):
with torch.no_grad():
logits = draft_model(cur)[:, -1, :]
probs = torch.softmax(logits, dim=-1)
next_id = torch.multinomial(probs, num_samples=1)
q = probs.gather(-1, next_id)
candidate_ids.append(next_id)
candidate_qs.append((probs, next_id)) # store full dist for residual
cur = torch.cat([cur, next_id], dim=1)
return torch.cat(candidate_ids, dim=1), candidate_qs
# === Step 3: Parallel verification on the main model ===
def main_verify(main_model, ids, candidate_ids):
full_seq = torch.cat([ids, candidate_ids], dim=1)
with torch.no_grad():
logits = main_model(full_seq)
K = candidate_ids.shape[1]
# We want the main model's distribution at the K positions where
# the draft proposed tokens. If candidate token t_i was proposed
# at position len(ids)+i, then the main model's distribution for
# that position is at logits[:, len(ids)+i-1, :], because the
# output at position j predicts the token at position j+1.
main_dists = []
start = ids.shape[1] - 1
for i in range(K):
p_logits = logits[:, start + i, :]
main_dists.append(torch.softmax(p_logits, dim=-1))
# The +1 distribution (after the last candidate) is the bonus
# distribution we sample from if all K candidates are accepted.
bonus_dist = torch.softmax(logits[:, start + K, :], dim=-1)
return main_dists, bonus_dist
# === Step 4: The acceptance rule and the residual distribution, walked through ===
def speculative_step(draft_model, main_model, ids, K=4):
candidate_ids, draft_dists = draft_propose(draft_model, ids, K)
main_dists, bonus_dist = main_verify(main_model, ids, candidate_ids)
accepted_ids = []
for i in range(K):
t = candidate_ids[:, i:i+1]
q_full, _ = draft_dists[i]
p_full = main_dists[i]
q_t = q_full.gather(-1, t)
p_t = p_full.gather(-1, t)
u = torch.rand_like(q_t)
if (u < torch.clamp(p_t / q_t, max=1.0)).all():
accepted_ids.append(t)
else:
# Residual distribution: max(0, p - q), normalized.
residual = torch.clamp(p_full - q_full, min=0.0)
residual = residual / residual.sum(dim=-1, keepdim=True)
replacement = torch.multinomial(residual, num_samples=1)
accepted_ids.append(replacement)
return torch.cat(accepted_ids, dim=1)
# All K accepted: append the bonus token from the main model.
bonus = torch.multinomial(bonus_dist, num_samples=1)
accepted_ids.append(bonus)
return torch.cat(accepted_ids, dim=1)
# === Step 5: Measure speedup against naive KV-cached generation ===
def measure_speculative_tps(draft, main, prompt, K=4, total=256):
ids = prompt.clone()
accepted_total = 0
rounds = 0
start = time.time()
while ids.shape[1] - prompt.shape[1] < total:
new_ids = speculative_step(draft, main, ids, K=K)
ids = torch.cat([ids, new_ids], dim=1)
accepted_total += new_ids.shape[1] - 1 # exclude the bonus/residual
rounds += 1
elapsed = time.time() - start
accept_rate = accepted_total / (rounds * K)
return total / elapsed, accept_rate
# === Step 6: Sweep K from 1 to 8 and find the sweet spot ===
results = {}
for K in [1, 2, 3, 4, 5, 6, 7, 8]:
tps, accept = measure_speculative_tps(draft, main, prompt, K=K, total=512)
results[K] = (tps, accept)
print(f"K={K}: {tps:.1f} tokens/sec, acceptance rate {accept:.3f}")
# === Step 7: Medusa — drafting from the main model itself ===
class MedusaHead(nn.Module):
def __init__(self, d_model, vocab_size, depth=1):
super().__init__()
# A small residual MLP on top of the main hidden state.
layers = []
for _ in range(depth):
layers.append(nn.Linear(d_model, d_model))
layers.append(nn.SiLU())
layers.append(nn.Linear(d_model, vocab_size))
self.net = nn.Sequential(*layers)
def forward(self, h):
return self.net(h)
The Steps
7 reference filesEach step is the code for one piece, in isolation — read them as the book introduces them. They assemble into build.py above.
Step 01Train A Tiny Draft Model On The Same Datastep_01_train-a-tiny-draft-model-on-the-same-data.py
"""
Project 14: Step 1 — Train a tiny draft model on the same data
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
def measure_tokens_per_second(model, prompt_ids, num_tokens=128):
model.eval()
start = time.time()
ids = prompt_ids.clone()
with torch.no_grad():
for _ in range(num_tokens):
logits = model(ids)[:, -1, :]
next_id = torch.argmax(logits, dim=-1, keepdim=True)
ids = torch.cat([ids, next_id], dim=1)
elapsed = time.time() - start
return num_tokens / elapsed
Step 02Naive Draft Generation Of K4 Tokensstep_02_naive-draft-generation-of-k4-tokens.py
"""
Project 14: Step 2 — Naive draft generation of K=4 tokens
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
def draft_propose(draft_model, ids, K=4):
candidate_ids = []
candidate_qs = []
cur = ids.clone()
for _ in range(K):
with torch.no_grad():
logits = draft_model(cur)[:, -1, :]
probs = torch.softmax(logits, dim=-1)
next_id = torch.multinomial(probs, num_samples=1)
q = probs.gather(-1, next_id)
candidate_ids.append(next_id)
candidate_qs.append((probs, next_id)) # store full dist for residual
cur = torch.cat([cur, next_id], dim=1)
return torch.cat(candidate_ids, dim=1), candidate_qs
Step 03Parallel Verification On The Main Modelstep_03_parallel-verification-on-the-main-model.py
"""
Project 14: Step 3 — Parallel verification on the main model
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
def main_verify(main_model, ids, candidate_ids):
full_seq = torch.cat([ids, candidate_ids], dim=1)
with torch.no_grad():
logits = main_model(full_seq)
K = candidate_ids.shape[1]
# We want the main model's distribution at the K positions where
# the draft proposed tokens. If candidate token t_i was proposed
# at position len(ids)+i, then the main model's distribution for
# that position is at logits[:, len(ids)+i-1, :], because the
# output at position j predicts the token at position j+1.
main_dists = []
start = ids.shape[1] - 1
for i in range(K):
p_logits = logits[:, start + i, :]
main_dists.append(torch.softmax(p_logits, dim=-1))
# The +1 distribution (after the last candidate) is the bonus
# distribution we sample from if all K candidates are accepted.
bonus_dist = torch.softmax(logits[:, start + K, :], dim=-1)
return main_dists, bonus_dist
Step 04The Acceptance Rule And The Residual Distribution Walked Throughstep_04_the-acceptance-rule-and-the-residual-distribution-walked-through.py
"""
Project 14: Step 4 — The acceptance rule and the residual distribution, walked through
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
def speculative_step(draft_model, main_model, ids, K=4):
candidate_ids, draft_dists = draft_propose(draft_model, ids, K)
main_dists, bonus_dist = main_verify(main_model, ids, candidate_ids)
accepted_ids = []
for i in range(K):
t = candidate_ids[:, i:i+1]
q_full, _ = draft_dists[i]
p_full = main_dists[i]
q_t = q_full.gather(-1, t)
p_t = p_full.gather(-1, t)
u = torch.rand_like(q_t)
if (u < torch.clamp(p_t / q_t, max=1.0)).all():
accepted_ids.append(t)
else:
# Residual distribution: max(0, p - q), normalized.
residual = torch.clamp(p_full - q_full, min=0.0)
residual = residual / residual.sum(dim=-1, keepdim=True)
replacement = torch.multinomial(residual, num_samples=1)
accepted_ids.append(replacement)
return torch.cat(accepted_ids, dim=1)
# All K accepted: append the bonus token from the main model.
bonus = torch.multinomial(bonus_dist, num_samples=1)
accepted_ids.append(bonus)
return torch.cat(accepted_ids, dim=1)
Step 05Measure Speedup Against Naive Kv Cached Generationstep_05_measure-speedup-against-naive-kv-cached-generation.py
"""
Project 14: Step 5 — Measure speedup against naive KV-cached generation
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
def measure_speculative_tps(draft, main, prompt, K=4, total=256):
ids = prompt.clone()
accepted_total = 0
rounds = 0
start = time.time()
while ids.shape[1] - prompt.shape[1] < total:
new_ids = speculative_step(draft, main, ids, K=K)
ids = torch.cat([ids, new_ids], dim=1)
accepted_total += new_ids.shape[1] - 1 # exclude the bonus/residual
rounds += 1
elapsed = time.time() - start
accept_rate = accepted_total / (rounds * K)
return total / elapsed, accept_rate
Step 06Sweep K From 1 To 8 And Find The Sweet Spotstep_06_sweep-k-from-1-to-8-and-find-the-sweet-spot.py
"""
Project 14: Step 6 — Sweep K from 1 to 8 and find the sweet spot
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
results = {}
for K in [1, 2, 3, 4, 5, 6, 7, 8]:
tps, accept = measure_speculative_tps(draft, main, prompt, K=K, total=512)
results[K] = (tps, accept)
print(f"K={K}: {tps:.1f} tokens/sec, acceptance rate {accept:.3f}")
Step 07Medusa Drafting From The Main Model Itselfstep_07_medusa-drafting-from-the-main-model-itself.py
"""
Project 14: Step 7 — Medusa — drafting from the main model itself
Pedagogical reference: this file shows the code for this step in isolation.
For the full assembled, runnable build, use build.py in this same folder.
"""
class MedusaHead(nn.Module):
def __init__(self, d_model, vocab_size, depth=1):
super().__init__()
# A small residual MLP on top of the main hidden state.
layers = []
for _ in range(depth):
layers.append(nn.Linear(d_model, d_model))
layers.append(nn.SiLU())
layers.append(nn.Linear(d_model, vocab_size))
self.net = nn.Sequential(*layers)
def forward(self, h):
return self.net(h)
Break It
break_it.py$python projects/14_speculative-decoding/break_it.py --tiny
"""
Project 14: BREAK IT experiment.
Deliberately sabotages one mechanism from build.py to show what happens
when it's removed. Compare outputs to the working version.
"""
mis_draft = train_draft_on_different_data(...)
tps_aligned, acc_aligned = measure_speculative_tps(
aligned_draft, main, prompt, K=4
)
tps_mismatched, acc_mismatched = measure_speculative_tps(
mis_draft, main, prompt, K=4
)
baseline_tps = measure_tokens_per_second(main, prompt)
print(f"Baseline: {baseline_tps:.1f} tps")
print(f"Aligned spec: {tps_aligned:.1f} tps "
f"(α={acc_aligned:.2f}, {tps_aligned/baseline_tps:.2f}x)")
print(f"Mismatched spec: {tps_mismatched:.1f} tps "
f"(α={acc_mismatched:.2f}, {tps_mismatched/baseline_tps:.2f}x)")
Read it in the book
This page ships the code. The full derivation — why the chain rule is mechanical not magical, what the computational graph buys you, the long-form reasoning — lives in Chapter 14 of Under the Hood.
Get the book →