Dustin: 草稿增强的稀疏验证用于高效长上下文生成与推测解码

# 草稿增强稀疏验证：利用推测解码实现高效长上下文生成  
来源：https://arxiv.org/html/2606.24957  

Jian-Jia Chen，Xiaolin Lin，Pei-Shuo Wang，Chi-Chih Chang，Chun-Che Yang，Ning-Chi Huang，Grace Li Zhang，Kai-Chiang Wu  

###### 摘要  

虽然推测解码提升了多批次长上下文大语言模型（LLM）的推理吞吐量，但其效率常受限于验证瓶颈——在此过程中，键值（KV）缓存的加载成为延迟主因。现有压缩方法在此场景下表现欠佳：静态淘汰因显著性偏移而损失精度，动态选择则在验证路径上引入高昂的计算开销。我们提出 Dustin，一个专为长上下文推测解码设计的稀疏验证框架。Dustin 融合草稿模型的前向信号与目标模型的历史注意力，在多步验证窗口内高保真地识别关键 token。为降低重计算延迟，该方法进一步采用稀疏估计方案，将重要性评分限制在少量注意力头上。在 PG-19 和 LongBench 上使用 Qwen2.5-72B 的评估表明，Dustin 在 32k 序列长度下实现了自注意力 27.85 倍的加速以及端到端解码 9.17 倍的加速，且精度损失可忽略不计。  

机器学习，ICML  

## 1 引言  

图 1：单次推测解码步骤的延迟分解。实验采用 32k 输入长度和批次大小 16。我们比较了经典推测解码（SD）、MagicDec（MDec）（Sadhukhan 等，2024）以及我们提出的 Dustin。x 轴标记 Target(Draft) 仅表示目标模型和草稿模型的具体配对。  

大语言模型（LLM）（Achiam 等，2023；Yang 等，2024；AI@Meta，2024）满足了日益增长的长上下文任务需求，但在自回归解码过程中面临严重的内存带宽瓶颈（Yuan 等，2024；Pope 等，2023）。随着上下文长度增长，KV 缓存的线性扩展使得内存占用达到数百 GB，内存访问成为延迟的主要因素（Kwon 等，2023；Dao，2023）。近期研究（Sun 等，2024；Sadhukhan 等，2024；Yang 等，2025）指出，在多批次长上下文场景中，推测解码能有效提升吞吐量，因为验证的计算成本低于加载完整 KV 缓存的巨大开销。然而，KV 缓存加载成本增长的问题依然存在。如图 1 所示，在 32k 输入长度和批次大小 16 的条件下，验证占解码延迟的 87.5%，这限制了端到端加速的潜力。集成 KV 缓存压缩策略是一项有前景的解决方案，但标准方法对于推测解码并非最优。静态淘汰方法（Xiao 等，2023；Zhang 等，2023）永久丢弃上下文，由于显著性偏移（Zhao 等，2025）——即高注意力 token 集合随时间变化的现象，会导致精度损失。相反，动态选择方法保留完整 KV 缓存，但必须在每一步重新计算 token 的重要性，从而在验证路径上增加计算量。其成本取决于重要性估计的方式：一种简单的估计器——具体化所有层和所有头的注意力分数——代价高昂得令人望而却步——我们在第 5.4.1 节中的分析表明，在超过 4k token 时，其开销超过了全缓存自注意力的延迟。诸如 Quest（Tang 等，2024）的页面级方案通过廉价的 min/max 键统计量避免了精确的 \(QK^\top\)，但仍会产生不可忽略的评分成本，并且如我们在第 5.4.1 节所示，其速度仍慢于我们的稀疏估计器。此外，我们发现，仅依赖历史注意力分数往往会导致显著的精度下降，因为验证阶段需要同时处理多个未来步骤。  

为应对这些挑战，我们提出 Dustin，一种面向多批次长上下文推测解码的稀疏验证方法。通过融合草稿模型的前向信号与目标模型的历史注意力，Dustin 以可忽略的精度损失识别关键 token。为最小化延迟，我们采用稀疏估计方案，将重要性评分限制在一部分注意力头上，从而在保证生成质量的同时实现高速验证。我们在 LongBench（Bai 等，2024）和 PG-19（Rae 等，2019）上，针对 Llama3（AI@Meta，2024）和 Qwen2.5（Yang 等，2024）系列评估了 Dustin，在 32k 上下文长度下实现了高达 9.17 倍的解码加速，且精度损失极小。总之，我们的主要贡献如下：  

- • 我们分析了仅基于历史注意力分数或草稿模型前向注意力分数来预测重要 token 的局限性。  
- • 我们设计了一种混合 token 选择策略，并结合搜索算法，以找到最小的注意力头集合，从而在保持精度的同时降低开销。  
- • 我们的方法 Dustin 将自注意力计算加速了 27.85 倍，在 PG-19 基准测试上，使用 Qwen2.5-72B，在 32k 输入长度和批次大小 16 的条件下，实现了解码阶段 9.17 倍的加速。  

## 2 背景与相关工作  

### 2.1 推测解码  

推测解码（SD）（Leviathan 等，2023；Chen 等，2023，2024；Miao 等，2024）通过采用更快的草稿模型来草拟多个 token，然后由目标模型并行验证，从而加速自回归生成。近期研究重新审视了 SD 在长上下文和多批次场景中的应用，这些场景下推理逐渐以内存流量为主导。TriForce（Sun 等，2024）通过分层推测提升了可扩展性。MagicDec（Sadhukhan 等，2024）表明，在长上下文、大批次设置中，SD 仍能通过将完整 KV 目标验证分摊到多个草拟 token 上，并采用稀疏 KV 草稿来减轻 KV 缓存瓶颈，从而带来收益。QuantSpec（Tiwari 等，2025）通过使用量化权重/KV 进行自推测来降低草稿开销，而 LongSpec（Yang 等，2025）则设计了一种面向长上下文的草稿模型，其 KV 缓存大小恒定，并配备位置索引和注意力聚合机制，以实现高效的长上下文推测解码。  

### 2.2 KV 缓存淘汰  

长上下文推理常受限于 KV 缓存内存和注意力成本，这促使了通过仅保留或访问部分缓存 token 来减少计算的方法。一系列常见工作采用**注意力引导的淘汰**，利用最近的注意力权重作为 token 重要性的代理，来决定保留哪些 KV 条目（Xiao 等，2023；Liu 等，2023；Zhang 等，2023；Li 等，2024；Cai 等，2024；Oren 等，2024；Lin 等，2025）。然而，高注意力的 token 集合在解码过程中会发生变化。SmallKV（Zhao 等，2025）将此称为“显著性偏移问题”，并在一个小模型的辅助下缓解了该问题。与不可逆的淘汰相反，Quest（Tang 等，2024）保留完整 KV 缓存，并在每个解码步骤执行**查询感知**的选择。它使用廉价的 min/max 键统计量（避免精确的 \(QK^\top\)）对 KV 页面进行评分，并仅关注排名靠前的页面，从而在不永久移除的情况下实现稀疏注意力。  

### 2.3 推测解码中的目标端 KV 缓存压缩  

作为压缩草稿端状态的补充，另一类工作通过仅使用目标 KV 缓存的一个稀疏子集来验证草拟 token，从而降低**目标端验证**成本。SpecAttn（Shah，2025）通过使用一个小型草稿模型估计 token 重要性，使目标模型能够以 token 级的稀疏 KV 访问进行验证，从而在长上下文中减少验证时的注意力成本。这一方向与我们关注加速长上下文推理解码中**目标模型验证**阶段的目标最为相关。  

## 3 观察  

虽然推测解码（SD）加速了大批次长序列的推理，但 KV 缓存加载成本不断增加的问题依然存在。因此，KV 缓存压缩仍然至关重要。先前的研究表明，由于解码过程中 token 重要性的动态变化，永久性 KV 缓存淘汰会因**显著性偏移问题**（Zhao 等，2025）而导致显著的信息丢失。验证阶段涉及同时处理多个草拟 token。本节评估两种主要来源，用于预测并为采样范围 \([i, i+w-1]\)（考虑未来的 \(w\) 个 token）选择最重要的 \(k\) 个 token：  

1. 从目标模型先前前向传播中提取的**历史注意力分数**。  
2. 草稿模型在推测未来 token 时生成的**前向注意力分数**。  

我们展示仅基于历史或前向注意力分数预测重要 token 的局限性，并提出一种结合两者优势的混合方法，以保持生成质量。  

### 3.1 评估框架：注意力恢复率  

我们使用**注意力恢复率（ARR）**来衡量 KV 选择策略 \(\pi\) 保留注意力的程度。在解码步骤 \(i\) 处，设 \(\mathcal{V}_i\) 为有效的 KV 缓存位置，\(A_{i,j}\) 表示在 \(j \in \mathcal{V}_i\) 上的归一化注意力权重，满足 \(\sum_{j \in \mathcal{V}_i} A_{i,j} = 1\)。给定策略选择的子集 \(K_i^\pi \subseteq \mathcal{V}_i\)，ARR 定义为：  

\[
\mathrm{ARR}_i(\pi) \triangleq \sum_{j \in K_i^\pi} A_{i,j}. \tag{1}
\]

##### 窗口 ARR。  
为与 SD 的 token 块验证对齐，我们计算前向窗口长度为 \(w\) 的平均 ARR：  

\[
\mathrm{ARR}_i^{(w)}(\pi) \triangleq \frac{1}{w} \sum_{s=0}^{w-1} \mathrm{ARR}_{i+s}(\pi). \tag{2}
\]

##### SD 最优（参考上界）。  
最大化 \(\mathrm{ARR}_i^{(w)}\) 的最优子集需要访问在选取时不可用的未来注意力权重。因此，一种最优策略在完全访问未来注意力的条件下选择 \(K_i\)，作为理论上的上界：  

\[
\bar{A}_{i,j}^{(w)} \triangleq \frac{1}{w} \sum_{s=0}^{w-1} A_{i+s,\,j}, \qquad K_i^{\pi_{\mathrm{orc}}} \triangleq \operatorname{TopK}_k\big(\bar{A}_{i,\cdot}^{(w)}\big), \tag{3}
\]

从而得到  

\[
\mathrm{ARR}_i^{(w)}(\pi_{\mathrm{orc}}) = \sum_{j \in K_i^{\pi_{\mathrm{orc}}}} \bar{A}_{i,j}^{(w)}. \tag{4}
\]

直观上，ARR 是所选 KV token 保留的原始注意力权重的比例：ARR 为 1 表示所选位置覆盖了所有注意力权重，而较低的 ARR 表明压缩导致更多注意力丢失。因此，ARR 提供了一个直接度量，即在固定 KV 缓存预算下，有多少注意力信息被保留。我们进一步在附录 B 中说明，ARR 与输出 logit 的 KL 散度呈强负相关，这表明 ARR 是稀疏前向保真度的有意义的代理。实验使用 Qwen2.5-Instruct 系列模型，数据集为 LongReward（Zhang 等，2025），其平均上下文长度为 13.5k token。Qwen2.5-32B 作为主要目标模型用于时间和逐层分析，而 0.5B 变体作为草稿模型。为清晰起见，在后续图中，“KV 缓存预算”指的是用于压缩上下文的固定容量 \(k\)。  

### 3.2 目标历史信号的时间衰减  

图 2：注意力恢复率分析（历史分数）。比较使用未来平均注意力（最优）与历史注意力分数在 Qwen2.5-32B 模型上的注意力恢复率。微小差距表明时间稳定性高。

本节研究**历史注意力分数**作为未来 token 重要性低开销代理的有效性。具体而言，分析评估了在 \(w=4\) 个验证步骤中 token 显著性是否保持一致，并考虑一种基于历史的策略，利用目标模型过去的注意力。上下文 token 按 \(A_{i-\delta, \cdot}\) 排序，其中 \(\delta\) 表示解码步骤中的回顾距离。保留前 \(k\) 个 token，并与 SD 最优（公式 3）进行比较。如图 2 所示，结果表明短时一致性很强。在 KV 缓存预算 \(k=512\) 时，使用最近的历史注意力（\(\delta=1\)）达到了与最优 ARR 相差 1.04% 以内的结果。这表明 \(A_{i-1,\cdot}\) 可作为近期相关性的可靠代理。然而，验证阶段需要同时处理多个草拟 token，导致精度随草稿深度增加而下降。纯历史策略的局限性在第 5.4.2 节的消融研究中得到了进一步量化。  

### 3.3 草稿前向信号的不一致性  

图 3：注意力恢复率分析（前向分数）。比较目标模型（7B-72B）使用自身未来注意力分数（最优）与由 0.5B 草稿模型预测的前向分数之间的注意力恢复率。14B/32B 处的显著差距揭示了跨模型大小