⏶4
NOSA：原生且可卸载的稀疏注意力
发表
由 
Shawn Huang 提交
Shawn Huang 提交作者: Yuxiang Huang, Chaojun Xiao, Xu Han, Zhiyuan Liu
Yuxiang Huang, Chaojun Xiao, Xu Han, Zhiyuan Liu摘要
AI 生成总结
NOSA 是一个可训练的稀疏注意力框架，通过实现高效的 KV 缓存卸载，在不影响性能的情况下提高解码吞吐量。可训练稀疏注意力已成为解决 LLM 在长上下文处理中的解码效率瓶颈的有希望的解决方案，在最小化影响任务性能的同时，显著节省内存访问。然而，现有的稀疏注意力方法尚未解决一个关键限制：键值 (KV) 缓存的大小未减小，这限制了 GPU 上的批处理大小并降低了解码吞吐量，尤其是在大规模批量推理中。在本文中，我们表明可训练稀疏注意力在相邻解码步骤的 token 选择中自然表现出强烈的局部性，从而可以在不改变底层注意力计算的情况下实现 KV 缓存卸载。然而，固有的局部性不足以实现高效卸载，因为所选 KV 对在 CPU 和 GPU 之间的传输继续主导了整体解码成本。基于此见解，我们提出了 NOSA，一个可训练稀疏注意力框架，旨在本地支持 KV 缓存卸载。NOSA 通过将 token 选择分解为查询感知和查询无关的组件来引入显式的局部性约束，从而减少 KV 传输，同时保留训练期间使用的注意力计算。我们使用 NOSA 预训练了一个 10 亿参数的模型，并进行了广泛的基准测试，表明它在保持接近无损性能的同时，与标准可训练稀疏注意力基线 (InfLLM-V2) 相比，解码吞吐量提高了高达 2.3 倍。
可训练稀疏注意力已成为解决 LLMs 长上下文处理中解码效率瓶颈的有前途的解决方案，可在最大程度地减少内存访问的同时，对任务性能的影响极小。然而，现有的稀疏注意力方法仍有一个关键限制未解决：键值（KV）缓存的大小未减小，这限制了 GPU 上的批量大小，并降低了解码吞吐量，尤其是在大规模批量推理中。在本文中，我们表明，可训练稀疏注意力在跨相邻解码步骤的 token 选择中自然表现出强大的局部性，从而可以在不改变底层注意力计算的情况下实现 KV 缓存卸载。然而，固有的局部性不足以实现有效的卸载，因为所选 KV 对在 CPU 和 GPU 之间的传输仍然占据了整体解码成本。基于这一见解，我们提出了 NOSA，一个旨在原生支持 KV 缓存卸载的可训练稀疏注意力框架。NOSA 通过将 token 选择分解为查询感知和查询无关组件来引入显式的局部性约束，从而在保留与训练期间相同的注意力计算的同时减少 KV 传输。我们使用 NOSA 预训练了一个 1B 参数模型，并进行了广泛的基准测试，结果表明，与 vanilla 可训练稀疏注意力基线（InfLLM-V2）相比，它在保持近乎无损性能的同时，解码吞吐量提高了 2.3 倍。