Speculative decoding with draft models, explained

Speculative decoding makes a large model generate several tokens per forward pass without changing what it generates. A small, fast draft model proposes the next $\gamma$ tokens one at a time; the large target model scores all of them in a single parallel pass; a rejection-sampling rule keeps the longest prefix the target agrees with and corrects the first disagreement. Because a decode step of the target is limited by memory bandwidth rather than compute, checking five tokens costs almost the same as generating one. When the draft guesses well, latency per token falls by 2 to 3 times, and the output distribution is provably identical to sampling from the target alone. It is a latency technique first: it pays off most at small batch sizes, and this article shows with numbers where it stops paying. For a fine-tuned target, the draft has to be aligned to the fine-tune, which is a small post-training job in its own right and the part most deployments get wrong.

Where it came from#

Two groups published the method independently. Leviathan, Kalman and Matias (2022) introduced speculative decoding and demonstrated 2 to 3 times faster T5-XXL inference with identical outputs. Chen, Borgeaud, Irving, Lespiau, Sifre and Jumper (2023) introduced speculative sampling with the same modified rejection scheme and reported 2 to 2.5 times faster decoding of the 70B Chinchilla model with a 4B draft in a distributed setup. Both build on blockwise parallel decoding (Stern, Shazeer and Uszkoreit, 2018), which predicted several future tokens with extra heads and kept the prefix a scoring model validated, but only for greedy decoding. The new ingredient in 2022 was the acceptance rule that makes the scheme exact for sampling at any temperature.

The idea has since branched. SpecInfer (Miao et al., 2023) verifies a tree of candidates instead of a chain. Medusa and EAGLE replace the separate draft model with light heads on the target's own hidden states. Prompt lookup and REST (He et al., 2023) draft by copying n-grams from the prompt or a datastore, with no model at all. Lookahead decoding (Fu et al., 2024) generates n-grams with Jacobi iterations of the target itself. This article covers the original two-model form, its math, and how to train a draft for your own fine-tune.

The bottleneck: one weight read per token#

At small batch sizes, a decode step must stream all $2N$ bytes of weights for an $N$-parameter 16-bit model to produce one token per sequence. On an H100-class GPU an 8B model takes about 6.5 ms per step at batch 1 in our step-time model, of which under 0.3 ms is arithmetic. Scoring $\gamma +1$ positions in one pass multiplies the arithmetic by $\gamma +1$ but reads the weights once, so as long as $b(\gamma +1)$ tokens stay below the compute ridge (about 221 tokens in our model; see the roofline article), verification costs the same as a normal decode step. Speculative decoding converts idle FLOPs into fewer serial steps.

The math: why it is exact, and how fast it is#

The acceptance rule

At some position, let $p(\cdot )$ be the target's next-token distribution and $q(\cdot )$ the draft's, both conditioned on the same prefix. The draft samples $x\sim q$. The target accepts it with probability $\min (1,p(x)/q(x))$. On rejection, it samples a replacement from the residual distribution

The probability of emitting $x$ is the probability of drafting and accepting it plus the probability of rejecting anything and resampling $x$. Writing $\beta ={\sum }_x\min (p(x),q(x))$ for the total acceptance probability,

because $1-\beta ={\sum }_x\max (0,p(x)-q(x))$, which cancels the denominator of $r$. The output is distributed exactly as the target's, whatever the draft. A bad draft only costs speed. Note also that $\beta =1-D_{\text{TV}}(p,q)$: the acceptance rate is one minus the total variation distance between the two models. With greedy decoding the rule degenerates to "accept if the draft's token equals the target's argmax".

A hand example. Three tokens A, B, C with $p=(0.6,0.3,0.1)$ and $q=(0.3,0.6,0.1)$. Then $\beta =0.3+0.3+0.1=0.7$. If the draft samples B (probability 0.6), it is accepted with probability $0.3/0.6=0.5$. The residual is $\max (p-q,0)=(0.3,0,0)$, normalised to "always A". So $\mathrm{Pr}[A]=0.3\times 1+0.3=0.6$ (drafted and accepted, plus rejected B resampled), $\mathrm{Pr}[B]=0.6\times 0.5=0.3$, $\mathrm{Pr}[C]=0.1$: exactly $p$.

Expected tokens per step and speedup

Leviathan et al. model the per-position acceptance as independent with mean $\alpha =\mathbb{E}[\beta ]$. The number of accepted draft tokens is then truncated-geometric, and every step also produces one token from the target (the correction on rejection, or a bonus token sampled from the target's last position if all $\gamma$ are accepted). The expected number of tokens per target pass is

If one draft step costs $c$ times a target step, and a verification pass costs the same as a target decode step, the expected speedup in wall-clock time is

Speedup requires $\alpha >c$ at least, and the optimal $\gamma$ grows with $\alpha$. For $c=0.05$ (a draft about 20 times cheaper), the optimum is $\gamma =3$ at $\alpha =0.5$ (speedup 1.63), $\gamma =6$ at $\alpha =0.7$ (2.35), $\gamma =8$ at $\alpha =0.8$ (3.09) and $\gamma =13$ at $\alpha =0.9$ (4.67). The curves are flat near the optimum (Figure 2), so a $\gamma$ of 4 to 6 is rarely far off.

What changes at larger batch sizes

The assumption that verification costs one decode step fails as the batch grows. Verifying $\gamma +1$ tokens for $b$ sequences processes $b(\gamma +1)$ tokens, and once that exceeds the ridge the pass becomes compute-bound. In general

where $t_T(b)$ is a target decode step, $t_D(b)$ a draft step and $t_V$ the verification pass. Plugging in our step-time model for an 8B target and a 1B draft ($c\approx 0.19$ at batch 1, larger than the weight ratio because fixed per-step overheads do not shrink with the model) gives Figure 3. With $\alpha =0.8$ and $\gamma =4$ the speedup is 1.90 at batch 1, still 1.81 at batch 64, then 1.49 at 128 and 1.25 at 256, where verification costs 1.76 times a decode step. With $\gamma =8$ the speedup falls below 1 at batch 256, and with $\alpha =0.6$ and $\gamma =4$ it does too (0.86). Su, Giannoula and Pekhimenko (2023) observed the same thing on real GPUs: the optimal speculation length depends on batch size. MagicDec (Sadhukhan et al., 2024) adds the counterpoint: with long contexts, KV reads keep large batches memory-bound, and speculation with a sparse-KV draft still delivered up to 2.51 times speedup for Llama-3.1-8B at batch sizes 32 to 256.

Toy runs: exactness and the formula#

To check the claims numerically, we built a 32-token bigram "language model" as the target (random logits with standard deviation 2) and a draft whose logits add Gaussian noise of standard deviation 1.2, giving a mean acceptance of $\alpha =0.638$ under the target's stationary distribution. The script is seeded and the numbers are its output.

• Exactness. Over 60,000 generations, the empirical distribution of the third generated token under speculative sampling ($\gamma =4$) was at total variation 0.0101 from the exact target marginal; plain target sampling with the same number of draws was at 0.0089. Both are sampling noise.
• Tokens per step. Measured tokens per target pass were 1.634, 2.038, 2.297, 2.472, 2.645 and 2.707 for $\gamma =1,2,3,4,6,8$, against 1.638, 2.044, 2.303, 2.469, 2.641 and 2.712 from the formula with $\alpha =0.638$. The independence assumption is not true here (acceptance depends on context), but for this model the formula with the averaged $\alpha$ is accurate to within 0.3%.

Draft models for a fine-tuned target#

A draft must share the target's tokenizer and vocabulary, and should be a small model from the same family. For an off-the-shelf model that is easy: use the 1B sibling of an 8B model. For a fine-tune it is the main trap. Your SFT or RL run moved the target's distribution toward your domain, format and style; the stock draft did not move, so $\alpha$ falls, sometimes from above 0.8 to near 0.5, and the speedup with it. Three bodies of work show how to fix this.

DistillSpec (Zhou et al., 2023) distills the target into the draft, with two findings: use on-policy data generated by the draft (the contexts the draft will actually be in), and choose the divergence to match the decoding strategy, since $\alpha =1-D_{\text{TV}}$ and neither forward nor reverse KL is the acceptance objective. They report 10 to 45% speedups over standard speculative decoding with an off-the-shelf draft. Online speculative decoding (Liu et al., 2023) keeps distilling the draft on live traffic and reports acceptance gains of 0.1 to 0.65 and 1.42 to 2.17 times lower latency. Hong et al. (2025) studied domain drafts for function calling, biology and Chinese and found offline distillation beats online by 11 to 25%, white-box (logit) distillation beats black-box (generated-text) distillation by 2 to 10%, and synthetic prompts reach 80 to 93% of the benefit of real user queries.

The distillation objectives, for draft $q_{\varphi }$, target $p$ and contexts $c$ drawn from a data distribution $\mathcal{D}$, are

and the TVD gradient with respect to the draft's logits $z$ is ${\nabla }_z{\mathcal{L}}_{\text{TVD}}=q\odot (g-⟨g,q⟩)$ with $g_x=-1[q(x)<p(x)]$: push up exactly the tokens the draft under-predicts. This is the same machinery as knowledge distillation and on-policy distillation, aimed at a different metric.

Toy run: re-aligning a stale draft

We simulated the fine-tuning gap with a 64-token bigram target. The base model has structured logits; the "fine-tuned" target adds a strong shift on half of the contexts (its domain). The draft is deliberately small: its logits are a rank-4 factorisation, so it cannot match the target exactly and the choice of objective matters. All drafts start from one distilled on the base model. Acceptance is measured under the fine-tuned target's stationary distribution, with expected tokens per step at $\gamma =4$.

Re-aligning the draft to the fine-tune raised tokens per step from 2.39 to 3.33 even with plain SFT on generated text, and optimising the acceptance metric directly (TVD) reached 3.96, a 66% gain over the stale draft. A capacity-limited draft must choose which mistakes to make, and the KL objectives choose differently from what acceptance rewards. Two caveats: greedy agreement did not rank the objectives the same way as sampling acceptance (so measure $\alpha$ at the temperature you serve), and an on-policy variant gave no extra gain here (0.879), which is expected in a bigram model with no exposure bias; DistillSpec found on-policy data matters for real drafts.

A practical recipe

1. Collect 10k to 100k prompts from your fine-tune's real traffic or a synthetic set in its domain.
2. Generate completions with the fine-tuned target at the serving temperature. Store the target's top-k logits per position (k = 20 to 50 is usually enough) for white-box distillation.
3. Initialise the draft from the small model of the same family, ideally after the same SFT data.
4. Train with forward KL or TVD on the stored logits, optionally mixing in draft-generated contexts relabelled by the target.
5. Measure acceptance length on held-out prompts at the serving temperature, then pick $\gamma$ with the formula and a load test.

Other interactions with post-trained models#

Tokenizers and templates. If the fine-tune added special tokens or a new chat template, the draft must have the same vocabulary and must have seen the template; a draft that never emits your tool-call delimiters is rejected at every one of them. LoRA adapters. With many adapters on one base, a single shared draft is aligned to none of them; either distill a draft per high-volume adapter or use a model-free drafter such as n-gram prompt lookup. Quantised targets. A 4-bit target has a cheaper decode step, which raises $c$ for the same draft; recompute the optimal $\gamma$. Reasoning models. Long, repetitive chains of thought are often highly predictable, so acceptance is high, but the batch of concurrent long generations is also large, which pushes toward the right side of Figure 3.

Tuning knobs#

Failure modes and how to spot them#

Silent slowdowns. Speculation with a misaligned draft or at high batch can be slower than plain decoding while every output stays correct. Always log mean acceptance length and compare tokens per second with speculation on and off at your real concurrency.

Acceptance drift. A new checkpoint, a new system prompt or a traffic shift lowers $\alpha$. Track acceptance per route or adapter over time; a drop after a model release means the draft needs re-distillation.

Numerical non-identity. The method is exact in distribution, but verification runs the target on a different token shape than plain decoding, so floating-point results can differ slightly; vLLM's documentation describes the guarantee as lossless up to hardware numerics. Do not expect bit-identical greedy outputs when switching speculation on.

Memory pressure. The draft's weights and KV cache take memory that would otherwise hold target KV, reducing maximum batch size.

When to use it#

Implementation sketch#

In vLLM, speculation is a server option. A distilled draft for your fine-tune:

from vllm import LLM, SamplingParams

llm = LLM(
    model="path/to/finetuned-8b",
    speculative_config={
        "model": "path/to/distilled-draft-1b",   # same tokenizer and chat template
        "num_speculative_tokens": 4,
    },
)
# Model-free alternative for copy-heavy workloads:
# speculative_config={"method": "ngram", "num_speculative_tokens": 4, "prompt_lookup_max": 4}

In Hugging Face Transformers, assisted generation runs the same algorithm with model.generate(**inputs, assistant_model=draft). The draft distillation step is a short PyTorch loop over stored target logits:

import torch, torch.nn.functional as F

def draft_loss(draft_logits, target_logits, mask, kind="tvd", T=1.0):
    """logits: (B, L, V) over the same tokens; mask: (B, L) completion positions."""
    q = F.softmax(draft_logits / T, -1)
    p = F.softmax(target_logits / T, -1)
    if kind == "fkl":
        per_tok = (p * (p.clamp_min(1e-9).log() - q.clamp_min(1e-9).log())).sum(-1)
    else:  # total variation: 1 - expected acceptance
        per_tok = 0.5 * (p - q).abs().sum(-1)
    return (per_tok * mask).sum() / mask.sum()

In POFE, a draft is trained as a distillation recipe against the fine-tuned checkpoint you are deploying, and the evaluation report includes measured acceptance length at the serving temperature next to the quality metrics.

Further reading#

Leviathan et al. (2022) and Chen et al. (2023) are the primary sources and both are short. Xia et al. (2024) survey the field and compare methods under one setup. Zhou et al. (2023) and Hong et al. (2025) are the references for training drafts. For heads that live inside the target instead of a separate model, continue with Medusa and EAGLE.

Frequently asked questions#

References#

1. Yaniv Leviathan, Matan Kalman, Yossi Matias (2022). Fast Inference from Transformers via Speculative Decoding. arxiv.org/abs/2211.17192
2. Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, John Jumper (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arxiv.org/abs/2302.01318
3. Mitchell Stern, Noam Shazeer, Jakob Uszkoreit (2018). Blockwise Parallel Decoding for Deep Autoregressive Models. arxiv.org/abs/1811.03115
4. Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, Rishabh Agarwal (2023). DistillSpec: Improving Speculative Decoding via Knowledge Distillation. arxiv.org/abs/2310.08461
5. Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Alvin Cheung, Zhijie Deng, Ion Stoica, Hao Zhang (2023). Online Speculative Decoding. arxiv.org/abs/2310.07177
6. Fenglu Hong, Ravi Raju, Jonathan Lingjie Li, Bo Li, Urmish Thakker, Avinash Ravichandran, Swayambhoo Jain, Changran Hu (2025). Training Domain Draft Models for Speculative Decoding: Best Practices and Insights. arxiv.org/abs/2503.07807
7. Qidong Su, Christina Giannoula, Gennady Pekhimenko (2023). The Synergy of Speculative Decoding and Batching in Serving Large Language Models. arxiv.org/abs/2310.18813
8. Ranajoy Sadhukhan, Jian Chen, Zhuoming Chen, Vashisth Tiwari, Ruihang Lai, Jinyuan Shi, Ian En-Hsu Yen, Avner May, et al. (2024). MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding. arxiv.org/abs/2408.11049
9. Minghao Yan, Saurabh Agarwal, Shivaram Venkataraman (2024). Decoding Speculative Decoding. arxiv.org/abs/2402.01528
10. Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, et al. (2023). SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification. arxiv.org/abs/2305.09781
11. Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D. Lee, Di He (2023). REST: Retrieval-Based Speculative Decoding. arxiv.org/abs/2311.08252
12. Yichao Fu, Peter Bailis, Ion Stoica, Hao Zhang (2024). Break the Sequential Dependency of LLM Inference Using Lookahead Decoding. arxiv.org/abs/2402.02057
13. Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wenjie Li, Zhifang Sui (2024). Unlocking Efficiency in Large Language Model Inference: A Comprehensive Survey of Speculative Decoding. arxiv.org/abs/2401.07851
14. vLLM project (2025). vLLM documentation: Speculative Decoding. docs.vllm.ai/en/latest/features/speculative_decoding/

Related reading