NVIDIA 的研究人员构建了一种新的 Transformer 变体。

只需对层结构做一个小改动，就能实现：

• 解码速度提升 1.7 倍
• 长链推理准确率提升 6.5 个百分点

在典型的 Transformer 架构中，每个注意力层都会计算 Q、K 和 V。

NVIDIA 的改进增加了一个第四投影，用于预测下一层需要什么。

为了理解他们为什么这样做，我们先看看当前 Transformer 架构在推理时实际发生了什么。

稀疏注意力是为了应对长上下文推理而提出的方法。现代设计不再关注每一个缓存的 token，而是按块对 KV 缓存进行评分，保留 top-k 的块，只关注这些块。

这减少了注意力计算和带宽，但仍然遗留了两个问题。

首先，KV 缓存仍然会随着每个生成的 token 而增长。

当上下文超过 100K 时，缓存不再能放入 GPU 内存，而被卸载到 CPU RAM 中。

现在每一层都必须先将其选中的 KV 块从 CPU 内存复制回 GPU。这个复制过程很慢，GPU 在等待时处于空闲状态，而且这种停顿会在每个解码步骤的每一层重复发生。

其次，选择步骤本身也不是没有代价的。

标准的选择器会使用 GQA 组（分组查询注意力，其中多个查询头共享一个 KV 头）中的每个查询头对每个候选块进行评分，然后对每个头的分数进行 softmax，并在组内求和。

在解码过程中，稀疏注意力本身开销很小，因为只有一个查询 token。

但开销大的部分是决定要关注哪些块，而且这一成本会随着上下文长度增长而增加。

这两个问题都源于当前稀疏注意力方法中的同一个设计：即注意力查询（query）驱动了块的选择。

选择需要查询向量 Q，而 Q 只在其所属层开始运行时才存在。到那时再预取数据已经来不及了。

查询还会将其多头布局带入选择过程，因此所有评分计算都只是为了做出一个 top-k 决策。

最近一篇来自 NVIDIA 和 MIT 的论文——SparDA，通过一个架构上的改动打破了这种耦合。

每个层现在输出四个投影，而不是三个：

↳ Q、K、V 和一个预测向量（Forecast）。

来自第 L 层的 Forecast 预测第 L+1 层需要哪些 KV 块。

第 L+1 层自身的查询会针对这些选定的块执行稀疏注意力。

这一个改动解决了两个问题。

由于下一层的块集合在当前层仍在计算时就已经知道，运行时会在一个独立的 CUDA 流上将这些块从 CPU 内存预取过来。

复制操作与当前层的计算重叠，因此 GPU 不再需要等待。

由于 Forecast 与注意力查询是解耦的，它不需要为每个查询头设置一个分数。

SparDA 在每个 GQA 组中使用一个 Forecast 头，从而消除了每个查询头都要进行评分循环的开销，并且完全跳过了 softmax 步骤。

DeepSeek 在 DSA 中做了类似的事情，他们使用一个小型的索引器来选择重要的 token，而不是由查询来完成这一工作。

SparDA 将同样的思想应用于块，并增加了 DSA 没有涉及的预取角度。

这一改动带来的成本很小。

Forecast 在一个 8B 模型上只增加了 33.5M 参数（0.41%），并且只训练这些投影，使用 KL 损失来匹配原始选择器的块分布。

在 MiniCPM4.1-8B 和 NOSA-8B 上，准确率与稀疏基线持平或略有提升，其中 NOSA-8B 在长链推理上提升了 6.5 个点。

预填充速度比稀疏卸载基线快 1.25 倍，解码速度快 1.7 倍。

还有一个额外的好处。

由于预取隐藏了卸载成本，大部分 KV 缓存可以放在 CPU RAM 中，释放出的 GPU 内存可以容纳更大的批次，从而将解码吞吐量提升至非卸载稀疏基线的 5.3 倍。

不过，这种前瞻性仅在解码阶段且配合 CPU 卸载时才有收益。在预填充阶段，所有键已经位于 GPU 上，因此收益仅来自更高效的选择过程。

这是论文链接：https://arxiv.org/abs/2606.04511

我写了一份关于 KV 缓存工作原理的入门级分析，详细解释了为什么模型要存储键和值、为什么缓存会随着每个 token 增长，以及有无 KV 缓存时 LLM 生成速度的对比。

请阅读下面的内容。

SparDA：面向高效长上下文 LLM 推理的稀疏解耦注意力

来源：https://arxiv.org/abs/2606.04511 查看 PDF (https://arxiv.org/pdf/2606.04511)

摘要：稀疏注意力减少了长上下文 LLM 推理中的计算和内存带宽。然而，仍然存在两个关键挑战：（1）KV 缓存容量仍随序列长度增长，卸载到 CPU 内存会引入 PCIe 传输瓶颈；（2）稀疏选择步骤本身保留了 O(T^2) 的复杂度，在长上下文时可能占据注意力计算的主导地位。我们提出了 SparDA，一种解耦的稀疏注意力架构，它在 Query、Key 和 Value 之外，为每一层引入第四个投影——Forecast。Forecast 预测下一层所需的 KV 块，从而实现前瞻性选择，将 CPU 到 GPU 的预取与当前层的执行重叠。由于 Forecast 与注意力查询解耦，我们的 GQA 实现中每个 GQA 组使用一个 Forecast 头，相比于原始的多头选择器减少了选择开销。SparDA 增加了不到 0.5% 的参数，并且仅通过匹配原始选择器的注意力分布来训练 Forecast 投影。在两个稀疏预训练的 8B 模型上，SparDA 准确率持平或略有提升，在与稀疏注意力卸载基线的对比中，预填充速度提升高达 1.25 倍，解码速度提升高达 1.7 倍。通过允许在单 GPU 上使用更大的批次大小，SparDA 进一步使解码吞吐量达到非卸载稀疏基线的 5.3 倍。我们的源代码可从以下网址获取：https://github.com/NVlabs/SparDA。

提交历史

来自：Yaosheng Fu [查看邮箱 (https://arxiv.org/show-email/481e6733/2606.04511)]
[v1] 2026年6月3日星期三 06:42:05 UTC (328 KB)