Search papers, labs, and topics across Lattice.
This paper introduces CompressKV, a novel framework for compressing key-value (KV) caches in long-context large language models (LLMs) by leveraging Semantic Retrieval Heads (SRHs) to prioritize the retention of semantically important tokens. By focusing on the unique functionalities of attention heads rather than applying a uniform heuristic across all heads, CompressKV significantly enhances the efficiency of KV cache usage. Experimental results reveal that CompressKV maintains over 97% of full-cache performance using only 3% of the KV cache in question-answering tasks, demonstrating a substantial improvement in resource allocation for LLM inference.
CompressKV achieves over 97% performance retention with just 3% of the KV cache, revolutionizing resource efficiency in long-context LLMs.
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
