Speculative decoding

Problem

On the latency and throughput page, one decode step of an 8B model at small batch reads every weight to produce one token per sequence, and the GPU's arithmetic units sit almost entirely idle. Running the same weights over several positions at once costs nearly the same time, because the bytes read do not change, only the FLOPs. The trouble is that autoregressive decoding needs token t before it can compute token t + 1, so there is nothing to put in those extra positions.

Speculative decoding fills them with guesses. The question is how to use guesses without changing what the model says. A sampled chatbot must still sample from the large model's distribution, not from a blend of it with a weaker model's, or the quality you evaluated is no longer the quality you serve. The decoding page showed one step of the rule; this page covers the whole algorithm.

Intuition

Think of a junior writer and an editor. The junior drafts the next few words quickly. The editor reads the whole draft at once and keeps the longest prefix they would plausibly have written themselves, crosses out the first word they disagree with, and writes their own word in its place. If the editor agrees with every drafted word, they also add the next word, since they have already read the whole draft. Each round therefore costs one read by the editor and produces at least one word, often several.

The probabilistic part is how the editor "disagrees". A drafted token that the target likes at least as much as the drafter did (p ≥ q) is always kept. A token the drafter liked more than the target is kept only with probability p/q, and on rejection the replacement is drawn from the tokens the drafter under-proposed. Those two pieces add up to exactly p, so nothing about the output changes except how many target passes it took.

Mechanics

One round

Let the target model give p(· | prefix) and the drafter q(· | prefix). With draft length γ:

1. Draft. Sample x₁ ~ q(· | prefix), x₂ ~ q(· | prefix, x₁), …, up to x_γ. This is γ sequential steps of the cheap model.
2. Verify. Run the target once over the γ + 1 new positions (the last emitted token and the γ drafts) to get p at every position, p(· | prefix, x₁, …, x_i−1) for i = 1 … γ + 1.
3. Accept left to right. For each i, draw u ~ Uniform(0, 1) and keep x_i if u < min(1, p(x_i)/q(x_i)). Stop at the first rejection.
4. Correct. On a rejection at position i, sample a replacement from the residual r = max(0, p − q) / Σ max(0, p − q), using that position's p and q, and end the round.
5. Bonus. If all γ drafts are accepted, sample one more token from p at position γ + 1, which the verification pass already computed.

Every round emits between 1 and γ + 1 tokens, so it is never slower in tokens per target pass than plain decoding; it only spends the drafter's time.

Why the output is exactly the target's

Fix one position with target p and draft q. The token emitted there is x if the drafter proposed x and it survived, or if any proposal was rejected and the residual drew x:

P(emit x) = q(x) · min(1, p(x)/q(x)) + (1 − α) · r(x)

The first term is min(p(x), q(x)). The probability of a rejection is 1 − α, where

α = Σₓ min(p(x), q(x)) = 1 − TV(p, q)

is the acceptance rate (TV is total variation distance). Since min(p, q) + max(0, p − q) = p, summing over x gives Σ max(0, p − q) = 1 − α, so (1 − α) · r(x) = max(0, p(x) − q(x)). The two terms add to min(p, q) + max(0, p − q) = p(x). The emitted token is distributed exactly as p. The next position is then conditioned on a prefix whose distribution is already the target's, so the argument repeats position by position. Rejected and discarded drafts never reach the output, which is why the bonus token, drawn from the target's own p, is safe too.

The rule needs q(x) > 0 whenever the drafter proposes x, which holds by construction. Temperature, top-p and any penalties must be applied to p and q before the test, because the rule preserves exactly the p it is given, no more.

The greedy case

At temperature 0, p puts all its mass on argmax p. Then min(1, p/q) is 1 when the draft equals the target's argmax and 0 otherwise, and the residual is the target's argmax. So greedy verification is: keep drafts while they match the target's argmax at each position, replace the first mismatch with the target's argmax, and add the bonus argmax if all match. The output equals plain greedy decoding token for token, up to floating point: the target evaluates γ + 1 positions in one batched call, which can round differently from one position at a time and flip a near-tie.

Tokens per round and speedup

If every position had the same acceptance rate α, the number of accepted drafts K would satisfy P(K ≥ k) = α^k for k ≤ γ, and a round emits K + 1 tokens, so

E[tokens per round] = Σ_k=0..γ α^k = (1 − α^γ+1) / (1 − α)

Let c be the cost of one drafter step relative to one target step. A round costs γc + 1 target-step units, and plain decoding emits one token per unit, so the speedup in wall-clock time is

speedup(α, γ, c) = (1 − α^γ+1) / ((1 − α)(γc + 1))

The numerator saturates at 1/(1 − α) while the denominator keeps growing with γ, so there is a best draft length: longer for a better-aligned drafter (higher α) and for a cheaper one (lower c). Two assumptions hide in this formula. It assumes verifying γ + 1 positions costs the same as one decode step, which holds only while decoding is bound by memory bandwidth. And real acceptance varies with context, so the i.i.d. formula is an approximation; the worked example measures the gap.

Verification attention and the cache

The verification pass feeds γ + 1 query positions into a model whose KV cache already holds n keys for the prefix. After appending the new keys there are n + γ + 1 keys, and query i (0-based among the new ones) may attend to keys 0 … n + i. That is a causal mask aligned to the lower-right corner of the score matrix. The upper-left alignment, which is_causal=True uses in scaled_dot_product_attention when the matrix is not square, would let the first new query see only key 0, which is wrong. PyTorch's causal_lower_right bias encodes the correct variant.

After verification the cache contains entries for all γ + 1 new positions, including drafts that were rejected or discarded. If k drafts were accepted, keep the first n + k + 1 entries (the old prefix, the last token, the accepted drafts) and truncate the rest. The newly sampled replacement or bonus token has no cache entry yet; it is the "last token" of the next round. The drafter's own cache needs the same rollback.

Where drafts come from

• A small model of the same family (the original papers): shares the tokenizer, cheap per step, but must be trained and served alongside and run sequentially for γ steps.
• Extra heads on the target. Medusa adds several decoding heads that each predict a token further ahead from the target's last hidden state. Multi-token-prediction (MTP) modules trained with the model, as in DeepSeek-V3 and MiMo-V2-Flash, can be reused as drafters at inference.
• Feature-level drafters. EAGLE trains a light autoregressive head on the target's hidden features instead of on tokens, which raises α; EAGLE-3 fuses features from several target layers.
• No model at all. Prompt lookup or n-gram drafting copies a continuation from earlier in the context when the last few tokens match. It is nearly free (c ≈ 0) and works well for editing, retrieval and code, where outputs repeat inputs.

Two extensions change the verification. Tree verification (SpecInfer) drafts several candidate branches and verifies the whole tree in one pass with a tree-shaped attention mask, raising expected accepted length at the cost of more positions. Block verification (Sun et al., ICLR 2025) decides on the drafted block jointly instead of token by token, which provably accepts at least as many tokens on average while preserving the target distribution.

Worked example

A three-token vocabulary {a, b, c} where each model's next-token distribution depends only on the previous token. These are exact fractions, computed by tested helpers in ml_models_docs.py:

Implementation

One sampled round in PyTorch for a batch of one, with models that take a cache and return logits for every new position. The cache objects are assumed to support crop(length):

import torch
from torch.nn.attention.bias import causal_lower_right

@torch.inference_mode()
def speculative_round(target, draft, last_tok, t_cache, d_cache, gamma, temp=1.0):
    n = t_cache.length                                     # cached prefix length
    drafts, q_rows, tok = [], [], last_tok
    for _ in range(gamma):                                 # gamma cheap sequential steps
        logits, d_cache = draft(tok, cache=d_cache)
        q = torch.softmax(logits[:, -1] / temp, dim=-1)
        tok = torch.multinomial(q, 1)
        drafts.append(tok); q_rows.append(q)
    new = torch.cat([last_tok] + drafts, dim=1)            # [1, gamma + 1]
    mask = causal_lower_right(gamma + 1, n + gamma + 1)    # queries see the whole prefix
    logits, t_cache = target(new, cache=t_cache, attn_mask=mask)
    p_rows = torch.softmax(logits[0] / temp, dim=-1)       # [gamma + 1, V]
    out = []
    for i, tok in enumerate(drafts):
        x = tok.item()
        p, q = p_rows[i], q_rows[i][0]
        if torch.rand(()) < torch.clamp(p[x] / q[x], max=1.0):
            out.append(x); continue
        residual = torch.clamp(p - q, min=0)
        out.append(torch.multinomial(residual / residual.sum(), 1).item())
        break
    else:
        out.append(torch.multinomial(p_rows[gamma], 1).item())   # bonus token
    accepted = len(out) - 1
    t_cache.crop(n + accepted + 1)                         # drop rejected drafts
    d_cache.crop(n + accepted + 1)
    return out, t_cache, d_cache

The drafter's cache in this sketch is one position behind on a full accept, because the last draft was never fed to it; real implementations run one more drafter step or feed it next round. Production servers do this for many sequences at once, where each sequence accepts a different number of tokens, which is what makes batched speculative decoding fiddly: caches of ragged lengths, and padding or paged blocks that can be truncated per sequence.

Tradeoffs

• Gains shrink as the batch grows. The speedup formula assumes verifying γ + 1 positions is as cheap as one step. At large batch, decoding approaches the compute ridge on the roofline, so each extra position costs real FLOPs and rejected drafts become wasted work. Speculation buys per-user latency (interactivity) most at low concurrency, and servers often shorten γ or switch it off as load rises.
• The drafter costs memory and engineering. A separate model needs its own weights, cache and serving path. Heads on the target (MTP, Medusa, EAGLE) share the target's features and avoid a second tokenizer, but must be trained per target model and retrained when it changes.
• Acceptance depends on the workload. Code, extraction and templated text accept long drafts; open-ended creative sampling at high temperature accepts few. Measure acceptance length on your own traffic, not on a benchmark average.
• Exactness is a property of the rule, not the system. Mismatched temperature or top-p between p and q, a drafter applying a penalty the target does not, or nondeterministic kernels each break "identical distribution" quietly.
• It is a latency tool before a cost tool. It spends extra FLOPs to cut wall-clock time per token. Mistral Large 4's report gives the scale a production stack sees: with its DSpark speculator and block verification, decode interactivity rose 90% over no speculation and 15% over an EAGLE-3 drafter, both at eight concurrent requests.
• Distillation helps the drafter. Training the drafter on the target's outputs, as on the distillation page, raises α directly, since α is one minus the total variation distance between the two.

Connections

References

• Leviathan, Kalman and Matias. Fast Inference from Transformers via Speculative Decoding. ICML, 2023.
• Chen et al. Accelerating Large Language Model Decoding with Speculative Sampling. 2023.
• Stern, Shazeer and Uszkoreit. Blockwise Parallel Decoding for Deep Autoregressive Models. NeurIPS, 2018.
• Miao et al. SpecInfer: Accelerating Large Language Model Serving with Tree-based Speculative Inference and Verification. ASPLOS, 2024.
• Cai et al. Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. ICML, 2024.
• Li, Wei, Zhang and Zhang. EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. ICML, 2024; EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test. NeurIPS, 2025.
• Sun, Mendlovic, Leviathan, Aharoni, Ro, Beirami and Suresh. Block Verification Accelerates Speculative Decoding. ICLR, 2025.
• DeepSeek-AI. DeepSeek-V3 Technical Report (multi-token prediction). 2024.
• Mistral AI. Mistral Large 4. 2026 (technical report), section on inference optimizations and Figure 16.
• PyTorch v2.15.0: causal_lower_right, scaled_dot_product_attention, torch.multinomial.