让大模型推理更快：通过智能稀疏注意力减少显存占用
NOSA: Native and Offloadable Sparse Attention
- 设计可训练的稀疏注意力机制，限制显存与显卡间数据传输量
- 在13亿到80亿参数模型上实现最高5.04倍推理加速
- 适合需要长文本生成和高吞吐量部署的研究者与工程师
更大的推理批次虽能提升解码吞吐量，但受限于GPU显存，而显存主要被键值（KV）缓存占用。此前无需训练的KV缓存卸载方法通过将冗余上下文保留在CPU上、仅加载稀疏子集进行注意力计算来缓解问题，但常因训练-推理稀疏模式不一致导致长文本生成质量下降。同时，可训练的稀疏注意力与高效卸载不兼容，因不受限的KV访问可能引发大量跨设备传输，抵消性能收益。为此，我们提出NOSA——一种原生为KV缓存卸载设计的可训练稀疏注意力机制，显式约束CPU-GPU间KV数据传输量，从而实现低通信开销与高解码吞吐。我们进一步构建了完全适配NOSA的推断系统NOSI。在13亿、30亿、80亿参数的大语言模型上的实验证明，NOSA在通用任务、长输入和长生成任务中均优于现有基线，相较FullAttn、InfLLMv2和ShadowKV分别提升解码吞吐达5.04倍、1.92倍和1.83倍。代码已开源。
原文摘要 · Abstract (English)
Decoding throughput improvements from larger inference batches are limited by GPU memory, which is largely consumed by the key-value (KV) cache. Prior training-free KV cache offloading alleviates this by keeping redundant context on the CPU and fetching only a sparse subset for attention, but it often degrades long-generation quality due to training-inference mismatch on sparse patterns. Meanwhile, trainable sparse attention is incompatible with efficient offloading, as unconstrained KV accesses may force large CPU-to-GPU transfers and erase throughput gains. To this end, we propose NOSA, a trainable sparse attention mechanism natively designed for KV cache offloading. NOSA explicitly constrains the volume of CPU-GPU KV transfers, thereby achieving low communication overhead and high decoding throughput. We further build NOSI, a KV cache offloading inference system that fully unlocks NOSA's efficiency. Empirical results on 1,3,8B LLMs demonstrate that NOSA outperforms KV cache offloading baselines on general, long-input, and long-generation tasks, while boosting decoding throughput by up to 5.04x, 1.92x, and 1.83x over FullAttn, InfLLMv2, and ShadowKV, respectively. We release our code at https://github.com/thunlp/NOSA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。