💬 NLP

Batch Speculative Decoding Done Right

This paper identifies that existing batch speculative decoding implementations fail to maintain output equivalence due to the ragged tensor problem, and proposes the first correct framework, EQSPEC, along with its optimized variant EXSPEC, which achieves up to 3x throughput improvement while guaranteeing algorithmic correctness across various model pairs.

Original authors: Ranran Haoran Zhang, Soumik Dey, Ashirbad Mishra, Hansi Wu, Binbin Li, Rui Zhang

Published 2026-02-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ranran Haoran Zhang, Soumik Dey, Ashirbad Mishra, Hansi Wu, Binbin Li, Rui Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a high-speed train station. Your goal is to get as many passengers (tokens) onto trains (the AI model's output) as possible, as quickly as possible.

The Problem: The "Guess and Check" Traffic Jam

Normally, an AI model writes one word at a time, checking its work before moving to the next. This is safe but slow.
Speculative Decoding is like hiring a fast, junior assistant (a "draft model") to guess the next 5 words for you. The main AI (the "target model") then quickly checks if those guesses are right. If they are, you save time. If not, you fix the mistake and try again.

This works great for a single passenger. But what happens when you have a group of passengers (a "batch") all trying to board at once?

The Disaster: The "Ragged Tensor" Nightmare

In the real world, different passengers accept different numbers of guesses.

Passenger A accepts all 5 guesses.
Passenger B only accepts the first 2, then says, "No, that's wrong!"
Passenger C accepts 4.

Now, you have a messy line. Some people are at position 5, some at 2, some at 4. In computer science, this is called a "Ragged Tensor" problem. The computer's brain (the GPU) loves neat, rectangular grids. It hates messy, jagged lines.

The Old Way: The "Broken" Solutions

Previous attempts to fix this were like trying to force a jagged line into a square box by breaking the rules:

The "Masking" Method: They tried to hide the differences, but it confused the train schedule. The AI started hallucinating, repeating words like "not not not not" or printing gibberish.
The "Rollback" Method: They forced everyone to stop at the shortest line. If Passenger B only accepted 2 words, everyone else had to throw away their extra 3 words and wait. This wasted all the speed gains.

The result? Existing batch systems were fast but produced garbage output. They were like a race car that drives 200 mph but crashes into a wall every 10 seconds.

The Solution: The "Perfectly Organized" Station

This paper introduces two new methods, EQSPEC and EXSPEC, to fix the station without breaking the rules.

1. EQSPEC: The "Strict Traffic Controller"

Think of EQSPEC as a very strict, organized traffic controller.

How it works: After every round of guessing, it stops the train. It takes everyone's tickets, shuffles them so everyone is perfectly aligned in a straight line, updates the schedule (Position IDs), and moves the luggage (KV-Cache) to the right spot.
The Cost: This shuffling takes time. It's like stopping the train to reorganize the seating chart. It adds about 40% overhead (extra work), but it guarantees the output is 100% correct. No gibberish, no crashes.
The Metaphor: It's like a librarian who stops the whole library to re-shelve books perfectly after every few pages are read. It's slow, but the books are always in the right place.

2. EXSPEC: The "Smart Grouping" System

EQSPEC is correct but a bit slow because it stops to reorganize every time. EXSPEC is the upgrade.

How it works: Instead of forcing everyone into one big, messy line, EXSPEC keeps a "pool" of passengers. It looks around and says, "Hey, Passenger A and Passenger C both accepted 4 words! Let's put just them on a train together."
The Magic: If everyone on a specific train accepted the same number of words, no reorganization is needed. The train leaves instantly.
The Metaphor: Imagine a bus system where the driver doesn't force everyone onto one bus. Instead, they group people by how far they are going. If a group of 5 people all want to go exactly 3 stops, they get on a bus that leaves immediately without any waiting or shuffling.

The Results: Fast AND Correct

The authors tested this on real AI models (like Vicuna and Qwen).

Old Methods: Fast, but produced nonsense (like a robot speaking in circles).
New Methods (EQSPEC/EXSPEC):
- Correctness: 95%+ of the time, the output is identical to the slow, perfect method. The tiny 5% difference is just due to the computer's math being slightly fuzzy (floating-point errors), not a logic error.
- Speed: At a batch size of 8, they are 3 times faster than running one by one, while still producing perfect text.

In a Nutshell

Previous batch speculative decoding was like a chaotic mosh pit: fast, but people got hurt (corrupted output).
This paper built a well-organized dance floor.

EQSPEC is the choreographer who stops the music to fix everyone's steps (guarantees correctness).
EXSPEC is the smart DJ who only plays songs where everyone knows the dance moves, so no one ever has to stop (maximizes speed).

They proved that you can have both speed and accuracy in AI, as long as you respect the rules of how the computer's memory works.

1. Problem Statement: The Failure of Existing Batch Speculative Decoding

Speculative decoding accelerates Large Language Model (LLM) inference by using a small "draft" model to propose multiple tokens, which a larger "target" model verifies in parallel. The fundamental requirement of speculative decoding is output equivalence: the generated token distribution must be identical to standard autoregressive generation.

The paper identifies a critical, systematic failure in all existing batch speculative decoding implementations (e.g., BSP, DSD, and HuggingFace's batched attempts):

The Ragged Tensor Problem: In a batch, different sequences often accept different numbers of draft tokens during verification. This creates "ragged" (irregular) tensors where sequence lengths vary within the same batch.
Desynchronization: Existing methods fail to properly synchronize Position IDs, Attention Masks, and KV-cache states across these ragged sequences.
Consequence: This leads to corrupted outputs, ranging from repetitive tokens to gibberish or <unk> symbols. The authors demonstrate that prior methods achieve high throughput but violate the defining criterion of valid speculative decoding, resulting in near-zero exact-match rates compared to standard generation.

2. Methodology

The authors propose a framework that prioritizes algorithmic correctness (output equivalence) while optimizing for throughput. They analyze three potential approaches to handling ragged tensors:

Masking: Fails because it creates non-contiguous position IDs, breaking RoPE and attention mechanisms.
Rollback: Truncates all sequences to the shortest accepted length. This wastes verified tokens from faster sequences, causing throughput to collapse as batch size increases.
Dynamic Padding: The only viable approach. It realigns sequences after every verification round to maintain rectangular tensors.

The paper introduces two complementary algorithms based on Dynamic Padding:

A. EQSPEC: Guaranteeing Algorithmic Correctness

EQSPEC is the first algorithm that formally guarantees output equivalence by enforcing Synchronization Invariants:

Rectangular Alignment (I1): Ensures that for every sequence $i$ , the sum of padding length ( $p_i$ ) and content length ( $c_i$ ) equals the batch length ( $L$ ) at every step.
Position-ID Contiguity (I2): Ensures position IDs are derived consistently from the attention mask (starting from 0 at the first content token), which is required for RoPE and other position-aware mechanisms.

Mechanism (Unpad-Append-Repad):
After each verification round, EQSPEC performs a three-step synchronization:

Unpad: Remove existing padding.
Append: Add accepted draft tokens and the bonus token.
Repad: Calculate a specific padding offset ( $\delta_i$ ) for each sequence to realign the batch to the new maximum length. This involves shifting KV-cache entries and recomputing position IDs and attention masks.

B. EXSPEC: Reducing Overhead via Cross-Batch Scheduling

The authors analyze the cost structure of EQSPEC and find that the realignment overhead grows superlinearly with batch size, consuming up to 40% of computation at batch size 8. This is an inherent cost of maintaining correctness across ragged tensors.

To mitigate this, they introduce EXSPEC, which expands the scheduling scope beyond a single static batch:

Sequence Pool: Maintains a pool of active sequences in their ragged states rather than forcing them into fixed batches immediately.
Dynamic Grouping: The scheduler scans the pool and dynamically groups sequences of identical lengths into batches.
Zero Overhead for Homogeneous Groups: If sequences in a batch accept the same number of tokens, the padding offset $\delta_i$ is zero. This eliminates the need for unpad-append-repad and KV-cache shifting, bypassing the alignment overhead entirely.
Fallback: Only when sequences cannot be grouped by length does the system fall back to the EQSPEC realignment process.

3. Key Contributions

Formalization of Invariants: The paper formally defines the synchronization invariants (Rectangular Alignment and Position-ID Contiguity) required for valid batch speculative decoding, proving that prior methods violate these.
EQSPEC Algorithm: The first implementation that guarantees output distribution equivalence by explicitly enforcing these invariants after every verification round.
EXSPEC Algorithm: A cross-batch scheduling mechanism that reduces alignment overhead by grouping same-length sequences, effectively turning a fixed synchronization cost into an optional one.
Comprehensive Analysis: A detailed breakdown of the "ragged tensor problem," distinguishing between batch-independent bugs (e.g., wrong sampling distributions) and batch-specific synchronization failures.

4. Experimental Results

The methods were evaluated on SpecBench using three model pairs (Vicuna-7B/68M, Qwen3-8B/0.6B, GLM-4-9B/0.6B) on NVIDIA A100 GPUs.

Output Equivalence (Correctness):
- EQSPEC and EXSPEC: Achieved ~95% exact-match rates with standard autoregressive decoding. The remaining ~5% divergence is attributed to floating-point non-determinism in GPU inference, not algorithmic errors.
- Prior Methods (BSP, DSD): Suffered catastrophic failures with near-zero exact-match rates (e.g., DSD at 0.0%, BSP at 1.9%), producing repetitive or nonsensical outputs.
Throughput (Performance):
- Scaling: At batch size 8, EXSPEC achieved up to a 3× throughput improvement over batch size 1.
- Overhead: Alignment overhead in standard batch speculative decoding grows superlinearly, reaching ~40% of total computation time at batch size 8. EXSPEC significantly reduces this by grouping sequences.
- Comparison: While production systems like vLLM and SGLang use continuous batching to avoid ragged tensors, they currently do not support external draft models for speculative decoding at scale (vLLM deprecated it due to memory overhead; SGLang is limited to EAGLE-family drafters).

5. Significance

This paper fundamentally corrects the trajectory of batch speculative decoding research and deployment:

Validity vs. Speed: It establishes that output equivalence is a binary requirement, not an optimization target. High throughput is meaningless if the output is corrupted.
The Cost of Correctness: It quantifies the "inherent cost" of maintaining correctness across ragged tensors (superlinear alignment overhead), a factor previously ignored or mishandled.
Practical Deployment: By introducing EXSPEC, the authors demonstrate that it is possible to combine batch parallelism with per-sequence speculation gains without sacrificing correctness, provided that intelligent scheduling (grouping by length) is used.
Open Source: The authors release the first verified reference implementation for batch speculative decoding, enabling future research and production systems to build on a correct foundation.

In summary, the paper argues that previous batch speculative decoding implementations were fundamentally broken due to the "ragged tensor problem." By formalizing the necessary synchronization invariants and introducing EQSPEC/EXSPEC, the authors provide the first method that is both algorithmically correct and scalable, achieving up to 3x speedup while maintaining the exact output distribution of the target model.