arXiv:2503.00392cs.LGcs.AI2025-03被引 7

提出渐进稀疏注意力机制,高效处理长文本上下文。

Progressive Sparse Attention: Algorithm and System Co-design for Efficient Attention in LLM Serving

  • 根据注意力分布动态调整每层的缓存预算,避免固定阈值带来的权衡。
  • 相比顶尖方法,缓存占用减少最多8.8倍,推理吞吐提升2.0倍。
  • 结合系统优化,适合高并发长文本生成场景的部署需求。

处理长上下文已成为现代大语言模型的关键能力。然而,服务长上下文大模型时,由于键值(KV)缓存的高内存开销,推理成本显著增加。现有工作采用动态稀疏注意力算法(DSAes)缓解缓存开销,但依赖于top-$k$缓存选择,导致准确率与效率之间的权衡:更大的$k$提升准确率但降低效率,更小的$k$则反之。为此,本文提出PSA(Progressive Sparse Attention),通过算法创新与系统协同设计,在保持高推理准确率的同时提升效率。PSA算法根据各令牌和层的真实注意力权重分布,自适应调整其KV缓存预算,而非依赖固定预算$k$,从而在保证精度的同时最小化缓存使用。为进一步提升执行效率,引入流水线迭代方案,减少计算过程中CPU-GPU交替与同步开销;同时实现统一的GPU内存管理,优化因各层内存需求不均导致的利用率问题。大量实验表明,相比最先进的DSAes和无稀疏注意力的系统,PSA将注意力计算的缓存占用分别减少最高2.4倍和8.8倍,端到端服务吞吐量提升最高1.4倍和2.0倍。

原文摘要 · Abstract (English)

Processing long contexts has become a critical capability for modern large language models (LLMs). However, serving long-context LLMs comes with significant inference costs due to the high memory overhead of the key-value (KV) cache. Existing work leverages dynamic sparse attention algorithms (DSAes) to mitigate the KV cache overhead, but these algorithms rely on top-$k$ KV cache selection, which results in a trade-off between accuracy and efficiency. A larger $k$ improves accuracy but decreases efficiency, while a smaller $k$ boosts efficiency but compromises accuracy. To overcome this trade-off, this paper presents PSA, a $\underline{P}$rogressive $\underline{S}$parse $\underline{A}$ttention mechanism that integrates algorithmic innovations with system co-design to achieve both high inference accuracy and improved efficiency in LLM serving. The PSA algorithm adaptively adjusts the KV cache budget of different tokens and layers according to their real attention weight distributions, rather than relying on a fixed budget $k$. This enables high accuracy while minimizing KV cache usage. To further enhance execution efficiency, we introduce a pipelined iteration scheme that reduces CPU-GPU interleaving and synchronization overhead during PSA computation. Additionally, we implement unified GPU memory management that optimizes PSA's memory utilization by accounting for uneven memory requirements across different model layers. Extensive experimental results demonstrate that PSA reduces KV cache usage for attention computation by up to 2.4$\times$ and 8.8$\times$, and increases end-to-end serving throughput by up to 1.4$\times$ and 2.0$\times$, compared to state-of-the-art DSAes and systems without sparse attention, respectively.

注意力机制长文本生成系统优化大模型推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。