🤖 machine learning

Vegas: Self-Speculative Decoding with Verification-Guided Sparse Attention

Vegas 是一种自我投机解码方法，它利用验证引导的稀疏注意力机制，将识别关键 KV 缓存条目作为验证过程的副产品，从而在与现有方法相比开销极小的情况下，提高草稿标记（draft token）的接受率和解码吞吐量。

原作者： Yikang Yue, Yuqi Xue, Jian Huang

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

原作者： Yikang Yue, Yuqi Xue, Jian Huang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你是一位顶级大厨（AI），正试图写一个非常长的故事。为了写出下一个句子，你需要记住之前写下的所有内容。在 AI 的世界里，这种“记忆”被称为 KV Cache。

随着故事变得越来越长，这个记忆库也会变得巨大。每当大厨想要写下一个词时，他们都必须扫描整个历史书来寻找最重要的线索。这种扫描非常耗时且耗能，会显著降低大厨的写作速度。这就是论文中所说的“记忆瓶颈”。

旧方法：猜测与检查

为了提高速度，之前的方法尝试了一种“草拟”策略。

草拟： 大厨利用一个“捷径”（只看历史书中的几页）快速猜测接下来的几个词。
检查： 然后，大厨停下来阅读整本历史书，以查看这些猜测是否正确。
结果： 如果猜对了，太棒了！如果猜错了，这些词就会被丢弃，大厨必须重新开始。

问题所在： 用于草拟的“捷径”往往是一个糟糕的猜测。大厨可能会猜出几个词，但因为捷径忽略了重要的上下文，导致“检查”阶段拒绝了大部分猜测。大厨花了大量时间进行检查，结果却白忙一场，把工作全扔了。

新方法：Vegas

这篇论文引入了 Vegas，一种更聪明的猜测游戏方式。其核心思想很简单：利用“检查”阶段来教“草拟”阶段如何更好地进行猜测。

以下是 Vegas 的运作方式，通过几个类比来解释：

1. “免费的神谕”（隐藏的线索）

在旧方法中，“检查”阶段仅仅是一个是非判断的守门员。但论文意识到，当大厨阅读整本历史书来验证猜测时，他们已经在计算哪些历史部分是最重要的了。

Vegas 的洞察： 为什么要丢弃那个计算过程呢？Vegas 将“检查”阶段视为一个免费的老师。它说：“嘿，既然你在检查时已经弄清楚了哪些历史页面最重要，那我们就把这个清单留着，用于下一次猜测吧！”

2. “Collect-2-Query” 技巧（不要过度思考）

你可能会想：“为了制作一份完美的重要页面清单，我需要检查草稿中的每一个词。”

问题： 检查每一个词会耗费太多时间，违背了提速的初衷。
Vegas 的解决方案： 作者们发现了一个聪明的捷径。你不需要检查每一个词。你只需要看草稿中的第一个词和最后一个词（即“奖励”词）。
类比： 想象一下，你试图根据第一场戏和最后一场戏来猜测一部电影的剧情。这两场戏通常捕捉到了整部电影最重要的主题。通过只看这两个“书签式”的场景，Vegas 几乎不费吹灰之力就能获得 95% 的准确度。这就是所谓的 “Collect-2-Query” 机制。

3. 结果：更快的烹饪速度

因为 Vegas 使用“检查”的结果来引导“草拟”，大厨的猜测变得更加准确。

旧方法： 猜测 5 个词，检查它们，结果只有 2 个被接受。
Vegas： 猜测 5 个词，检查它们，结果有 4 或 5 个被接受。

由于大厨在每一轮中接受的词更多，他们在不损失任何质量的前提下，能更快地完成故事。

总结

论文声称，通过使用这种“验证引导”的方法：

速度： 它使长文本生成速度比目前的标准方法快 1.15 倍至 2.81 倍。
质量： 它是“无损”的，这意味着故事的质量与大厨每次都阅读整本书时完全相同。
效率： 它通过聪明地选择哪些部分需要查看记忆，而不是盲目地查看一切，从而解决了“记忆瓶颈”问题。

简而言之，Vegas 将“检查”步骤从一项枯燥的苦差事变成了一堂有益的课程，让 AI 能够更快地创作出长篇、复杂的故事。

技术摘要：Vegas —— 基于验证引导稀疏注意力的自投机解码

1. 问题陈述

大语言模型（LLM）的长文本推理正日益受到内存带宽瓶颈的制约，其原因在于键值（KV）缓存随上下文长度呈线性增长。虽然稀疏注意力技术通过仅处理部分关键 KV 条目来减少内存访问，但往往会因丢弃上下文而导致输出质量下降。相反，结合稀疏注意力的自投机解码试图在速度与质量之间取得平衡，即利用稀疏注意力进行草拟（drafting），并使用全量注意力进行验证（verification）。然而，现有方法存在一个根本性的权衡：

草拟准确度 vs. 开销： 简单的查询无关方法（如滑动窗口）具有较低的选择开销，但草拟准确度较差。复杂的查询感知方法虽然提高了准确度，但识别关键 KV 条目需要高昂的计算开销，这往往抵消了带来的加速收益。
脱节的过程： 当前的方法将草拟阶段和验证阶段视为独立的进程。它们仅将验证阶段用于标记（token）的接受或拒绝，却忽略了验证过程中的全量注意力计算本质上已经计算出了每个 KV 条目的精确显著性（注意力权重/对数概率）。这个“免费的预言机（free oracle）”目前并未被用于指导后续的草拟阶段。

2. 方法论：Vegas

Vegas 是一个自投机解码框架，它通过协同设计草拟和验证阶段，利用验证结果来引导稀疏注意力。

核心机制：验证引导的稀疏注意力

Vegas 不再使用独立的算法来为草拟选择 KV 条目，而是将草拟阶段视为依赖于前一次迭代的验证阶段。

验证阶段： 模型执行标准的带有全量注意力的前向传播以验证草拟标记。在此过程中，它计算所有 KV 条目的注意力对数概率（logits）。
KV 选择： Vegas 将关键 KV 条目的识别作为此验证过程的副产品。这些条目随后被用于构建下一次草拟阶段的稀疏注意力掩码。
草拟阶段： 模型使用受限于所选关键 KV 条目的稀疏注意力来草拟一系列标记。

关键算法创新

A. 聚合注意力对数概率选择

为了避免过度拟合单个标记的注意力模式，Vegas 并不仅仅根据最后一次被接受的标记来选择 KV 条目。相反，它会聚合验证阶段中所有草拟标记（包括被接受和被丢弃的标记）的注意力对数概率。

原理： 被丢弃的标记通常也具有较高的接受概率，并且在语义上仍与目标生成保持一致。将它们纳入选择过程可以捕捉到后续标记链所需的更广泛的局部上下文，从而防止后期草拟标记的接受概率快速衰减。
公式化： 所选的前缀标记集 $T$ 最大化了所有草拟标记和查询头在所有对数概率上的累加和：
$T = \arg\max_{T \subseteq P, |T|=k} \sum_{t=1}^{\gamma+1} \sum_{h \in H} \sum_{i \in T} l^{(h)}_{ti}$

B. “Collect-2-Query” 机制

收集所有草拟标记的注意力对数概率可能会引入显著的内存带宽开销，尤其是在像 FlashAttention-3 这样经过优化的算子中。Vegas 引入了“Collect-2-Query”策略来最小化这种成本：

观察： 连续标记所选择的 KV 条目具有高度重叠。随着标记间位置距离的增加，这种重叠会逐渐减小。
策略： Vegas 不再收集全部 $\gamma+1$ 个标记的对数概率，而是仅提取第一个草拟标记和奖励标记（即在最后一个草拟标记之后生成的标记）的对数概率。这两个标记代表了草拟链中最大的位置距离。
影响： 这将对数概率收集的开销降低到了一个极小的常数水平（随 2 个标记而非 $\gamma$ 缩放），同时保持了与收集所有标记相当的草拟准确度。

C. 超参数微调

Vegas 采用系统性的三步微调策略，以平衡稀疏比例（所选 KV 条目的百分比）和草拟标记数量 ( $\gamma$ )：

确定接受标记数量趋于稳定的稀疏比例。
在选定的稀疏度下，通过遍历 $\gamma$ 来最大化解码吞吐量。
微调稀疏比例，以平衡草拟延迟与准确度。

3. 实现

Vegas 在 vLLM 推理框架内实现：

算子插桩（Kernel Instrumentation）： 作者修改了 FlashAttention-3 内核，使其在 Softmax 归一化之前将注意力对数概率写入全局内存（HBM），且不占用片上共享内存。
稀疏执行： 使用自定义钩子（hook）分发注意力算子。对于草拟阶段，它使用页面大小为 1 个标记的 PagedAttention 内核，将选定的标记索引作为页面索引传递，从而仅在关键条目上进行计算。
集成： Vegas 实现了 vLLM 的投机解码提议者（proposer）API，多次调用目标模型的带有稀疏注意力的前向传播。

4. 实验结果

实验在多种模型（Qwen3 系列、gpt-oss-20b）上进行，基准测试包括 AIME25、CodeElo 和 LongBench-v2（高达 128K 上下文），并在 NVIDIA H100 GPU 上运行。

吞吐量提升：
- 对比默认 vLLM： Vegas 在解码吞吐量上实现了 1.25×–2.81× 的加速。
- 对比最先进技术（SOTA）： Vegas 相比现有的基于稀疏注意力的自投机解码方法（MagicDec 和 SpecExtend）实现了 1.15×–1.29× 的加速。
草拟准确度： 与基准方法相比，Vegas 始终能获得更高的草拟准确度（每次迭代平均接受的标记数）。例如，在 Qwen3-8B 上，其表现优于 MagicDec-Stream 1.15×–1.23×。
开销效率： “Collect-2-Query” 机制使 KV 选择的开销保持在极低水平（Qwen3-8B 为 5.9%–9.4%），显著低于 MagicDec-Quest（21.7%）和 SpecExtend（11.2%–29.1%）等查询感知型基准方法。
长上下文扩展性： Vegas 的性能增益随上下文长度增加而提升。在 LongBench-v2（96K–120K tokens）上，Vegas 在基准方法表现挣扎的情况下仍能保持高接受率，比最具竞争力的基准方法高出 18%–29% 的吞吐量。

5. 重要性与主张

论文声称，Vegas 通过从根本上重新思考草拟与验证之间的关系，代表了无损 LLM 加速领域的重大进展。

协同设计： 它证明了将草拟和验证视为统一过程，可以让系统利用验证结果这一“免费预言机”来引导稀疏注意力，从而解决了传统的“草拟准确度 vs. 选择开销”的权衡问题。
无需训练： 不同于需要辅助草拟模型或模型重训练的方法，Vegas 是一种无需训练的方法，可以直接应用于现有的离线模型。
无损质量： 通过使用全量注意力验证所有草拟标记，Vegas 保证了最终输出分布与原始全量注意力解码完全一致，确保生成质量（如幻觉问题）不会下降。
可扩展性： 该方法对于内存带宽成为主要瓶颈的长上下文场景特别有效，为上下文长度扩展至数百万个 token 的高效推理提供了一条路径。