ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition

📄 arXiv: 2402.15220v4 📥 PDF

作者: Lu Ye, Ze Tao, Yong Huang, Yang Li

分类: cs.LG, cs.CL

发布日期: 2024-02-23 (更新: 2024-08-01)

备注: ACL 2024

💡 一句话要点

提出ChunkAttention以解决长序列自注意力延迟问题

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

关键词: 自注意力 大型语言模型 推理延迟 内存优化 前缀感知 多租户服务 性能提升

📋 核心要点

  1. 现有自注意力机制在处理长序列时存在显著的推理延迟，尤其是在多租户环境中。
  2. ChunkAttention通过前缀感知机制共享键/值张量，优化内存利用率并加速自注意力计算。
  3. 实验结果显示，ChunkAttention在性能上较现有方法提升3.2-4.8倍，显著提高了推理效率。

📝 摘要（中文）

自注意力是大型语言模型（LLM）的核心组成部分，但在处理长序列时会显著增加推理延迟。在多租户LLM服务场景中，利用多个请求共享系统提示的概率，可以优化自注意力的计算和内存操作成本。本文提出了ChunkAttention，这是一种前缀感知的自注意力模块，能够在多个请求中检测匹配的提示前缀，并在运行时共享其键/值张量，从而提高KV缓存的内存利用率。通过将单一的键/值张量分解为更小的块，并将其结构化为辅助前缀树，本文设计了一种高效的自注意力内核，并实现了两阶段分区算法，以提高在共享系统提示下自注意力计算的数据局部性。实验表明，ChunkAttention在系统提示长度为1024到4096时，能够将自注意力内核的速度提升3.2到4.8倍。

🔬 方法详解

问题定义：本文旨在解决长序列自注意力计算中的推理延迟问题，现有方法在多租户场景下的计算和内存开销较大。

核心思路：ChunkAttention通过识别多个请求中共享的提示前缀，优化键/值张量的存储和访问，从而提高内存利用率和计算效率。

技术框架：整体架构包括前缀树结构的KV缓存和高效的自注意力内核，采用两阶段分区算法以提升数据局部性。

关键创新：最重要的创新在于前缀感知的KV缓存设计和两阶段分区算法，这与传统的单一KV缓存方法有本质区别。

关键设计：在设计中，ChunkAttention将键/值张量分解为小块，并通过前缀树结构进行管理，确保在共享提示时的高效访问。

🖼️ 关键图片

img_0

📊 实验亮点

实验结果表明，ChunkAttention在自注意力内核的速度上相比于最先进的实现提升了3.2到4.8倍，尤其在系统提示长度为1024到4096时表现尤为突出，展示了其在长序列处理中的显著优势。

🎯 应用场景

ChunkAttention的研究成果具有广泛的应用潜力，特别是在需要处理长文本的自然语言处理任务中，如机器翻译、对话系统和文本生成等。其优化的自注意力机制能够显著提高多租户环境下的推理效率，降低计算成本，推动大规模语言模型的实际应用。

📄 摘要（原文）

Self-attention is an essential component of large language models (LLM) but a significant source of inference latency for long sequences. In multi-tenant LLM serving scenarios, the compute and memory operation cost of self-attention can be optimized by using the probability that multiple LLM requests have shared system prompts in prefixes. In this paper, we introduce ChunkAttention, a prefix-aware self-attention module that can detect matching prompt prefixes across multiple requests and share their key/value tensors in memory at runtime to improve the memory utilization of KV cache. This is achieved by breaking monolithic key/value tensors into smaller chunks and structuring them into the auxiliary prefix tree. Consequently, on top of the prefix-tree based KV cache, we design an efficient self-attention kernel, where a two-phase partition algorithm is implemented to improve the data locality during self-attention computation in the presence of shared system prompts. Experiments show that ChunkAttention can speed up the self-attention kernel by 3.2-4.8$\times$ compared to the state-of-the-art implementation, with the length of the system prompt ranging from 1024 to 4096.