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.
  • Build SMC-SD: keep a population of N draft particles and resample, don’t reject — reweight each K-token block by its importance weight and clone the survivors — so every round emits exactly K+1 tokens at a bounded O(1/N) approximation error: the module’s first deliberately approximate accelerator.
  • Make particles affordable with pointer exchange: page the KV cache, share the prompt’s pages across all N particles, and resample by exchanging page pointers — never copying KV tensors — with copy-on-write protecting shared pages.
  • Build SSD: overlap, don’t serialize — run the speculator and the verifier as concurrent processes on separate hardware, predict each round’s outcome and pre-speculate continuations into a speculation cache, so a hit hides drafting entirely and a miss falls back to lossless ordinary SD.

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-n 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.

The Systems Half: Per-Particle KV Caches with Pointer-Exchange Resampling

The statistics half you just built has a memory problem it never mentions. Every particle in smc_sd.py carries its own full prefix — N particles, N copies of the prompt’s KV cache. That is fine for a toy over step functions; it is not shippable. At N = 8 you would pay eight times the KV memory before generating a single token, and every resampling round would copy the survivors’ caches onto their clones. The systems half fixes both with one idea: page the cache, share the prefix, and resample by exchanging pointers — what the paper calls pointer exchange (§1), built on the PagedAttention (Kwon et al., 2023) plus RadixAttention (Zheng et al., 2024) lineage (§3.3).

The mapping onto machinery you already know is direct. A page is one fixed-size segment of what Module 08‘s KVCache stores as a single growable per-sequence block: instead of one private cache per particle, a shared pool owns fixed-size pages and each particle owns only a block table, a list of page ids. The prompt pages are allocated once; all N tables point at them, with a refcount counting the sharers. It is the same block-table idea as the paged-attention section’s PagedKVCache, pointed at a new job — making a particle population affordable. Resampling then moves no token data at all: it retains the survivors’ page ids — a page-metadata copy plus a refcount increment, an O(l_{seq}) metadata operation instead of an O(l_{seq} \times d_{model}) KV-tensor copy (§3.3, Observation 2) — and releases the old tables, freeing exactly the pages no survivor still needs.

The arithmetic is the book’s naive-upper-bound model, refined by prefix sharing — not the paper’s formula. Write L_p for the shared prompt length, L_s for each particle’s generated length, and c for the cache bytes per token. The naive layout, every particle holding its own full copy, costs

M_{naive} = N \cdot (L_p + L_s) \cdot c,

while the shared layout pays for the prompt once and only the generations per particle,

M_{shared} = (L_p + N \cdot L_s) \cdot c,

so the fractional cache saving is

reduction = \frac{(N-1) \cdot L_p}{N \cdot (L_p + L_s)}.

When the prompt dominates (L_p \gg L_s) the saving tends to (N-1)/N. The paper’s reported operating point sits in exactly that regime: a 72.3% cache reduction for Llama-1B→8B (Llama 3.2-1B→3.1-8B) at N = 8, K = 16, measured against naive per-particle KV caches with no prefix or page sharing — stated as inline text in the final sentence of §3.3, Observation 2.

smc_systems.py simulates exactly this layout: a KVPagePool of fixed-size pages with refcounts, and per-particle block tables that share the prompt. Build two particles over an 8-token prompt at 4 tokens per page — the prompt costs 2 pages, allocated once:

import sys

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

# Two particles, an 8-token prompt, 4 tokens per page: the prompt costs 2 pages, once.
pool = KVPagePool(page_size=4)
caches = ParticleCaches(pool, n_particles=2, prompt_tokens=[1, 2, 3, 4, 5, 6, 7, 8])
caches.extend(0, [9, 10, 11, 12])
caches.extend(1, [13, 14, 15, 16])

print("live pages (shared) :", caches.pages_used())
print("naive pages         :", naive_pages_total(8, 4, 2, 4))
print("reduction           :", round(caches.cache_reduction(), 3))
print("prompt refcounts    :", [pool.pages[pid]["refcount"] for pid in caches.tables[0][:2]])
live pages (shared) : 4
naive pages         : 6
reduction           : 0.333
prompt refcounts    : [2, 2]

Two prompt pages, each with refcount 2 — both particles point at the same page ids. Extending gives each particle one private page (the shared tail is full, so copy-on-write allocates), for 4 live pages against 6 naive: a reduction of 0.333. Now the paper’s regime shape — N = 8, a 256-token prompt, 64 generated tokens per particle, 16 per page — plus one pointer-exchange resample:

import torch
from smc_systems import KVPagePool, ParticleCaches, naive_pages_total

# The paper's regime shape: N = 8, a 256-token prompt, 64 generated tokens, 16 per page.
pool = KVPagePool(page_size=16)
caches = ParticleCaches(pool, n_particles=8, prompt_tokens=list(range(256)))
for i in range(8):
    caches.extend(i, [1000 * (i + 1) + j for j in range(64)])

print("naive pages  :", naive_pages_total(256, 64, 8, 16))
print("shared pages :", caches.pages_used())
print("reduction    :", round(caches.cache_reduction(), 3))

rng = torch.Generator().manual_seed(0)
ancestors = [int(a) for a in torch.multinomial(
    torch.full((8,), 1 / 8, dtype=torch.float64), 8, replacement=True, generator=rng)]
before = [list(s) for s in caches.sequences()]
freed = caches.pointer_exchange_resample(ancestors)
print("ancestors    :", ancestors)
print("freed pages  :", freed)
print("pages after  :", caches.pages_used())
print("pointer copy == naive deep copy :", all(
    caches.sequences()[i] == before[ancestors[i]] for i in range(8)))
naive pages  : 160
shared pages : 48
reduction    : 0.7
ancestors    : [7, 5, 3, 7, 5, 6, 1, 2]
freed pages  : 8
pages after  : 40
pointer copy == naive deep copy : True

160 naive pages against 48 shared: a reduction of 0.7 — the demo’s page-count model landing in the same regime as the reported 72.3%. The resample copied no sequences: it retained the ancestors’ page ids, freed the 8 pages owned only by extinct lineages (40 live afterwards), and every particle’s materialized tokens are exactly its ancestor’s — a pointer copy identical to a naive deep copy.

The widget’s curves are bridged from Python — the same demonstrate_pointer_exchange sweep the tests exercise: real ParticleCaches populations over N = 1 \dots 16 at shared-prefix fractions 0.5, 0.8, and 0.95, reading pages_used after per-particle extends.

from smc_systems import demonstrate_pointer_exchange

ojs_define(sysCacheData = demonstrate_pointer_exchange())

At the defaults — N = 8, prefix fraction 0.8 — the readout says naive 160 pages, shared 48, reduction 70.0%: the demo’s page model against the paper’s 72.3% in bytes. Drag N to 1 and both bars agree (reduction 0 — one particle shares with nobody); push the prefix dial to 0.95 and the reduction climbs to 83.1%, the prompt-dominated limit. The dashed line marks the paper’s N = 8 regime; the solid line is your dial.

WarningNever write into a shared page

A page with refcount > 1 belongs to several particles at once: writing one particle’s token into it would corrupt every sharer. extend enforces the copy-on-write rule — it fills the tail page in place only when that page has room and its refcount is 1, and allocates a fresh page otherwise. If you ever bypass extend and touch pool.pages[pid]["tokens"] directly, check the refcount first: shared pages are read-only.

The paper proves the per-round reweighting bound and leaves the resampling step’s error unbounded — measuring that gap is Exercise 10 below.

Speculative Speculative Decoding: Overlap, Don’t Serialize

Every method in this module has attacked one axis: the draft’s quality — better guesses (SD), self-drafting heads (Medusa), feature-space chains and dynamic trees (EAGLE, EAGLE-2), no draft at all (lookahead) — and SMC-SD rethought what rejection means. Speculative speculative decoding (SSD) (Kumar, Dao & May, Speculative Speculative Decoding, 2026; the optimized engine is called Saguaro) attacks the dependence every one of those methods still contains: round T{+}1’s speculation cannot start until round T’s verification returns. Look at the idle time: while the draft runs, the target’s GPUs wait; while the target verifies, the draft waits. Two machines, alternately half-idle — that serialization is what SSD removes.

SSD runs the speculator and the verifier as two concurrent processes on separate hardware — the draft gets its own GPU instead of sharing the target’s. While the verifier checks the current speculation s^T, the draft (a) predicts the likely verification outcome, (b) pre-speculates continuations for many possible outcomes in a single pass, and (c) stores them in a speculation cache. When the real outcome arrives, a cache hit returns the pre-made speculation immediately — drafting overhead vanishes. A cache miss falls back to synchronous speculation, which is lossless because it reduces to ordinary SD. The paper’s analogy is CPU speculative execution: while the current branch resolves, the CPU pre-executes the likely next branches so the pipeline never stalls — the gamble is free when right, and a flush is always safe. SSD gambles with draft compute instead of pipeline state, and the verifier’s accept rule makes every gamble lossless either way: the verifier never sees the cache, only drafts produced by the draft model from the true context, checked exactly as in SD, with one outcome sent per round (a single NCCL sync).

That completes the module’s arc in a deliberate way: SMC-SD rethought rejection (resample, don’t reject); SSD rethinks round structure (overlap, don’t serialize) — and it sits back on the exact, lossless line the module opened with. As with every engine section here, no kernels are simulated: latency is modeled with two explicit dials (T_p, T_b) against one normalized verifier pass, and the paper’s own reported numbers are quoted honestly at the end.

NoteKey Insight

Drafting overhead is a waiting problem, not a quality problem. SD’s serial round costs T_p + 1: a draft’s worth of target idle time, then a verify’s worth of draft idle time. SSD removes the waiting on two conditions — the draft lives on separate hardware, and it predicts outcomes well enough that most rounds hit the cache. The output never depends on those conditions: on a miss the fallback is ordinary SD, so a cache miss costs latency, never correctness.

The Math: One Round, Two Timelines

Set up the round exactly as in the SD section. A speculation s^T is K tokens drafted from q; the verifier returns the outcome v^T = (k, t^{*}): exactly k of the K tokens are accepted, followed by one more token t^{*} drawn from the residual r \propto (p - q)_{+} when k < K (the correction), or from p itself when all K are accepted — k = K — (the bonus). The acceptance machinery is unchanged, so both workhorses carry over:

\alpha \;=\; \sum_x \min(p(x), q(x)) \;=\; 1 - \mathrm{TV}(p, q), \qquad E \;=\; \frac{1 - \alpha^{K+1}}{1 - \alpha},

with E the expected tokens a round emits and P(k = K) = \alpha^{K} the share of rounds that end in full acceptance.

What SSD changes is the timeline, not the tokens. Normalize one verifier pass to one unit of time. Let T_p be the time the primary draft needs to prepare its speculation — in SSD, all the pre-speculated branches at once, which the paper’s sparse-mask branch decode keeps near a single draft pass — and let T_b be the time a cheaper backup draft needs to re-speculate a single continuation on demand. A serial SD round drafts, then verifies:

T_{\text{SD}} = T_p + 1.

An SSD round overlaps the two on separate hardware, so its duration depends on the round’s coin flip. On a hit the draft finished inside the verifier pass: the round costs \max(T_p, 1) — exactly 1 when T_p < 1, the regime where the draft is fully hidden. On a miss the verifier pass is spent and the backup must then draft and be verified: 1 + T_b. Per round:

T_{\text{SSD}} = p_{\text{hit}} \cdot \max(1, T_p) + (1 - p_{\text{hit}}) \cdot (1 + T_b),

and the rate comparison is expected tokens over round time,

\frac{E}{T_p + 1} \qquad \text{vs} \qquad \frac{E}{p_{\text{hit}} \cdot \max(1, T_p) + (1 - p_{\text{hit}}) \cdot (1 + T_b)} .

Two readings before any code. First, with T_p < 1 a hit round is exactly one verifier pass — the entire drafting cost vanishes; this is the whole point, and the paper’s App. B.1 confirms it is the dominant effect end to end (“most end-to-end speedup comes from hiding draft latency”). Second, since the backup is faster than the primary (T_b < T_p), even an all-miss SSD round, 1 + T_b, is no slower than SD’s serial round — SSD is never worse than SD in the toy, at any hit rate. As in the rest of the module the draft’s quality is untouched by the round structure, so the toy keeps E_{\text{hit}} = E_{\text{miss}} = E: every SSD gain is a denominator gain. That is a stated simplification, not a hidden one — richer drafts on the hit path would only help SSD.

The Prediction Problem: Which Outcomes Do You Arm?

The catch is in the word predicted: the draft must arm its branches before the verifier returns v^T = (k, t^{*}), and the outcome space is huge — K+1 accept lengths times V possible bonus tokens. The draft cannot pre-speculate all of them, so Saguaro frames the question as a budget allocation: with a budget B of pre-armed speculations, how many bonus-token guesses F_k should be armed for each accept length k, with

\sum_{k=0}^{K} F_k \;\le\; B ?

Each armed branch is a full K-token continuation, so the cache holds up to B \cdot K drafted tokens; the paper generates all of them in one forward pass with a custom sparse attention mask over the B branches (each branch attends its own continuation, all sharing the verified prefix), then serves them from the cache. The verifier never sees the cache — it receives exactly one speculation per round and checks it with the ordinary accept rule.

Where should the budget go? The outcome probabilities are the ones from the SD section:

P(k) = \alpha^{k}(1 - \alpha) \;\; (k < K), \qquad P(K) = \alpha^{K},

so for a good draft the mass sits at the extremes: rejections happen early, and the single largest atom is full acceptance. Saguaro arms the bonus-token guesses from the draft’s own top logits at each position — in the position-invariant toy, the top-F_k tokens of q; the paper adds a refinement the toy drops (the armed set excludes the token the draft just sampled — ssd.py’s banner states this). The cache hit rate is then

p_{\text{hit}} = \sum_{k=0}^{K-1} P(k) \cdot \Pr\!\big(t^{*} \in \mathrm{top}\text{-}F_k(q) \mid k\big) \;+\; P(K) \cdot \Pr\!\big(t^{*} \in \mathrm{top}\text{-}F_K(q)\big),

with the inner probabilities taken over the residual draw (k < K) or the target draw (k = K). Maximizing p_{\text{hit}} under the budget is a constrained optimization, and the paper proves (Theorem 12) that the optimal allocation is geometric: along the partial-acceptance positions the fan-out decays by a factor a_p^{1/(1+r)} per accepted token, and the all-accepted position carries an extra (1 - a_p)^{-1/(1+r)} factor — which makes it the largest single share whenever \alpha > 1/2:

F_k = F_0 \cdot a_p^{\,k/(1+r)} \;\; (k < K), \qquad F_K = F_0 \cdot a_p^{\,K/(1+r)} (1 - a_p)^{-1/(1+r)},

with F_0 set by \sum_k F_k = B and r the power-law exponent of the empirical miss curve 1 - p_{\text{hit}}(F) = F^{-r} (the paper measures miss rates falling as a power law in fan-out, Fig. 3). For \alpha near 1, the (1 - a_p)^{-1/(1+r)} term is huge: the budget leans into the deep positions and especially the all-accepted branch — the outcomes rounds actually land on. With a well-shaped cache the paper reports predicting the bonus token correctly up to 90% of the time.

There is one more lever, and it is subtle. The residual is (p - q)_{+} — it depends on the draft’s distribution. Saguaro sampling exploits that: by downweighting the draft’s own top tokens, the residual mass moves onto exactly the tokens whose continuations are cached, so corrections after a rejection land in the cache by construction. First the code, then that construction.

Code: SSD from Scratch

ssd.py implements the whole round model against the same toy convention as the rest of the module: fixed next-token distributions p and q over a small vocabulary, position-invariant — no kernels, no KV caches, no NCCL, just the algebra above plus a seeded round simulation. Every function below is the exact code the test suite exercises.

First the acceptance machinery — unchanged from the SD section, but now it feeds the outcome distribution and the hit rate:

import sys
import torch

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

# The section's sweep pair: a strong draft with NO ties, so "top-F" is
# unambiguous. The target p is flatter than the draft q, and the residual mass
# (tokens 1-3) sits exactly where q under-proposes.
p = torch.tensor([0.35, 0.30, 0.20, 0.15])
q = torch.tensor([0.50, 0.25, 0.15, 0.10])

alpha = acceptance_rate(p, q)
print("alpha      :", round(alpha, 4), "   (alpha = 1 - TV = 1 - 0.15)")
print("residual r :", [round(x, 4) for x in residual_distribution(p, q).tolist()])
print("E[K = 4]   :", round(expected_accepted_tokens(alpha, 4), 4), "tokens per SD round")
alpha      : 0.85    (alpha = 1 - TV = 1 - 0.15)
residual r : [0.0, 0.3333, 0.3333, 0.3333]
E[K = 4]   : 3.7086 tokens per SD round

\alpha = 0.85: a strong draft, but not a copy — \mathrm{TV} = 1 - \alpha = 0.15, so 15% of the mass differs. The residual puts the correction mass on tokens \{1, 2, 3\}, the tokens the target likes more than the draft does, so rejections are corrected from the tail of q. E = 3.71 tokens per round is what a serial SD round earns.

Now the outcome distribution the cache must cover — which accept lengths actually happen:

from ssd import outcome_probs

probs = outcome_probs(alpha, K=4)
print("P(exactly k accepted):", [round(x, 4) for x in probs])
print("sum                  :", round(sum(probs), 6))
P(exactly k accepted): [0.15, 0.1275, 0.1084, 0.0921, 0.522]
sum                  : 1.0

Rejections concentrate at the very first position (0.15) and full acceptance is the largest single outcome (\alpha^4 = 0.522); the long partial runs in between are the rarest rounds of all. A budget spent evenly across accept lengths is therefore spent mostly on outcomes that almost never happen — the geometric allocation exists because the outcome mass is this lopsided.

The budget machinery, the hit rates it buys, and the rates those hit rates produce:

from ssd import geometric_fan_out, hit_probability, tokens_per_time_ar, tokens_per_time_sd, tokens_per_time_ssd

K = 4
ar_tpt = tokens_per_time_ar(p)
print(f"AR  : {ar_tpt:.4f} tok/time   (autoregressive baseline: one pass per token)")
e_sd = expected_accepted_tokens(alpha, K)
sd_tpt = tokens_per_time_sd(alpha, K, t_p=0.5)
print(f"SD  : {sd_tpt:.4f} tok/time   (serial round = T_p + 1 = 1.5, E = {e_sd:.4f})")

flat5 = geometric_fan_out(alpha, K, budget=5)          # budget 5 forces one guess per k
uniform15 = [3, 3, 3, 3, 3]                            # fifteen branches, spent evenly
geo15 = geometric_fan_out(alpha, K, budget=15)         # fifteen branches, spent geometrically

ph_flat5 = hit_probability(p, q, flat5, K)
ph_uniform15 = hit_probability(p, q, uniform15, K)
ph_geo15 = hit_probability(p, q, geo15, K)
for name, fo, ph in [("flat      B =  5", flat5, ph_flat5),
                     ("uniform   B = 15", uniform15, ph_uniform15),
                     ("geometric B = 15", geo15, ph_geo15)]:
    ssd_tpt = tokens_per_time_ssd(ph, e_sd, t_p=0.5, t_b=0.25)
    print(f"{name}: fan-out {fo}   p_hit = {ph:.4f}   SSD {ssd_tpt:.4f} tok/time   ({ssd_tpt / sd_tpt:.2f}x SD)")
AR  : 1.0000 tok/time   (autoregressive baseline: one pass per token)
SD  : 2.4724 tok/time   (serial round = T_p + 1 = 1.5, E = 3.7086)
flat      B =  5: fan-out [1, 1, 1, 1, 1]   p_hit = 0.1827   SSD 3.0794 tok/time   (1.25x SD)
uniform   B = 15: fan-out [3, 3, 3, 3, 3]   p_hit = 0.7624   SSD 3.5007 tok/time   (1.42x SD)
geometric B = 15: fan-out [3, 3, 3, 2, 4]   p_hit = 0.8100   SSD 3.5404 tok/time   (1.43x SD)

One guess per accept length covers only the all-accepted branch’s most likely bonus, and it misses every rejection: p_{\text{hit}} \approx 0.18 and SSD earns about 3.08 tokens per unit time (1.25\times SD). Fifteen branches spent evenly — three guesses per accept length — jump the hit rate to \approx 0.76 and the rate to \approx 3.50 (1.42\times SD): the extra guesses sit on q’s tokens 1 and 2, the residual tokens inside the draft’s top-3, so two thirds of rejections now hit and the three guesses cover 85% of full-acceptance bonuses. The geometric split spends the same fifteen branches where the outcome mass actually is — one fewer on the rarest partial outcome (k = 3), the extra on the all-accepted branch k = K, which ends with the largest single share (four guesses) — so it clears the uniform row: p_{\text{hit}} \approx 0.81 and \approx 3.54 tok/time (1.43\times SD). That is Theorem 12’s point, which Exercise 9 checks against your own eyes.

Two structural facts about the numbers. The gain is bounded: even with every round a hit, the toy only reaches the 1.50\times ceiling, because a hit round is one verifier pass against SD’s 1.5 — the ceiling is exactly hiding the draft. And the floor is not SD:

for ph, tag in [(0.0, "all misses: the floor"), (1.0, "all hits: the ceiling")]:
    r = tokens_per_time_ssd(ph, e_sd, t_p=0.5, t_b=0.25)
    print(f"p_hit = {ph:.1f}  ({tag}): {r:.4f} tok/time   ({r / sd_tpt:.2f}x SD)")
p_hit = 0.0  (all misses: the floor): 2.9669 tok/time   (1.20x SD)
p_hit = 1.0  (all hits: the ceiling): 3.7086 tok/time   (1.50x SD)

At p_{\text{hit}} = 0 SSD still earns 2.97 (1.20\times SD): an all-miss round costs the backup’s 1 + T_b = 1.25, still less than the serial 1.5. Pre-arming pays whenever a hit is possible, and it pays more the better the draft predicts — but only the denominator moves.

The Worked Toy: Same Acceptance, Twice the Hits

The paper’s own construction (App. A.3.1) shows Saguaro sampling strictly speeding SSD up, on a vocabulary of four tokens where the draft is already nearly perfect:

p = (0.48,\ 0.48,\ 0.02,\ 0.02), \qquad q = (0.49,\ 0.49,\ 0.01,\ 0.01).

Here \alpha = 0.98: SD already earns E \approx 4.80 tokens per round and there is essentially no acceptance headroom left. The remaining lever is the hit rate, and it is a sampling lever, because the residual (p - q)_{+} is a function of the draft’s distribution. Saguaro sampling (Definition 14) downweights the draft’s own top-F tokens — exactly the tokens whose continuations the cache arms — before sampling:

q'(t) \;\propto\; \begin{cases} C \cdot q(t) & t \in \mathrm{top}\text{-}F(q), \\ q(t) & \text{otherwise,} \end{cases} \qquad C < 1,

renormalized. The draft samples the tail more often, and the residual — the mass where p exceeds q' — moves onto the downweighted top tokens: the correction tokens become the armed tokens by construction.

from ssd import saguaro_draft, acceptance_rate, residual_distribution

# The paper's App. A.3.1 pair (vocabulary of 4).
p_paper = torch.tensor([0.48, 0.48, 0.02, 0.02])
q_paper = torch.tensor([0.49, 0.49, 0.01, 0.01])

q_sag = saguaro_draft(p_paper, q_paper, F=2, C=47 / 147)
alpha_paper = acceptance_rate(p_paper, q_paper)
alpha_sag = acceptance_rate(p_paper, q_sag)
print("q'  (Saguaro, F = 2, C = 47/147):", [round(x, 4) for x in q_sag.tolist()])
print("alpha(p, q)  =", round(alpha_paper, 4), "   alpha(p, q') =", round(alpha_sag, 4))
print("residual(p, q)  =", [round(x, 4) for x in residual_distribution(p_paper, q_paper).tolist()])
print("residual(p, q') =", [round(x, 4) for x in residual_distribution(p_paper, q_sag).tolist()])
q'  (Saguaro, F = 2, C = 47/147): [0.47, 0.47, 0.03, 0.03]
alpha(p, q)  = 0.98    alpha(p, q') = 0.98
residual(p, q)  = [0.0, 0.0, 0.5, 0.5]
residual(p, q') = [0.5, 0.5, 0.0, 0.0]

Run it and check the two claims. First, acceptance is unchanged: q' = (0.47, 0.47, 0.03, 0.03) gives \sum_x \min(p, q') = 0.98 exactly like q — downweighting moved 0.04 of draft mass off the over-proposed pair (0.49 \rightarrow 0.47 each) onto the under-proposed tail (0.01 \rightarrow 0.03 each), and the min-sum did not move, so the round emits the same expected tokens on either draft. Second, the residual flips onto the armed set: under plain q the residual sits on tokens \{3, 4\} — corrections come from the tail, while the cache arms q’s top tokens \{1, 2\}; under q' the residual is (0.5, 0.5, 0, 0) — corrections come from exactly the armed tokens.

Why does that double the hit rate? Walk one rejection. Under plain q a rejection can only happen when the draft sampled its over-proposed token 1 or 2, and the correction is drawn from the tail \{3, 4\} — but the armed set is built from the draft’s own top guesses, which point at \{1, 2\}. In the module’s simplified top-F model no tail correction is ever armed; the paper’s model excludes the just-sampled token, so after a rejected token 1 the armed set is \{2, 3\} and only one of the two corrections (token 3) is covered — their per-rejection rate of 50%. Under q' the roles flip: a rejection can only happen on a sampled token 3 or 4, the residual is (0.5, 0.5, 0, 0) = \{1, 2\}, and the cache has armed exactly \{1, 2\} — every correction hits (100%). Same acceptance, and the per-rejection cache-hit rate goes 50% → 100% in the paper’s counting: “twice the hits” at zero acceptance cost. The paper proves this construction is not a curiosity: there exist draft/target pairs where Saguaro sampling is strictly faster end to end (their Theorem 19), precisely because “same E, higher p_{\text{hit}}” turns into a speedup through their Theorem 7.

What does doubling the hits buy in tokens per unit time? Feed the paper’s two hit rates through the round-time model with the toy’s dials (T_p = 0.5, T_b = 0.25):

from ssd import expected_accepted_tokens, tokens_per_time_sd, tokens_per_time_ssd

e_paper = expected_accepted_tokens(alpha_paper, 4)
sd_paper = tokens_per_time_sd(alpha_paper, 4, t_p=0.5)
print(f"E per SD round        : {e_paper:.4f} tokens   (alpha = {alpha_paper:.4f})")
print(f"SD rate (serial)      : {sd_paper:.4f} tok/time")
for ph, tag in [(0.5, "paper: plain draft, 50% hits"),
                (1.0, "paper: Saguaro draft, 100% hits")]:
    ssd_rate = tokens_per_time_ssd(ph, e_paper, t_p=0.5, t_b=0.25)
    print(f"SSD rate ({tag}): {ssd_rate:.4f} tok/time   ({ssd_rate / sd_paper:.2f}x SD)")
E per SD round        : 4.8040 tokens   (alpha = 0.9800)
SD rate (serial)      : 3.2026 tok/time
SSD rate (paper: plain draft, 50% hits): 4.2702 tok/time   (1.33x SD)
SSD rate (paper: Saguaro draft, 100% hits): 4.8040 tok/time   (1.50x SD)

The same E = 4.80 tokens per round on every line — the entire win is the denominator: hiding the draft on hits. At the plain draft’s 50% hit rate SSD earns 4.27 (1.33\times SD); at Saguaro’s 100% it earns 4.80 (1.50\times SD), the toy’s ceiling. This is the paper’s dominant effect made visible: Saguaro sampling does not make the draft agree with the target better — it makes the cache agree with the draft’s outcomes.

The Honest Numbers

How far does overlap go on real engines? The paper reports SSD end to end (App. B.3) with Llama-3.1-70B-Instruct as target on 4×H100 (tensor parallel) and Llama-3.2-1B as draft on its own H100, batch size 1, greedy decoding, across four datasets — tokens per second as AR / SD / SSD:

  • GSM8k: 54.7 / 188 / 301 — the headline up to 5.50× over AR;
  • HumanEval: 54.7 / 176 / 283 — 5.17× over AR;
  • Alpaca: 54.7 / 145 / 224 — 4.10× over AR;
  • UltraFeedback: 54.7 / 138 / 215 — 3.93× over AR;
  • Average: 54.7 / 161.8 / 255.8 — 1.58× over SD, 4.68× over AR.

With Qwen3-32B as target and a 0.6B draft: 88.8 / 136.8 / 203.8 tok/s on average — 1.49× over SD, 2.29× over AR. The paper’s abstract says “up to 5× faster than autoregressive decoding”, which is the GSM8k/HumanEval territory — the honest summary is up to ~5.5×, ~4.7× on average for Llama, ~2.3× for Qwen, ~1.5× over the authors’ own SD. And the “~30% faster than the strongest speculative decoding baselines” claim is measured against the best open-engine sweep, not against the authors’ SD: 255.8 tok/s vs ≈198–202 tok/s for the best open-engine SD configurations (vLLM v0.16 / SGLang v0.5.9, standalone and EAGLE-3) — about 1.29× — and the paper reports SSD still outperforming SD by ~20% at larger batch sizes despite doing more draft compute per sequence.

The section’s honest caveats are the paper’s own (Sec. 6), and they matter for what SSD is:

NoteKey Insight

SSD, like every SD-family method, is a latency win — the paper is explicit that it is “generally ineffective for throughput-bound workloads” such as large-scale RL or offline generation, where the batch is already large enough to keep the target busy. It also costs a separate draft device (five H100s here: four tensor-parallel plus one draft GPU), and adding more draft GPUs has diminishing returns. The gains ride on the cache-hit rate, which falls as temperature and batch size rise, and the custom sparse-mask branch decode dominates the critical path and needs KV-rollback bookkeeping on the target side. Composing SSD with EAGLE-3-style drafts has its own wrinkle: the draft would speculate round T{+}1 before round T’s activations exist, forcing it to self-condition on up to 2K tokens of its own output, which lowers acceptance on the latter K unless the draft is trained for it.

The toy’s clean 1.2×–1.5× band is the shape, not the size: real drafts are slower relative to the target than the toy’s T_p = 0.5, and real verification is more expensive than a toy pass, so hiding a real draft behind a real pass buys more. What transfers exactly is the logic — separate hardware, predict, arm geometrically, hide on hits, fall back losslessly.

Interactive Exploration

First, walk the round structure itself — SD’s serial draft→verify against SSD’s two overlapped lanes, and what a hit and a miss each do to the timeline:

The four idle labels on Step 0 are the whole motivation: SD alternates, and each machine spends half the round waiting. Step 1 runs the same round on two devices and the waiting is gone. Steps 2–4 show the price of the overlap — the draft must guess the outcome, so it arms a budget of branches, and the coin flip between Steps 3 and 4 decides whether the round costs 1 or 1 + T_b. Step 5 is the steady state: once the pipeline is full, speculation is always ready a round ahead and the verifier never idles.

Now drive the pre-arming dial itself. The data is bridged from Python — the same demonstrate_ssd() sweep the tests exercise: for each budget in a power-of-two sweep it computes the toy’s \alpha and expected tokens E, the geometric fan-out F_0 \dots F_K under that budget, the resulting hit probability, and the tokens-per-unit-time for AR, SD, and SSD (with T_p = 0.5, T_b = 0.25, so a hit round costs 1 and a miss round costs 1.25):

from ssd import demonstrate_ssd

ojs_define(ssdData = demonstrate_ssd())

At the middle budget the SSD curve sits above the flat SD line but below the all-hits ceiling — the gap is the hit rate the fan-out buys. Drag the dial: the SSD curve climbs as the budget doubles, and each doubling buys less — the first branches cover the outcomes that carry most of the probability mass, and later branches chase ever-rarer bonus tokens. The fan-out bars show where the budget went: always leaning into the deep accept lengths, with the all-accepted branch taking the largest single share. And the ratio readout never drops below 1 — even a starved cache keeps the toy’s miss-only floor of 1.2×, because a miss with the fast backup is still cheaper than a serial round.

TipTry This
  1. Read Step 0’s idle labels, then Step 1’s. The four idle zones of the serial round are what SSD deletes. Note what doesn’t change across the steps: the verifier pass, the accept rule, and the tokens emitted — overlap moves the draft in time, never the output.
  2. Slide the budget dial from 2 to 32. Watch p_{\text{hit}} climb and the SSD/SD ratio flatten: the first branches catch the probable outcomes, the last ones chase rare bonus tokens. Find where doubling the budget stops moving the ratio — that is the point of diminishing returns the paper’s Fig. 4 shows at high temperature.
  3. Starve the cache (B = 2). The fan-out cannot arm every accept length, yet the SSD line never dips below SD — the miss fallback (with T_b < T_p) is the floor, and it is lossless. A miss costs latency, not correctness; that is the design’s safety net.
  4. Tie it to the worked toy. The paper pair’s residual flip is what makes hits the lever when \alpha is already 0.98: same acceptance, twice the hits, and the entire gain in the denominator. Exercise 9 makes both effects quantitative.

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. Verify with the == greedy_decode anchor; treat it as a pure latency win, like the rest of m16.
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. Keep N large enough that (1 + χ²(p‖q))/N is small; verify on your own workload; never claim exactness.
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. Budget N and η for your worst-case draft divergence, not the average; watch the empirical error in the widget.
SSD: reading “lossless” as “always overlapped” Exactness says nothing about when a speculation exists: a hit hides drafting only when the draft really finishes before the verifier (T_p < 1), and temperature, batch size, and a slow draft all push rounds onto the miss path. Keep T_p < 1 on its own device and watch the hit rate — losslessness is a floor, not a speedup.
SSD: treating a cache miss as an error A miss is not a correctness event: the fallback speculates synchronously from the true outcome, which is ordinary SD — same output distribution, drafting latency just comes back. Arm the fan-out for the miss rate you can afford; never “repair” a miss by changing what gets verified.

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

Exercise 9: When does pre-arming pay?

All of SSD’s extra machinery — the separate device, the speculation cache, the fan-out — exists to buy cache hits, so it is worth being precise about when they pay. Verify the three claims behind this section: (a) the geometric fan-out dominates a uniform split at the same budget; (b) the worked toy’s numbers — acceptance unchanged at 0.98, the residual flipped onto the armed set, and the hit rate climbing; (c) SSD’s tokens-per-time never drops below SD’s in the toy.

import torch
from ssd import (acceptance_rate, residual_distribution, saguaro_draft,
                 geometric_fan_out, hit_probability, expected_accepted_tokens,
                 tokens_per_time_sd, tokens_per_time_ssd, simulate_ssd_rounds)

# (a) the same budget, two spendings — the section's sweep pair
p = torch.tensor([0.35, 0.30, 0.20, 0.15])
q = torch.tensor([0.50, 0.25, 0.15, 0.10])
alpha = acceptance_rate(p, q)

flat = [3, 3, 3, 3, 3]
geo = geometric_fan_out(alpha, K=4, budget=15)
print("uniform  :", flat, " p_hit =", round(hit_probability(p, q, flat, 4), 4))
print("geometric:", geo, " p_hit =", round(hit_probability(p, q, geo, 4), 4))

# (b) the worked toy: acceptance unchanged, residual now inside the armed set
p_paper = torch.tensor([0.48, 0.48, 0.02, 0.02])
q_paper = torch.tensor([0.49, 0.49, 0.01, 0.01])
q_sag = saguaro_draft(p_paper, q_paper, F=2, C=47 / 147)
print("alpha before/after :", round(acceptance_rate(p_paper, q_paper), 4),
      round(acceptance_rate(p_paper, q_sag), 4))
print("residual flips from", [round(x, 4) for x in residual_distribution(p_paper, q_paper).tolist()],
      "to", [round(x, 4) for x in residual_distribution(p_paper, q_sag).tolist()])
print("p_hit before/after :",
      round(hit_probability(p_paper, q_paper, [2, 2, 2, 2, 2], 4), 4),
      round(hit_probability(p_paper, q_sag, [2, 2, 2, 2, 2], 4), 4))

# (c) the round-time consequence across the whole hit-rate range
e_sd = expected_accepted_tokens(alpha, 4)
sd_tpt = tokens_per_time_sd(alpha, 4, t_p=0.5)
print("SD rate:", round(sd_tpt, 4))
for ph in [0.0, 0.25, 0.5, 0.75, 1.0]:
    r = tokens_per_time_ssd(ph, e_sd, t_p=0.5, t_b=0.25)
    print(f"p_hit={ph:.2f}: SSD rate = {r:.4f}   ({r / sd_tpt:.2f}x SD)")
print(simulate_ssd_rounds(p, q, K=4, budget=15, t_p=0.5, t_b=0.25,
                          n_rounds=2000, seed=0))

# Your task: (a) which line wins, and why does the flat split waste branches —
# which accept length is over-armed, and why is the all-accepted branch's bonus
# the easiest to cover? (b) confirm alpha = 0.98 on both drafts and the residual
# now lives inside the armed top-2 set; the paper's 50% -> 100% counts the same
# flip per rejection with their exclude-the-sampled-token rule (see the worked
# toy) — explain why the unconditional p_hit moves too, but less. (c) confirm
# SSD's rate stays at or above SD's for every p_hit in [0, 1] whenever the
# backup is no slower than the primary (T_b <= T_p), with equality only in the
# all-miss limit (p_hit = 0) when T_b = T_p — a miss costs exactly a serial SD
# round. Then compare the seeded simulation's tokens-per-time and hit/miss
# counts with the formula rows above: they should agree within the seed's
# tolerance (hits converge to p_hit, rates to the round-time model).
uniform  : [3, 3, 3, 3, 3]  p_hit = 0.7624
geometric: [3, 3, 3, 2, 4]  p_hit = 0.81
alpha before/after : 0.98 0.98
residual flips from [0.0, 0.0, 0.5, 0.5] to [0.5, 0.5, 0.0, 0.0]
p_hit before/after : 0.8855 0.9631
SD rate: 2.4724
p_hit=0.00: SSD rate = 2.9669   (1.20x SD)
p_hit=0.25: SSD rate = 3.1231   (1.26x SD)
p_hit=0.50: SSD rate = 3.2966   (1.33x SD)
p_hit=0.75: SSD rate = 3.4905   (1.41x SD)
p_hit=1.00: SSD rate = 3.7086   (1.50x SD)
{'tokens': 7437, 'time': 2093.5, 'tokens_per_time': 3.552424170050155, 'hits': 1626, 'misses': 374, 'p_hit': 0.813}

Exercise 10: Measure the resampling gap

The paper proves the per-round reweighting bound (Theorem 3.1) and stops there:

NoteOpen question

“We leave bounding the error of the resampling step of SMC-SD to future work.”

Your job is the teachable version of that gap: measure how much the resampling step itself contributes to the empirical error, on the toy smc_sd machinery. The scaffold below runs the same decode with resampling on (\eta = N/2) and off (\eta = 0.0, which never fires since \text{ESS} \ge 1):

import torch
from smc_sd import smc_sd_decode

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)
N, K = 8, 2

for eta in [0.0, N / 2]:
    rng = torch.Generator().manual_seed(0)
    out, trace = smc_sd_decode(lambda _: q, lambda _: p, [0], K, N, eta, 7, rng)
    fires = sum(t["resampled"] for t in trace)
    print(f"eta={eta:3.1f}: resamples={fires}, "
          f"min ESS={min(t['ess'] for t in trace):.2f}, emitted={out}")
eta=0.0: resamples=0, min ESS=2.15, emitted=[0, 6, 3, 1, 1, 1, 2, 2]
eta=4.0: resamples=2, min ESS=2.15, emitted=[0, 4, 2, 1, 1, 0, 2, 1]

Your task: sweep seeds and particle counts, and estimate each setting’s output distribution the honest way — per run, the fraction of the 7 generated tokens (after the prompt) equal to the target’s peak token (token 0, p(0) = 0.40), averaged over seeds — with resampling on and off. Report the two errors |\hat p(0) - 0.40| as a function of N: does the resampler shrink the error, add noise, or both at different N? Then say which half of the round your measurement still leaves unbounded, and why the quote above makes this an exercise and not a theorem.

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.
  11. SSD overlaps, it doesn’t serialize. The last serial link is the round structure: run the speculator and verifier as concurrent processes on separate hardware, and while the verifier checks s^T let the draft predict the outcome and pre-speculate continuations for many of them into a speculation cache (budget B, per-length fan-out F_k, geometric-optimal). A hit serves the next speculation the instant verification returns — drafting latency gone — and a miss falls back to synchronous speculation, i.e. ordinary SD, so the output stays exactly the target’s on every path. All the gain is latency: E is unchanged, and the method needs T_p < 1 and a separate draft device.
  12. Pages, not copies, make particles affordable. Sharing the prompt’s KV pages across all N particles and resampling by pointer exchange — a page-metadata copy plus a refcount increment, never a KV-tensor copy — cuts the cache from N full copies toward one shared prefix (the paper reports 72.3% off at N = 8, K = 16, §3.3). Shared pages are read-only: never write into a page with refcount > 1 — extend copy-on-write.

What’s Next

You have now built six approaches to the same wait: draft chains verified prefix-by-prefix, Medusa’s self-draft fan with tree attention, EAGLE-2’s dynamic draft tree, the deliberately approximate turn — SMC-SD’s resample-don’t-reject particles with pointer-exchange paging to make them affordable — and SSD, which overlaps speculation with verification on separate hardware. The lossless methods are pure-latency wins that leave the output untouched; SMC-SD trades a bounded error for a deterministic block. 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. The module’s closing section, speculative speculative decoding (SSD / Saguaro), attacks the one serial dependence the rest leave alone: drafting round T{+}1 waits on verifying round T. SSD overlaps the two on separate hardware and pre-arms a speculation cache with continuations for predicted outcomes — a hit hides drafting entirely, a miss falls back to ordinary SD, so the module ends back on the exact, lossless line it started from.

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: