Speculative Decoding

Ordinary autoregressive decoding is sequential, and small-batch decode is often memory-bound. Each new token requires a full forward pass through the model, reading every weight from memory, and most of the GPU's compute sits idle while that happens. We established this in Prefill vs. Decode and The One New Row.

Faster kernels, lower precision, and different hardware can reduce a decode step's cost. Speculation asks a different question: can one target pass advance several output positions?

Speculative decoding uses an accept/reject rule that preserves the target model's sampling distribution.

Prerequisites

Prefill vs. Decode (Post 5.2) explains the prefill/decode asymmetry, and The One New Row (Post 5.3) derives the weight traffic per decode step. The acceptance rule below compares probability distributions over the token vocabulary.

The core asymmetry

As in Post 5.2, prefill processes prompt positions in parallel. You hand the model 1,000 tokens and it processes all 1,000 in one forward pass, computing attention across all positions simultaneously. Processing 1,000 tokens performs far more arithmetic than processing one. The useful window for speculation is a small candidate block that can use otherwise idle compute without greatly increasing target-pass latency.

Verification can reuse that parallelism. If someone hands you $k$ candidate tokens and says "what would you have predicted at each of these positions?", the target model can answer that question for all $k$ positions in a single forward pass, at roughly the cost of processing a prompt of length $k$.

Ordinary autoregressive generation cannot do this: each sampled token becomes input for the next prediction. Verification already has the candidate sequence available.

Draft, then verify

The draft-and-verify setup uses two models:

1. A draft model that's small and fast. Something like a 1B parameter distilled version of your target, or even a separate lightweight model.
2. A target model that's the big one you actually want output from. Think 70B parameters.

The draft model generates $k$ candidate tokens autoregressively, each conditioned on the previous tokens. An illustrative pairing is Llama-2-7B drafting for Llama-2-70B, or a purpose-built 1B distilled model drafting for a 70B+ target. The draft must be fast enough that its extra work costs less than the target steps it saves.

The target evaluates the candidate block with causal attention. Logits at a position predict the next token: the first candidate is checked against the distribution after the accepted prefix, not against the logits produced after consuming that candidate. Implementations retain that prefix prediction or include its last unprocessed token in the verification call. The aligned predictions check all $k$ candidates and provide the next-position distribution for a bonus token.

Verification has a cost; for a small candidate block it can be close to a single decode pass in a bandwidth-limited regime. You're already loading all the weights for a decode step anyway, so you might as well use them on $k$ tokens instead of 1. The batch-one, weights-only FP16 example in Post 5.3 has roughly 1 FLOP per byte. Processing $k$ tokens during verification bumps the arithmetic intensity to $k$ FLOPs per byte, still well within the memory-bound regime, so the additional target cost can be modest. Cache traffic, batch size, candidate count, and kernel overhead determine the actual result.

The acceptance criterion

So the target model processes all $k$ draft tokens in one forward pass and produces logits at every position. At each position $i$, you now have two probability distributions:

• $q(x_i)$: the probability the draft model assigned to the token it chose
• $p(x_i)$: the probability the target model assigns to that same token

The rule works left to right, position by position:

If $p(x_i) \geq q(x_i)$: accept. The target assigns this token at least as much probability as the draft does.

If $p(x_i) < q(x_i)$: accept with probability $\frac{p(x_i)}{q(x_i)}$, reject otherwise. If you reject, you don't just throw the token away. You resample from an adjusted distribution:

This adjusted distribution captures the probability mass that the target model wants to assign but the draft model didn't. You sample one token from $p'$, use it as the replacement, and discard everything after the rejection point.

And then you start a new round. The draft model picks up from the last accepted (or resampled) token and guesses another $k$ tokens. After a rejection, discard the suffix from the draft and target caches as well as from the output. This cycle repeats until an accepted or sampled end-of-sequence token, or another stopping condition.

The correctness guarantee

Leviathan et al. and Chen et al. both independently proved that this accept/reject scheme produces the same sampling distribution as the target model in exact arithmetic, using the normalized draft and target distributions after their sampling transformations.

This is rejection sampling: the draft distribution $q$ is the proposal, and the acceptance rule corrects its mismatch with the target $p$. A poor proposal increases rejection without changing the resulting distribution.

For a token $x$, the accepted probability mass is $q(x)\min(1,p(x)/q(x))=\min(p(x),q(x))$. Write $\alpha=\sum_x\min(p(x),q(x))$. Rejection occurs with probability $1-\alpha$, and the residual distribution normalizes $\max(p-q,0)$ by that same quantity. Thus the total output mass is

Resampling directly from $p$ would instead give $\min(p(x),q(x))+(1-\alpha)p(x)$, which generally differs from $p(x)$. When $\alpha=1$, rejection never occurs and no residual distribution is needed.

Equal distributions do not imply the same token sequence for an arbitrary random seed. Finite-precision implementations can also introduce rounding differences.

How much faster?

The speedup depends on the acceptance rate $\alpha$, the probability that any given draft token gets accepted. If the draft model closely matches the target model's distribution, $\alpha$ is high and most candidates are accepted.

Under the simplifying assumption of a constant conditional acceptance probability $\alpha$ at each candidate position, the expected token yield per target pass is:

where $k$ is the number of draft tokens per round. At $\alpha=1$, use the limit $k+1$ (or the equivalent sum $1+\alpha+\cdots+\alpha^k$). Some examples:

• $\alpha = 0.8$, $k = 4$: about 3.4 tokens per step. That is token yield per target pass, not wall-clock speedup.
• $\alpha = 0.9$, $k = 5$: about 4.7 tokens per step.
• $\alpha = 0.5$, $k = 4$: about 1.9 tokens per step.

The $k+1$ in the exponent accounts for the bonus token. Even when a draft token is rejected, you resample one new token from the target model's adjusted distribution. If all $k$ drafts are accepted, sample the bonus token from the target distribution at the next position. An immediate rejection yields one replacement token, but still incurs the wasted draft work.

For wall-clock speedup, divide expected token yield by the round cost relative to an ordinary target step. If $c$ is draft cost per token and $v_k$ is verification cost, both in target-step units, the estimate is $E[N]/(kc+v_k)$. Low acceptance or expensive drafting can make this less than one.

Leviathan et al. reported 2–3× acceleration for their T5-XXL experiments.

A worked example

The following round uses constructed probabilities and a chosen coin-flip outcome.

Say we're mid-generation. The context so far is the beginning of a new sentence and the draft model proposes five tokens: "The", "cat", "sat", "on", "the".

The target model processes all five in one forward pass and produces its own probability distribution at each position. Going left to right:

1. "The" — target gives it $p = 0.85$, draft had $q = 0.70$. Since $p \geq q$, accept.
2. "cat" — target gives $p = 0.40$, draft had $q = 0.35$. Since $p \geq q$, accept.
3. "sat" — target gives $p = 0.30$, draft had $q = 0.25$. Since $p \geq q$, accept.
4. "on" — target gives $p = 0.10$, draft had $q = 0.30$. Since $p < q$, we flip a coin with probability $p/q = 0.33$. The coin comes up reject. We resample from $\text{norm}(\max(0, p - q))$ across the full vocabulary and get "down".
5. "the" — discarded. Everything after the rejection point gets thrown away.

Net result: 3 accepted tokens + 1 resampled token = 4 tokens from one target forward pass. The generation continues from "down", and the draft model starts a new round from there.

For "on," the rejection probability was $1 - p/q = 0.67$. The replacement comes from the residual distribution, not from the target’s top token or directly from $p$. That correction accounts for the probability mass already supplied by accepted draft tokens.

When it works and when it doesn't

Acceptance is determined by overlap between the draft and target distributions: at one position, $\alpha=\sum_x\min(p(x),q(x))$. Predictable continuations can help when both models concentrate on the same tokens, but low entropy alone does not ensure agreement. Measure acceptance on the actual prompts, sampling settings, and draft/target pair. A mismatched draft can remain poor even on syntactically constrained text.

Beyond draft models: Medusa

A separate draft model needs deployment and memory alongside the target. It must produce distributions over the same token vocabulary, run cheaply, and overlap enough with the target distribution to repay its cost.

Cai et al. proposed a different approach called Medusa. Instead of a separate draft model, you add extra prediction heads to the target model itself. Alongside the standard language-model head that predicts position $t+1$, you train additional lightweight heads that predict positions $t+2$, $t+3$, $t+4$, and so on.

Each Medusa head is a small feedforward network on the target model’s last hidden state. A subsequent pass verifies the proposed continuations.

The extra heads avoid deploying a separate draft model. Medusa-1 trains them while freezing the backbone; Medusa-2 also trains the backbone. Their training and acceptance costs differ.

Medusa also introduced tree-structured verification. Instead of proposing a single linear sequence of $k$ tokens, you generate a tree of candidates. At position $t+2$, you might have three candidates (from the Medusa head's top-3 predictions). At $t+3$, candidates from the next prediction head form further branches; the same head’s candidates can appear under multiple earlier choices. You verify the entire tree in one forward pass using carefully constructed attention masks. The best path through the tree gives you the longest accepted prefix. This can accept a longer path than a single chain because a single bad prediction at one position doesn't necessarily waste all subsequent positions.

The earlier work by Stern et al. explored a similar idea called blockwise parallel decoding. Medusa refined it with better training objectives and the tree attention mechanism.

Medusa’s extra heads predict future positions from a shared current hidden state. Their candidates do not acquire the same token-by-token conditioning as an autoregressive draft model.

Medusa’s typical acceptance policy can trade distribution fidelity for speed, so the exact rejection-sampling guarantee above does not automatically apply to every configuration.

Common misconceptions

"Speculative decoding is an approximation." The exact accept/reject algorithm preserves the target distribution in exact arithmetic. Use the actual post-temperature/filter proposal and target distributions in the acceptance calculation; heuristic acceptance variants need separate analysis.

"The draft model needs to be the same architecture." No. Any model works as long as it shares the tokenizer (or at least the token vocabulary). You could use a tiny 2-layer transformer as a draft for a giant 80-layer MoE model. The architecture is irrelevant because the acceptance criterion only compares probability distributions over tokens, not internal representations. What matters is that both models produce distributions over the same vocabulary.

"Speedup is always $k\times$." The maximum output yield is $k + 1$ tokens per target pass when every draft is accepted. Wall-clock speedup also depends on verification time and draft cost. And you have to account for the overhead of running the draft model. If the draft model's full $k$-step round (all 4 drafts together, not per draft) takes about 20% of the time of one target-model step, your effective speedup at $\alpha = 0.8, k = 4$ is roughly $3.4 / 1.2 \approx 2.8\times$, not $3.4\times$.

"It helps equally with throughput and latency." Speculative decoding primarily helps latency for individual requests. In high-throughput serving with large decode batches, the GPU utilization is already decent from batching, and the draft model competes for the same GPU resources. For throughput-optimized workloads with large batch sizes, the benefit is often minimal.

What's next

The small-batch weight traffic derived in Post 5.3 can also be amortized across independent requests.

The next post covers continuous batching, how inference engines schedule many requests at once by interleaving their prefill and decode steps, so queued work can enter at iteration boundaries.

Additional reading (and watching)

• Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023. First proof that draft-then-verify preserves the exact target distribution.

• Chen, C., et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. Independent derivation of speculative sampling with the same correctness guarantee.

• Cai, T., et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. Introduced multi-head prediction and tree-structured verification.

• Stern, M., et al. (2018). Blockwise Parallel Decoding for Deep Autoregressive Models. NeurIPS 2018. Early precursor to Medusa-style parallel prediction heads.

• Li, Y., et al. (2024). Eagle: Speculative Sampling Requires Rethinking Feature Uncertainty. Feature-level draft prediction using the target model's hidden states.

• Xia, H., et al. (2024). Unlocking Efficiency in Large Language Model Inference: A Comprehensive Survey of Speculative Decoding. Survey covering the full taxonomy of speculative methods as of early 2024.