渐进式稀疏注意力：面向 LLM 服务的高效注意力的算法与系统协同设计
Progressive Sparse Attention: Algorithm and System Co-design for Efficient Attention in LLM Serving
AI总结：
本文提出渐进式稀疏注意力机制PSA，通过算法与系统协同设计自适应调整KV缓存预算，并优化GPU内存与计算流水线，在保证准确率的同时大幅降低缓存使用并提升LLM服务吞吐量。
摘要：
处理长上下文已成为现代大型语言模型（LLMs）的一项关键能力。然而，由于键值（KV）缓存的高内存开销，提供长上下文 LLM 服务伴随着显著的推理成本。现有工作利用动态稀疏注意力算法（DSAes）来减轻 KV 缓存开销，但这些算法依赖于 top-$k$ KV 缓存选择，这导致了准确率与效率之间的权衡。较大的 $k$ 会提高准确率但降低效率，而较小的 $k$ 会提升效率但损害准确率。为了克服这一权衡，本文提出了 PSA，一种 $\underline{P}$rogressive $\underline{S}$parse $\underline{A}$ttention（渐进式稀疏注意力）机制，该机制将算法创新与系统协同设计相结合，以在 LLM 服务中同时实现高推理准确率和提升的效率。PSA 算法根据不同 token 和层的实际注意力权重分布，自适应地调整其 KV 缓存预算，而不是依赖固定的预算 $k$。这使得在最小化 KV 缓存使用的同时保持高准确率。为了进一步提高执行效率，我们引入了一种流水线迭代方案，减少了 PSA 计算期间的 CPU-GPU 交错和同步开销。此外，我们实现了统一的 GPU 内存管理，通过考虑不同模型层之间不均衡的内存需求来优化 PSA 的内存利用率。广泛的实验结果表明，与最先进的 DSAes 和没有稀疏注意力的系统相比，PSA 分别将注意力计算的 KV 缓存使用量减少了高达 2.4 倍和 8.8 倍，并将端到端服务吞吐量提高了高达 1.4 倍和 2.0 倍。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
