A small draft model makes Chinchilla 70B decode 2.5× faster without changing its distribution

One forward pass, one token

On sixteen TPU v4 accelerators, Chinchilla 70B emits one token every 14.1 milliseconds. Pair it with a 4-billion-parameter draft model and the modified rejection sampling below, and on the HumanEval coding benchmark it emits one every 5.73 milliseconds: a 2.46× speedup at batch size one, drawing from the same distribution the 70B model would have sampled alone.¹ The large model got no cheaper; each pass simply emitted more than one token.

A standard autoregressive decoder produces exactly one token per forward pass, because token t is sampled from a distribution that conditions on tokens 1 through t−1. A 200-token completion is 200 sequential passes through the whole network. Speculative decoding breaks that chain without changing the output distribution, and getting the acceptance rule subtly wrong is the classic way to ship a version that is fast and silently biased.

Decoding moves weights, not math

A single decode step at batch size one is dominated by the time to read the model's parameters out of high-bandwidth memory and into the compute units, not by the arithmetic those units then perform. Every weight is loaded to produce one token, and the multiplies that consume it finish well before the next weights arrive. The step is bound by memory bandwidth, the rate at which parameters stream onto the chip, and the arithmetic units sit mostly idle.

That idle arithmetic is the whole opportunity. Chen's paper puts it plainly: "the latency of parallel scoring of short continuations, generated by a faster but less powerful draft model, is comparable to that of sampling a single token from the larger target model."¹ Scoring five candidate tokens in one pass loads the target's weights once, the same load a single-token step already pays, and spends the wasted arithmetic checking all five. The extra tokens are close to free because the pass was never arithmetic-limited.

Draft small, verify in bulk

The structure is two models and two phases. A cheap draft model proposes a short continuation of γ tokens one at a time; the target model then scores all γ candidate positions, plus the one that follows, in a single batched pass. Both founding papers share the shape: Leviathan runs γ draft steps per target step,² Chen drafts K tokens and scores K+1 logit sets at once.¹

The draft loop is ordinary autoregressive sampling; the verify is one target call whose output covers every prefix length at once, so the γ extra positions cost a single pass.

import torch

def speculative_step(prefix, draft, target, gamma):
    # 1. DRAFT: sample gamma tokens autoregressively from the small model.
    draft_tokens, draft_probs = [], []
    ctx = prefix
    for _ in range(gamma):
        d = draft(ctx).softmax(-1)          # the draft's distribution over the vocab
        x = torch.multinomial(d, 1)         # one proposed token
        draft_tokens.append(x)
        draft_probs.append(d)
        ctx = torch.cat([ctx, x], dim=-1)

    # 2. VERIFY: one batched target pass scores every drafted position at once,
    #    so the gamma extra tokens cost a single forward pass, not gamma of them.
    target_probs = target(ctx).softmax(-1)  # gamma + 1 distributions

    return accept_or_resample(draft_tokens, draft_probs, target_probs, gamma)

The test that keeps the target's distribution

The correctness lives in one place: the rule that decides, token by token, whether to keep a draft or correct it. The two founding papers letter it in opposite ways, the surest route to convincing yourself a correct implementation is buggy. Leviathan writes p for the target and q for the draft; Chen swaps them.² The algebra is identical. Fix one and hold it: for the rest of this section, following Chen, q is the target we must preserve and p the draft.¹

For each drafted token x, draw a uniform random number r between 0 and 1, and accept x if r < min(1, q(x)/p(x)). Read the ratio directly: if q(x) ≥ p(x) the token is always kept, and if the target likes it less it is kept with probability exactly q(x)/p(x). On the first rejection, stop, discard the rest of the draft, and resample the replacement from the normalized residual (q − p)₊, meaning max(0, q(x) − p(x)) divided by its own sum.¹

Now watch the token that comes out of a single position. The draft proposes x with probability p(x), and that x survives the accept test with probability min(1, q(x)/p(x)). Multiply them: the probability a given x is proposed and then kept is

p(x) · min(1, q(x)/p(x)) = min(p(x), q(x)).

Sum that over the vocabulary and the total probability of keeping any draft is Σ min(p, q), the mass the two distributions share, so a rejection happens with probability 1 − Σ min(p, q). When it does, the replacement is drawn from (q − p)₊, whose unnormalized mass is Σ max(0, q − p). Because both distributions sum to one, that residual mass equals 1 − Σ min(p, q), the exact probability of reaching the rejection branch: the normalizer and the branch probability cancel. So the probability the emitted token is x is the keep path plus the correct path:

min(p(x), q(x)) + max(0, q(x) − p(x)) = q(x).

The right-hand side is the target distribution, unchanged. That identity is the whole guarantee, and it is Chen's Theorem 1: for any two numbers, the smaller one plus the amount by which the target exceeds the draft is the target.¹ The draft's quality never enters. A better draft is accepted more often and runs faster, a worse one is corrected more often, and the emitted token is drawn from q either way.

In code it is a few lines, and the one that carries the argument is the comparison against min(1, q(x)/p(x)). The listing names the distributions rather than lettering them.

def accept_or_resample(draft_tokens, draft_probs, target_probs, gamma):
    emitted = []
    for i in range(gamma):
        draft_p  = draft_probs[i]           # the draft's distribution at step i
        target_p = target_probs[i]          # the target's distribution at step i
        x = draft_tokens[i]
        r = torch.rand(1)
        # Accept x with probability min(1, target(x) / draft(x)).
        # Invariant this preserves: the emitted token is distributed exactly as
        # the target, because min(draft, target) + max(0, target - draft) = target.
        if r < min(1.0, (target_p[x] / draft_p[x]).item()):
            emitted.append(x)               # keep the drafted token unchanged
        else:
            residual = torch.clamp(target_p - draft_p, min=0)   # (target - draft)_+
            residual = residual / residual.sum()                # renormalize the leftover mass
            emitted.append(torch.multinomial(residual, 1))      # draw the correction
            return emitted                  # stop at the first rejection

    # All gamma accepted: the verify pass already scored one more position,
    # so a bonus token drawn from the target comes for free.
    emitted.append(torch.multinomial(target_probs[gamma], 1))
    return emitted

You need not trust the derivation. Reduce it to one position with a wrong draft and count the outputs: the frequencies converge to the target, never the draft.

import collections, torch

draft_p  = torch.tensor([0.10, 0.60, 0.30])   # a deliberately bad draft
target_p = torch.tensor([0.50, 0.20, 0.30])   # the distribution we must reproduce
counts = collections.Counter()

for _ in range(200_000):
    x = torch.multinomial(draft_p, 1)                     # draft proposes
    if torch.rand(1) < min(1.0, (target_p[x] / draft_p[x]).item()):
        counts[x.item()] += 1                             # accepted
    else:
        residual = torch.clamp(target_p - draft_p, min=0)
        counts[torch.multinomial(residual / residual.sum(), 1).item()] += 1

total = sum(counts.values())
print({k: round(counts[k] / total, 3) for k in sorted(counts)})
# -> {0: 0.5, 1: 0.2, 2: 0.3}, i.e. target_p, not draft_p

From acceptance rate to walltime

How much a step buys depends on how often the draft is right. Define the acceptance rate α as the expected fraction of drafted tokens the test keeps, which Leviathan shows equals the expected overlap mass E[Σ min(p, q)]: two models that agree on 70 percent of the probability mass accept about 70 percent of the time.² With γ draft tokens per step, the expected number of tokens a step emits is a geometric sum,

E[tokens per step] = (1 − α^γ+1) / (1 − α),

straight from Leviathan.² That figure is an upper bound: it counts tokens per target pass and ignores the draft's own cost. Let c be the cost ratio, one draft run over one target run. Theorem 3.8 folds it in, and the walltime improvement over plain autoregressive decoding is

(1 − α^γ+1) / ((1 − α)(γc + 1)).²

The two factors pull against each other. Raising γ lifts the numerator toward its ceiling of 1/(1 − α) but grows the γc term below, since every drafted token is a draft run whether it survives or not. The chart plots the first factor alone for two draft lengths; acceptance rate matters more than draft length, the curves fanning apart only once α is high.

Those are formula outputs. The measured numbers land lower, because c is real and no draft agrees perfectly. Leviathan ran T5-XXL with a 77-million-parameter T5-small draft, against the standard T5X implementation: on English-to-German translation, γ=7 gave 3.4× at an acceptance rate of 0.75 (greedy) and 2.6× at 0.62 (sampled); on CNN/DM summarization, 3.1× and 2.3×.² Chen's Chinchilla 70B, a 4-billion draft with γ=4 at batch size one on sixteen TPU v4 chips, hit 1.92× on XSum and 2.46× on HumanEval against a 14.1-millisecond per-token baseline.¹ The spread is the acceptance rate moving: code and summaries are more predictable to the draft than open translation.

Where the draft can come from

Everything above assumes a separate draft model, trained or found and served alongside the target. Three variants change where the proposal comes from, and one changes what "loss-free" means.

Medusa keeps a single model and bolts extra decoding heads onto the backbone's last hidden state, each predicting one more future token in parallel, then verifies candidate continuations together with a tree-structured attention pass.³ It reports over 2.2× with frozen-backbone heads and 2.3–3.6× with heads and backbone trained jointly, on Vicuna-class models. The catch is the acceptance rule. Medusa deliberately replaces rejection sampling with a "typical acceptance" threshold, keeping any candidate whose original-model probability clears min(ε, δ·exp(−H)), and states that it is "typically unnecessary to match the distribution of the original model."³ At temperature 0 this reverts to greedy decoding and is exact; above it, Medusa's sampling mode does not reproduce the target's distribution, so any blanket "loss-free" claim has to carve it out.

EAGLE keeps true rejection sampling and moves the draft down a level. It autoregresses over the target's second-to-top-layer features rather than tokens, feeding in a token sequence shifted one step ahead to resolve the uncertainty pure feature prediction leaves, then verifies with the same test, so it "maintain[s] the distribution of the generated text."⁴ On LLaMA2-Chat 70B it reports 2.7–3.5× latency and roughly double the throughput against vanilla decoding.⁴

The lightest variant has no draft model at all. Prompt-lookup decoding replaces it with string matching: it finds the most recent n-gram in the prompt or output so far and proposes the tokens that followed it last time.⁵ Verification is the same target-model test, so output is unchanged, and the gain concentrates where the output echoes the input. On summarization and context question-answering with Mistral-7B-Instruct on an A100, greedy, it reports a consistent 2.4× average, 2–4× across such tasks and near nothing when the output does not reuse the prompt.⁵

When to reach for it

Every speedup above carries the same fine print: batch size one, or low query rate, on a memory-bound decode. Raise the batch, and the target's pass amortizes its weight load across many sequences until it becomes arithmetic-bound. Now the spare compute is gone, and the draft and verify burn arithmetic the server needed for throughput. vLLM's own guidance puts speculative decoding in "memory-bound" workloads at "medium-to-low" query rates, and rates the gain high only in the latency-focused regime.⁶

A fleet running large batches for tokens per dollar is already arithmetic-bound, and there the method can go negative: it adds draft and rejected-token cost to a pass with no spare compute to donate. A low α caps the walltime formula near one however long the draft, and a high draft cost c lets the denominator eat the gain. It rewards one shape of workload: latency-bound, lightly batched, with a cheap draft that agrees often.

The decision reduces to three quantities you can estimate before writing any code: the acceptance rate between target and draft on real traffic, the overlap mass E[Σ min(p, q)]; the batch size, which decides whether the decode is memory-bound at all; and the draft cost c. Put them in the walltime factor (1 − α^γ+1) / ((1 − α)(γc + 1)), and if it clears one by a worthwhile margin, the test in Fig. 2 guarantees you pay nothing in output quality to collect it.²