LLMs Know What to Drop: Self-Attention Guided KV Cache Eviction for Efficient Long-Context Inference

📄 arXiv: 2503.08879v1 📥 PDF

作者: Guangtao Wang, Shubhangi Upasani, Chen Wu, Darshan Gandhi, Jonathan Li, Changran Hu, Bo Li, Urmish Thakker

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

发布日期: 2025-03-11

💡 一句话要点

提出SAGE-KV，利用自注意力指导KV缓存淘汰，提升长文本LLM推理效率。

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

关键词: 长文本推理 KV缓存 自注意力机制 模型压缩 内存效率

📋 核心要点

  1. 长文本LLM推理面临KV缓存过大和注意力计算复杂度高的挑战，导致内存瓶颈和推理延迟。
  2. SAGE-KV利用自注意力机制的稀疏性，在预填充后选择性地淘汰不重要的token和head，压缩KV缓存。
  3. 实验表明，SAGE-KV在保持准确率的同时，显著提升了内存效率，优于StreamLLM和Quest等现有方法。

📝 摘要（中文）

随着大型语言模型(LLMs)采用128K到1M tokens的上下文窗口，高效的长文本推理至关重要。然而，不断增长的键值(KV)缓存和注意力机制的高计算复杂度在内存使用和延迟方面造成了显著的瓶颈。本文发现，在各种长文本任务中，注意力表现出稀疏性，并且LLMs在预填充阶段之后隐式地“知道”哪些tokens可以在head级别被丢弃或淘汰。基于这一洞察，我们提出了一种简单而有效的KV淘汰缓存方法，即自注意力引导淘汰(SAGE-KV)，用于长文本推理。在预填充之后，我们的方法在token和head级别执行一次性的top-k选择，以压缩KV缓存，从而实现使用减少的缓存进行高效推理。在LongBench和三个长文本LLMs (Llama3.1-8B-Instruct-128k, Llama3-8B-Prolong-512k-Instruct, and Qwen2.5-7B-Instruct-128k)上的评估表明，SAGE-KV在保持与完整注意力相当的准确性的同时，显著提高了效率。具体而言，SAGE-KV实现了比静态KV缓存选择方法StreamLLM高4倍的内存效率，并且准确性更高，比动态KV缓存选择方法Quest高2倍的内存效率，并且准确性更好。

🔬 方法详解

问题定义：论文旨在解决长文本LLM推理过程中，由于KV缓存过大导致的内存效率低下和推理延迟问题。现有方法，如StreamLLM和Quest，在KV缓存选择方面存在不足，要么是静态选择导致精度损失，要么是动态选择计算开销大。

核心思路：论文的核心思路是利用LLM自注意力机制的内在特性，即注意力具有稀疏性，并且LLM本身能够识别哪些token和head对于后续推理是不重要的。通过在预填充阶段后进行一次性的重要性评估和淘汰，可以有效地压缩KV缓存，从而提高推理效率。

技术框架：SAGE-KV方法主要包含两个阶段：预填充阶段和KV缓存压缩阶段。在预填充阶段，模型正常处理输入序列，生成完整的KV缓存。在KV缓存压缩阶段，SAGE-KV基于自注意力权重，对每个token和head的重要性进行评估，并选择保留最重要的top-k个token和head，淘汰其余部分。压缩后的KV缓存用于后续的解码阶段。

关键创新：SAGE-KV的关键创新在于利用自注意力权重作为token和head重要性的度量标准。与现有方法相比，SAGE-KV无需额外的训练或微调，可以直接应用于现有的LLM。此外，SAGE-KV采用一次性的选择策略，避免了动态选择带来的额外计算开销。

关键设计：SAGE-KV的关键设计包括：(1) 使用平均注意力权重作为token重要性的度量；(2) 使用head的平均注意力权重作为head重要性的度量；(3) top-k选择的比例，即保留多少比例的token和head。这些参数需要根据具体的模型和任务进行调整，以达到最佳的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SAGE-KV在LongBench数据集上，与Llama3.1-8B-Instruct-128k, Llama3-8B-Prolong-512k-Instruct, and Qwen2.5-7B-Instruct-128k等模型结合使用时，能够在保持与完整注意力相当的准确率的前提下，实现比StreamLLM高4倍的内存效率，并且准确性更高，比Quest高2倍的内存效率，并且准确性更好。这证明了SAGE-KV在长文本LLM推理效率方面的显著优势。

🎯 应用场景

SAGE-KV可广泛应用于需要处理长文本的LLM应用场景，例如长文档摘要、代码生成、对话系统和知识库问答等。该方法能够显著降低内存需求，提高推理速度，使得在资源受限的设备上部署长文本LLM成为可能，并降低了大规模应用的成本。

📄 摘要（原文）

Efficient long-context inference is critical as large language models (LLMs) adopt context windows of ranging from 128K to 1M tokens. However, the growing key-value (KV) cache and the high computational complexity of attention create significant bottlenecks in memory usage and latency. In this paper, we find that attention in diverse long-context tasks exhibits sparsity, and LLMs implicitly "know" which tokens can be dropped or evicted at the head level after the pre-filling stage. Based on this insight, we propose Self-Attention Guided Eviction~(SAGE-KV), a simple and effective KV eviction cache method for long-context inference. After prefilling, our method performs a one-time top-k selection at both the token and head levels to compress the KV cache, enabling efficient inference with the reduced cache. Evaluations on LongBench and three long-context LLMs (Llama3.1-8B-Instruct-128k, Llama3-8B-Prolong-512k-Instruct, and Qwen2.5-7B-Instruct-128k) show that SAGE-KV maintains accuracy comparable to full attention while significantly improving efficiency. Specifically, SAGE-KV achieves 4x higher memory efficiency with improved accuracy over the static KV cache selection method StreamLLM, and 2x higher memory efficiency with better accuracy than the dynamic KV cache selection method Quest.