arXiv:2503.08879cs.CLcs.AI2025-03被引 18

让大模型自己决定哪些上下文可丢,提升长文本推理效率

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

  • 利用注意力机制自动筛选可丢弃的键值缓存,分头与分词两级压缩
  • 在128K到512K上下文任务中,内存效率提升4倍且准确率更高
  • 适合追求长文本推理速度与内存优化的研究者和开发者

随着大语言模型上下文窗口扩展至128K至1M token,高效长上下文推理成为关键挑战。当前键值(KV)缓存增长与注意力计算复杂度导致显著的内存与延迟瓶颈。本文发现,在多种长上下文任务中,注意力具有稀疏性,大模型在预填充阶段后隐式知道哪些头部层级的令牌可被丢弃。基于此,提出SAGE-KV:一种简单高效的KV缓存淘汰方法。预填充后,该方法在词元与注意力头层面进行一次性Top-k选择,压缩KV缓存以实现高效推理。在LongBench及三款长上下文LLM(Llama3.1-8B-Instruct-128k、Llama3-8B-Prolong-512k-Instruct、Qwen2.5-7B-Instruct-128k)上的评估显示,SAGE-KV在保持与全注意力相当精度的同时显著提升效率:相比静态淘汰方法StreamLLM,内存效率提升4倍且准确率更高;相比动态淘汰方法Quest,内存效率提升2倍且准确率更优。

原文摘要 · Abstract (English)

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.

长上下文注意力缓存优化推理加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。