Module 16: Speculative Decoding

Introduction

Generating text one token at a time (m08) is slow for a frustrating reason: it is memory-bound, not compute-bound. Every new token needs a full forward pass of the model, and each pass mostly waits on reading the billions of weights from memory — the GPU’s arithmetic units sit nearly idle. Producing 100 tokens means 100 serial passes, 100 round-trips through the whole model. You cannot parallelise across tokens, because token t+1 depends on token t.

Speculative decoding breaks that serialization without changing a single output token. The idea: keep a second, much smaller draft model that guesses the next few tokens cheaply. Then run the big target model once, in parallel, over all of those guesses to check them. A carefully designed accept/reject rule keeps a prefix of the guesses and fixes the first wrong one — and the tokens that come out are distributed exactly as if the target model had generated them one at a time. You get the target model’s output, but with far fewer of its expensive forward passes.

Why it matters for LLMs:

  • Free speed, identical quality. Unlike quantization (m14) or a smaller model, speculative decoding does not approximate the target — its output distribution is provably unchanged. It is a pure latency win.
  • It is how fast inference works. Every major serving stack (vLLM, TensorRT-LLM, llama.cpp) ships speculative decoding; typical speedups are 2–3×.
  • A beautiful use of rejection sampling. The correctness proof is three lines and shows the mechanism has to be exact — a rare case where the fast path and the correct path are the same path.

What You’ll Learn

After this module, you can:

  • Explain why autoregressive decoding is memory-bound and how guess-then-verify amortises the target model’s passes.
  • Derive the modified rejection sampling accept rule u < min(1, p/q) and the residual distribution (p − q)₊ used on rejection.
  • Prove that the emitted tokens are distributed exactly as the target’s — the identity min(p, q) + (p − q)₊ = p.
  • Compute the acceptance rate α = 1 − TV(p, q), the expected tokens per pass (1 − α^{γ+1})/(1 − α), and the wall-clock speedup.
  • Implement speculative_verify and a full speculative_decode loop, and wire two real GPTModels together as draft and target.

Prerequisites

This module requires familiarity with:

  • Module 08: Generation — autoregressive sampling, temperature, and the KV cache; speculative decoding accelerates exactly this loop.
  • Module 06: Transformer — the GPTModel we use as both the draft and the target.

Intuition: Guess, Then Verify

A human editor reads far faster than they write. Speculative decoding gives the big model that same asymmetry: verifying a proposed continuation costs one forward pass, while writing it from scratch would cost one pass per token.

The loop, once per iteration:

  1. Draft. The small model proposes γ tokens (say γ = 4), sampling them autoregressively. This is cheap — the draft is small.
  2. Verify in parallel. Feed the prompt plus all γ draft tokens to the big target model in a single forward pass. Because the tokens are already laid out, the target scores all γ+1 positions at once — it tells us the target’s next-token distribution as if it had reached each of those positions itself.
  3. Accept a prefix. Walk the draft tokens left to right. Accept each one with a probability set by how much the target agrees with the draft. Stop at the first rejection.
  4. Correct once. Replace the first rejected token with a single sample from a residual distribution (defined below). If every draft token was accepted, take one free bonus token from the target’s last position instead.

Each iteration costs one target pass (plus γ cheap draft passes) and emits anywhere from 1 token (draft was useless) to γ+1 tokens (draft was perfect). When the draft is good, most iterations emit several tokens per expensive pass — that is the whole speedup.

Step through one iteration:

NoteKey Insight

The target still does the same kind of work — a forward pass — but now a single pass can ratify several tokens instead of producing just one. Nothing about the target’s distribution is approximated; the draft only changes how many tokens each expensive pass yields.

The Math: Speculative Sampling

Let \(p(x)\) be the target model’s next-token distribution at some position and \(q(x)\) the draft’s distribution at that same position. The draft has sampled a token \(x \sim q\). We want a rule that decides whether to keep \(x\) such that the final token is distributed as \(p\) — even though \(x\) came from \(q\).

The accept rule. Draw \(u \sim \text{Uniform}(0,1)\) and

\[ \text{accept } x \quad \text{if} \quad u < \min\!\left(1, \frac{p(x)}{q(x)}\right). \]

If the target likes \(x\) at least as much as the draft did (\(p(x) \ge q(x)\)), the ratio is \(\ge 1\) and \(x\) is always accepted. If the target likes it less (\(p(x) < q(x)\)), accept it only a \(p(x)/q(x)\) fraction of the time — the draft over-proposed it, so we thin it out.

The correction. When \(x\) is rejected, we do not just resample from \(p\) — that would double-count the tokens the accept step already handles. Instead we sample from the residual distribution, the part of \(p\) that the accept step left unfilled:

\[ p'(x) = \frac{\big(p(x) - q(x)\big)_+}{\sum_{x'} \big(p(x') - q(x')\big)_+}, \qquad (z)_+ = \max(0, z). \]

The residual puts mass exactly where the target wants more than the draft gave. Both live in speculative.py as residual_distribution and the accept test inside speculative_verify.

Why the Output Is Exactly the Target’s

Here is the whole proof, for a single drafted token. The output token equals a value \(v\) in one of two disjoint ways:

  • Accepted: the draft sampled \(v\) (prob \(q(v)\)) and passed the test (prob \(\min(1, p(v)/q(v))\)). Contribution: \(q(v)\min(1, p(v)/q(v)) = \min(q(v), p(v))\).
  • Rejected then corrected: some token was rejected — that happens with total probability \(\sum_x q(x)\big(1 - \min(1, p(x)/q(x))\big) = \sum_x (q(x)-p(x))_+ = \text{TV}(p,q)\) — and then the residual produced \(v\), with prob \(p'(v)\). Since \(p'(v) \cdot \text{TV}(p,q) = (p(v) - q(v))_+\), the contribution is \((p(v) - q(v))_+\).

Add them:

\[ P(\text{output} = v) = \min\!\big(q(v), p(v)\big) + \big(p(v) - q(v)\big)_+ = p(v). \]

The last step is just algebra: if \(p(v) \ge q(v)\) then \(\min = q(v)\) and the residual adds the gap \(p(v)-q(v)\); if \(p(v) < q(v)\) then \(\min = p(v)\) and the residual adds nothing. Either way the total is \(p(v)\) — the emitted token is distributed exactly as the target model’s, no matter how bad the draft \(q\) is. The draft only affects speed, never correctness.

NoteKey Insight

min(p, q) + relu(p − q) == p is the identity that makes the whole method exact. Our test suite asserts it directly (test_marginal_identity_recovers_target_exactly) rather than relying on sampling — the correctness is algebraic, not statistical.

The Speedup: How Many Tokens per Pass

Whether speculative decoding actually helps comes down to one number: the acceptance rate

\[ \alpha = \mathbb{E}_{x \sim q}\!\left[\min\!\left(1, \tfrac{p(x)}{q(x)}\right)\right] = \sum_x \min\big(p(x), q(x)\big) = 1 - \text{TV}(p, q). \]

\(\alpha\) is high when the draft and target agree (small total-variation distance). With a constant \(\alpha\) and \(\gamma\) draft tokens, the number of tokens emitted per iteration is a truncated geometric series plus the guaranteed correction/bonus token:

\[ \mathbb{E}[\text{tokens}] = 1 + \alpha + \alpha^2 + \dots + \alpha^\gamma = \frac{1 - \alpha^{\gamma+1}}{1 - \alpha}. \]

At \(\alpha \to 1\) this is \(\gamma + 1\) (every guess accepted, plus the bonus); at \(\alpha \to 0\) it collapses to \(1\) (only the correction survives). If a draft step costs a fraction \(c\) of a target step, one iteration costs \(\gamma c + 1\) target-equivalents, so the wall-clock speedup over serial decoding is

\[ \text{speedup} = \frac{1 - \alpha^{\gamma+1}}{(1 - \alpha)\,(\gamma c + 1)}. \]

There is a sweet spot for \(\gamma\): too small wastes the parallel pass, too large wastes draft work on tokens that will be rejected anyway. Drive it yourself.

Interactive Exploration

Set the acceptance rate \(\alpha\), the block length \(\gamma\), and the draft cost ratio \(c\), and watch the expected tokens-per-pass and the speedup curve. The faint dots are the values computed by the tested demonstrate_speculative in Python (at \(c = 0.1\)) — the live curve passes right through them.

TipTry This
  1. A good draft (α ≈ 0.9). The curve climbs well above 1× and the optimal γ is a handful of tokens. This is the realistic regime.
  2. A weak draft (α ≈ 0.4). Even the best γ barely beats 1×, and large γ makes it worse — you pay draft cost for tokens that get rejected.
  3. Cheaper draft (drop c). A smaller draft (lower c) both raises the speedup and pushes the optimal γ larger, since each wasted guess costs less.

Code: Verify a Draft

The core is speculative_verify: given the target’s distributions, the draft’s distributions, and the tokens the draft sampled, it returns the accepted prefix plus one correction/bonus token. Everything here follows speculative.py.

import torch
from speculative import acceptance_rate, residual_distribution, speculative_verify

# Two toy distributions over a 4-token vocabulary at one position.
p = torch.tensor([0.5, 0.3, 0.15, 0.05])   # target
q = torch.tensor([0.4, 0.4, 0.15, 0.05])   # draft (over-proposes token 1)

print(f"acceptance rate α = {acceptance_rate(p, q):.3f}  (= 1 - TV)")
print(f"residual p'      = {residual_distribution(p, q).tolist()}")
acceptance rate α = 0.900  (= 1 - TV)
residual p'      = [1.0, 0.0, 0.0, 0.0]
# One block of gamma=2 draft tokens. target_probs has gamma+1 rows
# (two draft positions + the bonus position); here we reuse p for all.
target_probs = torch.stack([p, p, p])   # (3, 4)
draft_probs = torch.stack([q, q])       # (2, 4)
draft_tokens = torch.tensor([0, 1])     # the draft sampled tokens 0 then 1

gen = torch.Generator().manual_seed(0)
emitted, n_accepted = speculative_verify(
    target_probs, draft_probs, draft_tokens, generator=gen
)
print(f"draft proposed : {draft_tokens.tolist()}")
print(f"accepted       : {n_accepted} token(s)")
print(f"emitted        : {emitted.tolist()}  (accepted prefix + 1 correction/bonus)")
draft proposed : [0, 1]
accepted       : 1 token(s)
emitted        : [0, 0]  (accepted prefix + 1 correction/bonus)

The emitted length is always n_accepted + 1. The verification is the exact mechanism from the proof — and because the target token comes out \(\sim p\) regardless of q, you can swap in any draft you like.

Code: A Full Decode Loop

speculative_decode wraps the block verifier in the autoregressive loop: draft γ tokens, score them with the target, verify, append, repeat. It takes two step functions — anything mapping a token sequence to a next-token distribution — so it works with toy distributions or real models.

from speculative import speculative_decode

vocab = 6
target = torch.tensor([0.4, 0.25, 0.15, 0.1, 0.06, 0.04])
draft = torch.tensor([0.3, 0.3, 0.2, 0.1, 0.06, 0.04])

# Context-free step functions for a clean demo.
target_step = lambda seq: target
draft_step = lambda seq: draft

gen = torch.Generator().manual_seed(1)
out, trace = speculative_decode(
    torch.tensor([0]), target_step, draft_step,
    gamma=4, max_new_tokens=12, generator=gen,
)
print(f"generated {out.shape[0]} tokens: {out.tolist()}")
for i, step in enumerate(trace):
    print(f"  iter {i}: proposed {step['proposed']} → accepted {step['n_accepted']}, emitted {step['emitted']}")
generated 12 tokens: [0, 0, 2, 0, 0, 3, 2, 1, 0, 0, 0, 4]
  iter 0: proposed [0, 0, 2, 0] → accepted 4, emitted [0, 0, 2, 0, 0]
  iter 1: proposed [3, 2, 1, 1] → accepted 3, emitted [3, 2, 1, 0]
  iter 2: proposed [0, 0, 4, 1] → accepted 3, emitted [0, 0, 4, 0]

Notice how many tokens each iteration emits from a single target pass — that count, averaged, is the (1 − α^{γ+1})/(1 − α) from above.

Code: Two Real Models

Now wire two GPTModels (m06) together. model_step adapts a model into the step function speculative_decode expects. Using an untrained pair here just checks the plumbing; in practice the draft is a small distilled version of the target.

import importlib.util, sys
from pathlib import Path

# Load the transformer module (numeric-prefixed dir) via importlib.
_tpath = Path("../m06_transformer/transformer.py").resolve()
_spec = importlib.util.spec_from_file_location("transformer", _tpath)
_t = importlib.util.module_from_spec(_spec)
sys.modules["transformer"] = _t
_spec.loader.exec_module(_t)
GPTModel = _t.GPTModel

from speculative import model_step

torch.manual_seed(0)
big = GPTModel(vocab_size=32, embed_dim=32, num_heads=4, num_layers=2, max_seq_len=64)
small = GPTModel(vocab_size=32, embed_dim=16, num_heads=2, num_layers=1, max_seq_len=64)

prompt = torch.tensor([1, 5, 9, 2])
gen = torch.Generator().manual_seed(7)
out, trace = speculative_decode(
    prompt, model_step(big), model_step(small),
    gamma=4, max_new_tokens=16, generator=gen,
)
accepted = [s["n_accepted"] for s in trace]
print(f"iterations      : {len(trace)}")
print(f"accepted/iter   : {accepted}")
print(f"tokens produced : {out.shape[0]} from {len(trace)} target passes")
iterations      : 5
accepted/iter   : [1, 4, 4, 1, 4]
tokens produced : 16 from 5 target passes
NoteKey Insight

Point model_step at the same model for both draft and target and every token is accepted (test_end_to_end_self_draft_accepts_all): \(p = q \Rightarrow \alpha = 1\). That is the sanity check that the verifier is faithful — a model never disagrees with itself.

Common Pitfalls

Pitfall Why it bites Fix
Draft too weak Low α → almost every guess rejected → you pay draft cost for nothing. Use a draft that tracks the target (distilled/pruned), not a random one.
γ too large Past the optimum, extra drafts are usually rejected and wasted. Tune γ to the α/c regime (use the explorer); 3–8 is typical.
Resampling from p instead of the residual Double-counts accepted mass → output is not the target distribution. Reject → sample from (p − q)₊, never from p.
Tokenizer mismatch Draft and target must share a vocabulary, or token ids are meaningless. Same tokenizer for both models.
Assuming it changes quality It does not — same α gives the same distribution, only faster. Don’t “tune” α for quality; tune the draft for speed.
Temperature applied inconsistently p and q must use the same sampling settings, or the ratio is wrong. Apply identical temperature/top-p to both step functions.

Exercises

Exercise 1: Acceptance rate by hand

Compute \(\alpha = \sum_x \min(p(x), q(x))\) for a draft that is off by a little vs. a lot, and confirm it equals \(1 - \text{TV}\).

from speculative import acceptance_rate, total_variation

p = torch.tensor([0.5, 0.3, 0.2])
q_close = torch.tensor([0.45, 0.35, 0.20])
q_far = torch.tensor([0.1, 0.1, 0.8])

for name, q in [("close", q_close), ("far", q_far)]:
    a = acceptance_rate(p, q)
    print(f"{name:>5}: α={a:.3f}, 1-TV={1 - total_variation(p, q):.3f}")
close: α=0.950, 1-TV=0.950
  far: α=0.400, 1-TV=0.400

Exercise 2: Sweep the block length

Use expected_accepted_tokens to find, for a fixed α, the γ that maximises tokens per pass under a cost model. (Hint: more tokens is not always faster.)

from speculative import expected_accepted_tokens, speculative_speedup

alpha, c = 0.85, 0.15
for gamma in range(1, 11):
    tok = expected_accepted_tokens(alpha, gamma)
    spd = speculative_speedup(alpha, gamma, c)
    print(f"γ={gamma:2d}: {tok:.2f} tokens/pass, {spd:.2f}× speedup")

# Your task: which γ is best here? Try changing c to 0.05 and 0.4.
γ= 1: 1.85 tokens/pass, 1.61× speedup
γ= 2: 2.57 tokens/pass, 1.98× speedup
γ= 3: 3.19 tokens/pass, 2.20× speedup
γ= 4: 3.71 tokens/pass, 2.32× speedup
γ= 5: 4.15 tokens/pass, 2.37× speedup
γ= 6: 4.53 tokens/pass, 2.38× speedup
γ= 7: 4.85 tokens/pass, 2.37× speedup
γ= 8: 5.12 tokens/pass, 2.33× speedup
γ= 9: 5.35 tokens/pass, 2.28× speedup
γ=10: 5.55 tokens/pass, 2.22× speedup

Exercise 3: The residual, visualised

Pick a p and q, compute the residual, and verify min(p,q) + residual*(1-α) recovers p exactly — the correctness identity.

from speculative import residual_distribution, acceptance_rate

p = torch.tensor([0.5, 0.2, 0.2, 0.1])
q = torch.tensor([0.2, 0.4, 0.3, 0.1])

alpha = acceptance_rate(p, q)
recovered = torch.minimum(p, q) + residual_distribution(p, q) * (1 - alpha)
print(f"recovered = {recovered.tolist()}")
print(f"target    = {p.tolist()}")
print(f"exact?    = {torch.allclose(recovered, p, atol=1e-6)}")
recovered = [0.4999999403953552, 0.20000000298023224, 0.20000000298023224, 0.10000000149011612]
target    = [0.5, 0.20000000298023224, 0.20000000298023224, 0.10000000149011612]
exact?    = True

Summary

Key takeaways:

  1. Guess, then verify. A cheap draft proposes γ tokens; the target ratifies them in one parallel pass, emitting between 1 and γ+1 tokens per pass.
  2. The accept rule is u < min(1, p/q). Rejections are corrected by a sample from the residual (p − q)₊, never from p directly.
  3. The output is provably the target’s. min(p, q) + (p − q)₊ = p makes each emitted token exactly target-distributed, for any draft — speculative decoding is a pure latency win, not an approximation.
  4. Throughput is set by α and γ. With acceptance rate α = 1 − TV(p, q), each pass yields (1 − α^{γ+1})/(1 − α) tokens; the speedup peaks at an intermediate γ that depends on the draft’s cost.
  5. A better draft is the only lever. Since correctness is fixed, all the speed comes from making the draft agree with the target (higher α) as cheaply as possible (lower c).

What’s Next

Speculative decoding assumes a separate draft model. The frontier removes that cost by making the target draft itself — extra lightweight heads (Medusa) or a feature-space draft (EAGLE) predict several tokens ahead, verified with the same rejection rule over a tree of candidates. The next module in the fast- inference arc is continuous batching / paged attention (vLLM), which is how these accepted-token bursts are scheduled across many concurrent requests.

For now, revisit Module 08: Generation — every decoding strategy there (temperature, top-k, top-p) composes with speculative decoding unchanged, because the target distribution it samples from is exactly preserved.

Going Deeper

Core Papers:

Going Further: