🤖 AI Summary
This work addresses the high memory overhead of key-value (KV) caching in long-context large language model inference, which hinders deployment on resource-constrained devices. Existing compression methods often overlook functional heterogeneity among attention heads, leading to critical information loss. To overcome this limitation, the authors propose CompressKV, a novel framework that introduces the concept of “semantic retrieval attention heads” for the first time. By identifying such heads to guide KV cache retention and incorporating offline-estimated layer-wise pruning errors to dynamically allocate cache budgets, CompressKV achieves remarkable efficiency. On LongBench, it retains over 97% of the original performance with only 3% of the KV cache, and in Needle-in-a-Haystack tasks, it attains 90% accuracy using merely 0.7% of the original storage—significantly outperforming current compression strategies.
📝 Abstract
Long-context large language model (LLM) inference is increasingly constrained by the memory footprint and decoding cost of key-value (KV) caches, limiting sustainable deployment on resource-constrained hardware. Existing KV cache eviction methods typically apply heuristic token scoring over all heads in GQA-based LLMs. These methods ignore the different functionalities of attention heads, leading to the eviction of critical tokens and thus degrading the performance of LLMs. To address this issue, we propose CompressKV, a resource-efficient KV-cache compression framework for GQA-based LLMs. Instead of aggregating attention scores from all heads, CompressKV identifies Semantic Retrieval Heads (SRHs) that capture both the initial and final tokens of a prompt and semantically important mid-context evidence, and uses them to select tokens whose KV pairs should be retained. Furthermore, CompressKV allocates cache budgets across layers according to offline estimates of layer-wise eviction error. Experiments on LongBench and Needle-in-a-Haystack show that CompressKV consistently outperforms existing KV-cache eviction methods across memory budgets. Notably, it preserves over 97\% of full-cache performance using only 3\% of the KV cache on LongBench question-answering tasks and achieves 90\% accuracy with just 0.7\% KV storage on Needle-in-a-Haystack. These results demonstrate an improved resource--performance trade-off for long-context LLM inference. Our code is publicly available at: https://github.com/TUDa-HWAI/CompressKV