Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

TL;DR

This work tackles the memory and latency bottlenecks of long-context inference in LLMs by exploiting attention sparsity to prune the KV cache. It introduces Self-Attention Guided Eviction (SAGE-KV), a single-pass, token-level KV-cache compression method that selects top-k evicted KV entries using the last token's query after prefill. The approach achieves near full-attention accuracy with substantial memory savings, outperforming both static and dynamic prior KV strategies on LongBench across multiple long-context models. Its simple integration and demonstrated efficiency gains make long-context inference more practical for large-scale LLM deployments.

Abstract

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.

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

TL;DR

This work tackles the memory and latency bottlenecks of long-context inference in LLMs by exploiting attention sparsity to prune the KV cache. It introduces Self-Attention Guided Eviction (SAGE-KV), a single-pass, token-level KV-cache compression method that selects top-k evicted KV entries using the last token's query after prefill. The approach achieves near full-attention accuracy with substantial memory savings, outperforming both static and dynamic prior KV strategies on LongBench across multiple long-context models. Its simple integration and demonstrated efficiency gains make long-context inference more practical for large-scale LLM deployments.

Abstract

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.

Paper Structure

This paper contains 9 sections, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. Self-Attention Guided KV Cache Eviction
  3. Experiments and Results
  4. Experimental Setup
  5. Results
  6. Conclusion and Future Work
  7. Appendix
  8. StreamLLM with Relative and Absolute Positions
  9. Implementation Details of SAGE-KV with LLMs

Figures (3)

  • Figure 1: The illustration of SAGE-KV with the single-pass KV cache selection. The full KV cache consists of four parts: initial tokens, tokens for eviction, recent tokens, and the last token. To construct a reduced KV cache, we select the top-$k$ evicted tokens based on their attention scores with the last token and concatenate them with the initial and recent tokens. The updated cache is used for continuous token generation, with each new token (green) added to the recent tokens, updating the recent window for the next step.
  • Figure 2: Token Budget Analysis.
  • Figure 3: Implementation of Absolute and Relative Positioning in StreamLLM. Relative position assigns indices within StreamLLM’s sliding window, dynamically shifting as the window moves to maintain a bounded range. In contrast, absolute position assigns a fixed index to each token based on its original sequence, increasing continuously as new tokens are added.