Task-KV: Task-aware KV Cache Optimization via Semantic Differentiation of Attention Heads

TL;DR

This work addresses the KV cache memory bottleneck in long-context LLM inference by introducing Task-KV, a dynamic, task-aware KV cache allocation method. It identifies heterogeneous attention heads, far from the semantic center, as carrying diverse semantic information and assigns them the full KV budget, while non-heterogeneous heads rely on a small set of sinks, recent tokens, and uniquely chosen middle activations to preserve core reasoning. The approach hinges on a semantic separator that computes head vectors from a local observation window and a budget strategy that adapts by layer via parameters $\beta$ and $m$, along with a formula for middle activations $k$. Empirically, Task-KV outperforms token- and head-level baselines on LongBench and LooGLE across Llama-2-7B-Chat and Mistral-7B-v0.2-Instruct, achieving full-context-like performance with only 40–60% KV budgets and maintaining competitive latency and memory usage, demonstrating its practical impact for scalable long-context inference.

Abstract

KV cache is a widely used acceleration technique for large language models (LLMs) inference. However, its memory requirement grows rapidly with input length. Previous studies have reduced the size of KV cache by either removing the same number of unimportant tokens for all attention heads or by allocating differentiated KV cache budgets for pre-identified attention heads. However, due to the importance of attention heads varies across different tasks, the pre-identified attention heads fail to adapt effectively to various downstream tasks. To address this issue, we propose Task-KV, a method that leverages the semantic differentiation of attention heads to allocate differentiated KV cache budgets across various tasks. We demonstrate that attention heads far from the semantic center (called heterogeneous heads) make an significant contribution to task outputs and semantic understanding. In contrast, other attention heads play the role of aggregating important information and focusing reasoning. Task-KV allocates full KV cache budget to heterogeneous heads to preserve comprehensive semantic information, while reserving a small number of recent tokens and attention sinks for non-heterogeneous heads. Furthermore, we innovatively introduce middle activations to preserve key contextual information aggregated from non-heterogeneous heads. To dynamically perceive semantic differences among attention heads, we design a semantic separator to distinguish heterogeneous heads from non-heterogeneous ones based on their distances from the semantic center. Experimental results on multiple benchmarks and different model architectures demonstrate that Task-KV significantly outperforms existing baseline methods.

Figures (13)

• Figure 1: Illustration of Task-KV compared with existing KV cache compression methods. (a) Token-level methods allocate the same KV cache budget to each attention head. (b) Head-level methods pre-identify important attention heads, but the KV cache budget among these heads remains fixed regardless of the task. (c) Our method identifies important attention heads based on the specific task and dynamically adjusts the KV cache budget among attention heads according to task semantics.
• Figure 2: For each specific layer, we use PCA to reduce the semantic vectors of different attention heads to two dimensions for visualization, allowing us to observe the differences between the semantic vectors.
• Figure 3: Results of comparative experiments in which only heterogeneous or non-heterogeneous heads were retained in different layers.
• Figure 4: Distribution of heterogeneous heads across different tasks within the Llama-2-7B-Chat model
• Figure 5: Distribution of heterogeneous heads across different tasks within the Mistral-7B-v0.2-Instruct model