Title: RoPK: A Head-Level Key Cache Channel Pruning Method for Efficient Long-Context LLM Inference
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models, KV Cache Compression, KV Cache Pruning

Abstract: The substantial memory overhead of the Key-Value (KV) cache is a critical bottleneck in Large Language Models (LLMs), limiting context length and inference throughput. While prior compression techniques have targeted various dimensions of the cache, they typically assume a uniform channel allocation for all attention heads, which leads to inaccurate channel pruning and a significant drop in accuracy. This paper introduces a novel head-level key cache channel pruning method that allocates channel budgets based on a new head importance estimation algorithm derived from Rotary Position Embedding (RoPE). We posit that an effective attention head should leverage different RoPE frequencies to capture dependencies at corresponding distances—low frequencies for long-range and high frequencies for short-range. Consequently, we define head importance as the correlation score between the attention contribution of frequency dimensions and the relative textual distance. Based on this metric, our method allocates larger channel budgets to more important heads. Within each head, channels with lower attention contribution scores are pruned. This approach enables significant, non-uniform channel pruning in the key cache without sacrificing performance. Evaluations on LLaMA and Mistral models using LongBench, RULER, and Needle-in-a-Haystack benchmarks demonstrate the method's effectiveness.
With 60\% of the channels pruned, our method achieves almost no performance degradation on LongBench. At a more aggressive pruning rate of 70\%, it surpasses the state-of-the-art Think method by 2\%. Code is available in the supplementary material.

TL;DR: We propose a novel method that efficiently compresses the KV cache without sacrificing performance by allocating a different number of channels to attention heads of varying importance.