BudgetDraft：面向稀疏KV投机解码的接受感知多视图训练

BudgetDraft: 接受感知的多视角训练用于稀疏KV投机解码
来源: https://arxiv.org/html/2606.00144
Liang He¹,∗ Jingbo Wen² Qishi Zhan³ Yixiong Chen⁴ Kangning Cui⁵ Qizhen Lan⁶ Xilu Wang⁷,∗ ¹上海光学精密机械研究所 ²悉尼大学 ³马凯特大学 ⁴约翰霍普金斯大学 ⁵维克森林大学 ⁶德克萨斯大学休斯顿健康科学中心 ⁷萨里大学 [email protected], [email protected] ∗通讯作者

###### 摘要

投机解码通过使用草稿模型提出多个令牌，并由验证模型并行验证，从而加速自回归解码。在资源受限的部署中，草稿模型使用稀疏KV缓存，在固定KV预算下限制峰值GPU内存和端到端延迟，而验证模型则保留完整的KV缓存。中长上下文推理（4K–16K上下文长度）在实际应用中很常见。然而，朴素的稀疏/完整投机解码随着上下文长度增长而遭受稀疏/完整不匹配的问题，导致接受率迅速下降。我们提出BudgetDraft，一种用于中长推理中稀疏草稿的多视角稀疏训练方法。草稿模型在训练过程中暴露于多个采样的KV预算，并学习将每个稀疏视角与一个共享的完整缓存教师目标对齐。BudgetDraft将完整缓存分支上的接受感知损失与稀疏缓存分支上的多视角损失相结合，产生一个单一的预算鲁棒草稿模型，无需额外的推理时组件即可恢复不同稀疏级别下的接受率。在PG-19、LongBench和LWM上的实验结果表明，BudgetDraft在4K、8K和16K上下文长度下分别实现了高达6.55×、4.46×和2.10×的端到端加速比（相对于自回归解码），同时保持推理管道内存友好。

BudgetDraft: 接受感知的多视角训练用于稀疏KV投机解码

Liang He¹,∗ Jingbo Wen² Qishi Zhan³ Yixiong Chen⁴ Kangning Cui⁵ Qizhen Lan⁶ Xilu Wang⁷,∗ ¹上海光学精密机械研究所 ²悉尼大学 ³马凯特大学 ⁴约翰霍普金斯大学 ⁵维克森林大学 ⁶德克萨斯大学休斯顿健康科学中心 ⁷萨里大学 [email protected], [email protected] ∗通讯作者

## 1 引言

中长上下文推理（上下文长度范围为4K到16K个令牌）在文档摘要、多轮对话和检索增强生成等实际应用中越来越常见（Liu et al., 2025a; Liao et al., 2025）。自回归解码（AR）仍然是大语言模型（LLM）推理的主导范式，但其逐令牌的顺序生成导致高延迟和高成本（Chen et al., 2026）。投机解码（SD）通过使用一个小型草稿模型提出候选令牌，并由一个更大的验证模型来验证，从而缓解了这一瓶颈（Leviathan et al., 2023）。然而，其加速效果关键取决于草稿模型与验证模型之间的对齐程度。当草稿模型与验证模型对齐良好时，每个验证步骤可以接受多个令牌。

在部署中的中长推理场景下，维持这种对齐变得困难（Xiao et al., 2025）。一个核心限制是KV缓存（Zhou et al., 2024）。为了控制峰值GPU内存（VRAM）和端到端延迟，草稿模型通常在固定KV预算下使用稀疏KV缓存运行，只保留缓存键值对的一个子集，而验证模型则保留完整的KV缓存以保持输出质量（Li et al., 2024a）。这种稀疏/完整不匹配随着上下文长度增长而加剧，并可能急剧降低接受率。重要的是，这也使得投机解码对KV预算高度敏感，这在部署中是有问题的，因为不同设备和工作负载的内存预算各不相同（Cai et al., 2026）。

参见图1的说明：在GS数据集上，SD（稀疏/完整）的接受率崩溃，γ=5。随着上下文长度增加，接受率下降，并在KV设置B∈{256,512,1024,full}的8K–16K范围内接近零，其中full设置在本实验中对应2048个令牌。图1突出了这种失败模式。随着上下文长度增加，使用稀疏草稿模型/完整验证模型设置的SD表现出快速的接受率下降，在16K时接受率接近零，而相对于自回归解码的端到端加速比下降到≤1×。我们称之为16K崩溃。这促使我们将重点放在8K–16K范围，在这个范围内稀疏草稿在实际中具有吸引力，但接受率下降成为加速的主要障碍。图1还显示接受率对KV预算的响应是非单调的：对于朴素的稀疏草稿，较小的KV预算可能产生比较较大的预算更高的接受率（更多细节见附录D）。

先前的工作要么提出在全缓存或短上下文草稿下更强的草稿模型（Li et al., 2024c, 2026b），要么通过稀疏KV缓存和结构方法控制内存或缓解长上下文瓶颈，但代价是单模型目标或额外的推理时阶段（Li et al., 2024b; Xiao et al., 2024; Sun et al., 2024）。在这些路线中，草稿模型针对固定设置进行调整，当部署时KV预算变化时会变得脆弱（Gao et al., 2025）。我们则是在训练时解决稀疏/完整不匹配问题，同时保持推理简单，并提出了BudgetDraft用于稀疏草稿，重点关注预算鲁棒性。我们的核心思想是多视角稀疏训练：在训练过程中，草稿模型暴露于多个随机采样的KV预算，并学习在每个稀疏级别下与相同的完整KV教师目标对齐。这产生了一个预算鲁棒的草稿模型，在部署时面对不同的内存约束，能够在不同KV预算下保持稳定的接受率。我们的贡献如下：

- • 我们刻画了从4K到16K上下文长度中稀疏/完整接受率崩溃的现象，将16K确定为朴素稀疏投机解码的实际失败边界，并进一步揭示了在4K时的非单调预算效应，即较小的KV预算可能产生更高的接受率。
- • 我们提出了BudgetDraft，它将接受感知对齐与多预算稀疏视角训练相结合，产生一个预算不变的草稿模型，能够在所有稀疏级别下恢复接受率。
- • 我们证明，在单块A100 GPU上，BudgetDraft在4K时相对于自回归解码实现了高达6.55×的端到端加速，在8K时实现了4.46×的加速，在16K时实现了高达2.10×的加速。

## 2 相关工作

#### 投机解码。
投机解码通过使用一个小型草稿模型提出令牌，并由验证模型在一个前向传播中检查，从而加速自回归推理（Leviathan et al., 2023）。最近的工作通过更强的草稿机制提高了草稿质量（Li et al., 2024c）。EAGLE及其变体从目标模型的隐藏状态训练轻量级草稿头，在短上下文中实现了显著的加速。EAGLE-3增加了多层特征融合和测试时模拟（Li et al., 2026a）。然而，这些方法通常在2K上下文和全缓存草稿下进行评估（Liu et al., 2025b），并没有解决当稀疏KV缓存在更长上下文中约束草稿模型时的稀疏/完整不匹配问题（Yang et al., 2025）。

#### 用于中长上下文推理的稀疏KV缓存。
许多方法通过稀疏化来减少KV缓存内存，以进行中长推理（Zhang et al., 2023; Li et al., 2024a）。StreamingLLM保留注意力汇聚点和最近的令牌，以在有界缓存支持下进行长上下文生成（Xiao et al., 2024）。H2O使用基于注意力分数的驱逐策略，保留预算内的KV状态子集（Zhang et al., 2023）。这类工作的常见设计模式是在固定KV预算下进行预算驱逐，并采用令牌级或块级选择（Zhang et al., 2023; Li et al., 2024b; Feng et al., 2026）。然而，在实际部署中，KV预算并不是固定的：它会随GPU内存可用性、工作负载并发性和延迟约束而变化（Gao et al., 2025）。因此，针对单一预算调整的草稿模型在不同预算下部署时可能会很脆弱。这促使