Speculative decoding and the correctness proof

Here is the number that starts everything, taken from §0.4. Llama-3-8B in bf16 streams 7.50 B parameters per token — 15.01 GB — because the input embedding table is gathered rather than streamed. At the H100 SXM's 3.35 TB/s that is a hard floor of 4.48 ms per decode step at batch 1, or 223 tokens per second, and roughly 98% of that traffic is weights.

One token through 7.50 B matmul parameters requires about 15 GFLOPs. Dividing by peak tensor-core throughput gives an arithmetic floor, not an observed duty cycle. Low tensor-core utilization does not mean the GPU is idle: memory hardware may be serving useful traffic throughout the step.

Write the arithmetic intensity of a whole decode step as a function of $T$, the number of token positions in the step:

Idealized weight-only model. BF16 weights cost two bytes and each multiply-add counts as two FLOPs, so intensity is approximately the number of jointly processed positions. The resulting ridge near 295 positions assumes peak rates and perfect weight reuse. It suggests spare arithmetic capacity, but positions 2-295 are not free: attention, activations, logits, scheduling, kernel efficiency and repeated reads change actual verify latency.

The catch is causality. You cannot compute token $t+1$ until you have token $t$, so a plain decode loop can only ever put one position per sequence into the step. Speculative decoding breaks the dependency by guessing the next $\kappa$ tokens from something cheap, then checking all $\kappa+1$ of them in a single target pass. The guessing is easy. The hard part — the part this chapter is about — is checking them in a way that provably leaves the output distribution untouched.

One target step produces a probability distribution at every position you feed it, not just the last one. That is a property of causal attention, not a feature anyone added: run the model over [x₁ … x_t, x̃₁, x̃₂, x̃₃] and position $t$ tells you the true distribution for $x̃_1$, position $t{+}1$ tells you the true distribution for $x̃_2$ given $x̃_1$, and so on. If the guesses happen to be right, you have just validated four tokens for the price of one weight stream. If they are wrong, you find out at the first wrong one and throw away everything after it.

So the loop becomes: draft $\kappa$ tokens cheaply, verify all $\kappa+1$ positions in one target pass, keep a random-length prefix, repeat. Figure 1 puts real numbers on both timelines.

The equal-width verify block in this diagram is a hypothetical weight-streaming model, not a measured schedule. Verification can reuse weights across positions but does not move exactly the same total bytes as a one-token pass. Its 6.70 ms step and 2.53-token yield are illustrative model outputs; benchmark draft, verify and sampler costs separately.

Fix a prefix. Let $V$ be the vocabulary, $p(\cdot)$ the target model's next-token distribution, and $q(\cdot)$ the draft's. Both are probability distributions on $V$; nothing else is assumed about $q$. Let $\kappa$ be the number of drafted tokens. vLLM calls it num_speculative_tokens. SGLang's equivalent is speculative_num_steps — not speculative_num_draft_tokens, which is the verify-window length: for chain drafting SGLang overwrites it with speculative_num_steps + 1 whenever speculative_eagle_topk == 1 (python/sglang/srt/arg_groups/speculative_hook.py:L674-L683). That off-by-one is what makes SGLang's KV reservation in §6.2.9 come out at $\kappa+1$ rather than $\kappa$.

def spec_step(prefix, target, draft, kappa):
    # 1. DRAFT. kappa sequential cheap forward passes.
    xs, qs = [], []
    for i in range(kappa):
        q_i = draft(prefix + xs)             # distribution over V
        xs.append(sample(q_i)); qs.append(q_i)

    # 2. VERIFY. ONE target pass over all kappa+1 positions.
    #    ps[i] is p(. | prefix + xs[:i]) -- the true conditional,
    #    because causal attention gives a distribution at every position.
    ps = target(prefix + xs)                 # kappa+1 distributions

    # 3. ACCEPT / REJECT, left to right, stop at the first rejection.
    out = []
    for i in range(kappa):
        x = xs[i]
        u = uniform(0, 1)
        if u < min(1, ps[i][x] / qs[i][x]):
            out.append(x)                    # accept
        else:
            residual = relu(ps[i] - qs[i])   # elementwise max(0, .)
            out.append(sample(residual / residual.sum()))
            return out                       # DISCARD ps[i+1:] -- see below
    # 4. BONUS. All kappa accepted, so ps[kappa] is conditioned on a
    #    prefix that actually happened. Free extra token.
    out.append(sample(ps[kappa]))
    return out

Three things in that listing are load-bearing and are the three things implementations get wrong. First, the acceptance test uses the ratio $p_i(x)/q_i(x)$ at the drafted token only — not a threshold on $p$, not a top-$k$ membership test. Second, on rejection you must resample from the residual $\max(0, p-q)$, not from $p$. Third, on rejection you must discard every later position, even though the target already computed distributions for them, because those distributions were conditioned on a prefix that did not happen.

Steps 3 and 4 are the whole of rejection_sample in vLLM and of speculative_sampling_classic_kernel in SGLang. Everything else in both codebases is plumbing.

The claim is not "approximately $p$", not "$p$ in expectation over many steps", not "$p$ up to a bounded divergence". It is that the single token emitted at each position is drawn from exactly $p$, for any $q$ whatsoever. Here is the argument, with every symbol defined.

One position

Fix a position and drop the index. Let $X$ be the token this procedure emits. Two events partition the sample space: accept (the drafted token survives the coin flip) and reject. Compute the probability of emitting a specific token $x \in V$ under each.

Accept branch. To emit $x$ by acceptance, the draft must have proposed $x$ — probability $q(x)$ — and the coin must have landed under the acceptance ratio — probability $\min\!\big(1, p(x)/q(x)\big)$. These are independent given $x$, so

The rewriting is the first small miracle: multiplying a probability by a clamped ratio is just a minimum. If $q(x) = 0$ the token is never proposed and the left side is $0$, which is also $\min(0, p(x))$ — so the identity holds on all of $V$ with no case split.

Summing over $x$ gives the acceptance probability:

where $D_{\mathrm{TV}}(p,q) = \tfrac{1}{2}\sum_x |p(x)-q(x)|$ is total variation distance. The one-token acceptance rate is one minus total variation distance. This is a useful overlap metric, not literally every draft's training objective: cross entropy, feature losses and sequence objectives are different surrogates. See the draft architecture discussion.

Reject branch. On rejection we sample from the normalised residual

Now the second and larger miracle. Use the elementary identity $\max(0, a - b) = a - \min(a, b)$, valid for all reals:

The residual mass equals the rejection probability exactly. This is not a tuning choice; it is forced, and it is forced by the single fact that $p$ and $q$ both sum to one. It is why the algorithm has no free parameter and no correction factor.

Adding the branches

Every line is an identity. There is no inequality, no limit, no assumption on $q$ beyond being a distribution on $V$. A uniform draft, an n-gram table, a draft model trained on a different corpus — all of them emit exactly $p$. A bad draft costs throughput and nothing else. Speculative decoding is not a quality/speed trade-off; it is a speed/compute trade-off wearing a probabilistic disguise.

Chaining the positions

The per-position result composes by induction on the position index. Suppose positions $1..i-1$ were all accepted. Then the emitted prefix is exactly $\tilde{x}_1..\tilde{x}_{i-1}$, which is exactly the prefix the target conditioned on when it produced $p_i$ in the single verify pass. So $p_i$ is the correct conditional for the sequence that actually happened, and the one-position argument applies verbatim: position $i$ emits from $p_i$.

At the first rejection, position j emits from the normalized positive residual, not from p_j conditional on rejection. Combining accepted and rejected outcomes recovers p_j. But $p_{j+1}, \dots, p_{\kappa}$ were computed conditioned on $\tilde{x}_j$, and $\tilde{x}_j$ was just rejected. Those distributions are conditionals of a sequence that does not exist. This is why the tail must be discarded, and it is the single most common way a speculative implementation silently biases its output. If all $\kappa$ are accepted, the bonus distribution was conditioned on an accepted prefix. Its logits are already available, although drawing and committing the bonus still has overhead. Thus the maximum emitted length is $\kappa+1$, not $\kappa$; the expectation equals that maximum only with perfect acceptance and no earlier stop condition.

Set temperature to zero. Then $p = \delta_{a}$ where $a = \arg\max_x p(x)$, and with greedy drafting $q = \delta_{d}$ where $d$ is the draft's argmax. The acceptance ratio collapses:

• If $d = a$: $p(d)/q(d) = 1/1 = 1$, accept with probability 1.
• If $d \ne a$: $p(d) = 0$, accept with probability 0.
• Residual on rejection: $\max(0, \delta_a - \delta_d) = \delta_a$, normalised to $\delta_a$ — emit the target's argmax.

So the entire probabilistic machinery degenerates to "accept while the draft matches the target's argmax; on mismatch, emit the target's argmax". vLLM implements exactly that as a separate kernel rather than deriving it at runtime:

    rejected = False
    for pos in range(num_draft_tokens):
        if not rejected:
            draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
            target_argmax_id = tl.load(target_argmax_ptr + start_idx + pos).to(tl.int32)
            # ...
            else:
                token_id = target_argmax_id
                rejected = draft_token_id != target_argmax_id
            tl.store(
                output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos,
                token_id,
            )

Note the elegance of token_id = target_argmax_id being assigned unconditionally: on acceptance the argmax equals the draft token, so one store covers both branches.

Is greedy speculative decoding bit-identical to greedy plain decoding? Only under two conditions that production violates routinely. First, argmax ties: torch.argmax returns the lowest index among equal maxima, and a draft that picked a different tied token is rejected — the emitted token is still the target's argmax, so the distribution is unchanged, but a downstream diff against a non-speculative run needs the same tie-break. Second, and much worse: the logits the target produces at position $i$ of a $(\kappa{+}1)$-token verify pass are not bitwise the logits the same model produces for the same prefix in a 1-token decode step, because the GEMM shape differs, the split-K reduction order differs, and floating-point addition is not associative. Near a near-tie that flips the argmax and the two runs diverge forever. This is batch invariance, and it is §10.4's whole subject. SGLang carries the tensor-parallel version of the same hazard in a comment:

        # Sync sampling results across TP ranks: different GPUs may
        # produce slightly different target_probs due to floating-point
        # non-determinism in softmax/top_k/top_p, causing different
        # sampled tokens. Broadcast from rank 0 to ensure consistency.
        tp_group = (
            get_parallel().attn_tp_group
            if is_dp_attention_enabled()
            else get_tp_group()
        )
        if tp_group.world_size > 1:
            tp_group.broadcast(predict, src=0)
            tp_group.broadcast(accept_index, src=0)
            tp_group.broadcast(num_correct_drafts, src=0)

Without that broadcast, two ranks accept different numbers of tokens, their sequence lengths diverge, and the next collective deadlocks. The ROCm branch above it (L745-L758) does the same for the greedy path because "the per-rank draft tokens can differ".

Expected tokens per step

Model per-token acceptance as i.i.d. Bernoulli($\alpha$). This is an approximation — real acceptance decays with position because errors compound — and vLLM's synthetic-acceptance harness makes the position dependence explicit by taking per-position unconditional rates $r_i = \mathbb{P}(\text{first } i \text{ all accepted})$ and converting them to conditional rates:

def unconditional_to_conditional_rates(rates: list[float]) -> list[float]:
    """Convert per-position unconditional rates to per-position conditional
    rates for the early-terminating rejection loop (c_i = p_i / p_{i-1})."""
    return [p / q if q > 0.0 else 0.0 for p, q in zip(rates, [1.0, *rates[:-1]])]

The i.i.d. model is the case $r_i = \alpha^i$, under which all conditional rates equal $\alpha$. Let $n \in \{0,\dots,\kappa\}$ be the number of accepted draft tokens. The loop stops at the first rejection, so $\mathbb{P}(n \ge i) = \alpha^{i}$, and exactly one extra token is always emitted — the recovered token on rejection, or the bonus token on full acceptance. Therefore

Write $E$ for that quantity. Note $E \le \kappa+1$ always, with equality only at $\alpha = 1$. You can never emit more than $\kappa+1$ tokens from one verify pass.

Cost per step, and the naive speedup

Let $t_T$ be one target step and $t_D$ one draft step, and $c = t_D/t_T$. Chain drafting is $\kappa$ sequential draft passes, so the step costs $t_T(1 + \kappa c)$ and

Derive $c$ rather than guessing it. In the memory-bound regime time is bytes over bandwidth, so $c$ is the ratio of streamed weight bytes. Llama-3.2-1B ($L=16$, $d=2048$, $h=32$, $h_{kv}=8$, $d_h=64$, FFN 8192, tied embeddings) streams $16 \times 60.8\text{M} = 973$ M transformer parameters plus a $128{,}256 \times 2048 = 263$ M output head — 1.236 B parameters, 2.47 GB, and $2.47/3.35 = 0.74$ ms. Against Llama-3-8B's 4.48 ms that is $c = 0.165$. Derived.

Two readings. First, $\kappa$ has an interior optimum, and it moves right as $\alpha$ rises: the marginal token from lengthening the draft is worth $\alpha^{\kappa+1}$ and decays geometrically, while the marginal cost $c$ is constant. Doubling $\kappa$ past the optimum is a pure loss. Second, at $\alpha = 0.5$ with this draft the whole thing is barely worth the complexity, and at $\kappa=7$ it is a net loss. Solving $E = 1 + \kappa c$ for $\kappa = 3$, $c = 0.165$ gives a break-even acceptance rate of $\alpha^{*} \approx 0.34$: below that, a 1B draft for an 8B target makes the server slower. Derived.

When speculation loses

The naive formula assumes the verify pass is free, which is true only while the step is memory-bound. Reintroduce the roofline. Let $T^{*} = 295$ be the ridge in token positions (§0.4, re-derived at the top of this chapter) and $B$ the batch size. Normalise all times by the bandwidth floor $W/\beta$. A step containing $T$ positions costs $\max(1, T/T^{*})$. Plain decode puts $B$ positions in the step; verify puts $B(\kappa{+}1)$; the draft model, being a scaled copy, costs $c\max(1, B/T^{*})$ per pass. Then

Three regimes fall out. For $B \le T^{*}/(\kappa{+}1)$ both maxima are 1 and $S = E/(1+\kappa c)$ — the textbook formula, valid only here. For $B \ge T^{*}$ everything is compute-bound and the maxima cancel:

Under this cost model, E is at most kappa+1 and positive draft cost makes the compute-bound limit less than one. With zero draft cost and perfect acceptance, equality is possible. This is not a hardware-independent theorem about every compute-bound implementation: batching efficiency and kernel choices can change the relative costs. In the model's middle regime, setting S(B)=1 gives the crossing:

for $\kappa=3$, $c=0.165$, $\alpha=0.70$. Derived. Beyond about 150 concurrent sequences on one H100, this configuration is making Llama-3-8B slower, and by $B = 295$ it has given away 44% of the machine's throughput.

Two caveats: treating old KV-read traffic as proportional to batch times context and independent of draft width assumes ideal sharing across verification queries. New KV writes, activation/logit traffic and repeated reads remain. Long context can shift the crossover; no revised crossing follows without a full traffic/latency model. The draft also has its own roofline and synchronization costs, so a constant cost ratio is an approximation. In measured workloads, speculation often is a latency optimisation that borrows from throughput, and the loan comes due at high batch. Both engines expose the knob to repay it — vLLM's num_speculative_tokens_per_batch_size is an explicit batch-size schedule (vllm/config/speculative.py:L180-L186) resolved into a lookup table each step at vllm/v1/core/sched/scheduler.py:L1265-L1271.

vLLM: two Triton kernels and a ragged prefix sum

vLLM flattens the whole batch's draft tokens into one 1-D tensor and carries a prefix sum to find each request's slice. That is SpecDecodeMetadata in full:

@dataclass
class SpecDecodeMetadata:
    # [num_tokens]
    draft_token_ids: torch.Tensor
    # [batch_size]
    num_draft_tokens: list[int]
    # [batch_size]
    cu_num_draft_tokens: torch.Tensor
    # [batch_size]
    cu_num_sampled_tokens: torch.Tensor
    # [num_tokens]
    target_logits_indices: torch.Tensor
    # [batch_size]
    bonus_logits_indices: torch.Tensor
    # [num_tokens + batch_size]
    logits_indices: torch.Tensor

    def __post_init__(self):
        self.max_spec_len = max(self.num_draft_tokens)

Requests in the same batch may carry different draft lengths — including zero, for a request that just finished prefill. The split of logits_indices into target_logits_indices (positions that get the accept test) and bonus_logits_indices (the last position of each request) is what lets one forward pass serve a mixed batch. The accept test itself:

        if not rejected:
            draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
            uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
            if draft_token_id < 0:
                # -1 is used for padded draft token ids that should be rejected.
                accepted = False
            # ...
            else:
                if NO_DRAFT_PROBS:
                    draft_prob = 1
                else:
                    draft_prob = tl.load(
                        draft_probs_ptr
                        + (start_idx + pos) * vocab_size
                        + draft_token_id
                    )
                target_prob = tl.load(
                    target_probs_ptr + (start_idx + pos) * vocab_size + draft_token_id
                )
                # NOTE(woosuk): While the draft probability should never be 0,
                # we check it to avoid NaNs. If it happens to be 0, we reject.
                accepted = draft_prob > 0 and target_prob / draft_prob >= uniform_prob
            if accepted:
                token_id = draft_token_id
            else:
                rejected = True
                token_id = tl.load(recovered_token_ids_ptr + start_idx + pos)

target_prob / draft_prob >= uniform_prob implements a non-strict comparison. It agrees with the strict acceptance rule almost surely for ideal continuous RNG, but finite RNG can hit endpoints; the endpoint section below explains that distinction. NO_DRAFT_PROBS sets $q = 1$, the correct value when the draft is deterministic — which is vLLM's default, draft_sample_method = "greedy" (vllm/config/speculative.py:L289-L296) — and the only option for an n-gram proposer that has no distribution at all. A one-hot $q$ makes $\min(1, p/q) = p(x)$: accept the draft with its own target probability. The proof holds unchanged.

The residual sampler avoids a normalisation pass entirely by racing exponentials — the Gumbel-max trick in its $\mathrm{Exp}(1)$ form. It draws one $[\,\text{batch},\,\text{vocab}\,]$ tensor of exponentials per request, not per position, and takes an argmax of $\text{prob} \times 1/q$:

            draft_prob = tl.load(
                draft_probs_ptr + token_idx * vocab_size + vocab_offset,
                mask=vocab_mask,
                other=0.0,
            )
            target_prob = tl.load(
                target_probs_ptr + token_idx * vocab_size + vocab_offset,
                mask=vocab_mask,
                other=0.0,
            )
            prob = tl.maximum(target_prob - draft_prob, 0.0)
            # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because
            # `tl.argmax` will select the maximum value.

That is $\max(0, p - q)$ from the proof, verbatim, and the comment is the observation that argmax is scale-invariant so $Z$ never needs computing. Sharing one exponential row across all $\kappa$ positions of a request is safe because at most one of them is ever consumed — the first rejection ends the loop. In the NO_DRAFT_PROBS branch the kernel instead masks the drafted token out of the target row (vocab_offset != draft_token_id, L913-L918), which is exactly the normalised residual against a one-hot $q$.

SGLang: a Triton chain sampler behind a flag, and a CUDA tree sampler by default

Read the dispatch before the kernel. In its non-greedy branch — an all-greedy batch takes verify_tree_greedy instead — eagle_sample picks between exactly two implementations, and the choice is one boolean:

        use_rejection_sampling = get_spec().speculative_use_rejection_sampling

        sampling_fn = (
            chain_speculative_sampling_triton
            if use_rejection_sampling
            else tree_speculative_sampling_target_only
        )

speculative_use_rejection_sampling defaults to False (python/sglang/srt/server_args.py:L2192-L2196), so the Triton chain sampler quoted next is not what a default SGLang server runs — it is opt-in via --speculative-use-rejection-sampling, and only for topk 1. Read it anyway: it is the clearest in-tree statement of this chapter's algorithm, and its default-path counterpart is the CUDA kernel dissected after it. The verification loop:

    while (step < NUM_SLOTS) and (continue_verifying == 1):
        draft_token = tl.load(cand_ptr_base + step * stride_cand_s)
        # ...
        p = tl.load(TargetProbs + offset_prob)
        q = tl.load(DraftProbs + offset_draft)

        coin = tl.load(uni_ptr_base + (step - 1) * stride_uni_s)

        if coin * q < p:
            num_accept += 1
            cur_prob_row = step
            tl.store(Predicts + last_accepted_global_idx, draft_token)
            # ...
            step += 1
        else:
            continue_verifying = 0

    tl.store(AcceptTokenNum + pid, num_accept)

coin * q < p is the same test with the division moved to the other side — cheaper, and it removes the divide-by-zero that forces vLLM's draft_prob > 0 guard. Note also that cur_prob_row tracks the position whose distribution the residual will be built from: it advances on acceptance and freezes at the first rejection. That variable is the "discard the tail" rule.

The residual sampler is an explicit two-pass inverse CDF rather than a Gumbel race:

    all_drafts_accepted = continue_verifying
    coin_final = tl.load(UniformSamplesFinal + pid)
    norm_sum = 0.0
    # ...
    # Pass 1: Sum
    for v_start in range(0, VOCAB_SIZE, BLOCK_V):
        # ...
        if all_drafts_accepted:
            val = p_val
        else:
            q_ptr = dp_base_ptr_safe + v_offsets * stride_dp_v
            q_val = tl.load(q_ptr, mask=mask, other=0.0)
            # Treat NaN q (degenerate draft rows) as 0: residual falls back to p.
            q_val = tl.where(q_val == q_val, q_val, 0.0)
            diff = p_val - q_val
            val = tl.where(diff > 0.0, diff, 0.0)

        norm_sum += tl.sum(val)

The all_drafts_accepted branch is the bonus token: sample plain $p$ from the $(\kappa{+}1)$-th row. The other branch is $\max(0, p-q)$ with $Z$ accumulated explicitly, then a second pass walks the CDF to coin_final * norm_sum. The engineering difference is real: vLLM materialises a $[\,B,\,|V|\,]$ fp32 exponential tensor — 32.8 MB at $B=64$ with a 128,256 vocabulary — and reads the vocab once; SGLang allocates one scalar per request and reads the vocab twice. At 3.35 TB/s both land around 20 µs. Pick your poison.

vLLM's own worked example, left in the source as a comment block, is the clearest description of the ragged layout in either codebase. Five requests with 3, 0, 2, 0 and 1 draft tokens:

        # Inputs:
        # cu_num_scheduled_tokens:  [  4, 104, 107, 207, 209]
        # num_draft_tokens:         [  3,   0,   2,   0,   1]
        # Outputs:
        # cu_num_draft_tokens:      [  3,   3,   5,   5,   6]
        # logits_indices:           [  0,   1,   2,   3, 103, 104, 105, 106,
        #                            206, 207, 208]
        # target_logits_indices:    [  0,   1,   2,   5,   6,   9]
        # bonus_logits_indices:     [  3,   4,   7,   8,  10]

Request 0 contributed 4 positions (1 real + 3 drafts) to a flat 209-token batch; request 1 is mid-prefill with 100 scheduled tokens and no drafts, so it gets one logit index, 103, and it lands in bonus_logits_indices only. Eleven logit rows total: six get the accept test, five are bonus positions. The draft token ids are recovered by shifting the same index vector by one, draft_token_ids[target_logits_indices + 1] (L2985-L2986) — the token at position $i+1$ is what the model was asked to predict at position $i$.

The step ends where it began, in the scheduler, undoing its own optimism. Every scheduled token was already counted as computed at _update_after_schedule (L1398-L1404, "If some tokens (e.g. spec tokens) are rejected later, the number of computed tokens will be adjusted in update_from_output"), so rejection is a subtraction.

Here is the part that is absent from both papers. The target's verify pass runs attention over all $\kappa+1$ positions, which means it writes KV entries for all of them, including the ones that are about to be rejected. Those entries are wrong: they encode a continuation that did not happen. Leave them in place and the next step's attention reads them.

vLLM does not free anything. It rolls back the logical length and lets the next step overwrite the physical slots:

                num_draft_tokens = len(scheduled_spec_token_ids)
                num_sampled = self.num_sampled_tokens_per_step
                num_accepted = max(len(generated_token_ids) - num_sampled, 0)
                num_rejected = num_draft_tokens - num_accepted
                # Rejections roll back num_computed_tokens (and, under async
                # scheduling, num_output_placeholders, which covers the spec
                # tokens). A stale rejection count predates the preemption
                # rollback and must not apply.
                if not output_is_stale:
                    if request.num_computed_tokens > 0:
                        request.num_computed_tokens -= num_rejected
                    if request.num_output_placeholders > 0:
                        request.num_output_placeholders -= num_rejected

Cheap — a scalar subtraction — but it leaks into three other places. The prefix cache must never commit a block containing a draft token, because a rejected token that has been hashed into the cache will be served to a future request:

        # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens
        # + num_external_computed_tokens + num_new_tokens, but must exclude
        # "non-committable" tokens (e.g., draft tokens that could be rejected).
        # Therefore, we cap the number at `request.num_tokens`, ensuring only
        # "finalized" tokens are cached.
        num_tokens_to_cache = min(
            total_computed_tokens + num_new_tokens,
            request.num_tokens,
        )

Sliding-window block eviction has the mirror-image problem — free a block based on the optimistic length and a rollback will want it back — so the free boundary is computed from processed rather than scheduled tokens, "because ... rejected spec tokens can roll it back" (vllm/v1/core/kv_cache_manager.py:L504-L510). And the scheduler must reserve slots for drafting on every decode request every step, whether or not that request ends up speculating: draft_slots = spec.max_num_new_slots_for_drafting is subtracted from the input budget at vllm/v1/core/sched/scheduler.py:L505-L506 and again at L709.

SGLang pays the allocation cost more visibly. Every decode step reserves, per request:

    # Spec decoding allocates max(topk * num_steps, num_draft_tokens) per decode step.
    spec_steps = spec.speculative_num_steps or 1
    spec_topk = spec.speculative_eagle_topk or 1
    spec_tokens = max_speculative_num_draft_tokens()
    page_size = get_alloc_page_size()
    # ...
    if page_size == 1 or spec_topk == 1 or not spec_algo.has_draft_kv():
        return max(spec_steps * spec_topk, spec_tokens)
    # ...

def get_alloc_reserve_per_decode() -> int:
    """KV length reserved per request at each decode step.

    The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap
    mode; see eagle_utils.eagle_prepare_for_decode.
    """
    return 2 * get_alloc_len_per_decode()

At $\kappa = 3$, topk 1, that is $2 \times 4 = 8$ KV slots reserved per request per step where plain decode needs 1. For Llama-3-8B at 128 KiB of KV per token, a batch of 64 reserves 64 MB per step instead of 8 MB. It is returned when the step commits, but the pool must be sized for the peak, so speculation shrinks your maximum concurrency before it ever changes your latency. Derived from the cited per-token KV size in §0.4.

Then there is the case chain decoding does not have. With chain drafting the accepted tokens are a prefix of the allocated run, so committing is just advancing seq_lens. With a tree draft (topk > 1, covered in §6.4) the accepted path is scattered across tree-node slots, and SGLang physically compacts it:

def move_accept_tokens_to_target_kvcache(
    batch: ScheduleBatch,
    accept_index: torch.Tensor,
    num_correct_drafts: torch.Tensor,
    token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
):
    """
    Move accepted tokens (drafts + bonus) to the target KV cache.

    Args:
        batch: The batch to run.
        accept_index: The index of the accepted tokens (incl. bonus).
        num_correct_drafts: Per-req count of correct drafts (excludes bonus);
            seq_lens is advanced by ``num_correct_drafts + 1`` to cover the bonus slot.
    """

That ends in move_kv_cache(tgt_loc, src_loc) (python/sglang/srt/mem_cache/memory_pool.py:L2801-L2819), a per-layer gather-scatter over the K and V buffers. Each moved token costs 128 KiB read plus 128 KiB written for Llama-3-8B; at batch 64 with 2.5 accepted tokens that is 40 MB of traffic, about 12 µs — small against a 4.5 ms step, but a whole extra kernel launch, and the reason _finalize_accept_tree_path exists as a separate code path from the chain case (python/sglang/srt/speculative/eagle_worker_common.py:L406-L434).

The coin endpoints

Both engines have a bug-shaped hole at a different end of $[0,1]$, and both have patched it in the source with an explanatory comment. vLLM's accept test is >=, so a coin of exactly 0.0 accepts unconditionally — including a token the target assigns probability zero:

    # NOTE(woosuk): We deliberately use float64 instead of float32 here
    # because when using float32, there's a non-negligible chance that
    # uniform_prob is sampled to be exact 0.0 as reported in
    # https://github.com/pytorch/pytorch/issues/16706. Using float64
    # mitigates the issue.
    uniform_probs = torch.rand(
        (num_tokens,),
        dtype=torch.float64,
        device=device,
    )

SGLang's seeded path has the opposite failure, at 1.0:

    # The float32 cast rounds the top 129 uint32 hashes to exactly 1.0, but
    # the sampling kernels expect half-open [0, 1) coins: a 1.0 coin walks
    # past the last CDF bucket and can return a zero-probability token.
    # Clamp to the largest float32 below one; every other coin value is
    # untouched, so previously verified bitwise baselines stay intact.
    max_coin = 1.0 - 2**-24
    coins = (
        uniforms[:, :draft_token_num].to(torch.float32).clamp_(max=max_coin)
    ).contiguous()

Same class of bug, opposite endpoint, same symptom: a token with zero probability emitted at a rate of roughly one in ten million. Rare enough to survive a test suite, frequent enough to produce a bug report at scale.

Knobs that quietly void the proof

The guarantee holds only for the exact rule. SGLang exposes two multipliers that break it, both defaulting to the identity:

    speculative_accept_threshold_single: A[
        float,
        "Accept a draft token if its probability in the target model is greater than this threshold.",
        NS("spec"),
    ] = 1.0
    speculative_accept_threshold_acc: A[
        float,
        "The accept probability of a draft token is raised from its target probability p to min(1, p / threshold_acc).",
        NS("spec"),
    ] = 1.0
    speculative_use_rejection_sampling: A[
        bool,
        "Use rejection sampling for speculative decoding (requires topk=1).",
        NS("spec"),
    ] = False

Lowering threshold_acc below 1.0 multiplies every acceptance probability by $1/\text{threshold}$, raising $\alpha$ and biasing the output toward the draft. It buys throughput by spending the thing this chapter proved you did not have to spend. Use it knowing that, or not at all.

Note also that speculative_use_rejection_sampling is off by default and requires topk 1. With it off, eagle_sample passes draft_probs = torch.zeros_like(target_probs) into tree_speculative_sampling_target_only (python/sglang/srt/speculative/eagle_utils.py:L810-L816) — a target-only rule equivalent to treating $q$ as deterministic, the same simplification as vLLM's NO_DRAFT_PROBS.

That default kernel is readable at this SHA, despite shipping through the compiled sgl_kernel wheel: the CUDA lives at python/sglang/kernels/aot/csrc/speculative/speculative_sampling.cuh. Its accept test is a sibling walk, and it is worth seeing why that is still the same theorem:

      DType target_prob_single = target_probs[cur_prob_offset + draft_token_id];
      prob_acc += target_prob_single;

      if (coin <= prob_acc / threshold_acc || target_prob_single >= threshold_single) {
        // accept token
        prob_acc = 0.;
        cur_prob_offset = (bx * num_draft_tokens + cur_index) * d;
        coin = uniform_samples[bx * num_draft_tokens + cur_index];
        predicts[last_accepted_retrive_idx] = draft_token_id;
        ++num_accepted_tokens;
        accept_index[bx * num_speculative_tokens + num_accepted_tokens] = draft_index;
        last_accepted_retrive_idx = draft_index;
        break;
      } else {
        // FIXME: leverage draft probs
        draft_probs[cur_prob_offset + draft_token_id] = target_probs[cur_prob_offset + draft_token_id];
        cur_index = retrive_next_sibling[bx * num_draft_tokens + cur_index];
      }

With both thresholds at their default 1.0 the test is coin <= prob_acc, where prob_acc is the target mass of the siblings tried so far at this node. For a chain — one sibling — that is accept-with-probability $p(x)$, which is exactly $\min(1, p/q)$ for the one-hot $q = \delta_x$ the caller passed in. For a tree it is recursive rejection sampling: each rejected sibling has its target mass copied into draft_probs, so the residual max(target − draft, 0) computed afterwards excludes precisely the branches already refused. The final resample is that residual, and it appears in the kernel under confusingly swapped names — q_vec is loaded from target_probs and p_vec from draft_probs, so relu_q_minus_p at L119 is $\max(0, p - q)$ in this chapter's notation. The default tree path implements the exact rule.

Stale documentation, and a real cost

RejectionSampler's class docstring states that "spec decode does not support" top-$p$/top-$k$ (vllm/v1/sample/rejection_sampler.py:L38-L59), while apply_sampling_constraints applies both to the target logits at draft positions (vllm/v1/sample/rejection_sampler.py:L510-L564). As of a556f3f the code is ahead of the comment. Read the code. And note what the code costs:

    # NOTE(woosuk): `apply_top_k_top_p` uses sorting to calculate the mask,
    # which is slow for large vocab sizes. This may cause performance issues.
    return apply_top_k_top_p(logits, top_k, top_p)

Speculation multiplies the number of sampling rows by kappa+1. A full-sort fallback therefore sorts more rows, but the quoted comment does not prove the active runner uses that fallback. Inspect pivot/Triton/FlashInfer dispatch and profile the actual sampler (§6.1).

Ragged batches

Both engines launch one program per request and let it loop serially over its own draft tokens. vLLM launches both kernels over the full batch and each early-exits on the wrong rows — if not is_greedy: return in the greedy kernel (L731-L733) and if is_greedy: return in the random one (L791-L794) — so a mixed greedy/random batch pays two full launches with half the blocks idle in each. Grammar-invalidated drafts are padded to -1 rather than shortening the row (vllm/v1/core/sched/scheduler.py:L2313-L2317), which is why the kernel's first branch is if draft_token_id < 0: accepted = False. The uniform-random generator skips requests with zero drafts entirely, "important for reproducibility" (vllm/v1/sample/rejection_sampler.py:L650-L653) — a seeded run must not consume RNG for a request that did not speculate.

For each token, accepted mass is min(p,q); rejected mass is redistributed proportionally to positive p-q. Their sum must recover p. No division by q is needed where q=0 because such tokens are never proposed. Production also needs EOS/stop truncation, accepted-prefix KV commit, rejected-slot reclamation and identical transforms of the distributions used for proposal and verification. Matching seeds is not this distributional proof.

import numpy as np

for p, q in [
    ([0.2, 0.5, 0.3], [0.6, 0.1, 0.3]),
    ([0.2, 0.5, 0.3], [1.0, 0.0, 0.0]),
    ([0.2, 0.5, 0.3], [0.2, 0.5, 0.3]),
]:
    p, q = np.array(p), np.array(q)
    accepted = np.minimum(p, q)
    residual = np.maximum(p - q, 0)
    reject_probability = 1 - accepted.sum()
    if reject_probability > 1e-12:
        replacement = residual / residual.sum()
        output = accepted + reject_probability * replacement
    else:
        output = accepted
    np.testing.assert_allclose(output, p, atol=1e-12)
    tv = np.abs(p - q).sum() / 2
    assert np.isclose(accepted.sum(), 1 - tv)
print("Exact one-step output mass equals the target in all three cases.")

The single number to instrument is mean acceptance length, which vLLM logs directly:

        # Conventionally, mean acceptance length includes the bonus token
        mean_acceptance_length = 1 + (num_accepted_tokens / num_drafts)

        pos_matrix = np.array(self.accepted_tokens_per_pos_lists)
        acceptance_rates = np.sum(pos_matrix, axis=0) / num_drafts
        rates_str = ", ".join(f"{p:.3f}" for p in acceptance_rates)

That mean_acceptance_length is exactly the $E$ of this chapter, and acceptance_rates is the vector of unconditional per-position rates $r_i$ whose i.i.d. idealisation is $\alpha^i$. Run the shipped example and read them:

python examples/features/speculative_decoding/spec_decode_offline.py \
    --method ngram --num-spec-tokens 3 --temp 0 --output-len 256 --print-output

Then do the experiment this chapter argues for: hold everything fixed and sweep concurrency until the speedup crosses 1. Plot $r_i$ against $\alpha^i$ for the fitted $\alpha$ and see how badly the i.i.d. assumption holds. That is Lab 08, spec-decode acceptance. Nothing in this chapter was measured; the lab is where the numbers become real.

• Leviathan, Kalman & Matias, Fast Inference from Transformers via Speculative Decoding (arXiv:2211.17192). The original acceptance rule and the $\mathbb{E}[\text{tokens}]$ derivation. vLLM's RejectionSampler docstring names this paper as the specification it follows.
• Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling (arXiv:2302.01318). Same rule, independently derived, with the modified-rejection-sampling proof written out at length and results on Chinchilla 70B.
• Sun et al., block verification — vLLM ships this as rejection_sample_method="block" (vllm/config/speculative.py:L219-L226), which verifies the draft jointly rather than one token at a time and accepts strictly more often at the same exactness guarantee.
• vllm/v1/spec_decode/ at a556f3f lists ten proposer implementations behind the one sampler in this chapter. Where they come from is §6.3; how EAGLE and MTP draft from the target's own hidden states, and how tree verification generalises the chain rule proved here, is §6.4; the rest of SGLang's zoo is §6.6.
• The two formulas this chapter owns live in FORMULAS; the roofline they are argued against is §0.4; the numerical reason "exactly equivalent" needs an asterisk is §10.4.