💬 NLP

NOSA: Native and Offloadable Sparse Attention

本文介绍了 NOSA，一种专为 KV 缓存卸载（KV cache offloading）设计的可训练稀疏注意力机制，该机制通过约束 CPU-GPU 数据传输来解决内存效率与生成质量之间的权衡问题，从而在现有基准之上实现了显著的解码吞吐量提升。

原作者： Yuxiang Huang, Pengjie Wang, Jicheng Han, Weilin Zhao, Zhou Su, Ao Sun, Hongya Lyu, Hengyu Zhao, Yudong Wang, Chaojun Xiao, Xu Han, Zhiyuan Liu

发布于 2026-01-30

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Yuxiang Huang, Pengjie Wang, Jicheng Han, Weilin Zhao, Zhou Su, Ao Sun, Hongya Lyu, Hengyu Zhao, Yudong Wang, Chaojun Xiao, Xu Han, Zhiyuan Liu

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

想象一下，你正试图在一台极小的、高速运行的平板电脑上阅读一本长达 10 万页的巨型百科全书。这台平板电脑（你的 GPU）速度极快，但内存非常小。它一次只能同时打开几页书。而这本书的其余部分都放在房间另一头一个巨大的、缓慢的图书馆里（你的 CPU）。

每当你想要理解一个新的句子时，你的平板电脑就必须跑回图书馆，抓取你需要的特定页面，然后把它们带回来。如果你每读一个词都要跑回去一次，你就会把所有时间都花在奔跑上，而不是在阅读上。这就是当前 AI 模型在处理长文本时遇到的问题：它们陷入了数据传输的交通堵塞。

旧有的解决方案（以及它们为何失败）

1. “随机抓取”法（无需训练的卸载/Training-Free Offloading）：
一些之前的方法试图通过这种方式来解决问题：“让我们从图书馆里随机抓取几页可能重要的页面吧。”

问题在于： AI 在训练时是被教导要阅读“整本书”的，但在测试时却突然被要求只阅读随机的片段。这就像是教学生通过阅读整本教科书来准备期末考试，但考试时却只允许他们看随机的句子。学生会感到困惑、出错，尤其是在处理长篇故事时，表现会变差。

2. “智能抓取”法（可训练的稀疏注意力/Trainable Sparse Attention）：
其他方法尝试教会 AI 变得聪明：“学会精准挑选哪些页面是重要的！”

问题在于： 虽然 AI 学会了挑选正确的页面，但它并没有学会如何让去图书馆的旅程变得“高效”。它可能会挑选那些散落在建筑各处的页面。AI 仍然需要频繁地往返于图书馆，而旅途的速度（数据传输）成为了瓶颈，拖慢了整体进度。

新的解决方案：NOSA 和 NOSI

论文作者提出了一种名为 NOSA（原生且可卸载的稀疏注意力/Native and Offloadable Sparse Attention）及其配套系统 NOSI 的新系统。

可以将 NOSA 理解为一套新的 AI 学生行为准则：

“邻里”规则： 与其挑选随机的页面或散落在各处的页面，不如训练 AI 去挑选在书中位置“靠近”的页面。
类比： 想象你在读一本推理小说。如果你正在读第 50 页，你很有可能接下来需要第 49 页和第 51 页。你现在可能并不需要第 10,000 页。NOSA 教会 AI 坚守它的“邻里范围”。
益处： 因为 AI 挑选的页面是聚集在一起的，它可以一次性抓取图书馆书架上的整个“块”（chunk），而不是跑向十个不同的书架通道。这使得去图书馆的旅程变得更快、频率更低。

NOSI 是作者建造的用于运送这些“块”的高效超级运输车。

旧的卡车在移动这些特定“块”时既慢又不高效。
NOSI 卡车是专门为快速移动这些“邻里块”而定制的，它能以物理极限的速度进行运输，确保 AI 把时间花在阅读上，而不是等待卡车。

他们的发现

研究人员在不同规模（小型、中型和大型）的 AI 模型上进行了测试，发现：

质量更好： 由于 AI 被训练去挑选“邻里”页面，它不会像“随机抓取”法那样感到困惑。它能更好地理解长篇故事和复杂的推理任务。
速度极快： 通过减少去图书馆的次数并提高每次往返的效率，该系统变得异常迅速。
- 它比标准方法快了高达 5 倍。
- 它比之前最好的“智能抓取”方法快了近 2 倍。
没有妥协： 他们在实现这种速度的同时，并没有牺牲理解能力。AI 依然既聪明又快速。

简而言之

这篇论文介绍了一种通过关注信息的“邻里”而非散落的页面，来教会 AI 阅读长篇书籍的方法。这使得 AI 能更多地留在其快速内存中，并进行更少、更高效的慢速内存往返。其结果是，AI 可以比以前更快、更准确地处理海量文本，而无需更昂贵的硬件。

技术摘要：NOSA：原生且可卸载的稀疏注意力机制

问题陈述

大语言模型（LLMs）的解码吞吐量正日益受到 GPU 显存限制的约束，这主要是由于键值（KV）缓存的存储需求所致。虽然 KV 缓存卸载技术（将冗余上下文保留在 CPU 上并将其子集取回 GPU）缓解了显存压力，但现有方法面临着一个关键的权衡：

无需训练的卸载（Training-Free Offloading）： 像 ShadowKV 和 InfLLM 这样的方法通过选择无训练的稀疏子集来实现。这造成了训练-推理不匹配（training-inference mismatch），即稀疏模式在训练期间未被观测到，会导致性能显著下降，尤其是在长生成场景中，选择错误会不断累积。
可训练的稀疏注意力（Trainable Sparse Attention）： 学习稀疏模式的方法（如 InfLLMv2）避免了这种不匹配，但却与高效卸载不兼容。由于缺乏约束，这些方法可能会从 CPU 抓取任意的 KV 块，从而导致巨大的 PCIe 通信开销，抵消了吞吐量的提升。

目前的解决方案无法同时实现高解码效率（通过卸载）和高生成质量（通过训练一致性）。

方法论

作者提出了 NOSA（Native and Offloadable Sparse Attention，原生且可卸载的稀疏注意力机制），这是一种专门为 KV 缓存卸载设计的可训练稀疏注意力机制，并配合了专门的推理系统 NOSI。

1. 算法：NOSA

NOSA 在训练期间引入了局部性约束（locality constraint），以确保稀疏注意力模式在解码步骤中保持稳定，从而最大限度地减少 CPU-GPU 的数据传输。

双重选择机制（Dual-Selection Mechanism）： NOSA 将 KV 选择分解为两个部分：
- 查询感知选择（Query-Aware Selection, $\Gamma_q$ ）： 使用查询来检索特定的、可能较远的信息（保留召回率）。
- 查询无关选择（Query-Agnostic Selection, $\Gamma_e$ ）： 使用受 DMA 启发的驱逐头（eviction head）来分配独立于查询的重要性得分。该组件执行严格的驱逐策略：如果一个 token 在步骤 $t$ 未被选中，则它在未来的任何步骤 $t' > t$ 中都不能被选中。
局部性保证（Locality Guarantee）： 通过为查询无关选择分配固定的预算 $k_e$ ，NOSA 在数学上保证了最小局部性 $\gamma(t) \geq k_e/k$ 。这确保了在步骤 $t$ 中被选中的大多数 KV 块在步骤 $t-1$ 中也被选中，从而允许它们留在 GPU 缓存中。
块状实现（Block-Wise Implementation）： 为了最大化 PCIe 带宽利用率，NOSA 以 块（blocks） 而非单个 token 为单位进行操作。它采用两阶段池化策略（在子块内进行均值池化，随后在块级进行最大池化）来计算选择得分。
ED-DMA（Exp-Delayed DMA）： 为了处理驱逐头中的数值精度问题，NOSA 将指数运算推迟到注意力计算阶段，基于预指数得分进行块的选择。

2. 系统：NOSI

作者观察到，原生实现（如 Hugging Face Transformers）会引入过多的算子启动（kernel launch）开销，从而掩盖了真实的通信瓶颈。NOSI 是一个专门构建的自定义推理系统，旨在暴露并优化这些瓶颈：

算子融合（Kernel Fusion）： 将驱逐头计算、QKV 分割和池化操作融合进单个 Triton 算子中，以减少启动开销。
优化的内存布局： 将 KV 缓存组织成最小的块大小单元，并使用自定义 CUDA 算子高效计算用于块交换的集合差集（set-differences），避免了对连续内存的要求。
高带宽通信： 实现使用统一虚拟寻址（UVA）的自定义主机到设备（host-to-device）传输算子，以实现接近峰值的 PCIe 带宽，克服了标准 PyTorch 实现中的碎片化问题。

核心贡献

可行性分析： 论文指出，虽然可训练的稀疏注意力（如 InfLLMv2）具有内在的局部性，但如果没有显式约束，不足以克服 PCIe 带宽限制。
NOSA 算法： 一种新型的可训练稀疏注意力机制，通过通过强制执行基于查询无关驱逐策略的局部性约束，明确限定了 CPU-GPU KV 传输的体积。
NOSI 系统： 一个专为 NOSA 定制的推理系统，消除了软件开销，揭示了通信才是真正的瓶颈，并实现了高吞吐量的卸载。
全面评估： 在 1B、3B 和 8B 模型上的广泛基准测试表明，NOSA 优于现有的无需训练的卸载基准方法和其他可训练稀疏注意力方法。

实验结果

作者在 1B、3B 和 8B 模型上针对通用任务、长上下文输入和长生成推理任务对 NOSA 进行了评估。

性能： 在长上下文基准测试（LongBench, HELMET）中，NOSA 始终优于无需训练的卸载基准（ShadowKV, InfLLM, ArkVale）和可训练基准（InfLLMv2, DMA）。在无需训练的方法因训练-推理不匹配而导致性能显著下降的情况下，NOSA 保持了高召回率和推理能力。
吞吐量：
- 与 FullAttn 相比，NOSA 的解码吞吐量最高可提升 5.04倍（得益于卸载带来的更大 Batch Size）。
- 与 InfLLMv2 相比，其吞吐量最高可提升 1.92倍（得益于通过局部性约束减少的通信）。
- 与 ShadowKV 相比，其吞吐量最高可提升 1.83倍（得益于更高效的处理稀疏模式的能力和系统优化）。
长生成： 在长生成推理任务中，NOSA 保持了较低的困惑度（perplexity），而无需训练的卸载方法则表现出误差累积和困惑度爆炸。
消融实验： 研究证实，块状（block-wise） 选择模式对于 PCIe 效率至关重要，且 ED-DMA 驱逐头实现带来了最佳性能。

意义与主张

论文声称，NOSA 代表了从“无需训练”或“无约束可训练”向**算法-系统协同设计（algorithm-system co-design）**转变的必要趋势。通过在训练期间显式约束 KV 传输的体积，NOSA 弥合了内存效率需求（卸载）与生成质量需求（训练一致性）之间的鸿沟。

作者指出，他们的工作使得在显存受限的硬件上部署具有长上下文能力的 LLM 成为可能，且不会牺牲推理质量。他们强调，该方法通过提高解码效率降低了计算和能源成本，有助于促进 LLM 应用的更广泛部署。代码已开源，以鼓励对原生卸载机制的进一步研究。