Speculative decoding and other speedups

The big idea

When a single user is decoding, the accelerator spends most of each step waiting for weights to arrive from memory, and the arithmetic finishes almost instantly. So checking 5 tokens in one pass costs about the same time as producing 1. The question is where to get 5 good guesses.

Speculative decodingA decoding method in which a small, fast draft model proposes several tokens and the large target model checks them all in one parallel pass, keeping those it agrees with. The output distribution is the same as the target model's. answers it with a cheap assistant. Think of a senior writer with a fast junior. The junior drafts the next few words; the senior reads the draft in one glance, keeps the words they would have written anyway, crosses out the first wrong one and writes the correct word there. When the junior is usually right, the senior produces several words per glance instead of one, and the text reads exactly as if the senior had written every word.

That last point is what makes speculative decoding special among speedups: done correctly, it changes nothing about what the model outputs, only how fast.

How it works

Draft, then verify in parallel

Each round has two parts:

1. Draft. A small model (the draft) generates $\gamma$ tokens autoregressively, for example $\gamma =4$. In the original paper's experiments, drafts about two orders of magnitude smaller than the target usually worked best, and a draft step cost under 5 percent of a target step, so those 4 steps cost a small fraction of one target step.
2. Verify. The large model (the target) runs once over the context plus all $\gamma$ draft tokens. This is like a tiny prefill: it computes, in parallel, what the target would predict at each of the $\gamma +1$ positions. While decoding is memory-bound, this pass takes about as long as an ordinary single-token step.

Then the draft tokens are checked left to right. Each is either accepted or rejected; the first rejection ends the round and everything after it is thrown away. The target always contributes one token of its own at the end, either the correction at the rejected position or a bonus token after the last accepted one. So every round produces at least one token, as normal decoding would, and up to $\gamma +1$.

Why the output does not change

With greedy decoding the rule is simple: accept a draft token if it is the target's top choice. With sampling, a clever acceptance rule keeps the target's exact distribution. If the draft proposed token $x$ with probability $q(x)$ and the target gives it $p(x)$:

• accept with probability $\min \left(1,\frac{p(x)}{q(x)}\right)$;
• if rejected, sample a replacement from the leftover distribution, proportional to $\max (0,p(x)-q(x))$.

Tokens the draft under-rates are always kept; tokens it over-rates are kept only as often as the target would have chosen them, and the correction step fills in the missing probability. The proof that the result is distributed exactly as $p$ takes a few lines, and you can check it empirically in the code below.

How much faster?

The key number is the acceptance rateIn speculative decoding, the probability that the target model accepts a token proposed by the draft model. It equals the sum over tokens of min(p, q), so it is high when the two models agree. $\alpha$, the chance that a draft token is accepted. Assuming each token is accepted independently with probability $\alpha$, the expected number of tokens per target pass is

With $\alpha =0.8$ and $\gamma =4$: $(1-0.8^5)/0.2=(1-0.328)/0.2\approx 3.36$ tokens per pass. If the draft costs $c$ times a target pass, the wall-clock speedup is that divided by $(\gamma c+1)$; with $c=0.05$, $3.36/1.2=2.8$ times. Two things follow:

• Agreement matters most. At $\alpha =0.6$ the same setup gives only 2.3 tokens per pass, and a longer draft barely helps because the chain of acceptances breaks early.
• Longer drafts have diminishing returns. Each extra draft token costs $c$ but is only used if every token before it was accepted.

When it does not help

Speculative decoding spends extra compute (the draft, and target computation on rejected tokens) to save memory-bound steps. That trade only pays while the target has idle compute. At large batch sizes, where decoding becomes compute-bound ( How LLM inference works ), the extra work competes with other users' tokens, and the gain shrinks or turns into a loss. It also needs a draft that shares the target's tokenizer and agrees with it often, plus memory for the draft's weights and KV cache.

Try it

1. At the defaults (80 percent acceptance, 4 draft tokens, draft cost 5 percent), press Step a few times. How many tokens does each target pass produce, and how does that compare with the expected value?
2. Lower the acceptance rate to 60 percent. What happens to the best draft length, and why do long drafts stop paying?
3. Raise the draft cost to 30 percent. At what acceptance rate does speculative decoding stop being worth it at all?

A toy check that speculative sampling really keeps the target's distribution:

# Speculative sampling with toy models: the output follows the TARGET distribution exactly.
import random
from collections import Counter

V = ["the", "cat", "sat", "on", "mat"]
def target(prev):                      # big model: next-token probabilities given prev
    p = [1.0] * len(V); p[(V.index(prev) + 1) % len(V)] = 8.0
    return [x / sum(p) for x in p]
def draft(prev):                       # small model: similar, but less sure
    p = [1.0] * len(V); p[(V.index(prev) + 1) % len(V)] = 4.0
    return [x / sum(p) for x in p]

def speculative(prev, gamma, rnd):
    """One round: draft gamma tokens, verify them all with one target pass."""
    out, ctx = [], prev
    for _ in range(gamma):
        q, p = draft(ctx), target(ctx)  # in a real system p comes from ONE parallel pass
        x = rnd.choices(range(len(V)), q)[0]
        if rnd.random() < min(1.0, p[x] / q[x]):         # accept with prob min(1, p/q)
            out.append(V[x]); ctx = V[x]; continue
        resid = [max(0.0, a - b) for a, b in zip(p, q)]   # rejected: resample from p - q
        out.append(V[rnd.choices(range(len(V)), resid)[0]])
        return out
    out.append(V[rnd.choices(range(len(V)), target(ctx))[0]])  # all accepted: bonus token
    return out

rnd = random.Random(0)
passes, tokens, firsts = 0, 0, Counter()
while passes < 20000:
    out = speculative("cat", 4, rnd)
    passes += 1; tokens += len(out); firsts[out[0]] += 1

p = target("cat")
alpha = sum(min(a, b) for a, b in zip(p, draft("cat")))
print(f"acceptance rate per token: {alpha:.2f}")
print(f"tokens per target pass: {tokens / passes:.2f} "
      f"(formula {(1 - alpha ** 5) / (1 - alpha):.2f})")
print("token after 'cat'   target p   speculative output")
for i, w in enumerate(V):
    print(f"  {w:<6}{p[i]:>18.3f}{firsts[w] / passes:>20.3f}")

acceptance rate per token: 0.83
tokens per target pass: 3.57 (formula 3.59)
token after 'cat'   target p   speculative output
  the                0.083               0.082
  cat                0.083               0.086
  sat                0.667               0.667
  on                 0.083               0.082
  mat                0.083               0.083

The draft is a less confident copy of the target, so they agree on about 83 percent of tokens and each target pass yields about 3.6 tokens, matching the formula. The first token of each round, which is either an accepted draft token or the target's correction, appears with the target's probabilities to within sampling noise.

Other speedups

Multi-token prediction: the model drafts for itself

Instead of a separate draft model, a model can be given extra output heads that predict the 2nd, 3rd and 4th next tokens from the same hidden state. Those guesses are then verified the same way. Medusa adds such heads to an existing model and checks several candidate continuations at once with a tree-shaped attention mask. Models can also be trained from the start to predict several future tokens: one study found models trained with 4-token prediction up to 3 times faster at inference, with better scores on coding benchmarks as a side benefit. Self-drafting avoids running and maintaining a second model.

Prefix caching: skip prefill you already did

Many requests start with the same tokens: a long system prompt, tool definitions, a few-shot block, or the earlier turns of the same conversation. Their keys and values are identical every time. Prefix cachingKeeping the KV cache of a prompt prefix after a request finishes so that later requests starting with the same tokens can reuse it and skip that part of prefill. keeps those KV blocks in GPU memory after a request finishes, indexed by the tokens they cover, and lets later requests that start with the same tokens reuse them. With paged KV blocks ( Batching and PagedAttention ) this is natural: a cached prefix is just a set of shared blocks.

It cuts time to first token and prefill compute for repeated prefixes, but it does nothing for decode, so it helps most when prompts are long and answers short. Provider-side prompt caching ( Prompt caching and cost control ) is the same idea exposed as a price discount. To benefit, put the stable parts of the prompt first and the variable parts last.

FlashAttention: fewer trips to slow memory

Standard attention computes the full score matrix of every query against every key, writes it to GPU main memory, reads it back for the softmax, writes the result, and reads it again to multiply by the values. For long sequences that matrix is huge, and the reading and writing, not the arithmetic, dominate.

FlashAttentionAn exact attention algorithm that processes queries, keys and values in tiles that fit in fast on-chip memory, using a running softmax so the full score matrix is never written to main GPU memory. computes exactly the same result in tiles: it loads a block of queries and a block of keys and values into fast on-chip memory, computes their scores, and keeps a running softmax (a running maximum and sum) so it never stores the full matrix. Fewer bytes move between memory levels, so attention runs faster and uses memory that grows linearly, not quadratically, with sequence length. It is a kernel change, invisible to the model's outputs, and widely used serving engines ship attention kernels built on this idea.

Shrinking the KV cache with attention variants

Grouped-query and multi-query attention let several query heads share one key/value head, which shrinks the KV cache and the bytes each decode step reads ( Context windows and the KV cache ). Unlike the other speedups here, this is part of the model's architecture: it is chosen at training time, not switched on at serving time.

Going deeper

Choosing the draft length

The best $\gamma$ maximizes $\frac{1-{\alpha }^{\gamma +1}}{(1-\alpha )(\gamma c+1)}$. High agreement and a cheap draft favor longer drafts; low agreement favors 1 to 3. Acceptance also varies with the content and the prompt, so measure it on your own traffic rather than assuming a number.

Where the draft comes from

• A smaller model of the same family, sharing the tokenizer: simplest when one exists.
• Extra heads on the target (Medusa-style) or a trained multi-token head: no second model to serve.
• Text already in the context: when the output is likely to copy from the prompt, as in editing code or quoting a document, candidate tokens can be looked up from matching text in the prompt with no model at all. vLLM, for example, lists n-gram methods alongside draft models and prediction heads, noting that they give modest speedups without adding load at peak traffic.

Trade-offs

Common mistakes

• Turning on speculative decoding for a high-throughput batch workload and seeing no gain or a slowdown.
• Using a draft model with a different tokenizer or a very different style, so the acceptance rate is too low to pay for the draft.
• Believing speculative decoding trades quality for speed: with the correct acceptance rule the output distribution is unchanged.
• Putting variable content, such as a timestamp or user name, at the start of the prompt, which breaks prefix caching for everything after it.
• Expecting prefix caching to speed up long answers, when it only skips prefill.

In the wild

Check yourself

Recap