Chapter 20 · Section 20.5 · 90 min

Speculative decoding

Speculative decoding makes a large language model produce several tokens per forward pass without changing what it writes. A small, fast model guesses the next few tokens. The large model checks all the guesses in one pass, keeps the ones it agrees with, and adds one token of its own. When the guesses are good, each expensive pass yields three or four tokens instead of one. The text that comes out has exactly the distribution the large model would have produced on its own.

Two names hold for the whole section. The large model whose output we want is the target. The small model that guesses is the draft. The target is fixed, and the draft is ours to choose.

The trick works because checking a few tokens costs about as much as writing one. To write a token, a model must read every one of its weights from the GPU's memory. On a GPU that reading takes far longer than the arithmetic the token needs (§20.1). Here is how lopsided it is for a real model.

✏️ Where the time of one token goes

Take Llama 3.1 70B, with 70.6 billion parameters stored in bfloat16 (2 bytes each), on a machine of eight H100 GPUs.

Reading. The weights are $2\times 70.6\times10^9 \approx 141$ GB. Each H100 reads 3.35 TB per second from its memory, so eight of them read the weights in $141\times10^9/(8\times 3.35\times10^{12}) \approx 5.3$ ms.
Arithmetic. One token costs about two floating-point operations (FLOPs) per weight, one multiply and one add, so $2\times 70.6\times10^9 \approx 1.4\times10^{11}$ FLOPs. Eight H100s do $8\times 989\times10^{12}$ FLOPs per second, so that takes 0.018 ms.
The gap. The GPUs wait on memory for 5.3 ms and compute for 0.018 ms, a third of a percent of the step.

Feed the model five tokens instead of one. The arithmetic grows to 0.09 ms, but the weights are still read once, so the pass still takes about 5.3 ms. Five tokens are checked for the price of one.

Two questions decide whether this is worth anything. First, how can the target accept guesses without changing what it would have written, even when it samples at a temperature instead of always taking its top token? Second, how often must the guesses be right, and how cheap must the guesser be, for the scheme to pay? This section answers both exactly. Then it follows the idea into a busy server, where checking stops being free, into trees of guesses, and through the kinds of guessers that frontier serving systems use in 2026. The proof of exactness and the speedup formula are both standard interview questions on inference.

Checking several tokens costs about as much as writing one

Every speedup in this section is measured against the cost model of a decode step built in §20.1, so we restate it here. A decode step does two kinds of work. It reads bytes from the GPU's memory, and it does arithmetic. The GPU overlaps the two, so the step takes as long as the larger of them.

The bytes come in two kinds. The weights are read once per step, however many sequences share the step. Each sequence also reads its own KV cache, the stored keys and values of its earlier tokens, kept so that they are not recomputed (§6.5). The arithmetic is about two FLOPs per weight for every token processed, plus the attention of that token over its cache.

In symbols, a decode step for $B$ sequences whose mean context, the number of tokens already cached, is $\bar T$ takes

t_{\text{step}}(B) = \max(t_W + Bt_K,\; Bt_F) + t_0 .

Each term is a time.

$t_W = M_W/(n_{\text{gpu}}\beta)$ is the time to read the weights once. $M_W = s_WN$ is their size in bytes, $N$ parameters at $s_W$ bytes each. $\beta$ is one GPU's memory bandwidth in bytes per second, and $n_{\text{gpu}}$ is the number of GPUs that share one copy of the model.
$t_K = \bar Tm_{\text{KV}}/(n_{\text{gpu}}\beta)$ is the time to read one sequence's cache, where $m_{\text{KV}}$ is the bytes cached per token.
$t_F = (2N_{\text{act}} + 4Ld_{\text{attn}}\bar T)/(n_{\text{gpu}}\pi)$ is the time of one token's arithmetic. $N_{\text{act}}$ is the number of parameters multiplied per token, $L$ the number of layers, $d_{\text{attn}}$ the attention width (the number of query heads times the head size), and $\pi$ one GPU's peak rate in FLOPs per second.
$t_0$ is what every step pays regardless, for communication between GPUs and for launching kernels.

As in §20.1, $\pi$ and $\beta$ always mean these two hardware numbers in this chapter. Their ratio $I^\star = \pi/\beta$ is the ridge point, the number of FLOPs a kernel must do per byte it reads before arithmetic becomes the bottleneck. For an H100 in bfloat16 it is $989/3.35 \approx 295$ FLOPs per byte.

Now feed the model $k$ new tokens per sequence in one pass instead of one. This is the chunked forward pass of §6.5. Each new token attends to the whole cache and to the new tokens before it, through the causal mask of §4.4 aligned to its bottom-right corner. Every new token does a full token's arithmetic. The weights and the caches are still read once. So the pass takes

t_{\text{pass}}(B, k) = \max\big(t_W + Bt_K,\; kBt_F\big) + t_0 .

In plain words, checking $k$ tokens per sequence costs one decode step as long as the extra arithmetic still hides under the memory time, that is, while $kBt_F \le t_W + Bt_K$ . The pass also writes $kB$ new cache entries, kilobytes against gigabytes, which we drop as §20.1 did.

For one sequence with a short context, the condition reads $k \le t_W/t_F = s_WI^\star/2$ . In bfloat16, $s_W = 2$ , so that is about 295 tokens. §20.1 met the same number as the knee, the batch size at which plain decoding turns compute-bound. Both count the tokens processed per read of the weights. A weight read from memory does not care whether those tokens come from 295 different sequences or from one sequence checking 295 guesses.

Python · cell 1

In this model a pass over 256 tokens of one sequence costs exactly one decode step, for both models. Only at 512 does the arithmetic exceed the memory traffic.

Real kernels are less generous. A matrix multiply works in tiles of 64 or 128 rows, so a few extra tokens can cost a whole extra tile. Engines also record each step as a CUDA graph, a replayable sequence of kernel launches, for only a few fixed shapes, and pad every pass up to the next shape. So the cost of a verification pass that vLLM measures at startup is a staircase set by those shapes, not a smooth roofline (vLLM Team 2026). In practice the free range of a few hundred tokens is never the limit. Guesses far down a draft are rarely right, as the speedup formula below shows, so deployed chains check from 1 to about 16 tokens and trees of guesses have a few dozen nodes.

On several GPUs there is one more cost, the all-reduces of tensor parallelism (§20.1), which sum each layer's partial results across the GPUs. Their cost does not grow with $k$ either. Each carries the activations of a handful of tokens, so its time is set by the latency of a message, not by bytes. Chen et al. (2023) list exactly these three reasons why scoring a short continuation costs about as much as scoring one token. The linear layers are memory-bound at small batch. The attention cost is dominated by reading the cache, which does not grow with $k$ . And the all-reduces are latency-bound.

An interviewer will ask why a forward pass over five tokens costs about the same as a pass over one. The answer is this box. At small batch the step is set by bytes read, and checking more tokens adds arithmetic, not bytes.

Greedy speculation keeps the target's own choices

Two symbols run through the rest of the section. Given the tokens so far, $y_{1:t}$ , the target assigns a probability $p(x \mid y_{1:t})$ to every possible next token $x$ . That list of probabilities is its next-token distribution $p$ . The draft has its own next-token distribution, written $q$ .

Start with a target that decodes greedily, always taking its most likely token, $y_{t+1} = \arg\max_x p(x \mid y_{1:t})$ . This case needs no probability, and the whole idea is already visible in it. The draft guesses a few tokens. The target scores all of them in one pass and says, at every position, what it would have picked there. We keep the guesses up to the first disagreement and take the target's own pick at that position.

The example took four steps. The draft made its guesses. One target pass returned the target's own pick after every guess. The guesses were kept up to the first disagreement. Then the target's pick at that position was added. Written for any prefix and any number of guesses, these are the four steps of a round.

📐 A round of greedy speculative decoding

Let $y_{1:t}$ be the tokens so far and $\gamma \ge 1$ the draft length, the number of tokens the draft guesses per round. (The letter follows Leviathan et al. 2023, and it is not a discount factor.)

Draft. For $i = 1, \dots, \gamma$ , let $x_i = \arg\max_x q(x \mid y_{1:t}, x_{1:i-1})$ , one cheap draft step each.
Verify. Run the target once on the last committed token followed by the $\gamma$ drafted tokens. The last committed token is the target's own output from the previous round, so the target has not read it yet and its keys and values are not cached. The logits at these $\gamma + 1$ positions give the target's own choices $z_i = \arg\max_x p(x \mid y_{1:t}, x_{1:i-1})$ for $i = 1, \dots, \gamma + 1$ .
Accept. Let $J$ be the largest $j \le \gamma$ with $x_i = z_i$ for every $i \le j$ .
Emit $x_1, \dots, x_J$ and then $z_{J+1}$ . That last token is a correction if $J < \gamma$ and a free bonus token if $J = \gamma$ .

The bonus token is why a round where every guess is right emits $\gamma + 1$ tokens, one more than the draft wrote. The target pass computed its pick after the last guess anyway, so that token costs nothing extra. The question that remains is whether the output is still the target's own. It is, by a short induction over the positions of a round.

∂ Proof that greedy speculation reproduces greedy decoding

Let $g_1, g_2, \dots$ be the target's greedy continuation of $y_{1:t}$ . We show by induction that the round emits $g_1, \dots, g_{J+1}$ . For $i \le J$ , assume $x_1, \dots, x_{i-1} = g_1, \dots, g_{i-1}$ . Then $z_i$ is the target's argmax after $y_{1:t}, g_{1:i-1}$ , which is $g_i$ by definition. Also $x_i = z_i$ by the definition of $J$ , so $x_i = g_i$ . The last emitted token is $z_{J+1}$ , the target's argmax after $y_{1:t}, g_{1:J}$ , which is $g_{J+1}$ . The next round starts from $y_{1:t}, g_{1:J+1}$ , and the same argument applies. The round emits $J + 1 \ge 1$ tokens with $J \le \gamma$ . ∎

In plain words, keep the drafts for as long as they agree with what the target would have chosen anyway, and take one token from the target at the first disagreement. A round never emits fewer tokens per target pass than plain decoding, since it always emits at least the target's own next token. It can still cost more time, because the draft steps are not free, which the speedup formula below accounts for.

The guarantee holds in exact arithmetic. On a GPU, a pass over $\gamma + 1$ tokens runs kernels of different shapes from $\gamma + 1$ one-token passes. Kernels of different shapes add numbers in different orders, and floating-point addition is not associative, so the logits can differ in their last bits. At a near-tie that can flip an argmax. This is the batch-invariance problem of §6.4. So greedy speculation reproduces the greedy output of the same engine with the same verification kernels. At rare near-ties it can differ from a run without speculation, which matters if you test it by comparing strings.

Interview

Greedy speculative generation, with call counting

Implement greedy_speculative_generate(target_logits, draft_next, prompt, n_new, gamma), greedy speculative decoding that also counts how many times it calls the target.

The two models are plain functions over a vocabulary of V tokens.

target_logits(tokens) runs the target once on a list of tokens. It returns an array of shape (len(tokens), V) whose row i scores the token that comes after tokens[:i+1].
draft_next(tokens) returns the draft's logits for the token after tokens, an array of shape (V,).

Each round does three things. It drafts gamma tokens greedily, one draft_next call each. It makes exactly one target call, on the current tokens followed by the drafts. It keeps the longest prefix of drafts that match the target's argmaxes and appends the target's own argmax after that prefix.

Return two things, the prompt followed by exactly n_new new tokens, and the number of target calls you made. The tokens must equal plain greedy decoding with the target. The setup provides make_models(seed, noise), which returns a target, a draft, and a dict counting the calls to each, and plain_greedy for comparison.

import numpy as np

def greedy_speculative_generate(target_logits, draft_next, prompt, n_new, gamma):
    toks = list(prompt)
    for _ in range(n_new):                                     # placeholder: no speculation at all
        toks.append(int(np.argmax(target_logits(toks)[-1])))
    return toks, n_new

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

When the target samples, keep a guess with probability p over q

Most deployed models sample, and then agreement with the target's top token is the wrong test. The target does not have one right answer at each position, only a distribution $p$ . From here on, $p$ and $q$ mean the distributions each model really samples from, after temperature, truncation such as top- $p$ , and penalties have been applied (§6.4). We need a rule that turns a draft sampled from $q$ into a token distributed exactly as $p$ , and that keeps the draft as often as possible.

There is an easy exact rule, and it is poor. Sample $y \sim p$ from the verification pass, and keep the draft $x$ only if $x = y$ . The output is always $y$ , so it is distributed exactly as $p$ . But two independent samples agree only with probability $\sum_xp(x)q(x)$ . That is small whenever the distributions are spread out, even when they are identical. Take $p = (0.5, 0.3, 0.2)$ and $q = (0.2, 0.6, 0.2)$ over three tokens $a, b, c$ . They agree with probability $0.10 + 0.18 + 0.04 = 0.32$ .

The good rule treats the draft's tokens unequally. Where the draft proposes a token less often than the target wants it, keep every such proposal. Where it proposes a token too often, keep only the fraction $p(x)/q(x)$ of those proposals. When a draft is rejected, draw a replacement from the probability that the draft under-covers. That is the residual $\max(0, p - q)$ , normalized to sum to 1.

✏️ The rule by hand, three tokens

Take $p = (0.5, 0.3, 0.2)$ and $q = (0.2, 0.6, 0.2)$ over tokens $a, b, c$ .

Keep probabilities. $\min(1, p/q)$ is $1$ for $a$ , $0.5$ for $b$ , and $1$ for $c$ . Only $b$ is over-proposed.
Acceptance. A draft is kept with probability $\sum_x\min(p(x), q(x)) = 0.2 + 0.3 + 0.2 = 0.7$ .
Residual. $\max(0, p - q) = (0.3, 0, 0)$ , which normalizes to $(1, 0, 0)$ . Every rejection becomes an $a$ .
Output. $a$ comes out with probability $0.2 + 0.3\times 1 = 0.5$ (drafted and kept, or rejected and replaced). $b$ comes out with probability $0.6\times0.5 = 0.3$ (drafted, kept half the time). $c$ comes out with probability $0.2$ . That is exactly $p$ .

The naive rule kept the draft 32% of the time. This one keeps it 70% of the time.

Is 0.7 good, or could a cleverer rule keep the draft more often? To answer, we need a way to measure how far apart two distributions are.

In plain words, total variation is the largest difference in probability that the two distributions can give to any event, and one minus it is how much their histograms overlap. In the example, $\operatorname{TV} = \tfrac12(0.3 + 0.3 + 0) = 0.3$ and the overlap is 0.7.

∂ Proof of the two identities

Since $\min(a, b) = (a + b - \lvert a - b\rvert)/2$ , summing over $x$ gives $\sum_x\min(p, q) = (1 + 1 - 2\operatorname{TV})/2 = 1 - \operatorname{TV}$ . For the maximum, let $A^\star$ be the set of tokens with $p(x) > q(x)$ . Both distributions sum to 1, so $\sum_x(p - q) = 0$ , and the positive and negative parts of $p - q$ have the same total. Hence $p(A^\star) - q(A^\star) = \sum_x\max(0, p - q) = \tfrac12\sum_x\lvert p - q\rvert$ . Any other set $A$ can only lose positive terms or add negative ones, so $p(A) - q(A) \le p(A^\star) - q(A^\star)$ , and by symmetry the same bound holds for $q(A) - p(A)$ . ∎

Training usually measures distance with the KL divergence of §0.5, not with total variation. The two are tied together by Pinsker's inequality, $\operatorname{TV}(p, q) \le \sqrt{D_{\mathrm{KL}}(p\,\Vert\,q)/2}$ with the KL in nats (Cover & Thomas 2006, Lemma 11.6.1, which we use without proof). So a draft close to the target in KL is close in total variation. That is why distilling the target into the draft, training the draft to match the target's distributions, raises acceptance.

∴ The speculative sampling rule

Draw $x \sim q$ . Accept it with probability $\min\big(1, p(x)/q(x)\big)$ . If it is rejected, draw instead from the residual distribution

p_{\text{res}}(x) = \frac{\max\big(0,\, p(x) - q(x)\big)}{\sum_{x'}\max\big(0,\, p(x') - q(x')\big)} .

The token that comes out is distributed exactly as $p$ , whatever $q$ is, and the draft is accepted with probability

\alpha = \sum_x\min\big(p(x), q(x)\big) = 1 - \operatorname{TV}(p, q).

This $\alpha$ is the acceptance rate. It is not the link latency $\alpha_{\text{link}}$ of §20.1.

Proof. The output is $x$ in one of two disjoint ways. Either $x$ was drafted and accepted, with probability $q(x)\min(1, p(x)/q(x)) = \min(q(x), p(x))$ , or the draft was rejected and the residual produced $x$ . Summing the first over $x$ gives the acceptance probability $\alpha = \sum_x\min(p, q)$ . The rejection probability is therefore $1 - \alpha = \sum_x\big(p(x) - \min(p(x), q(x))\big) = \sum_x\max(0, p(x) - q(x))$ . That is exactly the residual's normalizer, so the second way has probability $(1 - \alpha)\,p_{\text{res}}(x) = \max(0, p(x) - q(x))$ . The total is $\min(p(x), q(x)) + \max(0, p(x) - q(x)) = p(x)$ , because $\min(a, b) + \max(0, a - b) = a$ for any two numbers. A token with $q(x) = 0$ is never drafted, so its ratio is never needed. When $p = q$ every draft is accepted and the residual is never needed either. ∎

In plain words, the probability thrown away by rejecting over-proposed tokens is exactly the probability that the draft under-covers elsewhere, and the residual hands it back there. §0.4 met this rule as a variant of rejection sampling, with the same acceptance rate $\alpha$ . Unlike plain rejection sampling, a rejection is never wasted, because one exact sample comes out either way. In an interview, this proof is the core of the topic. Expect to give it on a whiteboard in under five minutes.

So is the 0.7 of the example the best possible? Yes. No exact rule can keep a single draft more often than $1 - \operatorname{TV}$ .

Proof. For every token $x$ , the event that $X = x$ and $Y = x$ is contained both in the event $X = x$ and in the event $Y = x$ , so its probability is at most $\min(q(x), p(x))$ . Summing over $x$ gives the bound. The speculative rule emits $Y = X$ exactly when it accepts, which has probability $\sum_x\min(p, q)$ , and its $Y$ is distributed as $p$ by the previous theorem. ∎

In plain words, two random tokens with distributions $q$ and $p$ can agree at most as often as their histograms overlap. A joint distribution of two random variables with given distributions is called a coupling. One that makes them equal as often as possible is a maximal coupling, and speculative sampling is a maximal coupling of the draft and the target.

What this changes in practice is where to look for more speed. The verification rule cannot be improved, so higher acceptance from a single draft needs a draft closer to the target. The other route is several drafts per position, the token trees below. Sun et al. (2023) analyze that case as an optimal transport problem, the problem of moving one distribution onto another at least cost.

Three special cases are worth knowing by heart.

A greedy target. Then $p$ is one-hot at its argmax $z$ , all its probability on one token. A draft $x \ne z$ is rejected with certainty, the residual is one-hot at $z$ , and the rule reduces to greedy speculation.
A deterministic draft, for example one that always proposes its own argmax $\hat x$ . Then $q$ is one-hot at $\hat x$ , and the draft is accepted with probability $p(\hat x)$ . On rejection the residual is $p$ with $\hat x$ removed and renormalized. So a greedy draft is allowed even when the target samples, as long as the rule uses the $q$ the draft really sampled from.
A processed target. $p$ must be the distribution the target would really sample from at that position. That means the temperature, top- $p$ , and penalties of §6.4 are applied with the drafted prefix as context. Chen et al. (2023) apply nucleus sampling and temperature this way.

Temperature matters, and here both models sample at the same temperature $\tau$ . As $\tau$ goes to 0, both approach their argmax. So $\alpha$ tends to 1 where their argmaxes agree and to 0 where they do not. As $\tau$ grows, both spread toward the uniform distribution, and $\alpha$ tends to 1 again. In between, $\alpha$ can move either way, and on real models it is usually highest near 0. Leviathan et al. (2023) drafted for T5-XXL (11B parameters) with T5-small (77M) on English to German translation. They measured $\alpha = 0.75$ under argmax sampling and $0.62$ at temperature 1. With LaMDA 100M drafting for LaMDA 137B on dialog, they measured $0.61$ and $0.57$ . An interviewer may ask whether speculation helps more at temperature 0 or 1. Usually at 0, and these numbers are why.

Core

Acceptance rate and temperature

Implement two functions.

tv(p, q) returns the total variation distance between distributions stored along the last axis of two arrays of shape (..., V). It returns one value per row.
acceptance_rate(p_logits, q_logits, tau) returns the acceptance probability $\sum_x\min(p, q)$ , one value per row, when both models sample at temperature tau from the given logits. That is, $p$ is the softmax of p_logits / tau, and $q$ is the softmax of q_logits / tau. Treat tau = 0 as greedy decoding on both sides.

Then try logits of your own and watch how $\alpha$ moves with the temperature. It need not move in one direction.

import numpy as np

def tv(p, q):
    raise NotImplementedError

def acceptance_rate(p_logits, q_logits, tau):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

A whole round of speculative sampling

A real round drafts $\gamma$ tokens in a row and verifies all of them in one target pass. It applies the one-token rule to the drafts from left to right and stops at the first rejection. The target pass supplies the target's distribution at every drafted position at once, each conditioned on the drafts before it.

📐 A round of speculative sampling

Let $y_{1:t}$ be the tokens so far and $\gamma \ge 1$ the draft length.

Draft. For $i = 1, \dots, \gamma$ , compute $q_i = q(\cdot \mid y_{1:t}, x_{1:i-1})$ and sample $x_i \sim q_i$ .
Verify. One target pass over the drafts gives $p_i = p(\cdot \mid y_{1:t}, x_{1:i-1})$ for $i = 1, \dots, \gamma + 1$ .
Accept. Draw independent uniforms $u_1, \dots, u_\gamma$ on $(0, 1)$ . Let $J$ be the number of leading positions with $u_i < p_i(x_i)/q_i(x_i)$ . So $J = \gamma$ if every draft passes, and otherwise $J = i - 1$ for the first $i$ that fails.
Emit $x_1, \dots, x_J$ and one more token. If $J < \gamma$ , sample it from the residual of $p_{J+1}$ and $q_{J+1}$ . If $J = \gamma$ , sample a bonus token from $p_{\gamma+1}$ .

This is Algorithm 1 of Leviathan et al. (2023). It was found independently at DeepMind, where it is Algorithm 2 of Chen et al. (2023). Chen et al. call the target $q$ and the draft $p$ , the reverse of our letters, which is worth knowing when you read both papers.

∂ Proof that a whole round is exact

By the chain rule it is enough to show that each emitted token, given everything that happened before it was produced, has the target's conditional distribution given the tokens emitted so far. Look at the moment the procedure produces a token at some position of some round, with emitted prefix $w$ . Two cases are possible.

If the round has accepted every earlier draft and this position $i \le \gamma$ holds a draft, then $x_i$ was sampled from $q_i = q(\cdot \mid w)$ given the past, $u_i$ is a fresh uniform independent of the past, and $p_i = p(\cdot \mid w)$ . The token emitted here is $x_i$ if accepted and a residual sample otherwise. That is the single-token rule applied to $p_i$ and $q_i$ , so the token is distributed as $p(\cdot \mid w)$ by the theorem above.

If instead $i = \gamma + 1$ , the bonus token is drawn from $p_{\gamma+1} = p(\cdot \mid w)$ directly.

A round ends after its first token that is not an accepted draft, and the next round starts from the emitted prefix with fresh randomness, so these two cases cover every token. The drafts after a rejected position were sampled conditioned on the rejected token and are discarded, so they never enter the output. ∎

In plain words, each position of a round is a fresh use of the one-token rule, and the round stops at the first position where the rule did not keep the draft.

The proof relies on one requirement that deserves a name now, because the schemes later in this section that choose how much to verify must respect it. The decision whether to verify position $i$ at all may depend on the tokens emitted before $i$ . It must never depend on the drafted token $x_i$ itself, or on anything computed from it, such as the drafts after it. A rule with this property is called non-anticipating. The proof used it without saying so, when it treated the token at position $i$ as a plain application of the one-token rule.

Interview

Verify one round of speculative sampling

Implement speculative_verify(p, q, drafts, rng), the verification step of one round. Its inputs are the following.

p has shape (gamma + 1, V). Row i is the target's distribution after the prefix and the first i drafts.
q has shape (gamma, V). Row i is the distribution that draft i was sampled from.
drafts holds the gamma drafted token ids.

Return the emitted tokens as a list of ints, the accepted drafts followed by one more token.

Use the random generator the way an engine does. Call rng.random() once per verified position, and accept when the result is below $p_i(x_i)/q_i(x_i)$ . Call rng.choice(V, p=...) once for the extra token. It comes from the normalized residual after a rejection, or from p[gamma] when every draft is accepted. The tests pass in a generator that records its calls.

import numpy as np

def speculative_verify(p, q, drafts, rng):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

The simulation below checks the theorem on small models we can solve exactly. Both are bigram models over six tokens, models whose next-token distribution depends only on the previous token $s$ . The target is $p(\cdot\mid s)$ , and the draft $q(\cdot\mid s)$ is a perturbed copy of it. Every drafted position conditions on the drafts before it, as in a real model.

We generate the first two tokens of 3,000 sequences and compare their joint frequencies with the exact target probabilities. We do it twice. The first run uses the correct residual. The second has the most common bug, resampling from $p$ after a rejection.

Python · cell 2

✏️ A broken sampler looks healthy on a speed dashboard

With the residual, the 36 joint frequencies are within sampling noise of the target. To judge that we use Pearson's chi-square statistic, the sum over cells of the squared count error divided by the expected count. It is 37.6 here. An exact sampler gives about $35 \pm 8$ , since the statistic's mean is its 35 degrees of freedom. The total variation of 0.04 is the noise of 3,000 samples.

Resampling from $p$ moves the distribution by 0.20 in total variation and pushes the statistic to 586. Yet the number of tokens per round, the only thing a speed benchmark measures, is the same, 2.94 against 2.95. A broken speculative sampler looks perfectly healthy on a throughput dashboard. That is why the tests of Exercise 20.5c check distributions and not speeds, and why you should test your own engine the same way.

How many tokens a round produces, and the speedup

A round that accepts $J$ drafts emits $J + 1$ tokens, the accepted drafts plus the target's correction or bonus token. Start with a tiny case. Let the draft guess $\gamma = 2$ tokens, each accepted with probability 0.8 if the one before it was.

The first draft is rejected with probability 0.2. The round emits 1 token.
The first is accepted and the second rejected with probability $0.8\times0.2 = 0.16$ . The round emits 2 tokens.
Both are accepted with probability $0.8^2 = 0.64$ . The round emits 3 tokens, the two drafts and a bonus.

The expected count is $0.2\times1 + 0.16\times2 + 0.64\times3 = 2.44$ . The same number comes out as $1 + 0.8 + 0.64$ , one token always, plus the chance of reaching each draft. That second way of counting is the general rule. Write $a_i = \Pr(J \ge i)$ for the survival probability of position $i$ , the probability that the first $i$ drafts are all accepted.

∴ Expected tokens per round

The expected number of tokens a round emits is

\bar n = \mathbb{E}[J + 1] = 1 + \sum_{i=1}^{\gamma}a_i .

Suppose every draft is accepted with the same probability $\alpha$ , given that the drafts before it were, independently of everything else (the i.i.d. model). Then $a_i = \alpha^i$ , and

\bar n = 1 + \alpha + \alpha^2 + \dots + \alpha^\gamma = \frac{1 - \alpha^{\gamma+1}}{1 - \alpha} \qquad (\alpha < 1).

Under that model $J$ is a geometric variable capped at $\gamma$ , with $\Pr(J = i) = \alpha^i(1 - \alpha)$ for $i < \gamma$ and $\Pr(J = \gamma) = \alpha^\gamma$ .

∂ Proof of the expected number of tokens

$J = \sum_{i=1}^{\gamma}\mathbb{1}[J \ge i]$ , because $J$ counts the positions $1, \dots, J$ . Taking expectations gives $\mathbb{E}[J] = \sum_i\Pr(J \ge i) = \sum_ia_i$ . Under the i.i.d. model, surviving to position $i$ means $i$ independent acceptances, so $a_i = \alpha^i$ . Then $\Pr(J = i) = a_i - a_{i+1} = \alpha^i(1 - \alpha)$ for $i < \gamma$ , and $\Pr(J = \gamma) = a_\gamma = \alpha^\gamma$ . The geometric sum is $(1 - \alpha^{\gamma+1})/(1 - \alpha)$ . ∎

In plain words, each position contributes the probability of being reached, and the leading 1 is the token the target always adds. With $\alpha = 0.8$ and $\gamma = 4$ , a round emits $1 + 0.8 + 0.64 + 0.512 + 0.4096 = 3.36$ tokens on average.

The first formula is exact for any draft. The second is a model. In a real model the acceptance probability of position $i$ is $\alpha_i = \sum_x\min(p_i(x), q_i(x))$ , and it changes with the context. Leviathan et al. call this per-position value $\beta$ , define $\alpha$ as its expectation, and assume the values are i.i.d. to get the closed form. In real drafts acceptance usually decays along the block. So the survival form is the one to use, and the schedulers later in this section work with it directly.

Now put a price on a round. Let $c$ be the time of one draft step divided by the time of one target step, so a draft step costs $c$ target steps. Let the verification pass cost one target step, which the first part of this section showed is true at small batch.

∴ The speedup of speculative decoding

Under these assumptions a round costs $\gamma c + 1$ target steps and emits $\bar n$ tokens, so over many rounds the time per token falls by the factor

S = \frac{\bar n}{\gamma c + 1} = \frac{1 - \alpha^{\gamma+1}}{(1 - \alpha)(\gamma c + 1)} .

With $\gamma = 1$ , $S = (1 + \alpha)/(1 + c)$ , which exceeds 1 exactly when $\alpha > c$ . So any draft that is accepted more often than its relative cost helps, at the right $\gamma$ .
With a free draft, $c = 0$ , $S = \bar n < 1/(1 - \alpha)$ , the bound approached as $\gamma \to \infty$ .
A parallel drafter, one that proposes all $\gamma$ tokens in a single pass of relative cost $c$ , gives $S = \bar n/(1 + c)$ .
The total arithmetic per emitted token grows by the factor $(1 - \alpha)(\gamma\hat c + \gamma + 1)/(1 - \alpha^{\gamma+1})$ , where $\hat c$ is the draft's FLOPs per token over the target's. The reason is that the target does the arithmetic of $\gamma + 1$ tokens per round and keeps only $\bar n$ of them.

∂ Proof of the speedup formula

Rounds are independent and identically distributed under the model, so by the law of large numbers the long-run tokens per unit time is the expected tokens per round over the expected time per round. That is $\bar n/(\gamma c + 1)$ in units of target steps, against 1 token per target step without speculation. Item 1 substitutes $\gamma = 1$ and $\bar n = 1 + \alpha$ , and $(1 + \alpha)/(1 + c) > 1$ exactly when $\alpha > c$ . Item 2 bounds the geometric sum by its infinite version. Item 3 replaces the draft cost $\gamma c$ by $c$ . For item 4, a round does $\gamma\hat c + \gamma + 1$ target-token units of arithmetic and emits $\bar n$ tokens. These are Theorems 3.8 and 3.11 and Corollary 3.9 of Leviathan et al. (2023). ∎

In plain words, speculation trades arithmetic for fewer trips through memory. Every rejected draft is arithmetic thrown away. That costs nothing while the GPU is waiting on memory anyway, and real time once it is not. The best $\gamma$ is found by trying integers. It is finite whenever $c > 0$ , because $\bar n$ never exceeds $1/(1 - \alpha)$ while the cost grows linearly in $\gamma$ .

Python · cell 3

✏️ Choosing the draft length

Take $\alpha = 0.8$ and a draft that costs $c = 0.05$ of a target step. With $\gamma = 4$ a round emits $\bar n = 3.36$ tokens and costs $4\times0.05 + 1 = 1.2$ target steps, a speedup of $3.36/1.2 = 2.80$ . The best draft length is $\gamma = 8$ , at 3.09 times, against a ceiling of $1/(1 - 0.8) = 5$ for a free draft.

A draft four times as expensive, $c = 0.2$ , is best at $\gamma = 4$ , with 1.87 times.
A worse draft, $\alpha = 0.6$ at $c = 0.05$ , also peaks at $\gamma = 4$ , with 1.92 times, against a ceiling of 2.5.
The arithmetic bill is real. At $\alpha = 0.8$ with $\hat c = 0.05$ , $\gamma = 4$ does 1.55 times the FLOPs per token of plain decoding, and $\gamma = 8$ does 2.17 times, for a 10% gain in speed.

Core

Expected tokens, speedup, and the best draft length

Implement three functions.

expected_tokens(alpha, gamma) returns $\bar n$ , the expected number of tokens a round emits.
speedup(alpha, gamma, c, parallel=False, v=0.0) returns $\bar n$ divided by the cost of a round in target steps.
best_gamma(alpha, c, max_gamma, parallel=False, v=0.0) returns the gamma in 1..max_gamma with the largest speedup, and the smallest such gamma on ties.

alpha is either one number, the i.i.d. acceptance probability, or an array with at least gamma entries (max_gamma for best_gamma). In the array, entry i is the probability that draft i + 1 is accepted, given that the drafts before it were.

A round costs three things, in target steps. The draft costs gamma * c, or just c for a parallel drafter that proposes the whole block in one pass. The verification pass costs 1. And each verified draft adds v, the extra target cost per token once verification is no longer free, so gamma * v in all.

import numpy as np

def expected_tokens(alpha, gamma):
    return gamma * alpha                      # placeholder: forgets the bonus token and the compounding

def speedup(alpha, gamma, c, parallel=False, v=0.0):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

The closed form assumed i.i.d. acceptance. For the bigram models of the simulation we can compute the survival probabilities exactly and see how far off that assumption is. Let $\mathbf{A}$ be the $6\times 6$ acceptance matrix with entries $A_{sx} = \min(p(x\mid s), q(x\mid s))$ , the probability of drafting $x$ after $s$ and accepting it. Every sequence of the simulation starts from token 0, so let $\mathbf{e}_0$ be the one-hot vector of token 0. Entry $x$ of the row vector $\mathbf{e}_0^\top \mathbf{A}^i$ is the probability of accepting the first $i$ drafts and ending at token $x$ . So the survival probability $a_i$ is the sum of that row.

Python · cell 4

Measured and exact agree, 2.932 against 2.938. The i.i.d. formula, with $\alpha$ set to the first position's acceptance, gives 2.911. It is off by about 1%, because acceptance here depends on which token was accepted before. On real models the gap is larger, since acceptance decays along the draft. So check published speedups against measured survival curves, not a single $\alpha$ .

✏️ Published numbers through the formula

T5 translation. Leviathan et al. (2023) sped up T5-XXL (11B) on English to German translation with T5-small (77M) as the draft, at batch 1 on a TPU v4. Under argmax sampling they measured $\alpha = 0.75$ and, with $\gamma = 7$ , a 3.4 times speedup. The formula gives $\bar n = (1 - 0.75^8)/0.25 = 3.60$ tokens per round. With the draft cost they estimated from profiles, $c = 0.02$ , the predicted speedup is $3.60/(1 + 7\times0.02) = 3.16$ , and they tabulate 3.2 against the measured 3.4. At temperature 1, $\alpha = 0.62$ predicts 2.26, against 2.6 measured. Even a bigram table as the draft gave $\alpha = 0.2$ . A table lookup costs nothing next to the target, so $c = 0$ , and the speedup approaches the ceiling $1/(1 - \alpha) = 1.25$ , the 1.25 times they report.

Chinchilla. Chen et al. (2023) drafted for Chinchilla 70B with a 4B model of only 8 layers, shallow so that it runs fast on the same 16 TPU v4 chips. It takes 1.8 ms per token against the target's 14.1 ms, so $c \approx 0.128$ . With $\gamma = 4$ (their $K$ ) they measured 1.92 times on XSum summarization with nucleus sampling, 2.01 times with greedy decoding, and 2.46 times on HumanEval code. Solving $\bar n = S(4c + 1)$ for $\alpha$ gives 0.72, 0.75, and 0.85. This inversion is ours, and it slightly underestimates $\alpha$ , because their verification time also grew a little with $K$ .

DeepSeek-V3. DeepSeek-V3 drafts one token ( $\gamma = 1$ ) with its multi-token prediction module, an extra block trained during pretraining to predict the token after next (§8.1). The report says the drafted token is accepted 85 to 90% of the time and that generation runs 1.8 times as fast. The formula gives $\bar n = 1 + \alpha = 1.85$ to $1.9$ , so the module's cost and the verification overhead together take only a few percent.

Speculation in a busy server

Everything so far assumed that verifying $\gamma + 1$ tokens costs one decode step. That holds while the GPU waits on memory, and a busy server changes that. It decodes $B$ sequences together, so that one read of the weights serves all of them (§20.1). The batch already uses up some of the spare arithmetic.

With speculation, each round runs $\gamma$ draft steps and then one target pass over $(\gamma + 1)B$ tokens, which costs $t_{\text{pass}}(B, \gamma + 1)$ . Plain decoding costs $t_{\text{pass}}(B, 1)$ per token per sequence. Here is the scale. For a dense model with short contexts, plain decoding turns compute-bound near $B = 295$ on an H100. Checking 4 tokens per sequence puts $4B$ tokens through each weight read, so it reaches the same point near $B = 295/4 \approx 74$ . Past that point the extra tokens of a verification pass are no longer free.

∴ When verification is free in a batch

Take the step model without overhead, $t(B, k) = \max(t_W + Bt_K,\; kBt_F)$ , ignore the draft's cost, and let $S(B) = \bar n\,t(B, 1)/t(B, \gamma + 1)$ be the speedup at batch $B$ .

If $(\gamma + 1)Bt_F \le t_W + Bt_K$ , the verification pass is memory-bound and $S(B) = \bar n$ , the full benefit.
If $Bt_F \ge t_W + Bt_K$ , plain decoding is already compute-bound and $S(B) = \bar n/(\gamma + 1) \le 1$ , a loss unless every draft is accepted.
In between, $S(B) = \bar n\,(t_W + Bt_K)/\big((\gamma + 1)Bt_F\big)$ , which falls as $B$ grows.

Case 1 holds at every batch size exactly when $(\gamma + 1)t_F \le t_K$ , that is, when the mean context is at least

\bar T_c(\gamma + 1) = \frac{2N_{\text{act}}(\gamma + 1)}{I^\star m_{\text{KV}} - 4Ld_{\text{attn}}(\gamma + 1)},

provided the denominator is positive. Otherwise case 1 holds only for $B \le t_W/\big((\gamma + 1)t_F - t_K\big)$ .

∂ Proof of the batch regimes

Cases 1 to 3 substitute the branches of the maximum into $S(B)$ . In case 1 both passes take $t_W + Bt_K$ . In case 2 both take their compute branch, $Bt_F$ and $(\gamma + 1)Bt_F$ . In case 3 the baseline takes $t_W + Bt_K$ and verification takes $(\gamma + 1)Bt_F$ , and the ratio $(t_W + Bt_K)/B$ falls as $B$ grows.

For the last part, $(\gamma + 1)Bt_F \le t_W + Bt_K$ for all $B$ if $(\gamma + 1)t_F \le t_K$ , since $t_W > 0$ . Otherwise it fails for large $B$ , first at $B = t_W/((\gamma + 1)t_F - t_K)$ . Multiplying $(\gamma + 1)t_F \le t_K$ by $n_{\text{gpu}}\pi$ gives $(\gamma + 1)(2N_{\text{act}} + 4Ld_{\text{attn}}\bar T) \le \bar Tm_{\text{KV}}\pi/\beta$ . Collecting the terms in $\bar T$ , with $I^\star = \pi/\beta$ , gives the condition. ∎

In plain words, long contexts keep speculation alive. Each sequence's cache is read once per pass, however many of its tokens are checked. When contexts are long, the pass is dominated by bytes that verification does not multiply, and the extra arithmetic stays hidden at any batch. §20.1 called the same threshold the critical context, the context beyond which decode is memory-bound at every batch. Here it is computed with $\gamma + 1$ queries per sequence instead of one. Sadhukhan et al. (2024) made this argument, and their MagicDec system showed that speculation pays at batch 32 to 256 once sequences are long enough.

The cell puts numbers on it for Llama 3.1 8B on one H100. The drafter is one Transformer layer of 0.5B parameters, about the size of the drafters that read the target's features, described later. It costs about 6% of a target step at batch 1. The cell also checks the same quantities for latent attention and for a mixture of experts.

Python · cell 5

Three regimes show up.

At 64 tokens of context the speedup is flat at 2.49 up to batch 64 and falls below 1 by batch 256. At batch 256 a plain step is still memory-bound, but the verification pass over four tokens per sequence is compute-bound, so each round pays for the arithmetic of its rejected drafts. By batch 1,024 plain decoding is compute-bound too, and the speedup sits near $\bar n/(\gamma + 1) = 0.74$ , less the drafter's cost.
At 512 tokens the cache delays the fall, which the theorem places at $B = 104$ .
At 4,096 tokens the context is past $\bar T_c(4) = 1{,}755$ , so the verification pass stays memory-bound at every batch. The speedup even rises slightly with $B$ , because the drafter's cost shrinks relative to the target's as the target's caches grow.

Only 119 sequences of 4,096 tokens fit in one H100's 80 GB, so the right end of that curve needs more GPUs. The ratios do not depend on how many GPUs share the work. What this changes in practice is the default. Speculation is not only a low-load trick, and whether it pays at high load depends on the context length, the attention design, and whether the model is a mixture of experts.

🔬 Latent attention and mixtures of experts move the boundary

Latent attention (§12.5) caches one small latent vector per token and lets every head attend to it. In DeepSeek-V3 the latent has 576 numbers, 1,152 bytes in bfloat16. Each of the 128 query heads scores all 576 and then averages the first 512, about $4\times128\times544$ FLOPs per cached token. So one query does about 242 FLOPs per cached byte, close to the ridge point of 295. With two or more queries per cached latent, its attention alone is compute-bound. The ratio printed above, 1.22, is the largest number of queries per cached token for which DeepSeek-V3's attention stays memory-bound on an H100. So the denominator of $\bar T_c$ is negative for every $\gamma \ge 1$ , and verification is never free at large batch.

Mixtures of experts are the opposite case (§13.6). Each token is routed to $k$ of $E$ expert networks, the feed-forward blocks of a layer, and a decode step reads only the experts its tokens touch. With $n$ tokens routed uniformly at random, the expected number of experts touched is $E\big(1 - (1 - k/E)^n\big)$ . With $E = 256$ and $k = 8$ , checking four tokens per sequence at batch 1 touches 30.5 experts per layer instead of 8, almost four times the expert bytes. At batch 64 it touches nearly all 256 instead of 222. For an MoE, verification is expensive at small batch and nearly free at large batch, the reverse of a dense model.

The figure redraws the cell's curve for any context length and draft length, starting at 64 tokens of context. Slide the context up and watch the drop move to larger batches, until past 1,755 tokens verification is free at every batch. Hover over the curve to read the speedup at any batch.

✏️ Speculation at batch 64 in a production engine

The SGLang team measured a drafter called EAGLE-3 in their serving engine, for Llama 3.1 8B on one H100 on the MT-Bench chat benchmark, and the EAGLE-3 paper reports the results. The drafter is a one-layer model that reads the target's hidden features (described below), and here it proposed chains of 3 tokens.

Throughput relative to no speculation was 1.81 times at batch 2, 1.48 times at 16, and 1.38 times at 64.
The original EAGLE drafter, with lower acceptance, fell below 1 at batch 24, to 0.93 times.
At batch 1, SGLang went from 158 tokens per second without speculation to 244 with EAGLE-2 and 373 with EAGLE-3.

Sadhukhan et al. (2024) reported up to 2.51 times for Llama 3.1 8B at batch sizes from 32 to 256 on moderate to long sequences. Their draft is the target itself reading only a small sparse window of its cache, so the draft's cost does not grow with the context.

Choosing how much to verify at every step

A fixed $\gamma$ is right for one load and wrong for every other. When the GPU is memory-bound, a long, risky draft is free. When it is compute-bound, each verified token displaces real work for other requests. So in 2026 engines began choosing, at every step, how many drafted tokens of each request to verify. This is called adaptive verification.

It needs an estimate of how likely each draft is to survive. Cheng et al. (2026) give their drafter, DSpark, a small confidence head, an extra output that predicts, for each drafted position, the probability that the token is accepted given that the positions before it were. The running product of these confidences along request $r$ is that request's survival curve, $a_{r,1} \ge a_{r,2} \ge \dots$ .

Now count the gain and the cost. Suppose the step verifies $\ell_r$ drafts of request $r$ . Then it emits $\sum_r(1 + \sum_{i\le\ell_r}a_{r,i})$ tokens in expectation, and it costs a target pass over $\sum_r(1 + \ell_r)$ tokens. Each survival curve decreases, so the best way to spend $b$ draft slots is on the $b$ largest survivals over all requests. Those automatically form a prefix of each request's draft. The engine tries every $b$ and keeps the one with the most expected tokens per second, using a table of step costs it measured at startup (vLLM Team 2026). The cell runs that rule against our step model for a parallel drafter of 7 tokens, at three loads, with 512 tokens of context. It leaves out the drafter's own cost. That cost is the same for every $b$ , so it does not change the choice, but it makes the gain over no speculation look slightly larger than it is.

Python · cell 6

With 8 requests the pass is memory-bound, and the rule verifies all seven drafts. With 64 it already trims the unlikely tails, verifying 4.75 drafts per request on average. With 256 it verifies 1.35. At that load, always verifying all seven is slower than not speculating at all, 22,980 against 25,800 tokens per second. The adaptive rule reaches 51,209.

Real deployments behave the same way. vLLM ran adaptive verification for DeepSeek-V4-Pro on eight B300 GPUs, drafting 7 tokens per pass, from 1 to 256 concurrent requests (vLLM Team 2026). At every load it sat on the Pareto frontier of throughput against interactivity, the tokens per second each user sees. That means no fixed draft length was better on both at once. It acted like a long fixed draft at low load and a short one at high load. In that deployment the last of the 7 drafted tokens survived less than 10% of the time, and the first more than 70%.

In DeepSeek's own production serving, the earlier one-token draft verified 2 tokens per request, the draft and the bonus position. The scheduler expanded that to roughly 4 to 6 tokens per request under the moderate loads that dominate their traffic (Cheng et al. 2026).

There is a catch, and it is the non-anticipating requirement from the whole-round proof. DSpark's confidence for position $k + 1$ reads the token drafted at position $k$ . Suppose a search looks at all the survivals at once and then decides how far to verify. Then whether position $k$ is verified can depend on $x_k$ , through the confidences after it, and the drafted tokens influence their own verification.

The paper avoids this in two ways. Its scheduler admits the draft slots one at a time, in order of survival, and stops at the first slot that lowers the expected throughput. So the decision about a slot never looks at the slots after it. The production versions choose the number of slots from older confidences, two steps old in DeepSeek's serving and one step old in vLLM, which cannot depend on the current drafts. Then they fill those slots in order of the current survivals. Our cell is safe for a simpler reason, since its survivals are random numbers that do not depend on any draft. An interviewer who asks about adaptive verification wants to hear this requirement.

⚠️ Trimming a draft can break exactness

Deciding how much of a draft to verify is safe only if the decision about position $i$ uses information available before $x_i$ was drafted. Here is a counterexample built on the three-token example, $p = (0.5, 0.3, 0.2)$ and $q = (0.2, 0.6, 0.2)$ . Draft $x \sim q$ and verify it only when $x \ne b$ . When $x = b$ , skip verification and let the target sample the token from $p$ itself.

Each piece looks harmless, but the output becomes $(0.5, 0.18, 0.32)$ instead of $p$ . Follow the probability. Drafts of $a$ and $c$ are always accepted, so $c$ already gets its full 0.2 from them. The correct rule would turn the 0.6 of drafted $b$ 's into 0.3 of $b$ and 0.3 of $a$ . The fresh sample from $p$ instead turns it into 0.3 of $a$ , 0.18 of $b$ , and 0.12 of $c$ . So $c$ gets 0.12 more than it should, all taken from $b$ . The cause is that the event "verified" now depends on which token was drafted.

A confidence computed from the draft's state before $x_i$ is fine. A search that peeks at $x_i$ , or at confidences that depend on $x_i$ , and then decides whether to keep $x_i$ , is not. The DSpark paper's appendix gives the same kind of counterexample.

Python · cell 7

Verifying always and verifying never are both exact. Verifying depending on the drafted token is not, and the cell prints the same $(0.5, 0.18, 0.32)$ as the hand count.

Token trees check several guesses at once

A chain of drafts dies at its first wrong guess. When the draft is unsure between two continuations, it can propose both, and the target can check both in the same pass. Say the prefix is "The cat" and the draft cannot decide between "sat" and "is". It proposes both, and after "sat" it adds "on". The target scores the three guesses in one pass. If it would have written "is", the round still accepts one draft, where a chain of "sat on" would have accepted none.

A token tree holds several candidate continuations that share prefixes, like this one. Verifying it costs one pass over all its nodes, so a tree of $n$ nodes is free in the same sense as a chain of $n$ tokens.

The pass needs two changes from an ordinary chunked pass. Each node must attend only to the prefix and to the nodes on its own branch, never to a rival guess. In the example, "on" may see "sat" but not "is". And each node must get the position id it would have in its own branch, the index fed to the position encoding (§5.4). Both "sat" and "is" are the third token of their sentence, so both get the same position id.

📐 Token tree and tree attention

A token tree is a list of $n$ drafted tokens with a parent array, $\text{par}(i) \in \{-1, 0, \dots, i - 1\}$ , where $-1$ means that node $i$ follows the last committed token directly. Nodes are listed parents before children. The depth of a node is the number of its ancestors, so children of the prefix have depth 0. Verifying the tree runs the target once on the $n$ nodes appended to a prefix of $P$ cached tokens, with

the mask $\mathbf{M} \in \{0, 1\}^{n\times(P + n)}$ , where node $i$ may attend to every prefix position and to node $j$ exactly when $j = i$ or $j$ is an ancestor of $i$ , and
position ids $P + \operatorname{depth}(i)$ , so every node sits at the position it would have in its own branch.

Siblings never see each other. A chain, $\text{par}(i) = i - 1$ , gives the bottom-right causal mask of §4.4.

Here is a small tree. Its parents are $(-1, -1, 0, 0, 1, 2)$ . Nodes 0 and 1 are two guesses for the first token, nodes 2 and 3 continue the first guess, node 4 continues the second, and node 5 continues node 2. The node-to-node block of the mask and the depths are

\mathbf{M}_{:,\,P:} = \begin{pmatrix} 1&0&0&0&0&0\\ 0&1&0&0&0&0\\ 1&0&1&0&0&0\\ 1&0&0&1&0&0\\ 0&1&0&0&1&0\\ 1&0&1&0&0&1 \end{pmatrix}, \qquad \operatorname{depth} = (0, 0, 1, 1, 1, 2),

and the prefix block $\mathbf{M}_{:,\,:P}$ is all ones. One pass returns the target's distribution after every node, conditioned on exactly that node's branch.

The figure starts from this tree, with "The cat" as the prefix. Click a node to see its row of the mask. Pick what the target would write to see which branch a greedy target keeps, by the walk described after the exercise.

Interview

The tree attention mask

Implement tree_attention_mask(parents, prefix_len). The tree is given as a parent array, with -1 for nodes that follow the prefix directly, and parents are listed before their children. Return two arrays.

The boolean mask, of shape (n_nodes, prefix_len + n_nodes). In row i, node i sees every prefix position, itself, and its ancestors, and nothing else.
The int64 position ids, prefix_len + depth.

Tree-based speculative decoders build exactly this mask before the verification pass.

import numpy as np

def tree_attention_mask(parents, prefix_len):
    n = len(parents)
    mask = np.tril(np.ones((n, prefix_len + n), dtype=bool), k=prefix_len)   # placeholder: a plain causal chain
    pos = prefix_len + np.arange(n, dtype=np.int64)
    return mask, pos

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Verifying a tree under greedy decoding is a walk. Start at the prefix and move to the child that equals the target's choice there. Repeat from that child. Stop at the first node with no such child, and emit the target's choice there as the extra token.

Under sampling, a node has several children to try, and the one-token rule extends to them. Try the children one at a time. Each rejection removes from the target what that child's draft distribution covered, and the next child is tested against what is left.

∴ Recursive rejection at a node with several children

Let a node's $k$ children $s_1, \dots, s_k$ be drawn independently from $q$ , and let $R_1 = p$ . For $i = 1, \dots, k$ , accept $s_i$ with probability $\min(1, R_i(s_i)/q(s_i))$ and stop. If it is rejected, set $R_{i+1}(x) = \max(0, R_i(x) - q(x))/\sum_{x'}\max(0, R_i(x') - q(x'))$ and try the next child. If all $k$ are rejected, sample from $R_{k+1}$ . The emitted token is distributed exactly as $p$ , and some child is accepted with probability

1 - \prod_{i=1}^{k}\Big(1 - \sum_x\min\big(R_i(x), q(x)\big)\Big).

∂ Proof of recursive rejection

Show by backward induction on $i$ that the procedure started at try $i$ with target $R_i$ outputs a token distributed as $R_i$ . Started at try $k + 1$ , it samples from $R_{k+1}$ directly. Started at try $i \le k$ , it draws a fresh $s_i \sim q$ and accepts it with probability $\min(1, R_i(s_i)/q(s_i))$ , and otherwise continues from try $i + 1$ . By the induction hypothesis that continuation outputs a token distributed as $R_{i+1}$ , the residual of $R_i$ and $q$ . That is the single-token rule with target $R_i$ and residual $R_{i+1}$ , so the output is distributed as $R_i$ . Taking $i = 1$ gives $p$ .

Try $i$ is accepted with probability $\sum_x\min(R_i, q)$ given that it is reached, and the tries are independent given the targets, which gives the product. (If some $R_i = q$ , try $i$ is accepted with certainty and $R_{i+1}$ is never needed.) ∎

This is the multi-step speculative sampling of Miao et al. (2024), whose SpecInfer system introduced token trees with this mask. SpecInfer reported 1.5 to 2.8 times over existing serving systems for distributed inference. After a child is accepted, verification moves to that child's own children, drafted from $q$ conditioned on the branch. The argument for a chain applies level by level, so the whole tree is exact.

Chen et al. (2024), in their Sequoia method, draw a node's children without replacement instead. They zero each rejected token in $q$ and renormalize before the next try, so the same wrong token is never proposed twice. That matters at low temperature, where an i.i.d. draft keeps repeating its favorite.

The cell checks the theorem at one node over five tokens. It runs 100,000 trials for 1, 2, 4, and 8 children, and compares the output frequencies with $p$ and the chance of accepting some child with the formula.

Python · cell 8

The output stays at $p = (0.40, 0.25, 0.20, 0.10, 0.05)$ for every $k$ . The chance that some child survives climbs from 0.55 with one child to 0.86 with eight. It climbs slowly, because i.i.d. children keep re-proposing the draft's favorite token, which the target likes less than the draft does.

How the tree is shaped is where the methods differ.

A fixed product tree. Sometimes the guesses for different positions are made independently, as with the extra prediction heads described in the next part. Then the natural tree is the Cartesian product of the top few guesses at each position, fixed in advance.
A tree grown each round. Li et al. (2024) observed that their drafter's confidence scores approximate acceptance rates well. So their EAGLE-2 grows a different tree in every round, expanding the nodes whose product of confidences along the branch is largest. It gave 3.05 to 4.26 times over plain decoding, 20 to 40% faster than the same drafter with a fixed tree.
A tree optimized offline. Sequoia chooses the tree by dynamic programming to maximize the expected accepted length for a given number of nodes, and picks that number from the hardware.

After a tree is verified, the cache holds entries for every node, but only the accepted branch is kept. Its entries are scattered among discarded siblings, so they must be gathered into consecutive slots before the next round.

Where the guesses come from

Every method in this section uses the same verifier. They differ only in where the guesses come from. The guesser sets both the acceptance rate $\alpha$ and the relative cost $c$ , so it sets the speedup. Exactness always comes from the verifier, never from the drafter. Here are the families, from the oldest to the newest.

A smaller model of the same family. This is the original setting of Leviathan et al. (2023) and Chen et al. (2023). The draft must share the target's tokenizer, so that a drafted token id means the same string to both. Leviathan et al. found that drafts about two orders of magnitude smaller than the target worked best, with $\alpha$ between 0.5 and 0.9.

Since $\alpha = 1 - \operatorname{TV}(p, q)$ , the natural way to improve a draft is to distill the target into it, training the draft to match the target's outputs (§20.4). Zhou et al. (2024) found that two choices matter most. One is training on the draft's own generations, which is called on-policy data. The other is choosing the divergence to match the decoding strategy. Their method, DistillSpec, reported 10 to 45% speedups over speculative decoding with an undistilled draft.

Extra heads on the target. Cai et al. (2024) add $K$ small output heads on top of the target. Each reads the target's last hidden state $\mathbf{h}_t \in \R^d$ , the vector that the output layer turns into next-token logits. In their method, Medusa, head $k$ predicts the token at position $t + k + 1$ as

\operatorname{softmax}\big((\mathbf{h}_t + \operatorname{SiLU}(\mathbf{h}_t\mathbf{W}_1^{(k)}))\mathbf{W}_2^{(k)}\big),

written in this book's input-by-output layout. $\mathbf{W}_1^{(k)} \in \R^{d\times d}$ starts at zero, and $\mathbf{W}_2^{(k)} \in \R^{d\times V}$ starts as a copy of the target's output head. So each head begins as the ordinary next-token predictor. Each head guesses its position from $\mathbf{h}_t$ alone, so the guesses for different positions are independent of each other. That is why Medusa verifies a Cartesian product tree.

Medusa comes in two versions. Medusa-1 trains only the heads on a frozen target, and it is exact with the rejection rule, at over 2.2 times. Medusa-2 also fine-tunes the target, at 2.3 to 2.8 times, which changes the model being served. Medusa also proposed typical acceptance. It keeps a candidate when the target gives it probability above a threshold that depends on the entropy of the target's distribution. It is faster at high temperature, and it is not exact.

DeepSeek-V3's multi-token prediction module (§8.1) is the same idea built into pretraining. Unlike a Medusa head, it also reads the token just emitted. So the module is a one-token draft that ships with the model.

Drafting from the target's features. The EAGLE drafter of Li et al. (2024) is one Transformer layer. It reads the target's features, the hidden states that enter the target's output head, together with the token sequence shifted one step ahead. It predicts the next feature and turns it into a token with the target's own output head. Features are more regular than tokens. Feeding in the sampled token tells the drafter which of several possible continuations the feature should follow. For a 70B target the drafter has under 1B trainable parameters and trains in one to two days on four A100 GPUs. It gave 2.7 to 3.5 times on LLaMA2-Chat 70B.

EAGLE-3 (Li et al. 2025) dropped the loss on predicting the next feature. It predicts tokens directly, from a fusion of low, middle, and high layers of the target. It is trained with training-time test, which means the drafter is trained on inputs that include its own earlier predictions, exactly as it will see them at inference. With those changes, the speedup kept improving as the training data grew. EAGLE-3 reached 3.0 to 6.5 times over plain decoding, 20 to 40% faster than EAGLE-2.

The target drafting for itself. Elhoushi et al. (2024) train with layer dropout and an early-exit loss, so that the first layers alone make usable predictions. Their method, LayerSkip, then drafts with the early layers and verifies with the rest, reusing the early layers' computation. It reported up to 2.16 times on CNN/DM summarization.

MagicDec (Sadhukhan et al. 2024) keeps the whole model as the draft but lets it read only a sparse window of its cache, the recent tokens plus a few initial ones. That is in the spirit of sliding-window attention with sinks. It is what long-context, large-batch serving needs, since the draft's cost no longer grows with the context.

No draft model at all. When the output copies the input, as in summarizing, editing code, or answering from retrieved documents, the context itself is a draft. Prompt lookup decoding (Saxena 2023) and LLMA (Yang et al. 2023) find an earlier occurrence of the last few tokens and propose what followed it. LLMA reported over 2 times speedups with outputs identical to greedy decoding.

Lookahead decoding (Fu et al. 2024) builds its own pool of n-grams with Jacobi iteration. That method guesses all future positions at once, runs the model on them in parallel, replaces each guess with the model's prediction, and repeats. Lookahead does this in the same pass as verification and is exact. It reported up to 1.8 times on MT-Bench and 4 times with several GPUs on code completion.

Core

Prompt-lookup drafting

Implement ngram_draft(context, n, k), the drafter of prompt lookup decoding. context is a list of tokens, strings or ints.

Take the last n tokens of context, the suffix. Find its most recent earlier occurrence. That occurrence must start before the suffix itself does, so that the suffix never matches itself.
Propose the tokens that followed that occurrence, at most k of them.
If the suffix never occurred earlier, try the last n - 1 tokens, then n - 2, down to 1. Return an empty list if nothing matches.

For example, take a b c X a b c Y a b c with n = 3 and k = 2. The suffix a b c last occurred earlier starting at the second a. It was followed by Y a, so that is the proposal. The setup's tok splits code-like text into tokens for the tests.

def ngram_draft(context, n, k):
    return []

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Drafting a whole block in one pass. An autoregressive drafter, one that guesses a token at a time, pays $\gamma c$ per round, which forces it to be shallow. A drafter that proposes the whole block in one pass pays $c$ once, so it can afford to be deeper. Chen, Liang & Liu (2026) built such a drafter, DFlash, from a small block diffusion model, one that fills all the masked positions of a block in a single forward pass. It is conditioned on hidden features of the target, injected into the keys and values of every draft layer. Their five-layer drafter proposing 16 tokens took less time to draft, and had more drafts accepted, than EAGLE-3's one-layer drafter proposing 8. They reported over 6 times lossless acceleration and up to 2.5 times the speedup of EAGLE-3.

Parallel drafting has one weakness. Positions predicted in parallel do not see the tokens sampled at each other, so acceptance decays quickly along the block. Cheng et al. (2026) call this suffix decay. Their drafter, DSpark, keeps the parallel backbone and adds a light sequential step. To position $k$ 's logits it adds a bias $\mathbf{W}_1[x_{k-1}]\,\mathbf{W}_2$ . Here $\mathbf{W}_1 \in \R^{V\times 256}$ is looked up at the token just sampled at position $k - 1$ , and $\mathbf{W}_2 \in \R^{256\times V}$ . Together they form a rank-256 table of which token tends to follow which. DSpark also adds the confidence head that drives adaptive verification.

Its reported results are these.

On Qwen3 targets from 4B to 14B, DSpark's average accepted length was 26.7 to 30.9% longer than EAGLE-3's and 16.3 to 18.4% longer than DFlash's.
In DeepSeek-V4's production serving it raised per-user generation speed over the previous one-token multi-token prediction draft, at matched throughput. The gain was 60 to 85% for V4-Flash and 57 to 78% for V4-Pro.
A DSpark drafter for Kimi K3 (Zhao et al. 2026), with five layers and 5B parameters proposing 8 tokens, has an average acceptance length of 4.11 tokens per verification round across nine domains, and 6.42 on math.

Drafter	Guesses from	Extra weights	Training	Exact with the rule	Reported speedup
Small model of the same family	its own forward passes	a separate model	pretraining, ideally distillation	yes	1.9 to 3.4 times (T5-XXL, Chinchilla)
Medusa heads	the target's last hidden state	one layer per head	heads only (Medusa-1) or joint (Medusa-2)	Medusa-1 yes, typical acceptance no	2.2 to 2.8 times
Multi-token prediction (MTP) module	the target's last state and the token just emitted	one block	during pretraining	yes	1.8 times at $\gamma = 1$ (DeepSeek-V3)
EAGLE 1 to 3	the target's features and the shifted tokens	one layer, under 1B	one to two days on four A100s (EAGLE, 70B)	yes	2.7 to 6.5 times
Early exit (LayerSkip)	the target's first layers	none	layer dropout and early-exit loss	yes	up to 2.16 times
Prompt lookup, LLMA, lookahead	the context or a Jacobi n-gram pool	none	none	yes	up to 1.8 times on chat (lookahead), over 2 times on copy-heavy tasks (LLMA)
Parallel block (DFlash, DSpark)	a one-pass block drafter on target features	a few layers, up to 5B (Kimi K3)	on target hidden states, at scale	yes	over 6 times (DFlash), 60 to 85% over MTP in production (DSpark)

Every "yes" holds only with the rejection rule of this section, and up to floating-point differences between kernels. Every speedup is a batch-1 or low-load number on the paper's own hardware and tasks, so compare them with care. An interviewer who asks which drafter to use wants to hear that the answer depends on the batch, the context length, and whether you can train a drafter on the target's features.

Keeping the caches consistent after a rejection

The verification pass computed keys and values for every drafted position, so after a round the target's cache holds entries for all of them. After $J$ acceptances, the entries of the accepted drafts are correct, because they were computed on the right prefix. The entries of the $\gamma - J$ rejected positions must be discarded. The correction or bonus token has no entry yet, and the next round computes it as the first token of its verification pass.

How rollback works depends on the cache layout.

With a contiguous cache, rollback moves a length pointer back.
With the paged cache of §20.6, which stores each sequence in fixed-size blocks listed in a block table, it trims the table and returns blocks that became empty to the free list. The scheduler must also reserve blocks for $\gamma + 1$ tokens per sequence before the pass.
The draft's cache rolls back the same way. At the start of the next round the draft must also process the tokens it has not yet read, the correction or bonus token, and after a fully accepted round its own last draft too.
Trees need one more step, gathering the accepted branch's entries, which are scattered among discarded siblings, into consecutive slots.

Recurrent layers are harder. A state space or linear-attention layer folds every token into one fixed-size state, so a rejected token cannot be removed by truncation. The engine must save the state after each verified position or recompute from the last accepted one. That is a real cost for the hybrid models that mix attention with recurrent layers, which are becoming common.

⚠️ Four ways to lose exactness without noticing

Each of these keeps the speed and silently changes the output distribution.

Resampling from $p$ instead of the residual after a rejection biases the output toward the tokens the draft over-proposes, as the simulation showed.
Accepting a draft whenever the target ranks it among its top few, or Medusa's typical acceptance, is a different sampler, and a lossy one.
Using a $q$ in the ratio that is not the distribution the draft actually sampled from, for example the draft's logits before its own top- $p$ truncation, makes $p/q$ wrong.
Verifying a drafted position or not depending on the drafted token itself breaks the non-anticipating rule of the adaptive schedulers.

Test a speculative sampler by comparing output distributions on many samples, never by comparing speed or a handful of strings.

Speculative verification in PyTorch

Engines verify the whole batch at once, with no Python loop over sequences. The block below does it with gather, a cumulative product for the accepted prefix, and one multinomial call for the extra token. A row of zeros appended to $q$ makes the bonus case fall out of the residual formula, since $\max(0, p_{\gamma+1} - 0) = p_{\gamma+1}$ . The check draws drafts from $q$ for 200,000 sequences and compares the first two emitted tokens with the target rows.

🔥 In PyTorchfor a machine with PyTorch

Both marginals match the target rows to within 0.001. The mean accepted length matches $\sum_i\prod_{j\le i}\alpha_j$ , the survival formula with the three positions' own acceptance rates.

In practice you call an engine. Hugging Face generate takes a draft model as assistant_model, or drafts from the prompt with prompt_lookup_num_tokens. vLLM takes a speculative_config. Here the draft shares the target's tokenizer, so a drafted token id means the same string to both.

🧭 What to remember

Decode is memory-bound, so a target pass over $k$ tokens of one sequence costs about one decode step while $kBt_F \le t_W + Bt_K$ , up to a few hundred tokens at batch 1 in the roofline model. Speculation turns that free arithmetic into extra tokens per read of the weights.
The rule is to draft $x \sim q$ , accept with probability $\min(1, p(x)/q(x))$ , and on rejection sample from $\max(0, p - q)$ , normalized. The output is exactly $p$ , the acceptance probability is $\sum_x\min(p, q) = 1 - \operatorname{TV}(p, q)$ , and no exact rule can accept a single draft more often.
A round of $\gamma$ drafts emits $\bar n = 1 + \sum_i a_i$ tokens, $(1 - \alpha^{\gamma+1})/(1 - \alpha)$ under i.i.d. acceptance, and runs $\bar n/(\gamma c + 1)$ times faster when verification is free. With $\alpha = 0.8$ and $c = 0.05$ , $\gamma = 8$ gives 3.09 times, at about twice the arithmetic per token.
In a batch, verification is free only while the pass stays memory-bound. With short contexts at large batch the speedup falls to $\bar n/(\gamma + 1) < 1$ , and with contexts beyond $\bar T_c(\gamma + 1)$ it keeps its full benefit. Latent attention and mixtures of experts shift the boundary. Adaptive verification picks per step how many drafts to check from survival estimates and a cost table, and must never let the decision depend on the drafted token.
Token trees verify several candidates in one pass with an ancestor mask and depth-based positions, and recursive rejection keeps them exact. Separate draft models, Medusa heads, EAGLE drafters, MTP modules, early exit, n-gram and lookahead drafting, and parallel block drafters such as DFlash and DSpark differ only in where the guesses come from. Rejected tokens must be rolled back in both models' caches.
Next, §20.6 turns from more tokens per sequence to more sequences per step, with continuous batching, the paged cache that speculation's rollback relies on, and prefix sharing.

Leviathan, Kalman & Matias (2023). Fast Inference from Transformers via Speculative Decoding. The rejection rule and its proof, the acceptance rate as $1 - D_{LK}$ , the expected tokens, walltime, and arithmetic formulas, and the T5 and LaMDA measurements. Read it with Chen et al. (2023), Accelerating Large Language Model Decoding with Speculative Sampling, the same algorithm found independently at DeepMind, with the three reasons scoring a short continuation costs one step, and a 2 to 2.5 times speedup for Chinchilla 70B.
Stern, Shazeer & Uszkoreit (2018). Blockwise Parallel Decoding for Deep Autoregressive Models. The greedy draft-and-verify idea with extra prediction heads, years before the sampling version.
Cai et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. Drafting heads on the target, Cartesian token trees with tree attention, and typical acceptance.
Li et al. (2025). EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test. Multi-layer feature fusion and training-time test, with batched throughput measured in SGLang. Read it after EAGLE and EAGLE-2.
Miao et al. (2024). SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification. Token trees, the tree attention mask, and the multi-step speculative sampling that keeps trees exact. Sun et al. (2023), SpecTr: Fast Speculative Decoding via Optimal Transport, treat several drafts per position as an optimal transport problem.
Sadhukhan et al. (2024). MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding. Why speculation pays again at large batch once the cache dominates, and drafting with a sparse cache.
Chen, Liang & Liu (2026). DFlash: Block Diffusion for Flash Speculative Decoding. A block diffusion drafter that proposes a whole block in one pass from the target's features.
Cheng et al. (2026). DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation. Parallel block drafting with a sequential head, confidence-scheduled verification, and production results in DeepSeek-V4. The engine side is in vLLM's Adaptive Verification in vLLM.

Made for larger screens

Speculative decoding

Checking several tokens costs about as much as writing one

Greedy speculation keeps the target's own choices

Greedy speculative generation, with call counting

When the target samples, keep a guess with probability p over q

Acceptance rate and temperature

A whole round of speculative sampling

Verify one round of speculative sampling

How many tokens a round produces, and the speedup

Expected tokens, speedup, and the best draft length

Speculation in a busy server

Choosing how much to verify at every step

Token trees check several guesses at once

The tree attention mask

Where the guesses come from

Prompt-lookup drafting

Keeping the caches consistent after a rejection

Speculative verification in PyTorch

Read more

Watch next