论文
Carryover Drafting:复用被拒绝的隐藏状态改善推测解码
Carryover Drafting: Recycling Rejected States for Speculative Decoding
摘要
推测性解码通过并行验证多个起草的词元来加速 LLM 推理,从而允许单个目标前向传递接受多个词元。通过构造,验证计算接受和拒绝的词元的表示。然而,传统的草稿模型只保留已接受词元的表示,从而导致验证者在拒绝词元上花费的大量计算实际上被浪费了。我们发现,目标转发期间生成的这些丢弃的隐藏状态保留了有关未来词元的有用信息,可以改进后续草案。然而,实现这一机遇面临着两个独特的挑战。据推断,回收开销会增加起草延迟,从而减少因增加接受长度而获得的加速。在训练期间,标准并行草稿模型训练不会产生推理对齐的拒绝状态,而通过顺序采样轨迹获得它们会牺牲训练位置之间的并行性。我们引入了结转起草,它可以解决这两个挑战。结转将拒绝的目标隐藏状态回收为临时 KV 上下文,允许草稿模型有选择地关注它们。它重用了草稿模型的现有界面,并仅添加一个学习嵌入来区分拒绝状态和提交上下文。额外的 KV 上下文在每一轮起草中都会被替换,并保持其长度受一个提案块的限制。我们引入并行草稿-验证-草稿训练,使草稿模型暴露于推理对齐的拒绝状态,同时保持训练位置之间的并行性。在两个目标模型上使用 DFlash 和 DSpark 衍生的半自回归绘图器进行的实验表明,与相应的基线相比,这种简单的 Carryover 机制将平均接受长度提高了 6.5--14.7%,端到端 vLLM 加速提高了 7.9--14.4%,翻译时的加速增益达到 28.8%。