Speculative Decoding

ch15 · Speculative Decoding¶

The problem¶

Decode is serial by construction. Token n+1 depends on token n, so the model runs, produces one token, and runs again. ch03 · The Arithmetic of Inference showed why that is the worst possible ratio: a forward pass reads every weight in the model, and we are spending that entire read on a single token.

ch07 · Continuous Batching fixed this for many requests by putting them in the same pass. But one user, alone on a machine, still gets one token per full weight read. Batching cannot help them — there is nobody to batch with.

The idea¶

Something cheap proposes several tokens. The real model checks all of them in one forward pass, because verification is prefill-shaped and prefill is parallel. Proposals the model agrees with are free; the rest are thrown away.

The economics are simple and unusual: a round costs one target forward pass and yields somewhere between one and k+1 tokens. It can never yield zero, because even total disagreement leaves the target’s own token at the first mismatch. Speculation cannot produce fewer tokens per pass than not speculating.

What proposes the tokens¶

Two options, and the cheap one is more interesting than it sounds.

A draft model is the canonical approach: a smaller model of the same vocabulary, run autoregressively. It costs real compute per proposal, so it must be much faster than the target.

An n-gram drafter costs nothing at all — it looks for where the recent context appeared before and copies whatever followed:

class NgramDrafter(Drafter):
    """Proposes by finding where this context appeared before and copying what followed.

    No second model, no training, no memory to speak of. It works precisely when the output
    repeats its input — summarising a document, editing code, answering from a quoted context —
    which is a large slice of real traffic. Where nothing repeats it proposes nothing useful, and
    the acceptance rate simply falls to zero.
    """

    name = "ngram"

No model, no training, no memory. It works exactly when output repeats input — summarising, editing code, answering from a quoted document — which is a large share of real traffic. Where nothing repeats, acceptance falls to zero and it costs nothing to have tried.

Why the result is exactly the target model¶

This is the part worth slowing down for, because it is what separates speculation from a heuristic.

For greedy decoding the argument is trivial: accept a proposal only if it is what the target would have chosen. Same tokens, fewer passes.

For sampling it is subtler, and the rule looks arbitrary until you see why it works. Accept a proposed token x with probability min(1, p_target(x) / p_draft(x)). On rejection, do not resample from p_target. Sample from the normalised positive part of p_target - p_draft:

def accept_sampled(
    target_logits: torch.Tensor,
    draft_probs: torch.Tensor,
    proposed: list[int],
    params: SamplingParams,
    generator: torch.Generator | None = None,
) -> tuple[list[int], int]:
    """Accept proposals by rejection sampling, preserving the target distribution exactly.

    For each proposed token ``x``, accept with probability ``min(1, p_target(x) / p_draft(x))``.
    On rejection, resample from the normalised positive part of ``p_target - p_draft``.

    That residual is the whole trick, and it is worth seeing why it works. Proposing from
    ``p_draft`` and accepting at that ratio leaves token ``x`` under-sampled by exactly
    ``max(0, p_target(x) - p_draft(x))``. Drawing the replacement from precisely that shortfall
    puts the missing mass back where it belongs, so the combined procedure samples from
    ``p_target`` — not approximately, exactly.

    Drop the residual and resample from ``p_target`` instead and the result is subtly biased
    towards tokens the draft model liked. It still produces fluent text, which is what makes the
    bug hard to notice.
    """

Proposing from p_draft and accepting at that ratio leaves each token under-sampled by exactly max(0, p_target(x) - p_draft(x)). Drawing the replacement from precisely that shortfall puts the missing mass back where it belongs. The combined procedure samples from p_target — not approximately, exactly.

Resampling from p_target instead is the obvious-looking simplification, and it is wrong in a way that is almost impossible to notice: the drafted token gets its mass counted twice, once through acceptance and again through the correction.

The build¶

def speculative_generate(
    target: TinyGPT,
    drafter: Drafter,
    prompt_token_ids: list[int],
    params: SamplingParams,
    k: int = 4,
    generator: torch.Generator | None = None,
    stats: SpeculationStats | None = None,
) -> list[int]:
    """Generate with speculation, verifying each round in a single target forward pass.

    Deliberately written without a KV cache for the target. Caching across speculative rounds is
    fiddly — a rejected proposal has to be rolled back out of the cache — and it would obscure the
    part of this that is worth understanding. Chapter 15's measurement therefore counts *target
    forward passes*, which is the quantity speculation reduces, rather than wall-clock.
    """

One deliberate omission: the target has no KV cache across speculative rounds. Rolling a rejected proposal back out of a cache is fiddly and would bury the idea under bookkeeping. The measurement therefore counts target forward passes, which is the quantity speculation actually reduces.

The measurement¶

Does it pay?¶

At the largest k the engine produces eight tokens per target forward pass. Every configuration produces output identical to greedy decoding, which is the claim being checked rather than assumed.

Two conditions made this measurable, and both are worth naming.

The model had to be trained. Against random weights the target’s next token is arbitrary, so an n-gram drafter matches almost nothing and acceptance sits at zero — measuring the drafter against noise rather than measuring speculation. The first version of this benchmark did exactly that.

Acceptance is a property of the workload, not of the technique. This corpus repeats, so copying earlier text works extremely well. On text that never repeats, the same drafter proposes nothing useful. Any acceptance rate quoted without its workload is meaningless.

Is it the same model?¶

The correct rule lands within half a standard error of the target. The plausible-looking bug lands nearly seven standard errors away, at almost exactly the probability theory predicts for it.

That negative control is the point. A correctness argument you cannot falsify is not worth much, so the benchmark implements the wrong rule too and shows the test detects it. Without that, “our speculation is distribution-preserving” is a claim about code nobody checked.

Note also what the test had to avoid. Measured on the trained model, the drafted token has probability 0.9998 — a distribution that close to a point mass cannot reveal a sampling bias, because every rule returns that token essentially always. The distribution test runs on the untrained model for that reason. A measurement has to be taken where it can discriminate, and choosing those conditions honestly is different from choosing the conditions that flatter you.

When it does not pay¶

Speculation spends extra compute to save serialisation. At batch one, that compute was idle anyway. As the batch grows, ch07 · Continuous Batching is already amortising the weight read across many sequences, the device is no longer waiting on a dependency, and the extra verification work competes with real work. The technique stops helping and starts costing.

This makes speculation a low-load optimisation, which is an unusual shape: it improves exactly the regime where the server has capacity to spare, and stops when it does not. That is still valuable — single-user interactive latency is what most people judge a deployment by — but it is not a throughput technique, and deploying it as one disappoints.

The cost¶

• Extra compute per token, wasted whenever a proposal is rejected.

• The win evaporates at high batch, and can become a loss.

• A draft model is a second model to load, version, and keep in memory.

• Acceptance is workload-dependent and can collapse without warning when traffic shifts.

• The correctness argument is easy to get subtly wrong, and the failure mode is fluent, plausible text drawn from the wrong distribution. Only a distributional test catches it.

• KV cache interaction is genuinely awkward — rolling back rejected proposals is where real implementations get complicated.

Key takeaways¶

• Verification is prefill-shaped and therefore parallel: k proposals cost one target forward pass.

• A round always yields at least one token, so speculation never reduces tokens per pass.

• The residual acceptance rule makes the output distribution exactly the target’s. Resampling from p_target on rejection double-counts the drafted token and biases the result invisibly.

• Test the correctness claim with a negative control. An argument you cannot falsify is not evidence.

• Acceptance depends on the workload. Quote it with the trace or not at all.

• Speculation helps at low batch and hurts at high batch. It is a latency technique, not a throughput one.

Looking ahead¶

Speculation constrains which tokens appear by proposing them. ch16 · Constrained and Structured Decoding constrains which tokens are allowed — grammars and JSON schemas enforced during decoding — and runs into the same question from the other side: what does restricting the distribution do to the output, and what does enforcing it cost per step?

Further reading¶

Leviathan et al. and Chen et al. (Appendix E) introduced speculative decoding independently; both contain the proof sketched above, and it is short enough to read in full now that you have implemented it. Medusa and EAGLE replace the separate draft model with extra heads on the target, removing the second-model problem. Prompt lookup — the n-gram drafter here — is folklore rather than a paper, and is the one to try first because it costs nothing.