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.
  • Build self-drafting (Medusa): extra decoding heads that draft from the target’s own hidden state, a tree of candidates verified in one pass with a tree-attention mask, and a greedy verification that is provably exact.
  • Build EAGLE: draft in feature space with a single autoregressive head that reuses the target’s frozen embedding and LM head, feed each drafted token back to make the draft a chain rather than a fan, and see why the “token one step ahead” resolves feature uncertainty — reusing the same verification to stay exact.
  • Build EAGLE-2: give every draft node a value (the product of the head’s confidences along its path), expand only the top-value frontier and rerank to a budget — a dynamic tree that is ancestor-closed and provably optimal for free, spending its budget far better than any fixed shape.
  • Build lookahead / Jacobi decoding: reframe greedy decoding as a fixed-point system, solve it in parallel with no draft model, no extra heads, and no training, and reuse the trajectory’s n-grams to accept many tokens per pass — all while emitting the exact greedy sequence.

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.

Self-Drafting: Medusa

Everything so far assumes a separate draft model — a second network you must find, host, and keep in lock-step with the target’s vocabulary. That is a real cost. Medusa (Cai et al., 2024) removes it: the target model drafts itself.

The idea is disarmingly simple. A single forward pass already computes the last hidden state h_t — the vector the language-model head turns into the next-token distribution. Medusa bolts a few extra lightweight heads onto that same h_t, each trained to predict a token further ahead. One pass, one weight load, and now you have not one guess but K+1 of them — the base head’s next token plus K speculative continuations — for almost free.

Intuition: the model already knows a bit about the next few tokens

When a language model writes the cat sat on the, its hidden state does not only “know” that mat comes next — it has a strong hunch about . after that, and a weaker one about the token after that. The base head throws all of that away and reads off a single token. Medusa adds heads that read those fainter hunches off the same hidden state, so the model proposes several tokens from one pass instead of paying a full pass per token.

NoteKey Insight

A separate draft model needs its own forward pass to produce each guess. Medusa’s heads share the target’s single pass — the extra tokens cost only a couple of small matrix multiplies, so a “cheap draft” is essentially free.

The Math: what a Medusa head is

A Medusa head is a one-layer residual block followed by a vocabulary projection. For head k reading hidden state h_t:

p_t^{(k)} = \text{softmax}\!\left( W_2^{(k)} \left( \text{SiLU}(W_1^{(k)} h_t) + h_t \right) \right)

with W_1^{(k)} \in \mathbb{R}^{d \times d} and W_2^{(k)} \in \mathbb{R}^{d \times V}. The clever part is the initialization: set W_1^{(k)} = 0 and W_2^{(k)} to the base model’s LM-head weights. Then

\text{SiLU}(0 \cdot h_t) + h_t = h_t \quad\Rightarrow\quad p_t^{(k)} = \text{softmax}(W_2^{(k)} h_t),

so every head begins as an exact copy of the language-model head. Training a head is just teaching it to shift its gaze k tokens forward from that starting point — which is why Medusa heads train fast and the backbone can stay frozen (the “Medusa-1” recipe).

Code: Bolt Heads onto the Model

medusa.py builds MedusaHead and MedusaModel from scratch. A head is exactly the formula above:

import torch
from medusa import MedusaHead

head = MedusaHead(dim=8, vocab_size=5)
print("W1 is zero at init:", bool(torch.count_nonzero(head.w1.weight) == 0))

# With W1 = 0, SiLU(0)+h = h, so the head is a pure read-out of h through W2.
h = torch.randn(2, 8)
print("head(h) == W2(h):", torch.allclose(head(h), head.w2(h), atol=1e-6))
W1 is zero at init: True
head(h) == W2(h): True

MedusaModel wraps a trained GPTModel, reads its normalised last hidden state, and copies the LM head into every Medusa head — so an untrained head reproduces the base next-token logits exactly:

import importlib.util, sys
from pathlib import Path

_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 medusa import MedusaModel

torch.manual_seed(0)
gpt = GPTModel(vocab_size=16, embed_dim=16, num_heads=2, num_layers=1, max_seq_len=32)
medusa = MedusaModel(gpt, num_heads=3)

base_logits, head_logits = medusa(torch.tensor([[1, 5, 9, 2, 7]]))
print("base logits shape:", tuple(base_logits.shape))
print("num heads:", len(head_logits))
# Each head starts identical to the base head (the init trick).
print("head 0 == base at init:", torch.allclose(head_logits[0], base_logits, atol=1e-5))
base logits shape: (1, 5, 16)
num heads: 3
head 0 == base at init: True

The Candidate Tree

Each head emits a distribution, so its top few tokens are all plausible. Rather than commit to one continuation, Medusa keeps the top-w tokens per head and arranges them into a tree: depth-1 nodes are the base head’s guesses, each of which branches into head 1’s guesses, and so on. A single forward pass then scores every node — the trick is a tree attention mask that lets each node attend only to its own ancestors (its prefix), never to a sibling on a different branch.

from medusa import build_candidate_tree, tree_attention_mask

# Base head proposes {7, 8}; the next head proposes {3, 4} under each.
tree = build_candidate_tree([[7, 8], [3, 4]])
print("tokens :", tree["tokens"])
print("parents:", tree["parents"])
print("paths  :", tree["paths"])   # every root-to-leaf continuation

mask = tree_attention_mask(tree["parents"])
print("\ntree attention mask (node i may attend column j):")
print(mask.int())
tokens : [7, 8, 3, 4, 3, 4]
parents: [-1, -1, 0, 0, 1, 1]
paths  : [[7, 3], [7, 4], [8, 3], [8, 4]]

tree attention mask (node i may attend column j):
tensor([[1, 0, 0, 0, 0, 0],
        [0, 1, 0, 0, 0, 0],
        [1, 0, 1, 0, 0, 0],
        [1, 0, 0, 1, 0, 0],
        [0, 1, 0, 0, 1, 0],
        [0, 1, 0, 0, 0, 1]], dtype=torch.int32)

Row i of the mask has a 1 only on node i and its ancestors — so node “3 under 7” sees 7, 3 but never 8 or “3 under 8”. That is what makes the parallel scan equivalent to having decoded each branch on its own.

Verify the Tree — and Why It’s Exact

Verification is greedy and local. Starting at the root, follow the child whose token equals the token the target itself would greedily emit next; repeat down the tree until no child matches. The accepted path is ratified, plus one bonus token (the target’s greedy next given the accepted prefix), so at least one token always comes out — exactly like a plain greedy step.

from medusa import verify_tree_greedy, greedy_autoregressive, greedy_next_ids_from_map

# A toy target as a next-token map (greedy next depends on the last token).
chain = {0: 1, 1: 2, 2: 3, 3: 4, 4: 0, 5: 5, 6: 6, 7: 7}
tree = build_candidate_tree([[1, 5], [2], [3], [7]])   # right until the last head (7 wrong)
node_next = greedy_next_ids_from_map(chain, tree["tokens"])

out = verify_tree_greedy(tree, root_next=chain[0], node_next=node_next)
print("accepted   :", out["accepted_tokens"])   # 1, 2, 3 accepted; 4th head guessed wrong
print("emitted    :", out["emitted"])           # accepted + one bonus token
print("plain greedy:", greedy_autoregressive(chain, last_token=0, steps=out["accepted_len"] + 1))
print("exact match :", out["emitted"] == greedy_autoregressive(chain, 0, out["accepted_len"] + 1))
accepted   : [1, 2, 3]
emitted    : [1, 2, 3, 4]
plain greedy: [1, 2, 3, 4]
exact match : True
NoteKey Insight

Every accepted token is the target’s greedy next token given its prefix, and so is the bonus. Therefore greedy Medusa verification emits exactly what plain greedy decoding of the target would — token for token. Self-drafting is the same pure-latency trade as classic speculative decoding: test_medusa.py proves the identity across 25 random targets and tree shapes. (Medusa’s sampling mode swaps this greedy check for typical acceptance — accept a token when the target’s probability clears min(ε, δ·exp(−H)) — which trades a little exactness for a higher acceptance rate; the greedy case is the clean, provable core.)

The self-draft in action

demonstrate_medusa overfits a tiny GPT on a short sequence, trains three heads to look ahead, then verifies a tree against the model itself. The heads learn to predict 2–4 tokens out, and the verified emit matches plain greedy decoding exactly.

# `med` was computed in the (hidden) bridge cell just above — reuse it.
print("per-head accuracy :", [round(a, 2) for a in med["head_accuracy"]])
print("verified emit     :", med["emitted"])
print("plain greedy      :", med["greedy_reference"])
print("exact match       :", med["matches_greedy"])
per-head accuracy : [0.93, 1.0, 1.0]
verified emit     : [7, 5, 9, 4, 8]
plain greedy      : [7, 5, 9, 4, 8]
exact match       : True

Watch the tree get verified

The tree below is the real one drafted by the trained heads. Step through the verification: accepted nodes light up along the winning path, and the emitted tokens (accepted + bonus) appear beneath.

TipTry This
  1. Step to the end. The emitted tokens (orange = accepted draft, green = the free bonus) match med_greedy — self-drafting changed nothing but speed.
  2. Widen the root. In demonstrate_medusa, the tree uses a 2-wide root then a chain. A wider tree holds more candidates per pass — more chances to accept a longer prefix, at the cost of a bigger verification.

How Much Faster? The Head-Accuracy Explorer

Speed comes entirely from how often each head’s top guess is accepted. For a linear chain of heads with per-depth accept probabilities p_k, a draft token counts only if every shallower head was also right, so

\mathbb{E}[\text{tokens/step}] = 1 + \sum_{k=1}^{K} \prod_{j=1}^{k} p_j.

Drag the per-head accuracies (they decay with depth in practice — predicting further ahead is harder) and watch the expected tokens per pass. The dashed line is the tested expected_accepted_length; the faint marker sits at the trained demo’s measured head accuracies.

TipTry This
  1. Flatten the decay. Set all three to 0.9: the sum climbs toward 4 tokens per pass. Perfect heads would give exactly K+1.
  2. Kill the deep heads. Drop p_2, p_3 to near 0: you fall back to ~1 token per pass — plain decoding. A head only helps if the shallower ones usually hit.
  3. This is why EAGLE exists. EAGLE drafts in feature space and feeds each accepted token back, raising the deep-head accept rates that this curve shows matter most.

EAGLE: Draft the Feature, Not the Token

Medusa’s heads share one strength and one weakness. The strength: they read the target’s own hidden state, so a draft costs almost nothing. The weakness is right there in the acceptance curve you just drove — the deep heads, the ones that reach furthest, accept least. Head 3 predicts the token three ahead from the same h_t that head 1 used, having never seen what head 1 or head 2 chose. The drafted positions are a fan: independent guesses from one point, blind to each other. A language model’s tokens are anything but independent, so the fan frays with depth.

EAGLE (Li et al., 2024) keeps self-drafting but turns the fan into a chain. Two changes make it work:

  1. Draft the feature, not the token. Instead of guessing a discrete token a few steps ahead, EAGLE predicts the next feature — the second-to-top-layer hidden state, the very vector the LM head consumes. “Autoregression at the feature level is more straightforward than at the token level”: features are smooth and continuous, so a tiny network can extrapolate them.
  2. Feed each drafted token back. The head is autoregressive: to draft step k it takes the feature it just predicted and the token that feature produced, so every drafted position is conditioned on the one before it — a chain.

The draft head is deliberately tiny, and it borrows almost everything from the target. It reuses the target’s own embedding table and its own LM head — frozen, never retrained. All it learns is a single Autoregression Head: one fully-connected layer that fuses [feature, next-token embedding] from 2d back down to d, followed by one decoder layer. The draft only has to get the feature right; the target’s head turns it into a token.

Intuition: a chain, not a fan

Watch one draft unroll. The context runs through the target once, giving a feature f and — free — the target’s own next token x_1. From there the Autoregression Head takes over: fuse f with the embedding of x_1, predict the next feature \hat f_1, read a token x_2 off the reused LM head, then feed x_2 back in to draft the next step. Each link knows the one before it.

NoteKey Insight

Medusa’s K heads all read the same h_t: the drafted positions are conditionally independent, so deep guesses are blind. EAGLE’s one head runs a chain — draft k is conditioned on the token drafted at k-1 — and it drafts the smooth feature rather than a discrete token, reusing the target’s frozen LM head to turn each feature into a token. Same one-pass-cheap idea, a far longer reach.

The Math: the Autoregression Head

Write f_i for the feature at position i (the LM head’s input), E for the target’s embedding table, and x_{i+1} for the token taken one step ahead. The head advances the feature by one step:

z_i = W_{\text{fc}} \begin{bmatrix} f_i \\ E\,x_{i+1} \end{bmatrix}, \qquad \hat f_{i+1} = \text{Decoder}\big(z_{\le i}\big), \qquad p_{i+1} = \text{softmax}\big(W_{\text{lm}}\,\hat f_{i+1}\big)

with W_{\text{fc}} \in \mathbb{R}^{d \times 2d} the fusion layer, \text{Decoder} a single causal transformer block, and W_{\text{lm}} the target’s own LM head (frozen). Only W_{\text{fc}} and the decoder layer are trained. Notice x_{i+1} in the input: that is the whole difference from a fan. The next feature is a function of the token actually taken — which is exactly what makes the chain coherent, and exactly what a token-blind head cannot see.

Feature Uncertainty: Why the Token Comes Back In

Why feed the token in at all — can’t the head predict f_{i+1} from f_i alone? No, and the reason is the paper’s title: feature uncertainty. The next feature depends on which token gets sampled. After a prefix ending “I”, a model might sample “am” or “always” — two different tokens, two different continuations, two different next features. Predicting f_{i+1} from f_i alone asks one input to map to two answers: ill-posed. Feed in the sampled token and the ambiguity is gone — each (f_i, x_{i+1}) pair has a single correct next feature.

We can make that exact with no training at all. Give one feature two legitimate next features (one per token). A token-blind predictor is a function of f alone, so its single output is stuck between the two targets — it must err on at least one branch. A token-aware predictor is a function of (f, \text{token}) and hits both exactly:

from eagle import feature_uncertainty_demo

fu = feature_uncertainty_demo()
print(f"token-BLIND  worst-branch error : {fu['blind_worst']:.4f}   (forced to err: {fu['blind_must_err']})")
print(f"token-AWARE  worst-branch error : {fu['aware_worst']:.1e}   (exact on both: {fu['aware_is_exact']})")
token-BLIND  worst-branch error : 0.1596   (forced to err: True)
token-AWARE  worst-branch error : 0.0e+00   (exact on both: True)
from eagle import demonstrate_eagle

_demo = demonstrate_eagle(draft_length=8)
ojs_define(
    eagleContext = _demo["context"],
    eagleDraft = _demo["draft_tokens"],
    eagleTrue = _demo["true_continuation"],
    eagleAccepted = _demo["accepted_len"],
    eaglePerStep = _demo["accept_rate_per_step"],
    fuBlindErr = float(fu["blind_worst"]),
)

Drag the token the model samples. A token-blind predictor (grey) is pinned to the midpoint — its error to whichever branch you pick is the gap you see. The token-aware predictor (gold) snaps onto the branch your token selects, error zero.

NoteKey Insight

The token fed “one step ahead” is not a detail — it is what makes feature-level drafting possible. Without it, the head is asked to be a function that maps one feature to several futures. With it, the map is single-valued, and a small network can learn it well enough to draft many tokens deep.

Code: A Draft Head from Scratch

eagle.py builds EagleDraftHead and EagleDrafter. The head is exactly the fusion-plus-decoder above; the drafter reads the feature the LM head consumes and reuses the target’s frozen embedding and head. First, the feature is the LM head’s input — passing it back through the head reproduces the model’s own logits:

import torch
from eagle import EagleDrafter, GPTModel

torch.manual_seed(0)
gpt = GPTModel(vocab_size=16, embed_dim=16, num_heads=2, num_layers=1, max_seq_len=32)
drafter = EagleDrafter(gpt, num_heads=2)

ids = torch.tensor([[1, 5, 9, 2, 7]])
feats = drafter.features(ids)                 # the second-to-top-layer feature
print("feature == LM-head input:", torch.allclose(drafter.lm_head(feats), gpt(ids), atol=1e-5))
print("draft reuses target embed:", torch.allclose(drafter.token_embed(ids), gpt.token_embedding(ids)))
feature == LM-head input: True
draft reuses target embed: True

Now unroll a draft chain. The first token is the target’s own greedy next (free and exact); every later token is drafted by the head, each conditioned on the previous:

draft = drafter.draft_chain(ids, length=5)
print("drafted chain:", draft["draft_tokens"], " (first token is the target's own next)")
drafted chain: [11, 9, 5, 1, 8]  (first token is the target's own next)

It’s Exact — the Same Verification

An untrained head drafts nonsense, and that is fine: correctness never depended on the draft. A chain is just a width-1 candidate tree, so the same greedy verification from the Medusa section ratifies it — the emitted tokens are exactly the target’s greedy decode, always. Here we train the tiny head (only the fusion layer and one decoder block) to advance the features, draft a chain, and verify it:

from eagle import demonstrate_eagle

demo = demonstrate_eagle(draft_length=8)
print("context           :", demo["context"])
print("drafted chain     :", demo["draft_tokens"])
print("true continuation :", demo["true_continuation"])
print("emitted (verified):", demo["emitted"])
print("greedy reference  :", demo["greedy_reference"])
print("exact match       :", demo["matches_greedy"], "  ← output is the target's, unchanged")
print("accepted in one pass:", demo["accepted_len"], "drafted +1 bonus")
print("per-step accept   :", [round(p, 2) for p in demo["accept_rate_per_step"]])
context           : [7, 2, 13]
drafted chain     : [5, 9, 0, 14, 3, 11, 6, 1]
true continuation : [5, 9, 0, 14, 3, 11, 6, 1]
emitted (verified): [5, 9, 0, 14, 3, 11, 6, 1, 15]
greedy reference  : [5, 9, 0, 14, 3, 11, 6, 1, 15]
exact match       : True   ← output is the target's, unchanged
accepted in one pass: 8 drafted +1 bonus
per-step accept   : [1.0, 1.0, 0.93, 0.92, 0.92, 0.91, 0.9, 0.89]

Because each step is conditioned on the last, the trained chain reaches the whole continuation — accepting many tokens in a single verification pass, and emitting exactly what greedy decoding would. That is the EAGLE payoff: Medusa’s fan frayed with depth; the chain holds.

TipTry This
  1. Shorten the chain. Call demonstrate_eagle(draft_length=3) and watch the accepted count — a shorter draft can’t accept more than it drafts.
  2. Break the feedback. In draft_chain, the token x is fed back each step. Imagine pinning it to a constant instead: the head would draft blind, and the chain would collapse to the fan it replaced.
  3. Read the acceptance decay. The per-step accept rates drift down with depth — the same shape as Medusa’s, but starting far higher because the chain is conditioned. This is exactly the curve EAGLE-2 attacks.

EAGLE-2: A Dynamic Draft Tree

EAGLE drafts one chain; Medusa drafts one fixed fan. Both decide the draft’s shape in advance — and that is the last thing left to fix, because which draft tokens survive verification is context-dependent, not a function of depth. Deep inside an easy, predictable span the chain should run long; at a single surprising token even depth 1 will be rejected, and every node you drafted past it is wasted. A shape fixed ahead of time cannot tell those two situations apart. The accept-vs-depth curve you just drove is an average over both — and an average is exactly what a dynamic drafter refuses to settle for.

EAGLE-2 (Li et al., 2024) makes the draft a context-aware dynamic tree, grown fresh at every step. It rests on one empirical fact about the head you just built:

the draft head is well-calibrated — its softmax confidence c_j closely approximates the token’s true acceptance rate p_j.

That is a gift: the drafter already knows, for free, where its guesses are likely to be accepted. EAGLE-2 just listens to it — pouring draft budget into the branches the head is confident about and starving the ones it isn’t.

NoteKey Insight

Give every node a value: the product of confidences along its path from the root, V_i = \prod_{j \in \text{path}} c_j. That single scalar is the head’s own estimate that the whole prefix ending at node i survives verification — deep-but-shaky paths score low, short-and-sure ones score high. The entire dynamic tree is built by chasing high value.

The Math: Value, Expand, Rerank

A draft node t_i sits at the end of a path from the root; its value is the running product of the draft confidences along that path,

V_i \;=\; \prod_{t_j \in \text{Path}(\text{root} \to t_i)} p_j \;\approx\; \prod_{t_j \in \text{Path}(\text{root} \to t_i)} c_j ,

so a child’s value is its parent’s, scaled by one more confidence: V_{\text{child}} = V_{\text{parent}}\cdot c_{\text{child}}. Two phases turn that number into a tree:

  1. Expand. Grow the tree one layer at a time. At each layer, run the draft model only on the k highest-value nodes of the frontier, giving each its top-b children. Promising branches deepen; dead ones are never touched again.
  2. Rerank. From every node produced, keep the m with the highest value (ties broken toward shallower nodes) and verify only those. That is the tree the target ratifies in one pass.
NoteKey Insight

Value-reranking is exact, not a heuristic, because of three facts that fall out of c \in [0,1]:

  • Monotone. V_{\text{child}} = V_{\text{parent}}\cdot c \le V_{\text{parent}}: value never grows as you go down.
  • Connected for free. So the m highest-value nodes are ancestor-closed — if a node is kept, its parent outscores it and is kept too. The top-m set is always a valid tree; no repair pass, no bookkeeping.
  • Optimal. Those same m nodes are the maximum-total-value set of that size — and since they form a connected subtree, the best size-m draft tree there is. The greedy pick is also the optimum.

Watch the Tree Grow, Then Prune

Below is a real dynamic tree, built by training a tiny EAGLE head (exactly as in the last section) and running expand_draft_tree on it. Node size is value — big nodes are confident prefixes. Step through the expansion (layers deepen only under the highest-value frontier), then the rerank that keeps the top few and drops the rest:

from eagle2 import demonstrate_eagle2

# A small, legible tree (depth 3, fan 2) for the widgets — same machinery,
# fewer nodes than the headline demo below so every node is readable.
_e2 = demonstrate_eagle2(depth=3, branch=2, expand_k=2, keep_m=4)
ojs_define(
    e2Tokens = _e2["exp_tokens"],
    e2Parents = _e2["exp_parents"],
    e2Depths = _e2["exp_depths"],
    e2Values = _e2["exp_values"],
    e2Kept = _e2["kept"],
    e2FixedValues = _e2["fixed_node_values"],
    e2KeepM = _e2["kept_node_count"],
)

Code: The Dynamic Tree from Scratch

The whole algorithm is two functions plus a rule for reading confidences off a drafter. expand_draft_tree grows the tree, deepening only the top-value frontier; rerank_draft_tree keeps the top-m by value. Neither touches a model — they take a path oracle (the only thing that knows the drafter), so every property above is exactly checkable:

from eagle2 import expand_draft_tree, rerank_draft_tree, values_are_monotone, is_ancestor_closed

# A toy oracle: at any path, offer two candidates whose confidence decays with depth.
def oracle(path):
    base = 0.7 - 0.1 * len(path)
    return [(len(path) * 10 + 1, base), (len(path) * 10 + 2, base - 0.3)]

exp = expand_draft_tree(oracle, depth=4, branch=2, expand_k=2)
print("value is monotone down every path :", values_are_monotone(exp))

rr = rerank_draft_tree(exp, keep_m=5)
print("top-5 kept set is a valid tree     :", is_ancestor_closed(exp["parents"], rr["kept"]))
print("kept tokens (breadth-first)        :", rr["tokens"])
value is monotone down every path : True
top-5 kept set is a valid tree     : True
kept tokens (breadth-first)        : [1, 2, 11, 12, 11]

The reductions are exact too: a fan-out of 1 is EAGLE’s chain, and equal confidences (nothing to prefer) recover Medusa’s dense Cartesian tree — the two fixed shapes fall out as special cases of the dynamic one:

from eagle2 import uniform_oracle
from medusa import build_candidate_tree

chain = expand_draft_tree(oracle, depth=5, branch=1, expand_k=1)
print("branch=1 → a straight chain, depths:", chain["depths"])

cands = [[7, 8], [3, 4], [9, 1]]
dense = expand_draft_tree(uniform_oracle(cands), depth=3, branch=2, expand_k=999)
print("uniform confidences → Medusa tree  :", dense["tokens"] == build_candidate_tree(cands)["tokens"])
branch=1 → a straight chain, depths: [1, 2, 3, 4, 5]
uniform confidences → Medusa tree  : True

It’s Exact, and It Spends Budget Better

Now the payoff on a trained drafter. demonstrate_eagle2 expands a context-aware tree, reranks it to a fixed verification budget, and checks it two ways: the emitted tokens are exactly the target’s greedy decode (lossless, as always), and the dynamic tree carries far more accumulated value than a fixed-shape tree of the same node count — so more of the budget lands on tokens that get accepted:

from eagle2 import demonstrate_eagle2

demo = demonstrate_eagle2()   # depth 5, fan 3, verify budget = 8 nodes
print(f"explored {demo['expanded_node_count']} nodes (a full tree would be {demo['full_cartesian_count']}), kept {demo['kept_node_count']}")
print("emitted (verified):", demo["emitted"])
print("greedy reference  :", demo["greedy_reference"])
print("exact match       :", demo["matches_greedy"], "  ← output is the target's, unchanged")
print(f"value mass — dynamic {demo['dynamic_value_mass']} vs fixed {demo['fixed_value_mass']}"
      f"  ({demo['dynamic_value_mass'] / demo['fixed_value_mass']:.2f}× the useful budget)")
print("dynamic tree is provably optimal   :", demo["dynamic_is_optimal"])
explored 39 nodes (a full tree would be 243), kept 8
emitted (verified): [5, 9, 0, 14, 3, 11]
greedy reference  : [5, 9, 0, 14, 3, 11]
exact match       : True   ← output is the target's, unchanged
value mass — dynamic 3.9868 vs fixed 1.9408  (2.05× the useful budget)
dynamic tree is provably optimal   : True

At the same eight-node budget, the dynamic tree packs about twice the acceptance-value of the fixed one — and it is provably the best size-8 tree there is, all from a scalar the drafter was already computing.

The Budget Dial

How much does the budget matter? Slide the number of nodes you are willing to verify and watch the two strategies’ accumulated value. The dynamic tree (gold) dominates the fixed one (grey) at every budget — the gap is widest when the budget is tight, exactly when spending it well matters most, and closes only when the budget is large enough to keep everything:

TipTry This
  1. Starve the budget. Set the dial to 2. The fixed tree spends its second node on a low-value sibling; the dynamic tree spends it on the confident child one level down — the gap is nearly 2\times.
  2. Open it up. Slide to the maximum. The curves meet: with room to keep every node, shape no longer matters — dynamic wins only under a real budget.
  3. Change the drafter. In the widget bridge, raise expand_k: the expansion explores wider before reranking, so the kept tree can only get better (or tie).

Where It Went: EAGLE-3

EAGLE-2 changed the draft’s shape; EAGLE-3 (2025) changes what is drafted at all. It drops EAGLE’s feature-regression constraint entirely — a trick it calls training-time test — and fuses low-, mid-, and high-level features instead of only the top one, predicting tokens directly. Freed from having to match a single feature, it keeps improving as training data grows, reaching up to 6.5×.

The through-line is the one you built across all four techniques: draft cheaply, spend the budget where it pays off, and let the target’s exact verification guarantee the output never changes.

Speculative decoding needs a second model. Medusa needs extra trained heads. Both keep the same shape: something guesses, the target verifies. Lookahead decoding (Fu, Bailis, Stoica & Zhang, 2024) removes even that — no draft model, no extra heads, no training. The model drafts for itself, out of a purely algebraic observation about what greedy decoding actually is.

Here is the observation. Greedy decoding of the next m tokens is not really a loop — it is the solution of a system of equations. Token y_i must equal the model’s argmax given everything before it:

y_i = \arg\max P(\,\cdot \mid \text{prefix},\, y_1, \dots, y_{i-1}\,), \qquad i = 1, \dots, m.

Autoregressive decoding solves this system one equation at a time, top to bottom — that is the sequential dependency. But a system of equations can also be solved by iteration in parallel: guess the whole block [y_1, \dots, y_m], then refine every position at once. That is Jacobi iteration, and one refinement is exactly one forward pass.

NoteKey Insight

A forward pass over a length-T sequence already computes the next-token logits at every position, not just the last — autoregressive decoding throws all but the last one away. Jacobi decoding uses them: feed a guessed block, read the argmax at every position, and you have refined all m guesses in a single pass.

The Math: A Triangular System, Solved in Parallel

One Jacobi step takes the current block of guesses g_0, \dots, g_{B-1} and computes, for every position k in parallel,

a_k = \arg\max P(\,\cdot \mid \text{prefix},\, g_0, \dots, g_{k-1}\,).

Two facts turn this into an exact, never-slower decoder:

  1. a_0 is always right. It depends only on the verified prefix, not on any guess — so it equals what autoregressive decoding would emit next. Every pass locks in at least one true token. Lookahead therefore never takes more passes than plain decoding; the block of m tokens is solved in at most m passes.
  2. Accept only on a match. a_k is the correct greedy token exactly when every earlier guess already matched (g_j = a_j for all j < k) — because then the context that produced a_k was the true prefix. So we accept the longest leading run of matches (always \ge 1), append it to the verified prefix, and carry the still-unconverged refinements forward as the next block’s guesses.

Because a token is emitted only when it equals greedy’s own argmax on a correct prefix, the output is token-for-token identical to greedy decoding. Lookahead is a pure latency win — the same proof of exactness as speculative decoding, reached without a draft distribution at all.

The stepper below walks the Jacobi passes on the simplest case — a stream that repeats one token — so you can watch a whole block lock in at once on the second pass:

Code: Jacobi Decoding From Scratch

lookahead.py builds the whole method against a plain oracle — a function from a context to its greedy next token — so the core is exact for any model. jacobi_step is one parallel refinement; accepted_prefix_length counts what a pass locks in; and jacobi_decode runs the loop with the carry.

import importlib.util, sys
from pathlib import Path

_lpath = Path("lookahead.py")
_spec = importlib.util.spec_from_file_location("lookahead", _lpath)
_la = importlib.util.module_from_spec(_spec)
sys.modules["lookahead"] = _la
_spec.loader.exec_module(_la)

greedy_decode = _la.greedy_decode
jacobi_decode = _la.jacobi_decode
lookahead_decode = _la.lookahead_decode
model_oracle = _la.model_oracle

# A repeated-token stream: the easy case where a whole block converges at once.
constant = lambda ctx: 5
_, greedy_passes = greedy_decode(constant, [0], 20)
gen, jacobi_passes, _ = jacobi_decode(constant, [0], 20, block_size=5)

print(f"greedy:  {greedy_passes} forward passes")
print(f"jacobi:  {jacobi_passes} forward passes  (block_size=5)")
print(f"same output? {gen == greedy_decode(constant, [0], 20)[0]}")
greedy:  20 forward passes
jacobi:  5 forward passes  (block_size=5)
same output? True

Twenty tokens in five passes instead of twenty — and the output is unchanged. Set block_size=1 and Jacobi reduces to plain autoregressive decoding exactly; that is the sanity floor.

The N-gram Pool: This Is Lookahead

Jacobi alone only wins when the block converges — a repeated token, as above. On real text the block rarely converges by itself. Lookahead decoding adds the piece that makes it pay off everywhere text repeats: it harvests the n-grams the verified stream has already produced into a small pool, and when a token recurs it re-proposes that token’s known continuation as the next guess block. Where generation loops — a cycle, boilerplate, a copied identifier, a for i in range(...) — the pooled n-gram matches on the very next pass and the whole block ratifies at once.

The window has two dials, matching the paper: W (the block width — how far ahead we look) and N (the n-gram size — how far back each harvested n-gram reaches).

# A tiny GPT overfit on a periodic pattern → greedy decodes a cycle, exactly the
# regime where the n-gram pool pays off. (Faithful: the same code drives a real model.)
import torch, torch.nn.functional as F

_tpath = Path("..") / "m06_transformer" / "transformer.py"
_tspec = importlib.util.spec_from_file_location("m06_transformer_transformer", _tpath)
_t = importlib.util.module_from_spec(_tspec)
sys.modules["m06_transformer_transformer"] = _t
_tspec.loader.exec_module(_t)
GPTModel = _t.GPTModel

torch.manual_seed(0)
_period = [5, 8, 2, 8]
_data = torch.tensor([(_period * 16)[:60]])
_model = GPTModel(vocab_size=16, embed_dim=32, num_heads=2, num_layers=2, max_seq_len=64)
_opt = torch.optim.Adam(_model.parameters(), lr=3e-3)
_model.train()
for _ in range(220):
    _logits = _model(_data[:, :-1])
    _loss = F.cross_entropy(_logits.reshape(-1, 16), _data[:, 1:].reshape(-1))
    _opt.zero_grad(); _loss.backward(); _opt.step()

demo = _la.demonstrate_lookahead(model_oracle(_model), [5, 8], n_new=24, block_size=5, ngram_n=4)

# Bridge the numbers to the widgets.
_grid_block = 5
jacobi_grid = {
    "trace": jacobi_decode(constant, [0], 12, block_size=_grid_block)[2],
    "block_size": _grid_block,
}
demo_passes = {
    "greedy": demo["greedy"]["passes"],
    "jacobi": demo["jacobi"]["passes"],
    "lookahead": demo["lookahead"]["passes"],
    "n_new": demo["n_new"],
}
ojs_define(jacobiGrid=jacobi_grid, demoPasses=demo_passes)
identical output across all three?  True
    greedy: 24 passes   1.00 tokens/pass
    jacobi: 24 passes   1.00 tokens/pass
 lookahead:  8 passes   3.00 tokens/pass

Same 24 tokens every time. Greedy and plain Jacobi both take 24 passes here — the cycle never converges inside a single block. Lookahead re-proposes the cycle from its pool and finishes in a third of the passes, for free.

Interactive Exploration

Drive the comparison. The bars are the forward-pass counts for the three decoders on the periodic run above — fewer is faster, and all three emit the identical sequence.

TipTry This
  1. Break the loop. In the code cell, change _period to a non-repeating pattern and re-run: the pool stops hitting and lookahead collapses back toward Jacobi — the speedup lives entirely in repetition.
  2. Shrink the window. Drop block_size to 2: fewer positions can ratify per pass, so the pass count climbs. This is the W dial from the paper.
  3. Confirm the floor. Set block_size=1: Jacobi and lookahead both become plain autoregressive decoding — exactly n_new passes, byte-identical output.
NoteKey Insight

Lookahead trades FLOPs for steps: each pass processes a wider window, but there are fewer passes. Because decoding is memory-bandwidth-bound, those extra parallel FLOPs are nearly free on real hardware — which is why the paper reports 1.5–2.3× on a single GPU with no draft model, no heads, and no training, and up to ~4× with lookahead parallelism across GPUs.

SMC Speculative Decoding: Resample, Don’t Reject

Lookahead removed the draft entirely; SMC-SD goes the other way and exploits the draft — even when the draft is wrong. Here is the problem it solves. Plain speculative decoding accepts the draft block only up to the first mismatch and then discards everything after it. The accepted-run length is stochastic: when draft and target diverge — exactly the cases where the draft is not a distilled copy — the block truncates to one or two tokens and the speedup evaporates. Rejection throws the draft’s guesses away wholesale.

SMC-SD (Emara et al., Faster LLM Inference via Sequential Monte Carlo, 2026) keeps the draft’s tokens and reweights them instead. It maintains a population of N draft particles — N “what-if” continuations of the prompt — and each round treats the draft blocks as importance-sampling proposals: every particle extends by exactly K+1 tokens (its K drafts plus one target continuation), with a weight saying how much the target agrees with the block. A sequential-Monte-Carlo resampling step prunes degenerate weight vectors, so the population keeps exploring where the target actually wants to go. Nothing is ever rolled back, and every round is the same fixed-size, vectorized operation.

That swap buys the paper’s headline — a deliberate reversal of the guarantee that has held for every method in this module:

“approximation quality is stochastic and the speed-up factor is deterministic.”

The output is no longer exactly target-distributed; it is approximately so, with a per-round error bound that shrinks as 1/N. In exchange, the speedup stops depending on the draft’s luck: SMC-SD emits exactly K+1 tokens per round, every round, no matter how bad the draft is. It is the first deliberately approximate accelerator in m16’s lossless line.

NoteKey Insight

Rejection throws the draft away at the first mismatch — throughput collapses exactly when draft and target diverge. SMC-SD reweights the draft instead: a bad block is not discarded, it is down-weighted, and the particle population is resampled toward the survivors. Verification stops being a pass/fail gate and becomes a fixed-size vectorized reweighting — no rollback, ever.

The Math: Reweight, Don’t Reject

At one position, with target distribution p and draft distribution q, the importance weight of a drafted token x \sim q is the likelihood ratio

w(x) = \frac{p(x)}{q(x)}.

A ratio of 1 means the models agree; >1 means the target likes x more than the draft did — the particle deserves more weight, not rejection; <1 means the draft over-proposed it. Along a whole K-token block d = d_1 \dots d_K the weights multiply, giving the block weight

W^{(n)} = \prod_{j=1}^{K} \frac{p\!\left(d_j^{(n)} \mid x^{(n)} d_{<j}^{(n)}\right)} {q\!\left(d_j^{(n)} \mid x^{(n)} d_{<j}^{(n)}\right)},

computed in log space as \exp\!\big(\sum_j \log p(d_j) - \log q(d_j)\big) — the target scores all K positions in one forward pass, exactly like SD’s verification, but the score is a continuous weight instead of a 0/1 accept test. (This section develops importance sampling from scratch; nothing here assumes Module 08 taught it — only its next-token sampling.)

With weights w^{(1)}, \dots, w^{(N)}, the particles define an empirical distribution over continuations,

\hat p_N(x) = \sum_{n=1}^{N} \frac{w^{(n)}}{\sum_{m} w^{(m)}}\, \mathbb{1}\{x^{(n)} = x\},

the self-normalized importance-resampling estimator (Eq. 5 of the paper). How healthy is the weight vector? The effective sample size,

\text{ESS} = \frac{\left(\sum_n w^{(n)}\right)^2}{\sum_n \left(w^{(n)}\right)^2}, \qquad 1 \le \text{ESS} \le N.

Equal weights give \text{ESS} = N (every particle contributes); one dominating weight drives it toward 1 (the population has collapsed onto a single trajectory). One Algorithm-1 round is then:

  1. Draft. Each particle independently drafts K tokens from q(\cdot \mid x^{(n)}).
  2. Score. One target pass scores every draft position of every particle.
  3. Reweight. w^{(n)} \leftarrow w^{(n)} \cdot W^{(n)} — multiply by the block likelihood ratio.
  4. Extend. Each particle samples x^{+} \sim p(\cdot \mid x^{(n)} d^{(n)}) and becomes x^{(n)} d^{(n)} x^{+} — exactly K{+}1 tokens, nothing rolled back.
  5. Resample. If \text{ESS} < \eta, draw ancestors a_n \sim \text{Cat}(\bar{w}^{(1..N)}) i.i.d., copy their particles, and reset all weights to 1/N.

Resampling is faithful in expectation: \mathbb{E}[\#\text{copies of particle } i] = N \bar{w}_i — heavy particles spawn clones, light ones die out, and the weights restart uniform.

The bound (Theorem 3.1). Writing \chi^2(p \| q) = \mathbb{E}_{q}\!\left[(W - 1)^2\right] for the \chi^2-divergence between draft and target, the estimator’s L¹ bias is

O\!\left(\sqrt{\frac{1 + \chi^2(p \| q)}{N}}\right),

and both the L² bias and the MSE are O\!\left((1 + \chi^2(p \| q))/N\right). Two readings: more particles always help (the 1/N factor), and the draft’s divergence sets the constant — a bad draft does not break SMC-SD, it just makes the approximation looser at a fixed N.

The speed (Eq. 3). When decoding is memory-bandwidth-bound, drafting and scoring N particles in parallel is nearly free — idle compute buys the acceleration — and the roofline speedup over serial decoding is

S^{(\text{mem})} = \frac{K+1}{\rho K + 1},

with \rho the draft:target cost ratio. The numerator is deterministic — K+1 tokens per round, always — which is the reversal made quantitative.

NoteHonest caveat

The paper proves the per-round reweighting bound above, but explicitly leaves one step out: “We leave bounding the error of the resampling step of SMC-SD to future work.” The theory covers everything except the resampling round’s contribution; in practice the method stays within a few percent of the target’s accuracy, but the full end-to-end proof is open.

Code: SMC-SD from Scratch

smc_sd.py implements Algorithm 1 against the same toy step-function convention as the rest of the module: a step function maps a prefix (a list of token ids) to a next-token distribution. Every piece below is the exact code the test suite exercises.

First, the effective sample size — the quantity that decides whether a round resamples:

import sys

# Lesson cells run with cwd = this module's folder; make smc_sd importable.
sys.path.insert(0, ".")
from smc_sd import effective_sample_size

# ESS = (sum w)^2 / sum w^2. Equal weights give N; one heavy weight gives ~1.
print("equal weights [1,1,1,1] :", effective_sample_size([1, 1, 1, 1]))
print("one heavy    [10,1,1,1] :", round(effective_sample_size([10, 1, 1, 1]), 3))
equal weights [1,1,1,1] : 4.0
one heavy    [10,1,1,1] : 1.641

Four equal weights give \text{ESS} = 4 = N — all particles count. Replacing one weight with 10 spreads the vector and collapses the ESS to 13^2/103 \approx 1.64: the population is effectively one particle. That is the degeneracy \eta guards against.

Now one full Algorithm-1 round. The target p is peaked (token 0 gets 40% of the mass); the draft is uniform — deliberately as wrong as a draft can be. With N = 8 particles, K = 2 drafts, and the standard threshold \eta = N/2 = 4, watch the weights spread and the ESS trip the resampler:

import torch
from smc_sd import smc_sd_round

rng = torch.Generator().manual_seed(0)
p = torch.tensor([0.40, 0.25, 0.15, 0.08, 0.05, 0.03, 0.02, 0.02])
q = torch.full((8,), 0.125)  # uniform draft: no agreement with p at all

N, K, eta = 8, 2, 4.0  # resample when ESS < N/2
parts, ws, tr = smc_sd_round(
    [[] for _ in range(N)], [1.0] * N, lambda _: q, lambda _: p, K, eta, rng
)

print("draft blocks (K=2 each) :", tr["draft_blocks"])
print("block ratios            :", [round(r, 2) for r in tr["block_ratios"]])
print("ESS                     :", round(tr["ess"], 2))
print("resampled (ESS < 4.0)?  :", tr["resampled"])
print("weights after round     :", [round(w, 3) for w in ws])
print("lengths                 :", [len(x) for x in parts], "  <- K+1 = 3 tokens, nothing rolled back")
draft blocks (K=2 each) : [[6, 3], [0, 1], [4, 2], [2, 3], [4, 5], [4, 0], [6, 4], [0, 6]]
block ratios            : [0.1, 6.4, 0.48, 0.77, 0.1, 1.28, 0.06, 0.51]
ESS                     : 2.15
resampled (ESS < 4.0)?  : True
weights after round     : [0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125]
lengths                 : [3, 3, 3, 3, 3, 3, 3, 3]   <- K+1 = 3 tokens, nothing rolled back

Particle 1’s draft [0, 1] drew the target’s two favorite tokens, so its block ratio is 6.4 — the target endorses that guess, and the weight grows. The ESS collapses to 2.15 < \eta = 4, so the round resamples: the heavy particle spawns clones, the light ones die out, and the printed weights are the reset 1/N = 0.125. Every particle nevertheless came out exactly 3 tokens longer — the K drafts plus one x^{+} — no truncation, no correction pass.

Run the full loop over the same pair — 7 new tokens, three rounds, resampling firing twice:

from smc_sd import smc_sd_decode

rng = torch.Generator().manual_seed(3)
out, trace = smc_sd_decode(lambda _: q, lambda _: p, [0], K, N, eta, 7, rng)

print("emitted sequence :", out)
for i, t in enumerate(trace):
    print(f"round {i+1}: ESS={t['ess']:5.2f}  resampled={t['resampled']}  lengths={t['lengths']}")
emitted sequence : [0, 7, 0, 1, 1, 0, 0, 0]
round 1: ESS= 5.77  resampled=False  lengths=[4, 4, 4, 4, 4, 4, 4, 4]
round 2: ESS= 1.44  resampled=True  lengths=[7, 7, 7, 7, 7, 7, 7, 7]
round 3: ESS= 2.98  resampled=True  lengths=[10, 10, 10, 10, 10, 10, 10, 10]

The output — drawn from the terminal normalized weights — is dominated by token 0, the target’s peak, and every round shows the same lengths across all 8 particles: 4, 7, 10, i.e. exactly 1 + 3 \cdot \text{round} tokens. A uniform draft gives plain SD an acceptance rate of \alpha = \sum_x \min(p(x), q(x)) = 0.575, i.e. an expected \frac{1 - \alpha^{3}}{1 - \alpha} \approx 1.9 tokens per pass, with most passes wasted; SMC-SD instead advances the whole population 3 tokens per round, twice triggering a resample that keeps the heavy particles in the lead. That is the fixed-size, no-rollback contract: speed is deterministic, quality is what N and \eta buy.

Two Exact Anchors

Two settings reduce SMC-SD to machinery you already trust. The tests pin both — anchor, identity, assert, True.

Anchor 1: a perfect draft is exact. If q \equiv p, every block ratio is 1, the weights never spread, \text{ESS} \equiv N, and resampling never fires — each particle is just a target sampler, exactly like SD at \alpha = 1:

rng = torch.Generator().manual_seed(1)
out, trace = smc_sd_decode(lambda _: p, lambda _: p, [0], K=2, N=8, eta=4.0, max_new_tokens=6, rng=rng)

all_ones = all(abs(r - 1.0) < 1e-9 for t in trace for r in t["block_ratios"])
ess_is_n = all(abs(t["ess"] - 8.0) < 1e-9 for t in trace)
never_fires = all(not t["resampled"] for t in trace)

print("every block ratio == 1.0 :", all_ones)
print("ESS == N on every round  :", ess_is_n)
print("resampling never fires   :", never_fires)
every block ratio == 1.0 : True
ESS == N on every round  : True
resampling never fires   : True

Anchor 2: K = 0 is plain autoregressive sampling. With no draft tokens the block weight is the empty product (1), x^{+} is drawn from p(\cdot \mid x), and each round extends every particle by exactly one token — SMC-SD reduces to Module 08’s sampling loop, N independent target samples per round:

rng = torch.Generator().manual_seed(2)
out, trace = smc_sd_decode(lambda _: q, lambda _: p, [0], K=0, N=8, eta=4.0, max_new_tokens=5, rng=rng)

grew_by_one = all(t["lengths"] == [1 + i + 1] * 8 for i, t in enumerate(trace))
weights_one = all(w == 1.0 for t in trace for w in t["weights"])
never_fires = all(not t["resampled"] for t in trace)

print("each round adds exactly 1 token :", grew_by_one)
print("weights stay 1.0                :", weights_one)
print("resampling never fires          :", never_fires)
each round adds exactly 1 token : True
weights stay 1.0                : True
resampling never fires          : True

Both anchors are structural: the first shows the approximation error is exactly zero when the draft is the target, the second shows the resampler is a no-op when there is nothing to reweight. Everything in between is the deliberate trade the next section prices.

The Cost of Approximate

Read the Theorem-3.1 bound once more:

\text{bias} \sim O\!\left(\sqrt{\frac{1 + \chi^2(p \| q)}{N}}\right), \qquad \text{MSE} \sim O\!\left(\frac{1 + \chi^2(p \| q)}{N}\right).

The direction is what makes SMC-SD practical: error is a function of N and of the draft’s divergence, not of the draft’s luck. Plain SD’s per-pass yield depends on where the first rejection lands — a coin flip every pass; SMC-SD’s error depends on how many particles you can afford and how far q is from p — both quantities you can measure and budget ahead of time. That is the reversal, again: approximation quality is stochastic (the estimator is a random variable; the bound holds in expectation), and the speed-up factor is deterministic (K+1 tokens every round).

What the numbers look like (reported in the paper, Llama + Qwen families on GSM8K, MATH500, AlpacaEval, DS1000): up to 2.5× throughput vs an optimized SD implementation within 3% of the target’s accuracy; on multi-GPU setups, 2.36× over state-of-the-art SD and 5.2× over autoregressive decoding (Llama 1B→70B, 4 H100s). Resampling itself is nearly free: duplicated particles are realized by KV-cache pointer exchange (no data movement) plus shared-prefix caching (72.3% cache reduction for Llama-1B→8B at N=8, K=16).

The cost is real, and it is not free lunch:

WarningApproximate, on purpose

The estimator is self-normalized: \hat p_N divides by \sum_m w^{(m)}, so it is biased for any finite N (the unbiasedness of Lemma B.1 holds only for the unnormalized estimator, in expectation). Worse, N = 1 degenerates completely — one particle, weight always reset to 1 after resampling, and the output trusts the draft’s trajectory without correction. SMC-SD is the first method in this module where the output is not provably the target’s; the error is small and bounded, but it is nonzero, and it grows as the draft drifts.

The engineering answer to that caveat is the same as everywhere else in the book: keep N large enough that \frac{1 + \chi^2(p \| q)}{N} is small, keep \eta near the standard N/2, and verify on your own workload. The widget below lets you watch both sides of the trade — the fixed speedup and the growing error — at once.

Interactive Exploration

First, step through one Algorithm-1 round and watch what replaces rejection: the draft blocks are scored, the weights spread, the ESS gauge trips, and the population is resampled — all at a fixed K+1 tokens per particle.

TipTry This
  1. Read the weight bars at Step 3. Particles 0 and 2 hold most of the mass after reweighting — the others drafted blocks the target disliked. Step 4’s ESS gauge says the same thing in one number.
  2. Notice what never happens. Across all six steps no token is rejected or rolled back: each particle only ever gains the K drafts and one x^{+}.
  3. Contrast with the SD walkthrough at the top of the module. There, rejection stopped the block at the first mismatch; here the same divergence is absorbed as a weight, and the round completes at full K+1 length.

Now drive the particles-vs-rejection comparison itself. The data is bridged from Python — the same demonstrate_smc_sd sweep the tests exercise: for each \delta \in \{0, 0.1, \dots, 1.0\} it builds the interpolated draft q_\delta = (1-\delta)\,p + \delta\,q_{\text{far}}, computes SD’s expected accepted tokens \frac{1 - \alpha^{K+1}}{1 - \alpha} (which collapses as the draft drifts), SMC-SD’s fixed K+1 per round, and the seed-averaged empirical error of the SMC estimator.

from smc_sd import demonstrate_smc_sd

ojs_define(smcData = demonstrate_smc_sd())

At \delta = 0.5 — the dial’s start — the draft is half-way to q_{\text{far}}: SD’s expected accepted run has already shrunk from 5.0 to 4.10 tokens per pass, SMC-SD still emits its fixed 5.0, and the empirical error of \hat p_N sits at \approx 0.054 with \chi^2(p \| q_\delta) = 0.057. Drag \delta to 1.0: SD slides to 3.36 tokens while SMC-SD’s block never moves — and the error curve is the price of the fixed block.

TipTry This
  1. Slide \delta from 0 to 1. SD’s expected accepted run collapses (5.0 → 3.36 tokens) exactly as the draft drifts; SMC-SD’s K+1 = 5 line never moves — the deterministic speed-up.
  2. Watch the error, not just the lines. The dashed curve is the cost of being approximate: at \delta = 1.0 the empirical error is \approx 0.063 and \chi^2(p \| q_\delta) = 0.209 — the bound’s constant. More particles shrink the error; divergence sets how large it can be.
  3. Read the bars. The wide bars are the interpolated draft q_\delta at the dial, the narrow bars the target p. Resampling is what keeps SMC-SD’s population near p’s peak (token 0) even as the draft spreads its mass elsewhere.

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.
Medusa: a node attends the wrong prefix If the tree mask lets a node see a sibling branch, the parallel scan no longer equals decoding that branch alone. Build the mask so each node attends only its ancestors (tree_attention_mask).
Medusa: a dense tree that explodes A full Cartesian tree grows as the product of the head widths — huge verification cost for little gain. Prune to a sparse tree of the highest-probability paths (Medusa’s tuned tree).
Medusa: deep heads left untrained The furthest heads have the lowest accept rate, so an untrained one just wastes a tree slot. Train all heads (cheap, backbone frozen); weight the loss toward the reachable depths.
EAGLE: feature drift with depth The head consumes its own predicted features when drafting, so small feature errors compound and a chain drafted too deep collapses. Train the head free-running (feed its predictions forward), not only teacher-forced; keep the chain to the depth its acceptance still pays for.
EAGLE: relearning the LM head Giving the draft head its own vocabulary projection wastes parameters and drifts from the target. Reuse the target’s frozen embedding and LM head; train only the fusion layer + decoder.
EAGLE: dropping the fed-back token Predicting the next feature from the feature alone is ill-posed (feature uncertainty) — the chain becomes a blind fan. Always fuse in the embedding of the token one step ahead; that is what makes the map single-valued.
EAGLE-2: ranking by raw confidence, not path value A deep node with a locally-high confidence can still sit on an unlikely path; ranking by cⱼ alone breaks the acceptance estimate. Rank by the path product Vᵢ = ∏ cⱼ — that is what makes value monotone and the kept set a valid tree.
EAGLE-2: reranking without the monotone tie-break If ties don’t favour shallower nodes, a kept node’s ancestor can be dropped and the “tree” is disconnected. Sort by (value desc, depth asc); monotonicity then guarantees ancestor-closedness with no repair pass.
EAGLE-2: treating acceptance as depth-only Fixing the tree shape in advance (Medusa/EAGLE-1) mis-spends budget — acceptance is context-dependent. Grow the tree per step from the drafter’s own confidence; spend budget where it’s likely to be accepted.
Lookahead: expecting a win on non-repetitive text The n-gram pool only pays off where the stream repeats; on prose with no recurrence it collapses to plain Jacobi (~1 token/pass). Expect the big wins on code and structured/boilerplate output; measure, don’t assume.
Lookahead: “the init must be a good guess” It feels like the initial block should be accurate. It doesn’t — correctness holds for any init; the init only affects how many passes. Any fill works (repeat the last token); tune block_size/ngram_n for speed, never for correctness.
Lookahead: reading it as approximate Because it “guesses”, it looks lossy. It is not — only greedy-matching tokens are accepted, so the output is byte-identical to greedy.
SMC-SD: reading the weights as a distribution The self-normalized p̂_N is an estimator, not the target: it is biased for any finite N, and N=1 degenerates to a single trajectory that trusts the draft uncorrected.
SMC-SD: assuming quality is deterministic too The reversal cuts both ways — the speedup is deterministic (K+1 tokens every round), but the quality is stochastic and loosens as the draft drifts.

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

Exercise 4: Medusa’s expected speedup

Use expected_accepted_length to see how per-head accuracy sets the speedup. Head accuracy decays with depth in practice — try a geometric decay and find where extra heads stop paying off.

from medusa import expected_accepted_length

for decay in [0.9, 0.7, 0.5]:
    probs = [decay ** k for k in range(1, 6)]  # 5 heads, accuracy decaying with depth
    tokens = expected_accepted_length(probs)
    print(f"decay={decay}: per-head p={[round(p, 2) for p in probs]} -> {tokens:.2f} tokens/pass")

# Your task: with decay=0.7, how much does the 5th head add over the 4th?
decay=0.9: per-head p=[0.9, 0.81, 0.73, 0.66, 0.59] -> 3.72 tokens/pass
decay=0.7: per-head p=[0.7, 0.49, 0.34, 0.24, 0.17] -> 2.19 tokens/pass
decay=0.5: per-head p=[0.5, 0.25, 0.12, 0.06, 0.03] -> 1.64 tokens/pass

Exercise 5: Where lookahead first beats Jacobi

Jacobi alone helps only when a block converges. Sweep block_size on a short cycle and find the smallest window at which lookahead’s n-gram pool starts saving passes over plain Jacobi.

from lookahead import jacobi_decode, lookahead_decode

cycle = lambda ctx: {3: 1, 1: 2, 2: 3}[ctx[-1]]  # the 3-cycle 1,2,3,1,2,3,...
for B in range(1, 6):
    _, jp, _ = jacobi_decode(cycle, [3], 30, block_size=B)
    _, lp, _ = lookahead_decode(cycle, [3], 30, block_size=B, ngram_n=3)
    print(f"block_size={B}: jacobi {jp:2d} passes, lookahead {lp:2d} passes")

# Your task: at which block_size does lookahead first undercut Jacobi, and why not B=1?
block_size=1: jacobi 30 passes, lookahead 30 passes
block_size=2: jacobi 30 passes, lookahead 17 passes
block_size=3: jacobi 30 passes, lookahead 12 passes
block_size=4: jacobi 30 passes, lookahead 10 passes
block_size=5: jacobi 30 passes, lookahead 10 passes

Exercise 6: EAGLE’s chain is exact for any draft

The whole point of every technique here is that verification, not the draft, sets correctness. Confirm it for EAGLE: draft a deliberately wrong chain and check the emitted tokens still equal the target’s greedy decode.

import torch
from eagle import EagleDrafter, GPTModel, verify_chain_greedy, target_greedy_next

torch.manual_seed(1)
gpt = GPTModel(vocab_size=16, embed_dim=16, num_heads=2, num_layers=1, max_seq_len=32)
ctx = torch.tensor([[2, 7, 1]])

# A garbage draft — nothing to do with the model.
bad_draft = [0, 0, 0, 0]
res = verify_chain_greedy(gpt, ctx, bad_draft)
print("emitted          :", res["emitted"])
print("greedy reference :", res["greedy_reference"])
print("still exact?     :", res["matches_greedy"], "(and it emitted", len(res["emitted"]), "token via the bonus)")

# Your task: how many tokens come out of a wrong draft, and why is it never zero?
emitted          : [12]
greedy reference : [12]
still exact?     : True (and it emitted 1 token via the bonus)

Exercise 7: Value is monotone, so reranking is a valid tree

EAGLE-2’s whole “connected for free” claim rests on one inequality. Confirm it, then confirm the consequence: the top-m nodes by value never orphan an ancestor.

from eagle2 import expand_draft_tree, rerank_draft_tree, values_are_monotone, is_ancestor_closed

def oracle(path):
    base = 0.7 - 0.1 * len(path)
    return [(len(path) * 10 + 1, base), (len(path) * 10 + 2, base - 0.25)]

exp = expand_draft_tree(oracle, depth=4, branch=2, expand_k=3)
print("V_child <= V_parent on every edge :", values_are_monotone(exp))
for m in range(1, len(exp["tokens"]) + 1):
    kept = rerank_draft_tree(exp, keep_m=m)["kept"]
    assert is_ancestor_closed(exp["parents"], kept)
print("every top-m set is a valid tree   : True")

# Your task: hand-build a tree where a node's confidence is > its parent's confidence.
# Value still can't increase downward — why? (Hint: c <= 1, and V is a running product.)
V_child <= V_parent on every edge : True
every top-m set is a valid tree   : True

Exercise 8: When does the resampler first fire?

Sweep the draft from exact to far and find the divergence at which ESS first dips below \eta — the point where reweighting alone stops being enough and the population must be cloned. (Hint: resampling fires when the weight vector spreads; a peaked target amplifies the spread.)

import torch
from smc_sd import smc_sd_round, effective_sample_size

rng = torch.Generator().manual_seed(0)
p = torch.tensor([0.40, 0.25, 0.15, 0.08, 0.05, 0.03, 0.02, 0.02])
q_far = torch.full((8,), 0.125)
N, K, eta = 8, 2, 4.0

for delta in [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]:
    q = (1 - delta) * p + delta * q_far
    _, _, tr = smc_sd_round([[] for _ in range(N)], [1.0] * N,
                            lambda _: q, lambda _: p, K, eta, rng)
    print(f"δ={delta:.1f}: ESS={tr['ess']:5.2f}  resampled={tr['resampled']}")

# Your task: which δ first trips the resampler, and why does a more peaked p
# make it trip earlier? (Hint: ESS measures weight spread, and w = p/q spreads
# as q flattens the mass p concentrates.)
δ=0.0: ESS= 8.00  resampled=False
δ=0.2: ESS= 7.38  resampled=False
δ=0.4: ESS= 7.34  resampled=False
δ=0.6: ESS= 4.08  resampled=False
δ=0.8: ESS= 2.83  resampled=True
δ=1.0: ESS= 2.57  resampled=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).
  6. Self-drafting removes the second model. Medusa’s extra heads read the target’s own hidden state to draft K tokens from one pass; a tree-attention mask verifies a whole tree of candidates at once, and greedy verification is provably exact — same output, fewer passes, no draft model to host.
  7. EAGLE turns the fan into a chain. Medusa’s heads draft independently and fray with depth. EAGLE drafts the smooth feature with one autoregressive head, feeds each drafted token back so every step is conditioned on the last, and reuses the target’s frozen embedding and LM head. Feeding the token “one step ahead” resolves feature uncertainty; the same verification keeps it exact.
  8. EAGLE-2 makes the tree dynamic. The draft head is calibrated, so a node’s value — the product of confidences along its path — estimates its acceptance. Expand the top-value frontier, rerank to a budget: monotone value makes the kept set a valid tree and the maximum-value one, for free. Same lossless output, but the budget lands where acceptance is likely — ~2× a fixed tree’s useful mass at the same node count (3.05–4.26× end to end).
  9. Lookahead removes the draft entirely. Greedy decoding is a fixed-point system; Jacobi iteration solves it in parallel, one forward pass per step, locking in \ge 1 true token each pass — so it is exact and never slower, with no draft model, no heads, and no training. Reusing the trajectory’s n-grams ratifies whole blocks wherever the stream repeats, buying 1.5–2.3× for free.
  10. SMC-SD reweights, it doesn’t reject. A population of N draft particles extends by exactly K+1 tokens every round — no rollback, ever — with the block’s importance weight W = ∏ p(dⱼ)/q(dⱼ) deciding who survives; when ESS = (Σw)²/Σw² falls below η, ancestor resampling clones the heavy particles and resets the weights. The error is bounded (O(1/N) with a constant set by χ²(p‖q)), the speedup is deterministic, and the output is only approximately the target’s — the first deliberate trade in this module.

What’s Next

You have now built all four faces of speculative decoding: a separate draft model, Medusa’s self-draft heads, EAGLE’s feature-space chain with EAGLE-2’s dynamic draft tree on top, and lookahead’s draft-free Jacobi solving. All are pure-latency wins that leave the output untouched. The remaining frontier here is tree-verified sampling with typical acceptance, threading the dynamic tree through a real GPTModel generation loop, and EAGLE-3’s multi-level feature fusion — deeper drafts, still exact — and the module’s first deliberately approximate accelerator, SMC-SD: importance-weighted resampling over draft particles replaces rejection, buying a deterministic K+1 tokens per round at the price of a bounded approximation error.

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: