Batch Speculative Decoding Done Right

2025-10-26

AI Summary

In batch speculative decoding, asynchronous draft lengths produce ragged tensors, disrupting alignment of position IDs, attention masks, and KV caches—thereby violating output equivalence (i.e., failing to reproduce standard autoregressive outputs). This work first systematically formalizes synchronization constraints and proposes two novel methods: EQSPEC, which guarantees strict output equivalence, and EXSPEC, which trades off efficiency for controlled equivalence degradation. Both leverage dynamic equal-length grouping and a sliding-pool mechanism to achieve sequence-level synchronization and efficient KV cache reuse—without requiring custom kernels. Evaluated on SpecBench with batch size 8, our approach achieves up to 3× throughput improvement, maintains 95% output equivalence, and reduces alignment overhead by 40%.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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$ imes$ 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.

Problem

• Addresses ragged tensor issues in batch speculative decoding implementations
• Ensures output equivalence between speculative and autoregressive generation methods
• Reduces realignment overhead while maintaining high throughput scaling

Innovation

• Dynamic grouping of sequences to reduce realignment overhead
• Sliding pool method maintains speculative speedups per sequence
• Ensures output equivalence while scaling batch speculative decoding