NOSA：原生稀疏注意力
计算与语言
2026-01-30 v2 人工智能
机器学习
摘要
较大的推理批量可获得的解码吞吐量提升受限于GPU内存，而该内存主要被键值(KV)缓存消耗。以前的训练-free KV cache offloading通过保持冗余上下文在CPU上并仅获取稀疏子集用于注意力计算，缓解了这一问题，但由于稀疏模式的训练-推理不匹配，往往会降低长序列生成质量。与此同时，可训练稀疏注意力与高效offloading不兼容，因为不可约束的KV访问可能强制大量CPU到GPU传输，抵消吞吐量提升的优势。为了解决这一问题，我们提出了NOSA，一种原生设计用于KV cache offloading的可训练稀疏注意力机制。NOSA显式地约束CPU-GPU KV传输的体积，从而实现低通信开销和高解码吞吐量。我们进一步构建NOSI，一个完全释放NOSA效率的KV cache offloading推理系统。实验结果表明，在1.3.8B规模的LLMs上，NOSA在通用、长输入和长生成任务上均优于KV cache offloading基线方法，分别相对于FullAttn、InfLLMv2和ShadowKV提高了5.04倍、1.92倍和1.83倍的解码吞吐量。我们在https://github.com/thunlp/NOSA上发布了代码。
引用
@article{arxiv.2510.13602,
title = {NOSA: Native and Offloadable Sparse Attention},
author = {Yuxiang Huang and Pengjie Wang and Jicheng Han and Weilin Zhao and Zhou Su and Ao Sun and Hongya Lyu and Hengyu Zhao and Yudong Wang and Chaojun Xiao and Xu Han and Zhiyuan Liu},
journal= {arXiv preprint arXiv:2510.13602},
year = {2026}
}
备注
Preprint