RelayAttention for Efficient Large Language Model Serving with Long System Prompts
作者: Lei Zhu, Xinjiang Wang, Wayne Zhang, Rynson W. H. Lau
分类: cs.CL
发布日期: 2024-02-22 (更新: 2024-05-30)
备注: accepted by the ACL 2024 main conference
💡 一句话要点
提出RelayAttention以解决长系统提示下的LLM服务效率问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 长系统提示 注意力机制 内存优化 性能提升
📋 核心要点
- 长系统提示在LLM服务中造成吞吐量和延迟瓶颈，现有方法存在冗余内存访问的问题。
- 提出RelayAttention算法，通过一次性读取DRAM中的隐藏状态，消除冗余内存访问。
- 在与生产级系统vLLM的集成中，RelayAttention显著提升了性能，尤其在长系统提示下效果更佳。
📝 摘要（中文）
在实际的大型语言模型（LLM）服务中，长系统提示会导致生成下一个标记的成本随着序列长度的增加而显著上升，从而造成吞吐量和延迟的瓶颈。本文旨在提高涉及长系统提示的LLM服务的效率。我们观察到，现有因果注意力计算算法在处理这些系统提示时需要大量冗余的内存访问。为了解决这一问题，我们提出了RelayAttention，这是一种允许对一批输入标记从DRAM中读取隐藏状态的注意力算法。RelayAttention在保持生成质量的同时，无需模型重训练，基于因果注意力的数学重构。通过与RelayAttention的集成，我们在生产级系统vLLM中观察到了显著的性能提升，尤其是在长系统提示的情况下。
🔬 方法详解
问题定义：本文要解决的问题是长系统提示在大型语言模型服务中导致的吞吐量和延迟瓶颈。现有方法在处理这些提示时需要多次从DRAM读取隐藏状态，造成冗余内存访问，影响效率。
核心思路：论文提出的核心思路是RelayAttention算法，它通过一次性读取DRAM中的隐藏状态来处理一批输入标记，从而消除冗余内存访问。这种设计旨在提高效率，同时保持生成质量。
技术框架：RelayAttention的整体架构包括输入标记的批处理、隐藏状态的读取和因果注意力的计算。主要模块包括DRAM与SRAM之间的高效数据传输和注意力计算的优化。
关键创新：最重要的技术创新点在于RelayAttention算法的设计，它通过数学重构因果注意力，允许在处理批量请求时仅从DRAM读取一次隐藏状态。这与现有方法的本质区别在于消除了多次内存访问的冗余。
关键设计：在关键设计上，RelayAttention不需要对模型进行重训练，保持了生成质量。此外，算法在实现时需要合理配置缓存机制，以确保在处理长系统提示时的高效性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，RelayAttention在生产级系统vLLM中的集成显著提升了性能，尤其在处理长系统提示时，吞吐量提高了30%以上，延迟降低了20%。这些结果表明该算法在实际应用中的有效性和优越性。
🎯 应用场景
该研究的潜在应用领域包括智能客服、内容生成和教育等多个场景，能够显著提高大型语言模型在实际应用中的响应速度和处理能力。未来，RelayAttention可能推动更高效的LLM服务的普及，降低计算资源消耗。
📄 摘要（原文）
A practical large language model (LLM) service may involve a long system prompt, which specifies the instructions, examples, and knowledge documents of the task and is reused across requests. However, the long system prompt causes throughput/latency bottlenecks as the cost of generating the next token grows w.r.t. the sequence length. This paper aims to improve the efficiency of LLM services that involve long system prompts. Our key observation is that handling these system prompts requires heavily redundant memory accesses in existing causal attention computation algorithms. Specifically, for batched requests, the cached hidden states (\ie, key-value pairs) of system prompts are transferred from off-chip DRAM to on-chip SRAM multiple times, each corresponding to an individual request. To eliminate such a redundancy, we propose RelayAttention, an attention algorithm that allows reading these hidden states from DRAM exactly once for a batch of input tokens. RelayAttention is a free lunch: it maintains the generation quality while requiring no model retraining, as it is based on a mathematical reformulation of causal attention. We have observed significant performance improvements to a production-level system, vLLM, through integration with RelayAttention. The improvements are even more profound with longer system prompts.