Accelerating Large Language Model Decoding with Speculative Sampling

A large language model generates text one token at a time, and for big models at small batch sizes each step mostly waits on reading weights from memory. Scoring several tokens in one pass costs about as much as generating one. DeepMind's speculative sampling uses that slack: a small draft model writes a short continuation of K tokens, the large target model scores all K positions in a single call, and a rejection-sampling rule decides how many draft tokens to keep. The authors note that this was developed independently of, and at the same time as, Leviathan et al.'s speculative decoding at Google; the core algorithm is the same.

The acceptance rule keeps each draft token with probability min(1, q/p), where q is the target's probability for that token and p the draft's. At the first rejection, a replacement is sampled from the normalized positive part of q − p and the rest of the draft is discarded. If all K survive, the target's last set of logits gives one more token for free. The output then follows the target model's distribution exactly, up to numerical precision, for nucleus sampling, top-k and temperature as well as greedy decoding.

This paper's contribution is making it work for a large model served across many chips. The obvious draft, a smaller model in the same family, is fastest on fewer chips than the target; a 7B model is fastest on 4 TPU v4s and gets slower on 16. So DeepMind trained a wide, shallow draft: 4 billion parameters in only 8 layers, which runs at 1.8 ms per token on the same 16 chips that serve Chinchilla 70B at 14.1 ms per token. With four draft tokens per step, decoding was 1.9 to 2.5 times faster, and the samples scored the same on summarization and code benchmarks.

Speculative sampling is one of the few latency optimizations that cannot change answers, and this paper adds two practical lessons that the original speculative decoding paper does not.

Pick the draft for the serving topology. A small model in the same family can be a poor draft if the target runs across many GPUs: the small model is fastest on fewer devices and wastes the rest. A wider model with fewer layers needs less communication per step. When a customer serves a large model with tensor parallelism, ask whether the draft was chosen for that layout.

More draft tokens is not free. Latency per token plateaued or got worse past three or four draft tokens, and the spread of time per response grew with K. For products with a P90 or P99 latency target, choose K for the tail, not only the mean.

Gains depend on the content. Code (HumanEval) sped up more than summarization, because code has predictable runs a draft can guess. Measure acceptance on the customer's real traffic before promising a speedup.