Speculative Decoding

The core mechanism involves a two-stage process:

• Drafting Phase: A small, fast draft model (e.g., a distilled version of the target) runs autoregressively for k steps to propose a candidate sequence of tokens.
• Verification Phase: The larger, more accurate target model processes this entire candidate sequence in a single, parallel forward pass. It scores the draft tokens against its own predictions, accepting them in a contiguous block up to the first mismatch. This allows the target model to generate multiple tokens per its own expensive forward pass, breaking the sequential bottleneck.

The primary goal is to reduce end-to-end latency (wall-clock time). Speedup is achieved when the time saved by verifying k tokens in parallel outweighs the time cost of running the draft model.

• Key Metric: The acceptance rate—the average number of draft tokens accepted per verification step. A high acceptance rate is critical for gains.
• Theoretical Limit: Speedup is bounded by the draft model's latency and the verification cost. It is most effective when the target model is significantly slower than the draft model.
• Real-World Impact: Can achieve 2-3x latency reduction for large models like Llama-70B when paired with a well-chosen draft model, without altering the final output distribution.

A defining feature of standard speculative decoding is that it is mathematically lossless. The output distribution is identical to what the target model would produce running standard autoregressive decoding.

• Verification Guarantee: The target model's parallel verification acts as a rejection sampling mechanism. Only tokens that match the target model's own probability distribution are accepted.
• No Quality Degradation: Unlike quantization or pruning, there is no approximation error or change in the model's behavior. This makes it ideal for production systems where output fidelity is paramount.
• Deterministic Output: Given the same random seed, the speculative decoding run will produce the exact same sequence as the standard autoregressive run.

Successful deployment depends on a strategic draft-target pairing.

• Draft Model Traits: Must be small, fast, and aligned. Common choices include:
  • A distilled version of the target model.
  • A shallower or narrower architecture (e.g., fewer layers).
  • A n-gram or statistical model in some implementations.
• Alignment is Critical: The draft model must have a high probability of predicting the same next token as the target model. Poor alignment leads to low acceptance rates and negates the speedup.
• Training: Draft models are often specifically trained on the target model's outputs to maximize prediction alignment.

Speculative decoding trades increased memory bandwidth and compute intensity for reduced latency.

• Memory Bandwidth: The verification phase requires loading the full target model parameters and the KV cache for a batch size equal to the draft length k+1. This can be high but is amortized over multiple accepted tokens.
• Compute: The draft model adds extra FLOPs. The verification phase is a single forward pass but on a batch of k+1 sequences.
• Optimal k: The lookahead length k must be tuned. A larger k increases the potential reward (more tokens accepted) but also the cost of a failed draft (more wasted verification compute). The optimal k is hardware and model-dependent.

The core idea has inspired several variants and adjacent optimizations:

• Medusa & Self-Speculative Decoding: Uses multiple prediction heads attached to the same target model to draft future tokens, eliminating the need for a separate draft model.
• Speculative Sampling: The original algorithm for lossless acceleration.
• Lossy Variants: Techniques like EAGLE modify the verification step for greater speedup at a potential, controlled cost to output quality.
• Integration with Other Optimizations: Often combined with PagedAttention (for efficient KV cache management of draft sequences) and continuous batching to maximize GPU utilization during verification.

The technique operates in a two-stage loop. First, a small draft model (e.g., a distilled version of the target) runs autoregressively for γ steps to propose a candidate sequence. Second, the large target model processes this entire candidate sequence in a single, parallel forward pass. It scores the draft tokens and, upon finding a mismatch, rejects all subsequent tokens and generates a correction. This replaces γ serial target model calls with one parallel verification, achieving speedups of 2-3x.

Medusa is a prominent framework that adds multiple prediction heads (a 'medusa head') on top of the target model's backbone, enabling it to act as its own draft model. This self-speculation eliminates the need for a separate model. Key components include:

• Multiple Decoding Heads: Predict several future tokens in parallel.
• Tree Attention: Efficiently verifies a tree of candidate token sequences.
• Typical Acceptance Rates: Often achieves 2-4 accepted tokens per verification step, significantly reducing latency.

This is the standard algorithm for token acceptance. After the target model verifies the draft:

1. It generates adjusted probability distributions.
2. For each position, the draft token is accepted with probability min(1, p_target / p_draft).
3. If a token is rejected, a new token is resampled from the adjusted distribution. This ensures the output distribution is mathematically identical to that of the original target model, preserving quality. The speedup is directly tied to the draft model's accuracy and the number of speculative steps (γ).

Leading inference engines integrate speculative decoding to maximize hardware utilization.

• vLLM: Leverages its PagedAttention mechanism to efficiently manage the KV caches for both draft and target models within its continuous batching scheduler, minimizing memory fragmentation.
• TensorRT-LLM: Provides optimized kernels for the draft and verification passes, fusing operations and leveraging in-flight batching to maintain high GPU utilization during the variable-length acceptance phase.

Choosing the draft model is critical for performance. Common strategies include:

• Distilled Model: A smaller version of the target, trained to mimic its output distributions.
• N-gram or Lookup Model: Extremely fast, shallow models based on statistical frequency.
• Structured Drafting: Using a lossless method like BiLD, which employs a different, faster model family (e.g., a non-autoregressive model) to generate candidates without distributional mismatch. The goal is to maximize the acceptance rate while minimizing the draft model's computational overhead.

Speculative decoding introduces specific engineering trade-offs:

• Memory Overhead: Must maintain two sets of model weights and KV caches in GPU memory.
• Draft Model Accuracy: Speedup collapses if the draft model is poor; the rejection rate must be low.
• Context Length: The technique is most effective for the decode phase. Very long context prefill phases remain a bottleneck.
• Hardware Saturation: On memory-bound systems, the parallel verification pass may not fully hide the draft latency, diminishing returns.

An inference scheduling technique that dynamically groups incoming requests into a single, continuously executing batch. Unlike static batching, it allows:

• New requests to join an ongoing batch.
• Completed sequences to exit the batch independently.
• Maximized GPU utilization by eliminating idle time, directly complementing speculative decoding's goal of higher throughput.

The key-value cache is a memory buffer in transformer models that stores computed key and value tensors from previous tokens during autoregressive decoding. It is fundamental to efficient inference because:

• It prevents recomputation of attention for past tokens.
• Its memory footprint grows with context length, becoming a primary bottleneck.
• Speculative decoding reduces the number of serial decoding steps, thereby reducing the rate of KV cache writes for the large target model.

A memory management algorithm, central to the vLLM inference engine, that organizes the KV cache into non-contiguous, fixed-size blocks. It enables:

• Efficient memory sharing for the KV cache across sequences in a batch.
• Near-zero internal fragmentation, allowing larger batch sizes.
• Essential for practical deployment of speculative decoding, as it efficiently manages the variable-length KV caches of both draft and target models.

Attention variants designed to reduce the memory and bandwidth cost of the KV cache.

• Multi-Query Attention (MQA): All query heads share a single key and value head.
• Grouped-Query Attention (GQA): A tunable hybrid where groups of query heads share a key/value head.
• Both directly reduce the size of the KV cache that must be stored and accessed each decoding step, a critical optimization that works in parallel with speculative decoding's algorithmic speedup.

The two distinct computational phases of transformer inference:

• Prefill Phase: The initial, parallel processing of the entire input prompt to compute the initial KV cache. This phase is compute-bound.
• Decode Phase: The token-by-token autoregressive generation, which is heavily memory-bound due to sequential reads of the growing KV cache.
• Speculative decoding specifically accelerates the decode phase by reducing its sequential depth.

A training technique where a smaller student model is trained to mimic the behavior of a larger teacher model. It relates to speculative decoding in two key ways:

• The small draft model in speculative decoding is often a distilled version of the large target model.
• While distillation creates a permanently smaller, faster model, speculative decoding uses the small model transiently as a draft engine, retaining the full quality of the large model for verification.