批处理式推测解码的正确实现
计算与语言
2026-02-17 v3 人工智能
摘要
推测解码必须产生与标准自回归生成完全相同的输出分布——这种输出等效性不是优化目标,而是有效推测解码的定义标准。我们展示了所有现有的批处理式推测解码实现都违反了这一基本要求,导致产生损坏的输出,从重复标记到胶水一样的文本。这些失败源于“ragged tensor”问题:同一批次中的序列接受不同的草稿标记数量,导致位置 ID、注意力掩码和 KV 缓存状态不同步。我们提出了第一个真正的批处理式推测解码框架。我们(1)形式化了有效批处理式推测解码必须满足的同步不变量,(2)present 了第一种保证输出等效性的 EQSPEC 算法,并分析其成本结构表明对齐开销呈超线性增长,最多消耗 40% 的计算资源,(3)引入 EXSPEC,通过跨批次调度动态分组相同长度的序列来降低这种开销。在 SpecBench 上测试 Vicuna-7B/68M、Qwen3-8B/0.6B 和 GLM-4-9B/0.6B 等模型对,所提方法在 batch size=8 时实现了最高可达 3 倍的吞吐量提升,同时保持算法正确性。我们的 methods 实现 95% 的解码等效性,剩余的差异可归因于 GPU 推理中的浮点误差非确定性,而非先前方法导致几乎零等效性的同步失效。我们的代码已公开于 https://github.com/eBay/spec_dec。
引用
@article{arxiv.2510.22876,
title = {Batch Speculative Decoding Done Right},
author = {Ranran Haoran Zhang and Soumik Dey and Ashirbad Mishra and Hansi Wu and Binbin Li and Rui Zhang},
journal= {arXiv preprint arXiv:2510.22876},
year = {2026}
}