バッチ投機的デコーディングの正しい実装
Batch Speculative Decoding Done Right
October 26, 2025
著者: Ranran Haoran Zhang, Soumik Dey, Ashirbad Mishra, Hansi Wu, Binbin Li, Rui Zhang
cs.AI
要旨
推測デコードは、小さなドラフトモデルを使用してターゲットモデルが並列で検証する複数のトークンを提案することで、LLM推論を高速化する。このアイデアをバッチ処理に拡張することは実運用において不可欠だが、不規則テンソル問題を引き起こす。同一バッチ内のシーケンスが異なる数のドラフトトークンを受け入れるため、右揃えが崩れ、位置ID、アテンションマスク、KVキャッシュ状態が破損する。既存の複数のバッチ実装が、推測デコードが標準的な自己回帰生成と同一のトークンシーケンスを生成しなければならないという基本要件である出力等価性に違反していることを示す。これらの違反は、不規則テンソル問題の不適切な処理が原因で発生する。これに対し、我々は (1) 正確性を保証する同期要件を特徴付け、(2) オーバーヘッドの40%を占めるリアラインメント処理を明らかにする正確性優先のバッチ推測デコードEQSPECを提示し、(3) リアラインメントのオーバーヘッドを削減しつつシーケンス単位の推測的加速を維持するため、スライディングプールを管理し動的に同じ長さのグループを形成するEXSPECを提案する。SpecBenchデータセットにおいて、Vicuna-7B/68M、Qwen3-8B/0.6B、GLM-4-9B/0.6Bのターゲット/ドラフトモデルペアを用いた実験では、バッチサイズ1と比較してバッチサイズ8で最大3倍のスループット向上を達成し、バッチサイズ8まで効率的にスケーリングしながら95%の出力等価性を維持した。本手法はカーネルカスタマイズを必要とせず、既存の推論スタックに容易に統合可能である。コードはhttps://github.com/eBay/spec_decで公開されている。
English
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 3times 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.