chaowen.← Notes

AI/ML systems / inference serving deep dive

2026-09-30

Speculative decoding: accelerate LLM generation without changing its answers

Autoregressive generation makes a large model wait on itself: each token needs the previous token's result. Speculative decoding uses a cheaper draft model to propose several tokens, then verifies them in parallel with the target model. A rejection-sampling correction makes the output distribution exactly match ordinary target sampling—not merely look similar.

1. The serial bottleneck

A transformer can process many prompt positions in parallel during prefill, but generation is different. To sample token xt+1, the model needs the hidden state for the prefix ending at xt. The next forward pass therefore depends on the previous pass. For a memory-bandwidth-bound decoder, each step moves model weights through accelerators to emit just one token. Adding more FLOPs does not automatically remove this serial dependency.

Speculative decoding attacks latency rather than changing the model's goal: do several candidate steps with a cheap model, then spend one target-model pass checking a whole candidate block. Transformer causal attention lets the target compute logits at all candidate positions in one pass. This increases parallel work per target invocation; it does not make the target's arithmetic disappear.

Key distinction: the draft model proposes; the target model retains authority. Exactness comes from the verification-and-correction rule, not from the draft being especially accurate.

2. Draft, verify, correct

Let target distribution at position i be pi(x), conditioned on the already accepted prefix. Let draft distribution be qi(x). The draft samples a block y1:k sequentially from q. The target evaluates the prefix plus that block in one causal forward pass, producing all pi.

  1. For each proposal yi, accept with probability min(1, p_i(y_i) / q_i(y_i)).
  2. At the first rejection, sample a replacement token from the positive residual distribution normalize(max(0, p_i - q_i)).
  3. Discard all later draft proposals, because they were conditioned on a rejected token.
  4. If every proposal is accepted, sample one bonus token from the target distribution after the verified block.
prefix ── draft q: y1 → y2 → y3 → y4 └──── target p: logits at y1, y2, y3, y4 (one block pass) accept y1, accept y2, reject y3 emit y1,y2, corrected token from residual(p3 − q3) discard y4

The algorithm can emit multiple tokens per target pass. It may emit only one when the first proposal is rejected. Draft and target vocabularies and tokenization must align for this basic construction; production systems must ensure the distributions refer to the same token events.

3. Why this preserves the target distribution

At one position, a proposed token x is accepted with probability a(x)=min(1,p(x)/q(x)). The probability mass emitted by accepted proposals at x is therefore q(x)a(x)=min(q(x),p(x)). The unfilled target mass is p(x)-min(q(x),p(x)) = max(0,p(x)-q(x)).

Summing the accepted mass and residual mass gives min(q,p)+max(0,p-q)=p at each token. The residual total is the probability of rejection, so normalizing it gives precisely the conditional distribution of the target outcomes not already accounted for by acceptance. Thus the emitted token has distribution p. Apply the same argument conditionally at each prefix; induction gives the same joint sequence distribution as sampling the target autoregressively.

This is exact in probability, not necessarily bit-for-bit deterministic. Random number generators, floating-point kernels, temperature/top-p transformations, and tie handling must be applied consistently to the distributions being compared. If the implementation accepts solely when a token matches target argmax, it is a different algorithm and generally does not preserve sampling behavior.

4. Data flow and performance trade-offs

Execution path

  1. Keep the target KV cache for the accepted prefix. Run the draft autoregressively for a chosen lookahead length, saving its proposal tokens and the necessary draft-side state.
  2. Append proposals to the target input. A single causal target pass produces a logit vector for every proposal position; causal masking ensures each position sees only its valid prefix.
  3. Compute acceptance ratios using the exact transformed distributions. Stop at the earliest reject, draw residual correction, and commit only the accepted prefix plus correction.
  4. Update cache state. Target KV entries for unaccepted suffix tokens must be discarded or truncated; the draft cache must likewise be reconciled with the corrected token.

For a batch of requests, each request can accept a different number of proposals. Continuous-batching engines need ragged verification lengths, per-request cache rollback, and fair scheduling; padding every request to the maximum lookahead wastes compute and memory.

When it wins

Let draft cost for k proposals be D(k), target block verification cost be V(k), and expected accepted output per cycle be E[A] (including correction or bonus token). A rough throughput proxy is E[A] / (D(k)+V(k)), compared with one target token per target step. This omits queueing, memory pressure, and batching, but explains the central trade-off: longer lookahead amortizes target invocation only when acceptance remains high and draft cost stays small.

Acceptance is not simply “draft quality.” It depends on p and q after temperature, top-k/top-p truncation, repetition penalties, and any other decoding transform. A draft proposal assigned tiny q probability but substantial p probability can have high acceptance; conversely a confident draft can be rejected. Measure accepted tokens per target pass, latency percentiles, draft accelerator occupancy, target utilization, and cache bytes—not acceptance percentage alone.

Cache and memory costs

Verifying a block temporarily needs target KV entries for candidate positions. Larger lookahead can raise peak memory, trigger smaller batches, or cause cache eviction. The target pass may also become compute-bound, changing the benefit relative to single-token decode. If the draft shares hardware with the target, it can steal the bandwidth speculation was meant to save. A small, colocated draft is not automatically free.

5. Reproducible toy lab: inspect the exact acceptance law

Goal: verify by enumeration that accepted mass plus residual correction equals the target distribution. Prerequisite: Python 3.8+ on a local machine. This tiny script is not a neural model benchmark; it isolates the probability rule.

p = [0.50, 0.30, 0.20]
q = [0.20, 0.50, 0.30]

accepted = [min(pi, qi) for pi, qi in zip(p, q)]
residual = [max(0.0, pi - qi) for pi, qi in zip(p, q)]
reject_probability = sum(residual)
corrected = [x / reject_probability for x in residual]

print("target:        ", p)
print("accepted mass: ", accepted)
print("reject mass:   ", residual, "total=", reject_probability)
print("corrected law: ", corrected)
print("reconstructed: ", [a + r for a, r in zip(accepted, residual)])

Expected: accepted mass [0.2, 0.3, 0.2]; residual [0.3, 0, 0]; rejection probability 0.3; conditional correction [1, 0, 0]; reconstructed distribution exactly [0.5, 0.3, 0.2]. In a full sampler, the residual is sampled only after a rejected proposal; this output prints the conditional law for inspection.

Variation 1: set q = p. Every proposal is accepted with probability one, residual mass is zero, and the bonus token still comes from p. Do not divide by zero: correction is never called because rejection cannot occur.

Variation 2: try a draft assigning zero probability to a target-supported token. Residual correction still supplies that token's mass. In code, avoid evaluating p/q for q=0; such a token cannot have been sampled from q, so it cannot be the rejected proposal. Robust implementations handle zero probabilities explicitly.

6. Production failures, debugging, and security

For a production comparison, replay identical prompt and sampling settings with speculation on and off. Compare empirical token frequencies on a controlled small vocabulary, exact generated-token log probabilities under the target, latency distributions, cache occupancy, and rejection positions. A quality test that only compares fluent-looking outputs will not detect a subtle distribution bug. A speed test without the target-only baseline may hide a draft that consumes more time than it saves.

Speculative decoding is an optimization with a proof obligation. The target distribution is preserved only if proposal probabilities, target probabilities, rejection correction, and cache state all describe the same prefix and tokenization. The most revealing operational metric is useful accepted output per unit of total compute and memory—not raw proposal count.

Primary sources

  1. Leviathan, Kalman, and Matias, “Fast Inference from Transformers via Speculative Decoding” (2023) — algorithm and exact distribution argument.
  2. Chen et al., “Accelerating Large Language Model Decoding with Speculative Sampling” (2023) — speculative sampling formulation and evaluation.
  3. vLLM source repository — production inference engine; inspect the version-specific speculative decoding implementation and tests before relying on an API detail.
  4. vLLM documentation: Speculative Decoding — deployment modes, configuration, and caveats; documentation may evolve.