论文

BudgetDraft:稀疏 KV 推测解码 的验收感知多视图训练

BudgetDraft: Acceptance-Aware Multi-View Training for Sparse-KV Speculative Decoding

模型推理投机采样

摘要

推测解码 通过使用起草者提出多个词元供验证者并行验证来加速自回归解码。在资源受限的部署中,起草者使用稀疏的 KV 缓存来限制固定 KV 预算下的峰值 GPU 内存和端到端延迟,而验证者则保留完整的 KV 缓存。中长上下文推理(4K--16K 上下文长度)在实际应用中很常见。然而,随着上下文长度的增长,朴素的稀疏/完整 推测解码 会遭受稀疏/完整不匹配的问题,导致接受率迅速下降。我们提出了 BudgetDraft,一种用于中长推理中稀疏绘图的多视图稀疏训练方法。起草者在训练期间会接触到多个采样的 KV 预算,并学习将每个稀疏视图与一个共享的全缓存教师目标对齐。 BudgetDraft 将全缓存分支上的接受感知损失与稀疏缓存分支上的多视图损失相结合,生成一个预算稳健的起草器,无需额外的推理时间组件即可跨稀疏级别恢复接受。 PG-19、LongBench 和 LWM 上的实验结果表明,BudgetDraft 在 4K、8K 和 16K 上下文长度下与 AR 相比实现了高达 6.55 倍、4.46 倍和 2.10 倍的端到端加速,同时保持推理管道内存友好。