Speculative Decoding & Multi-Head Speculation: High-Throughput Autoregressive Serving

Architecture

Autoregressive transformer inference at low batch sizes is fundamentally constrained by GPU High-Bandwidth Memory (HBM) bandwidth rather than compute capability. During each token generation step, the entire model weight matrix must be loaded from HBM into on-chip SRAM to process a single token vector, yielding an arithmetic intensity near 1.0 FLOP/Byte.

Speculative Decoding alters this equation by decomposing generation into a fast speculation phase and a parallel verification phase:

1. Speculative Draft Phase: A lightweight mechanism (such as a compact draft model or parallel speculative heads) speculatively generates a sequence or tree of \gamma candidate tokens.
2. Parallel Verification Phase: The large target model evaluates all \gamma candidate tokens simultaneously in a single forward execution. Because transformer forward passes over sequence lengths \le 64 are compute-underutilized at batch size 1, verifying \gamma tokens requires approximately the same wall-clock time as generating a single token autoregressively.
3. Lossless Acceptance & Resampling: A modified rejection sampling kernel determines which candidate tokens match the target model's probability distribution, committing accepted tokens to the KV-cache and falling back to a normalized residual distribution upon the first rejection.

Mathematical Foundations: Lossless Rejection Sampling & Residual Resampling

The defining property of speculative decoding (Leviathan et al., 2022; Chen et al., 2023) is that it is mathematically lossless: the final sequence distribution is identically equal to direct autoregressive sampling from the target model M_{\text{target}}.

1. Rejection Sampling Criterion

Let x_1, \dots, x_\gamma denote candidate tokens generated sequentially by the draft model distribution q(x) = P_{\text{draft}}(x \mid x_{<t}).

For each candidate token x_i, the target model calculates probability distribution p(x) = P_{\text{target}}(x \mid x_{<t}). The acceptance probability \alpha(x_i) is evaluated as:

A standard uniform random variate u_i \sim \mathcal{U}(0, 1) is sampled:

• If u_i \le \alpha(x_i), candidate token x_i is accepted and appended to the generated sequence.
• If u_i > \alpha(x_i), candidate token x_i is rejected. Speculation terminates at index i, all subsequent candidates x_{i+1}, \dots, x_\gamma are discarded, and an alternate token is sampled from the residual distribution.

2. The Residual Resampling Distribution

When token x_i is rejected, drawing a replacement token directly from p(x) would bias the generation, because x_i was already rejected according to q(x).

To guarantee exact distribution alignment, the replacement token must be sampled from the normalized residual distribution p'(x):

3. Mathematical Proof of Exact Distribution Recovery

We prove that the probability of emitting any token x \in \mathcal{V} at step i under speculative decoding equals p(x) exactly:

1. Probability of drafting and accepting x:

2. Total probability of rejection across the vocabulary \beta:

3. Probability of rejecting and drawing x from the residual distribution p'(x):

4. Combining both terms:

The emitted token distribution matches the target model distribution with zero divergence (D_{\text{KL}}(P_{\text{emit}} \parallel P_{\text{target}}) = 0).

Multi-Head Speculation: Medusa & Eagle Tree Attention

While two-model speculative decoding pairs a small model (e.g., LLaMA-3-8B) with a large model (LLaMA-3-70B), maintaining two distinct model weight graphs complicates serving clusters.

Medusa and Eagle replace the standalone draft model with speculative prediction heads built directly on top of the target model's representations.

1. Architectural Taxonomy

• Medusa Multi-Head Decoding: Adds K single-layer MLP heads to the target model's final hidden states. Head k \in \{1, \dots, K\} predicts logits for offset t + k + 1 without running a recurrent draft backbone.
• Eagle Feature-Level Speculation: Passes both the target model's top hidden state and token embeddings into a lightweight single-layer transformer decoder, capturing autoregressive dependency between speculative heads.

2. Tree Attention Mask Construction

Generating candidates from independent heads creates multiple divergent continuation paths. Rather than verifying paths sequentially, modern serving runtimes assemble candidates into a directed tree verified in a single forward pass.

To ensure causal validity, tokens must attend only to their ancestors in the tree. This is enforced via a 2D Tree Attention Mask:

The custom attention mask ensures zero cross-talk between independent branches, enabling parallel verification of up to 64 speculative paths in a single CUDA kernel invocation.

Analytical & Empirical Speedup Curves and Memory Bandwidth Scaling

The speedup of speculative decoding depends on three core parameters:

1. \gamma: The number of speculative candidate tokens evaluated per cycle.
2. \alpha: The empirical acceptance rate (\alpha \in [0.0, 1.0]).
3. c: The cost ratio c = \frac{T_{\text{draft}}}{T_{\text{target}}}, representing the latency of drafting relative to target verification.

1. Analytical Speedup Derivation

Under independent acceptance probability \alpha, the expected number of accepted draft tokens \mathbb{E}[N] per cycle follows a truncated geometric sum:

Including the bonus token emitted by target verification or residual resampling, the expected total tokens emitted per cycle \mathbb{E}[\tau] is:

The latency of a speculative cycle is T_{\text{spec}} = \gamma \cdot T_{\text{draft}} + T_{\text{target}} = T_{\text{target}}(\gamma \cdot c + 1).

The analytical speedup factor S over standard autoregressive decoding is:

2. GPU Memory Bandwidth Reduction on NVIDIA H100 SXM

Consider serving a 70B parameter model in FP16 (140 GB weight footprint) on an 8x H100 SXM cluster (NVLink 4, 3,350 GB/s HBM3 per GPU).

| Serving Paradigm | Draft Lookahead \gamma | Mean Acceptance \alpha | Tokens Emitted / Cycle | HBM Traffic / Cycle | Effective HBM Traffic / Token | Empirical Throughput | |---|---|---|---|---|---|---| | Standard Autoregressive | 0 | N/A | 1.0 token | 140.0 GB | 140.0 GB/token | 23.8 tok/s/req | | Draft Model (70B + 8B) | 5 | 0.72 | 3.54 tokens | 140 GB + 5(16 GB) = 220 GB | 62.1 GB/token | 53.6 tok/s/req (2.25x) | | Medusa Heads (4 Heads) | 4 | 0.68 | 3.12 tokens | 140 GB + 0.8 GB = 140.8 GB | 45.1 GB/token | 68.4 tok/s/req (2.87x) | | Eagle Tree Speculation | 16 (Tree) | 0.81 | 4.88 tokens | 140 GB + 2.4 GB = 142.4 GB | 29.2 GB/token | 98.2 tok/s/req (4.12x) |

By reusing the 140 GB target model weight read across 4.88 accepted tokens, Eagle tree speculation reduces the memory bus traffic from 140 GB down to 29.2 GB per emitted token, achieving a 79.1% reduction in memory bandwidth consumption.

Decisions

| Decision | Required evidence | Review trigger | |---|---|---| | Deploy Speculative Decoding when serving low concurrency workloads (Batch Size \le 8) | Memory profiling confirming > 85\% memory-bandwidth bound execution (Roofline intensity < 20 FLOPs/Byte). | Workload concurrency scaling beyond 32 concurrent requests per GPU. | | Choose Medusa / Eagle Speculative Heads over Standalone Draft Models | VRAM headroom constrained to under 8 GB surplus memory per GPU; desire zero inter-process draft communication overhead. | Frequent model architecture updates where fine-tuning speculative heads adds CI/CD maintenance friction. | | Select Speculation Window \gamma = 4 or \gamma = 5 for draft models, and \gamma = 16 to 32 for tree attention | Benchmark sweeps identifying peak of S = \frac{1 - \alpha^{\gamma+1}}{(1 - \alpha)(\gamma \cdot c + 1)} on production query logs. | Acceptance rate dropping below \alpha = 0.55, triggering speculation latency penalties. | | Implement Dynamic Speculation Window Adaptation | Runtime telemetry tracking rolling acceptance rate \alpha; automatically throttling \gamma \to 0 when entropy spikes. | Generating high-entropy code or math tokens where draft speculation accuracy collapses. |

Alternatives and trade-offs

• Prompt Lookup Decoding (N-gram Speculation): Highly efficient for summarization, code editing, and document question-answering where candidate tokens already exist in the input prompt. Requires zero extra VRAM or training, but fails completely on novel generative tasks.
• Standalone Draft Models: Robust across varied generative distributions, but requires careful vocabulary alignment and introduces inter-model scheduling complexity in multi-tenant inference engines.
• Medusa / Eagle Speculative Heads: The industry standard for production serving runtimes (vLLM, SGLang, TensorRT-LLM). Delivers highest speedup with minimal VRAM overhead, though requiring fine-tuning alongside the base model.

Failure modes

• Domain Mismatch Collapse (\alpha < 0.4): If the draft model is trained on general internet text while the target model serves medical jurisprudence or formal proofs, the acceptance rate drops below 40%. The overhead of running the draft model (c \cdot \gamma) exceeds the gain from accepted tokens, resulting in negative speedup (speculation penalty).
• KV-Cache Fragmentation & Rollback Thrashing: Every rejected speculative token requires rolling back KV-cache pointers. Under non-paged memory allocation, rapid allocation and deallocation of candidate blocks causes memory fragmentation and allocator lock contention. Resolved by integrating with PagedAttention virtual block tables.
• Batch Size Inversion: At high batch sizes (e.g., B \ge 64), standard autoregressive decoding becomes compute-bound rather than memory-bound. Running speculative verification on B \times \gamma tokens saturates Tensor Core capacity, increasing P99 latency and degrading overall serving throughput.
• Tokenizer Discrepancy: If a draft model uses a different vocabulary, BPE merge table, or special token scheme than the target model, candidate tokens cannot be mapped into the target model's embedding matrix without costly string detokenization and retokenization.

Operational checklist

• [ ] Draft and target models share identical tokenizers, vocabulary sizes, and special token mappings.
• [ ] Runtime engine implements dynamic speculative window throttling: \gamma is dynamically scaled down when rolling \alpha < 0.50.
• [ ] KV-cache management utilizes PagedAttention block mapping to support zero-copy rollbacks upon token rejection.
• [ ] Speculative tree attention mask kernels are compiled with FlashAttention / CUTLASS for fused forward execution.
• [ ] Telemetry monitors the Speculative Acceptance Rate (\alpha), Empirical Speedup Factor (S), and HBM Bandwidth Savings (\%).
• [ ] Concurrency-aware autoscaler dynamically disables speculative decoding when batch size exceeds the compute-bound threshold (B \ge 16).

Connected practice

• Labs: /labs/speculative-decoding
• System breakdowns: /systems/inside-vllm, /systems/inside-vllm-paged-attention
• Paper breakdowns: /papers/speculative-decoding-leviathan-paper

Sources

• speculative-decoding-leviathan-2022
• vllm-v0102-model-runner