一句话总结

NVIDIA 等人提出 SparDA，一种解耦稀疏注意力架构，引入 Forecast 投影以实现前瞻 KV 块预取与计算重叠，其中每个 GQA 组使用一个 Forecast 头可降低选择开销，增加不到 0.5% 的参数，仅训练 Forecast 投影以匹配原始选择器的注意力分布，在两个稀疏预训练的 8B 模型上，与稀疏注意力卸载基线相比，实现了高达 $1.25\times$1.25× 的预填充加速、$1.7\times$1.7× 的解码加速，与非卸载稀疏基线相比，解码吞吐量提升高达 $5.3\times$5.3×。

核心贡献

• SparDA 是一种解耦稀疏注意力架构，每层添加一个 Forecast 投影来预测下一层所需的 KV 块，从而实现前瞻预取，将 CPU 到 GPU 的数据传输与当前层计算重叠。
• 该设计每个 GQA 组使用一个 Forecast 头，相比多头选择器降低了稀疏选择开销，同时增加不到 0.5% 的参数，并仅通过匹配原始选择器的注意力分布来训练 Forecast 投影。
• 在两个稀疏预训练的 8B 模型上，SparDA 相比稀疏注意力卸载基线实现了高达 1.25 倍的预填充加速和 1.7 倍的解码加速，并通过支持更大的 batch size，相比非卸载稀疏基线实现了高达 5.3 倍的解码吞吐量提升。

引言

长上下文 LLM 服务对基础设施造成压力：预填充注意力计算量大，解码期间内存带宽紧张，KV 缓存容量需求不断增长。稀疏注意力减轻了计算和带宽压力，但容量压力仍未解决，而将 KV 缓存卸载到 CPU 内存并进行选择性获取会引入 PCIe 延迟瓶颈。先前的前瞻预取方法（如 InfiniGen）试图通过使用原始隐藏状态作为未来注意力的代理来隐藏这一延迟，但当层间相似性减弱时，这些方法变得不准确。同时，稀疏选择本身会产生 O(T²) 的开销，随着上下文长度的增加可能成为主导，现有的优化通常以精度换取速度。作者提出了 SparDA，它通过训练轻量级的 Forecast 投影直接预测下一层的 KV 块选择，将稀疏选择与注意力解耦，从而实现准确的异步预取和紧凑的索引器，在不重新训练基础模型的情况下降低选择开销。

方法

作者在 InfLLM-V2 的基础上构建 SparDA，InfLLM-V2 是一种块稀疏注意力架构。InfLLM-V2 将注意力分解为初始 token、局部滑动窗口和 top-$k$k 选择的块。它将选择注意力和滑动窗口注意力合并到一个具有共享 KV 投影的稀疏注意力模块中。为了计算块级相关性分数，它采用了一个无参数的、由粗到精的压缩流水线，其中键被平均池化为重叠的压缩表示。

为了实现与层执行重叠的 KV 缓存预取以隐藏延迟，并降低稀疏选择的开销，作者引入了一个架构上的改变：每层增加一个名为 Forecast $\mathbf{F}_l$Fl​ 的投影。该投影与标准的 $\mathbf{Q}_l$Ql​、$\mathbf{K}_l$Kl​ 和 $\mathbf{V}_l$Vl​ 一起产生，以驱动一个紧凑的稀疏选择索引器。通过将稀疏选择与注意力查询解耦，Forecast 可以提前一层计算。

如下图所示：

在基线 InfLLM-V2 中，查询 $\mathbf{Q}_l$Ql​ 在同一层内同时驱动 top-$k$k 选择和稀疏注意力，将选择置于注意力的关键路径上。SparDA 解耦了这些角色。每层的线性投影 $\phi_l$ϕl​ 产生额外的 Forecast $\mathbf{F}_l$Fl​： $(\mathbf{Q}_l, \mathbf{K}_l, \mathbf{V}_l, \mathbf{F}_l) = \phi_l(\mathbf{X}_l).$(Ql​,Kl​,Vl​,Fl​)=ϕl​(Xl​). Forecast $\mathbf{F}_l$Fl​ 用于 top-$k$k 稀疏选择，而 $\mathbf{Q}_l$Ql​ 仅保留用于稀疏注意力计算。具体来说，$\mathbf{F}_l$Fl​ 与下一层的压缩键 $\widetilde{\mathbf{K}}_{l+1}$Kl+1​ 进行评分，得到的 top-$k$k 块与初始块和局部块合并，形成层 $l+1$l+1 的注意力集合。在预填充期间，Forecast 用每个 GQA 组使用一个 Forecast 头的索引器取代了昂贵的多头选择器。在解码期间，$\mathbf{F}_l$Fl​ 提前预测下一层的块，以便从 CPU 内存异步预取选定的 KV 块，有效地将 PCIe 传输隐藏在当前层计算之后。

在训练方面，作者将 SparDA 添加到现有的稀疏预训练模型中，仅训练 Forecast 投影，增加的参数量可忽略不计。受 DeepSeek DSA 的启发，Forecast 索引器通过 top-$k$k 限制和重新归一化后的 Kullback-Leibler 散度来匹配目标块注意力分数。目标和预测分数定义为： $\mathbf{S}_{l,m}^{\mathrm{tgt}} = \sum_{h=1}^{G} \operatorname{softmax}\left(\mathbf{Q}_{l,m,h} \widetilde{\mathbf{K}}_{l,m}^{\mathrm{tgt}\top} / \tau\right), \quad \mathbf{S}_{l,m}^{\mathrm{pred}} = \operatorname{softmax}\left(\mathbf{F}_{l-1,m} \widetilde{\mathbf{K}}_{l,m}^{\mathrm{pred}\top} / \tau\right),$Sl,mtgt​=∑h=1G​softmax(Ql,m,h​Kl,mtgt⊤​/τ),Sl,mpred​=softmax(Fl−1,m​Kl,mpred⊤​/τ), 其中预测分支使用来自前一层索引器的 Forecast 直接为每个 KV 头生成一个分数。训练损失计算 top-$k$k 分区分布上的 KL 散度，使索引器关注选定集合内的相对排序，同时通过一个剩余桶约束非选定块上的总质量。

为了提供细粒度的训练监督，作者在目标压缩键上使用了比推理时配置更小的核大小和步长。这种更细粒度的压缩提供了更高分辨率的监督信号，使索引器能够学习更清晰的选择。然后在计算 KL 损失之前将目标分数张量下采样到标准网格。

在推理时，SparDA 通过异步预取流水线利用 Forecast 的早期可用性。作者采用了一个基于统一虚拟寻址的持久化 Triton 内核，在专用的 CUDA 流上从固定 CPU 内存中获取选定的 KV 块。该内核维护一小组固定的 GPU 线程块，在单次启动中持续处理块传输任务，从而减少开销并限制对主计算流的干扰。为了平衡预取吞吐量和层执行速度，SparDA 采用了一种批次自适应的 CTA 分配启发式方法。随着 batch size 的增长，该方法分配更多的协作线程阵列，以最大化底层硬件的整体吞吐量。

实验

评估在 NVIDIA H100 和 A100 GPU 上，使用长上下文基准，将 SparDA 与密集、稀疏和 InfiniGen 基线在两个 8B 模型上进行比较，测量准确性和效率。SparDA 学习的 Forecast 索引器相比稀疏基线提高了准确性，尤其是在推理和长度泛化方面，同时减少了块选择开销以加速预填充和解码。其前瞻预取设计将 KV 缓存传输与计算重叠，在长上下文下带来显著的解码吞吐量提升。相比之下，无需训练的 InfiniGen 方法准确性损失严重，并且由于 CPU 端的收集操作而更慢。

SparDA 在两个模型上都提高了相比 Sparse 基线的平均准确性，在推理和 RULER 基准上提升最大。尽管密集注意力在 HELMET 和 RULER 上仍然领先，SparDA 在 RULER 上缩小了差距，并在推理任务上超越了密集注意力。InfiniGen 在所有指标上始终不如 Sparse。SparDA 的整体平均准确性高于 Sparse，在 MiniCPM4.1-8B 上提升适中，在 NOSA-8B 上提升更大。推理准确性受益最大，SparDA 在两个模型上甚至超过了密集注意力。RULER 分数随着 SparDA 提高，缩小了与密集注意力的差距，而 LongBench 基本保持不变。HELMET 准确性在 MiniCPM4.1-8B 上略有下降，但在 NOSA-8B 上有所提高。InfiniGen 在所有基准上落后于 Sparse，在 HELMET 和 RULER 上下降最大。

SparDA 在 MiniCPM4.1-8B 和 NOSA-8B 上，在每个评估的序列长度下，都一致地提高了相比 Sparse 基线的 RULER 准确性。在 NOSA-8B 上优势最为明显，差距从 32K 到 128K 稳步扩大，而 MiniCPM4.1-8B 显示出更多变化但始终为正的差距。这表明学习的 Forecast 索引器的泛化能力至少与无需训练的选择器相当，在某些情况下更好。在两个模型上，SparDA 在从 32K 到 128K 的所有序列长度下，RULER 准确性都优于 Sparse。在 NOSA-8B 上，准确性差距随上下文长度单调增长，在 128K 时达到 +4.3 个百分点。

密集注意力在短序列（32K）时实现最高的预填充吞吐量，但其二次方扩展导致随着上下文长度增加而急剧下降。SparDA 在更长上下文下超越所有方法，在 MiniCPM4.1-8B 上从 64K 开始领先，在 NOSA-8B 上从 96K 开始领先，相比 Sparse 和 Dense 的加速比随序列长度增长。这一改进源于更高效的块选择步骤，而 KV 缓存卸载本身引入的预填充开销可忽略不计。密集注意力在 32K 时领先，但其吞吐量随序列长度急剧下降，落后于稀疏方法。SparDA 在 MiniCPM4.1-8B 上从 64K 起，在 NOSA-8B 上从 96K 起，实现了最高的预填充吞吐量。在 MiniCPM4.1-8B 上 128K 时，SparDA 相比 Sparse 实现了 1.25 倍加速，相比 Dense 实现了 2.11 倍加速。加速来自一个比基线多头选择器更适应序列长度增长的块选择机制。KV 缓存卸载对预填充吞吐量的影响可忽略不计，这从 Sparse 有无卸载的几乎相同的数字可以看出。所有稀疏方法在 NOSA-8B 上都更慢，因为其查询无关的驱逐头增加了额外工作。

SparDA 的卸载设计允许解码吞吐量扩展到比非卸载基线大得多的 batch size，而非卸载基线在长上下文下 batch size 超过 16 时内存不足。在 32K 上下文下，SparDA 在 MiniCPM4.1-8B 上 batch size 128 时超过 Sparse，而 InfiniGen 始终更慢，因为其 CPU 端的 KV 收集成为瓶颈。最大的相对收益出现在中等 batch size 下，此时预取重叠最有效，SparDA 在 128K 时相比 Sparse 实现了高达 1.69 倍的加速。在 32K 下，SparDA 在 MiniCPM4.1-8B 上 batch size 128 时达到 1358.2 tok/s，超过 Sparse 的 1225.1 tok/s，而 Dense 无卸载在 batch size 16 后内存不足。卸载使 SparDA 能够运行高达 128 的 batch size，而非卸载基线在长上下文下 batch size 超过 16 时 OOM，从而相比 Dense 无卸载实现了高达 9.21 倍的吞吐量提升。InfiniGen 的解码吞吐量要低得多（例如，batch size 128、32K 时为 222.4 tok/s），因为它在将 KV 块传输到 GPU 之前先在 CPU 上收集。最大的加速出现在中等 batch size（8–16），此时 KV 预取与层执行的重叠大致平衡。

评估将 SparDA 与稀疏和密集基线以及 InfiniGen 在两个模型（MiniCPM4.1-8B 和 NOSA-8B）上进行比较，使用准确性基准、预填充吞吐量和解码吞吐量。SparDA 始终相比稀疏基线提高了准确性，在推理和 RULER 任务上提升最大，甚至在推理上超越了密集注意力，同时缩小了在 RULER 上的差距。在吞吐量方面，SparDA 由于高效的块选择在长上下文下实现了最高的预填充速度，其 KV 缓存卸载允许在更大的 batch size 下解码，相比非卸载方法带来了显著的加速。InfiniGen 在所有设置下准确性和吞吐量均落后。