SPIRe: Boosting LLM Inference Throughput with Speculative Decoding

TL;DR

This work tackles the cost-per-token challenge in autoregressive decoding by reframing speculative decoding as a throughput problem. It introduces SPIRe, a draft model that combines Sliding window KV cache, Pruned Initialization, and a Feedback Transformer with a distillation objective to maximize throughput multipliers for large batch sizes and long contexts. Through an implementation-agnostic evaluation framework and targeted ablations, SPIRe achieves substantial throughput gains over standard speculative decoding and the MagicDec baseline, with favorable cost-to-benefit trade-offs. The approach demonstrates practical impact for production LLM serving by enabling efficient exploration of draft-model designs across varying batch sizes and context lengths while providing a concrete evaluation methodology and analysis of trade-offs and extrapolations toward larger models like Llama 3 70B.

Abstract

Speculative decoding (SD) has been shown to reduce the latency of autoregressive decoding (AD) by 2-3x for small batch sizes. However, increasing throughput and therefore reducing the cost per token requires decoding with large batch sizes. Recent work shows that SD can accelerate decoding with large batch sizes too if the context is sufficiently long and the draft model's KV cache is sparse. We introduce SPIRe, a draft model that combines static sparse attention, pruned initialization, and feedback memory to increase the modeled throughput of speculative decoding by over 100% compared to speculation with a much smaller draft model and by over 35% compared to the strong baseline of sparse self-speculation. Our approach is particularly effective when context lengths vary significantly across requests.

Figures (6)

• Figure 1: Speedup of generating tokens in a batch of size $B=64$ given a context of length $L$. A round of speculation is the generation and subsequent verification of each block of $k=4$ draft tokens, the average generation length$\tau$ is the average number of tokens generated per round of speculation, and speedup is $\tau$ divided by how much longer a round of speculation takes compared to autoregressively decoding one token.
• Figure 2: The matrix $QK^\top$ in the 1st and 4th forward passes on a training sequence of length $L=12$ with a maximum speculation depth of $k=4$. In general we perform $k$ forward passes on every $k$-th prefix of a sequence, in a process similar to batched autoregressive decoding with teacher forcing.
• Figure 3: With a sufficiently long context, verification may never be compute-bound even in the limit as batch size tends to infinity. For context lengths where verification is compute bound for some batch size $B_\text{crit}$, there's no benefit in terms of throughput (and some cost in terms of latency) to increasing batch size above $B_\text{crit}$.
• Figure 4: Speedup of generating tokens in a batch of size $B$ given a context of length $L$. Using the highlighted values, we get that SPIRe increases the modeled throughput of speculative decoding by 100% compared to vanilla speculative decoding and by 35% compared to MagicDec.
• Figure 5: Ablations of the techniques used to train our draft model SPIRe. In order of importance, the techniques are 1) sparse KV cache, 2) attending to target model activations, 3) distillation loss, 4) initializing by pruning the target model, 5) feedback memory, and 6) $\text{MixedLoss}$.