Principle:Pytorch Serve Speculative Decoding Inference
| Knowledge Sources | |
|---|---|
| Domains | NLP, Optimization |
| Last Updated | 2026-02-13 18:52 GMT |
Overview
Speculative Decoding Inference is the principle of accelerating large language model text generation by using a smaller, faster draft model to propose candidate tokens that are then verified in parallel by the larger target model.
Description
In standard autoregressive decoding, a large language model generates tokens one at a time, with each forward pass producing a single token. This sequential process is heavily memory-bandwidth bound rather than compute-bound, meaning the expensive hardware is underutilized.
Speculative decoding addresses this by introducing a two-stage pipeline:
- Draft phase — A small, fast draft model (e.g., a distilled or smaller variant) generates K candidate tokens autoregressively. Because the draft model is lightweight, this step is fast.
- Verification phase — The large target model processes all K candidate tokens in a single forward pass (exploiting the parallelism of the transformer architecture). It computes the probability distribution at each position and accepts or rejects draft tokens based on a modified rejection sampling scheme.
The key guarantee is that the output distribution is mathematically identical to what the target model would produce on its own — speculative decoding provides a lossless speedup.
# Simplified speculative decoding loop
def speculative_decode(draft_model, target_model, prompt, K=5):
"""Generate tokens using speculative decoding."""
tokens = prompt
while not is_complete(tokens):
# Draft model proposes K candidate tokens
draft_tokens, draft_probs = draft_model.generate(tokens, num_tokens=K)
# Target model verifies all K tokens in one forward pass
target_probs = target_model.forward(tokens + draft_tokens)
# Accept/reject using modified rejection sampling
accepted = 0
for i in range(K):
ratio = target_probs[i] / draft_probs[i]
if random.random() < min(1, ratio):
accepted += 1
else:
break
tokens = tokens + draft_tokens[:accepted]
# Sample correction token from adjusted distribution if rejected
if accepted < K:
tokens = tokens + sample_residual(target_probs[accepted], draft_probs[accepted])
return tokens
Usage
Apply Speculative Decoding Inference when:
- Serving large language models where autoregressive generation latency is the primary bottleneck.
- The deployment environment has a suitable smaller draft model available (e.g., a distilled variant of the same model family).
- Lossless acceleration is required — the output quality must be identical to standard decoding.
- The target model's forward pass has spare compute capacity that can be exploited by processing multiple tokens simultaneously.
Theoretical Basis
Speculative decoding is grounded in modified rejection sampling. Given the draft model distribution q(x) and the target model distribution p(x), a proposed token x is accepted with probability:
min(1, p(x) / q(x))
If the token is rejected, a correction token is sampled from the residual distribution:
norm(max(0, p(x) - q(x)))
This scheme guarantees that the final output follows the exact distribution p(x) of the target model, regardless of the quality of the draft model. The expected speedup depends on the acceptance rate, which in turn depends on how closely the draft model approximates the target model. When the draft model is a good approximation, most proposed tokens are accepted, yielding speedups of 2-3x or more while preserving output fidelity.
The theoretical expected number of tokens generated per verification step is:
E[tokens] = (1 - alpha^(K+1)) / (1 - alpha)
where alpha is the average acceptance probability and K is the number of draft tokens proposed per step.