Batch Speculative Decoding Done Right

TL;DR

This work tackles batch speculative decoding for LLMs, where ragged-tensor alignment across sequences can break output equivalence with standard autoregressive generation. It analyzes the design space and proves that simple masking or rollback approaches fail to preserve correctness, leaving dynamic padding as the viable path. The authors introduce EqSpec to enforce minimal synchronization invariants during realignment and EXSpec to schedule same-length sequences across a sliding window, effectively bypassing realignment when possible. Across Vicuna-7B/68M, Qwen3-8B/0.6B, and GLM-4-9B/0.6B targets, the approach achieves approximately $95\%$ exact-output equivalence while delivering up to $3\times$ throughput at batch size $8$, with no custom kernels and clean integration into existing stacks. These results demonstrate a practical, correctness-preserving route to scale batch speculative decoding for production systems and inform future designs of inference engines.

Abstract

Speculative decoding speeds up LLM inference by using a small draft model to propose multiple tokens that a target model verifies in parallel. Extending this idea to batches is essential for production serving, but it introduces the ragged tensor problem: sequences in the same batch accept different numbers of draft tokens, breaking right-alignment and corrupting position IDs, attention masks, and KV-cache state. We show that several existing batch implementations violate output equivalence-the fundamental requirement that speculative decoding must produce identical token sequences to standard autoregressive generation. These violations occur precisely due to improper handling of the ragged tensor problem. In response, we (1) characterize the synchronization requirements that guarantee correctness, (2) present a correctness-first batch speculative decoding EQSPEC that exposes realignment as consuming 40% of overhead, and (3) introduce EXSPEC, which maintains a sliding pool of sequences and dynamically forms same-length groups, to reduce the realignment overhead while preserving per-sequence speculative speedups. On the SpecBench dataset, across Vicuna-7B/68M, Qwen3-8B/0.6B, and GLM-4-9B/0.6B target/draft pairs, our approach achieves up to 3$\times$ throughput improvement at batch size 8 compared to batch size 1, with efficient scaling through batch size 8, while maintaining 95% output equivalence. Our method requires no custom kernels and integrates cleanly with existing inference stacks. Our code is available at https://github.com/eBay/spec_dec.

Figures (6)

• Figure 1: Batch speculative decoding on Vicuna-7B/68M: Existing methods achieve high throughput but violate the fundamental requirement of output equivalence by producing corrupted outputs. Our approach maintains perfect correctness while still achieving competitive performance.
• Figure 2: The ragged tensor problem in batch speculative decoding. Differing numbers of accepted draft tokens across sequences in the same batch lead to ragged-shaped input IDs tensors and KV Cache that disrupt subsequent batch operations.
• Figure 3: EqSpec synchronizes via unpad–append–repad.
• Figure 4: EXSpec pools ragged sequences by length, avoiding realignment; only unmatched sequences need syncing, turning fixed overhead into optional cost.
• Figure 5: Decomposing batch speculative decoding performance. (a) Batch scaling efficiency normalized to BS=1, isolating GPU parallelism from per-sequence speculation; EXSpec initially exceeds the No-Ragged-Scaling upper bound before degrading due to alignment overhead. (b) Alignment overhead grows super-linearly with batch size, consuming up to 38% of inference time, validating that $c_{\text{overhead}}(B)$ dominates at scale. (c) Cross-batch grouping rates on Multi30k for random vs. uniform-length sequences, showing that length homogeneity transforms grouping effectiveness.