Speculative decoding

Setup#

pip install numpy

The acceptance rule is the whole algorithm, and it fits in six lines. Running it a few hundred thousand times against a known target distribution is how you convince yourself the guarantee is real.

The algorithm, and a proof by counting#

import numpy as np

VOCAB = ["the", "train", "is", "late", "again", "today"]
TARGET = np.array([0.40, 0.25, 0.15, 0.10, 0.06, 0.04])   # what the big model wants
DRAFT  = np.array([0.35, 0.30, 0.20, 0.10, 0.03, 0.02])   # what the small model guesses

rng = np.random.default_rng(42)


def speculative_step(p, q, rng):
    """One token. p = target distribution, q = draft distribution.
    Returns (token, accepted?) and provably samples from p."""
    x = rng.choice(len(q), p=q)                    # the draft's guess, free
    if rng.random() < min(1.0, p[x] / q[x]):       # the acceptance test
        return x, True
    residual = np.maximum(p - q, 0.0)              # what the draft under-supplied
    residual /= residual.sum()
    return rng.choice(len(p), p=residual), False   # resample, still one target call


N = 200_000
tokens = np.empty(N, dtype=int)
accepted = 0
for i in range(N):
    tokens[i], ok = speculative_step(TARGET, DRAFT, rng)
    accepted += ok

empirical = np.bincount(tokens, minlength=len(VOCAB)) / N
print(f"{'token':<8}{'target p':>10}{'draft q':>10}{'observed':>10}")
for w, p, q, e in zip(VOCAB, TARGET, DRAFT, empirical):
    print(f"{w:<8}{p:>10.3f}{q:>10.3f}{e:>10.3f}")
print(f"\nlargest gap between target and observed: {np.abs(empirical - TARGET).max():.4f}")
print(f"acceptance rate: {accepted / N:.3f}")
print(f"1 - total variation distance between p and q: "
      f"{1 - 0.5 * np.abs(TARGET - DRAFT).sum():.3f}")

# ---- how many tokens does one verification pass actually buy? ----
GAMMA = 4


def one_round(rng):
    """Draft GAMMA tokens, verify them all in one target pass."""
    for i in range(GAMMA):
        _, ok = speculative_step(TARGET, DRAFT, rng)
        if not ok:
            return i + 1                # i accepted, plus the corrected one
    return GAMMA + 1                    # all accepted, plus one free bonus token


rounds = 100_000
produced = sum(one_round(rng) for _ in range(rounds))
alpha = accepted / N
expected = (1 - alpha ** (GAMMA + 1)) / (1 - alpha)
print(f"\ndrafting {GAMMA} tokens per round, {rounds:,} rounds")
print(f"  measured tokens per target pass : {produced / rounds:.3f}")
print(f"  formula (1-a^(g+1))/(1-a)       : {expected:.3f}")

token     target p   draft q  observed
the          0.400     0.350     0.400
train        0.250     0.300     0.251
is           0.150     0.200     0.150
late         0.100     0.100     0.100
again        0.060     0.030     0.059
today        0.040     0.020     0.040

largest gap between target and observed: 0.0012
acceptance rate: 0.901
1 - total variation distance between p and q: 0.900

drafting 4 tokens per round, 100,000 rounds
  measured tokens per target pass : 4.096
  formula (1-a^(g+1))/(1-a)       : 4.100

Seeded, so these exact numbers reproduce. Checked on NumPy 1.26 and 2.4, on Python 3.10 and 3.13.

Reading the output#

The observed column matches the target column, not the draft. The draft wanted train at 0.300; the target wanted 0.250; the output landed on 0.251. The draft wanted today at 0.020, the target at 0.040, and the output gave 0.040. The draft's bias has been removed exactly. Largest deviation across the whole vocabulary is 0.0012, which is sampling noise at 200,000 draws.

This is the entire claim of the method, verified by counting rather than believed from a paper.

Acceptance rate 0.901 against a predicted 0.900. That is not a coincidence. The expected acceptance rate equals $1$ minus the total variation distance between the two distributions. The paper proves it; this run confirms it to three decimals. It also tells you what to optimise: pick a draft model whose distribution is close to the target's, and the entire speed-up follows.

4.096 tokens per target pass, against a formula prediction of 4.100. Four expensive passes' worth of work for one. That is the speed-up, in the only unit that matters.

The +1 in return GAMMA + 1 is free and easy to miss. When the target verifies four draft tokens, its forward pass also produced logits for the position after the fourth. That token needs no verification, so a fully-accepted round yields five tokens. Implementations that drop it lose a measurable fraction of the benefit.

The residual distribution, which is where people go wrong#

On rejection, the naive move is to sample from the target distribution p. That is wrong, and it biases the result.

The reason is that acceptance already consumed part of p. Tokens the draft over-proposes were sometimes accepted, so the remaining probability that still needs delivering is p - q where that is positive, and zero elsewhere:

residual = np.maximum(p - q, 0.0)
residual /= residual.sum()

Delete the maximum and use plain p and the whole guarantee vanishes, silently. Nothing crashes. Your output distribution is now somewhere between the draft and the target, and no test will tell you.

Greedy is the easy special case#

At temperature zero the rule collapses to string comparison:

# accept each draft token while it equals the target's argmax at that position
take = 0
for j, d in enumerate(draft):
    if target_argmax[j] == d:
        take += 1
    else:
        break

Accept the matching prefix, take the target's own token at the first mismatch, discard the rest. This is what most production implementations use, and it is exactly lossless because greedy decoding has no distribution to preserve.

Using it for real#

# vLLM: a smaller model from the same family as the drafter
from vllm import LLM
llm = LLM(
    model="Qwen/Qwen3-32B",
    speculative_config={"method": "draft_model",
                        "model": "Qwen/Qwen3-0.6B",
                        "num_speculative_tokens": 5},
)

# HuggingFace Transformers: assisted generation
out = model.generate(**inputs, assistant_model=small_model, max_new_tokens=64)

No output blocks for these two. Both need model downloads, and the vLLM path needs a GPU, so any output printed here would be invented rather than observed. The vLLM speculative-decoding API has changed several times. speculative_config is the current shape and it takes a method key. Check your installed version's docs rather than a blog post.

Common mistakes#

Mismatched tokenizers. Draft and target must share a vocabulary, or every id means something different. Use a smaller model from the same family. vLLM has added a cross-vocabulary mode, and at the time of writing it supports greedy draft sampling only.

Drafting too many tokens. The optimal num_speculative_tokens depends on the acceptance rate. Every rejected draft token is wasted work, and beyond roughly 3 to 8 the wasted drafts outweigh the saved passes on most workloads. Measure your acceptance rate first.

Expecting a gain under heavy load. Speculation spends spare arithmetic capacity. At batch size 1 there is plenty. At batch size 64 the GPU is already busy and speculation can make total throughput worse while still improving one user's latency. It is a latency optimisation, not a throughput one.

Forgetting the draft model's memory. It needs weights and a KV cache of its own, taken out of the same GPU. On a tight memory budget this can cost you more concurrency than the speed is worth.

Benchmarking on the wrong text. Acceptance is high on boilerplate and low on creative prose. A speed-up measured on repetitive prompts does not transfer.

Try it yourself#

Set DRAFT = TARGET and re-run. Acceptance goes to 1.0 and tokens per pass reaches GAMMA + 1, the theoretical ceiling. Then set DRAFT to a uniform distribution and watch acceptance and speed-up collapse together. Plot tokens-per-pass against acceptance rate and you have drawn the curve that decides whether this technique is worth deploying for you.

What to learn next#

• Medusa, EAGLE and self-speculation — removing the second model entirely.
• vLLM — where you actually switch this on.
• Knowledge distillation — how a good draft model gets trained in the first place.