LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention

📄 arXiv: 2502.14866v2 📥 PDF

作者: Shang Yang, Junxian Guo, Haotian Tang, Qinghao Hu, Guangxuan Xiao, Jiaming Tang, Yujun Lin, Zhijian Liu, Yao Lu, Song Han

分类: cs.CL, cs.AI, cs.DC, cs.LG, cs.PF

发布日期: 2025-02-20 (更新: 2025-04-21)

备注: Accepted by MLSys 2025. Code available at: https://github.com/mit-han-lab/omniserve

🔗 代码/项目: GITHUB


💡 一句话要点

LServe:通过统一稀疏注意力加速长序列LLM服务。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 长序列LLM 稀疏注意力 LLM服务 KV缓存 模型加速

📋 核心要点

  1. 现有LLM服务在长序列处理中面临预填充阶段的二次计算复杂度和解码阶段KV缓存的大内存占用问题。
  2. LServe通过统一的混合稀疏注意力框架,结合静态和动态稀疏模式,跳过不重要token的计算,加速长序列处理。
  3. 实验表明,LServe在保持长上下文准确性的前提下,显著加速LLM的预填充和解码过程,优于vLLM。

📝 摘要(中文)

大型语言模型(LLM)在处理长序列和复杂推理任务方面展现了卓越的潜力,但由于预填充阶段注意力计算的二次复杂度以及解码阶段KV缓存的大内存占用,高效地服务这些模型仍然具有挑战性。为了解决这些问题,我们引入了LServe,一个通过混合稀疏注意力加速长序列LLM服务的有效系统。该方法将用于预填充和解码注意力的不同硬件友好的结构化稀疏模式统一到一个框架中,其中对不太重要的token的计算以块方式跳过。LServe展示了长上下文LLM注意力中静态和动态稀疏的兼容性。这种设计通过结合这些优化实现了倍增的加速。具体而言,我们在预填充和解码阶段将一半的注意力头转换为几乎免费的流式头。此外,我们发现只需要恒定数量的KV页面即可保持长上下文和推理能力,而与上下文长度无关。然后,我们设计了一种分层KV页面选择策略,该策略基于以查询为中心的相似性动态修剪KV页面。平均而言,LServe在保持长上下文准确性的前提下,比vLLM加速LLM预填充高达2.9倍,解码加速1.3-2.1倍。代码已在https://github.com/mit-han-lab/omniserve发布。

🔬 方法详解

问题定义:现有的大型语言模型在处理长序列时,计算复杂度高,尤其是在预填充阶段,注意力的计算复杂度是序列长度的平方级别。此外,解码阶段的KV缓存需要大量的内存,限制了模型能够处理的上下文长度。这些问题导致长序列LLM服务的效率低下。

核心思路:LServe的核心思路是通过引入混合稀疏注意力机制,降低计算复杂度和内存占用。它将不同的硬件友好的结构化稀疏模式统一到一个框架中,从而能够同时优化预填充和解码阶段。通过跳过对不太重要的token的计算,减少了计算量。此外,通过动态地选择和修剪KV页面,降低了内存需求。

技术框架:LServe的整体框架包括以下几个主要模块:1) 混合稀疏注意力模块:该模块实现了统一的稀疏注意力机制,支持静态和动态稀疏模式。2) 流式注意力头:将部分注意力头转换为几乎免费的流式头,进一步降低计算成本。3) 分层KV页面选择策略:该策略基于查询相似性动态地选择和修剪KV页面,减少内存占用。

关键创新:LServe的关键创新在于统一了静态和动态稀疏注意力模式,并将其应用于长序列LLM服务中。与传统的稠密注意力机制相比,LServe通过跳过不重要的计算,显著降低了计算复杂度。此外,LServe的分层KV页面选择策略能够有效地减少内存占用,使其能够处理更长的上下文。

关键设计:LServe的关键设计包括:1) 混合稀疏注意力模式的具体选择,例如block-wise的稀疏模式。2) 流式注意力头的比例,论文中提到将一半的注意力头转换为流式头。3) 分层KV页面选择策略的相似度度量方式和阈值设定。4) KV页面数量的控制,论文中提到只需要恒定数量的KV页面即可保持长上下文和推理能力。这些参数的具体设置会影响LServe的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

LServe在长序列LLM服务中取得了显著的性能提升。实验结果表明,LServe比vLLM加速LLM预填充高达2.9倍,解码加速1.3-2.1倍,同时保持了长上下文的准确性。这些结果验证了LServe的有效性,并表明其在长序列LLM服务方面具有很大的潜力。

🎯 应用场景

LServe可应用于需要处理长序列的各种场景,例如长文本摘要、机器翻译、对话系统、代码生成等。通过提高LLM服务的效率,LServe可以降低计算成本,并支持更大规模的应用部署。该研究对于推动LLM在实际应用中的普及具有重要意义。

📄 摘要(原文)

Large language models (LLMs) have shown remarkable potential in processing long sequences and complex reasoning tasks, yet efficiently serving these models remains challenging due to the quadratic computational complexity of attention in the prefilling stage and the large memory footprint of the KV cache in the decoding stage. To address these issues, we introduce LServe, an efficient system that accelerates long-sequence LLM serving via hybrid sparse attention. This method unifies different hardware-friendly, structured sparsity patterns for both prefilling and decoding attention into a single framework, where computations on less important tokens are skipped block-wise. LServe demonstrates the compatibility of static and dynamic sparsity in long-context LLM attention. This design enables multiplicative speedups by combining these optimizations. Specifically, we convert half of the attention heads to nearly free streaming heads in both the prefilling and decoding stages. Additionally, we find that only a constant number of KV pages is required to preserve long-context and reasoning capabilities, irrespective of context length. We then design a hierarchical KV page selection policy that dynamically prunes KV pages based on query-centric similarity. On average, LServe accelerates LLM prefilling by up to 2.9x and decoding by 1.3-2.1x over vLLM, maintaining long-context accuracy. Code is released at https://github.com/mit-han-lab/omniserve.