CompressKV: Semantic Retrieval Heads Know What Tokens are Not Important Before Generation

ICLR 2026 Conference Submission16649 Authors

Keywords: KV Cache Compression, Efficient LLM Inference

TL;DR: CompressKV identifies semantic retrieval heads that effectively retrieve critical tokens for KV cache eviction and employs a layer‑adaptive KV cache allocation strategy.

Abstract:

Recent advances in large language models (LLMs) have significantly boosted long-context processing. However, the increasing key-value (KV) cache size poses critical challenges to memory and execution efficiency. Most KV cache compression methods rely on heuristic token eviction using all attention heads in Grouped Query Attention (GQA)-based LLMs. This method ignores the different functionalities of attention heads, leading to the eviction of critical tokens and thus degrades the performance of LLMs.

To address the issue above, instead of using all the attention heads in GQA-based LLMs to determine important tokens as in the previous work, we first identify the attention heads in each layer that are not only capable to retrieve the initial and final tokens of a prompt, but also capable of retrieving important tokens within the text and attending to their surrounding semantic context. Afterwards, we exploit such heads to determine the important tokens and retain their corresponding KV cache pairs. Furthermore, we analyze the cache eviction error of each layer individually and introduce a layer-adaptive KV cache allocation strategy. Experimental results demonstrate the the proposed framework CompressKV consistently outperforms state-of-the-art approaches under various memory budgets on LongBench and Needle-in-a-Haystack benchmarks. Notably, it retains over 97% of full‑cache performance using only 3% of KV cache on LongBench’s question‑answering tasks and achieves 90% of accuracy with just 0.7% of KV storage on Needle-in-a-Haystack benchmark.

Primary Area: generative models

Submission Number: 16649