HeadKV-R2: Head-Level KV-Cache Compression

• The paper introduces a head-level KV-cache compression method that retains only 1.5% of the original cache while maintaining around 97% of full-KV performance on contextual QA tasks.
• It employs a retrieval-and-reasoning-aware scoring mechanism to allocate unequal KV budgets across individual attention heads, enhancing efficiency in long-context inference.
• The method uses a SnapKV-style selector to retain key tokens within each head and demonstrates clear performance advantages over token-level or layer-level compression approaches.

HeadKV-R2 is a training-free, head-level KV-cache compression method for long-context LLM inference, introduced in “Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning” (Fu et al., 2024). Its central premise is that the appropriate granularity for KV-cache compression in long-context generation is the individual attention head rather than the token or the layer. The method targets contextual question answering over long inputs, especially settings that require both retrieval of supporting evidence and reasoning over that evidence. Relative to the earlier HeadKV-R variant, HeadKV-R2 replaces retrieval-only head scoring with a retrieval-and-reasoning-aware importance estimate, then uses those scores to allocate unequal KV budgets across heads while retaining tokens within each head using a SnapKV-style selector. In the reported experiments, it retains just $1.5\%$ of the KV cache while achieving about $97\%$ of full-KV performance on contextual QA at KV size $128$ (Fu et al., 2024).

1. Problem formulation and motivation

HeadKV-R2 is motivated by the KV-cache bottleneck in autoregressive decoding. For layer $l$ and head $h$, with input sequence

$$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$

the head-specific projections are

$$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$

Without compression, each head stores $n$ key vectors and $n$ value vectors, so per head and per layer the cache storage is proportional to $2nd_h$, and across all layers and heads the total KV-cache size scales as $97\%$0 (Fu et al., 2024). This linear growth in sequence length is the practical constraint that makes long-context inference memory-bound.

The method begins from the claim that token-level and layer-level compression are not sufficient for long-context contextual QA. Prior methods discussed in the paper largely operate through token eviction, fixed or dynamic per-layer budgets, pyramidal layer schedules, or local-window heuristics. HeadKV-R2 instead argues that attention heads play distinct functional roles and therefore should not receive uniform cache budgets. The paper ties this claim to contextual QA tasks that require two capabilities simultaneously: retrieval of relevant evidence from a long context and reasoning over that evidence to produce the answer (Fu et al., 2024).

The compressed representation is still expressed head-wise. For each head, the full cache

$97\%$1

is reduced to

$97\%$2

with $97\%$3. HeadKV-R2 does not prune heads from the model. All heads remain active, but different heads are allowed to retain different amounts of KV history (Fu et al., 2024).

2. Method architecture and position within HeadKV

HeadKV is the general framework for head-level KV-cache compression. Its mechanism is a hybrid of two decisions: a head-level budget allocation step and a per-head token-selection step. During prefill, the model processes the full input and forms K/V states. Instead of keeping all $97\%$4 KV entries for all heads, HeadKV assigns each head its own retained size and then selects which entries survive within that head. Heads judged more important receive larger budgets; less important heads receive smaller budgets (Fu et al., 2024).

Within this framework, HeadKV-R and HeadKV-R2 differ only in how head importance is estimated. HeadKV-R uses a prior retrieval-head distribution, whereas HeadKV-R2 uses a new retrieval-reasoning distribution designed to reflect contextual reasoning ability. Thus HeadKV-R2 is not a separate compression mechanism so much as the reasoning-aware scoring variant of HeadKV (Fu et al., 2024).

The full HeadKV-R2 pipeline has four stages. First, an offline head-importance estimation stage constructs retrieval-reasoning “needle-in-a-haystack” examples and scores heads by their attention to the correct answer tokens. Second, those scores $97\%$5 are normalized and treated as distributional weights for global allocation. Third, each head begins from a fixed initial budget $97\%$6, part of that budget is pooled globally, and the pool is redistributed proportionally to $97\%$7. Fourth, within each head, the actual retained KV entries are selected using the SnapKV mechanism: preserve the last $97\%$8 instruction tokens as a local window, use attention from that window to score earlier tokens, and keep the highest-scoring entries up to the head’s allocated budget (Fu et al., 2024).

Two common misconceptions are addressed directly by the paper. HeadKV-R2 is not head pruning, because it does not remove heads from the network. It is also training-free only in the sense that it introduces no finetuning of model parameters; it still requires an offline importance-estimation stage based on constructed calibration examples (Fu et al., 2024).

3. Retrieval-reasoning head importance estimation

The defining innovation of HeadKV-R2 is its retrieval-reasoning scoring rule. The earlier retrieval-only criterion used in HeadKV-R assigns credit to a head only when the maximally attended token falls on an inserted answer token: $97\%$9 The paper criticizes this rule on two grounds: it captures only retrieval-and-paste behavior, and it produces an overly sparse head-importance distribution, with about $128$0 of heads receiving zero score under the strict retrieval-head criterion (Fu et al., 2024).

HeadKV-R2 replaces that criterion with retrieval-reasoning examples in which the inserted answer sequence is decomposed as

$128$1

where $128$2 is a reasoning step, $128$3 is a wrong answer to the refined question, and $128$4 is the correct answer. The prompt is designed so that the model must use the reasoning information $128$5 to select $128$6 rather than the distractor $128$7. This operationalizes “contextual reasoning ability” as an attention-based proxy: heads are deemed important if they place strong attention on the correct answer sequence in tasks that require reasoning to determine what should be retrieved (Fu et al., 2024).

The resulting HeadKV-R2 score is

$128$8

This differs from the retrieval-only score in two explicit ways. First, it scores attention over the entire correct answer $128$9, not only a single copied token. Second, it uses attention-weighted multi-rank scoring rather than a binary top-1 indicator. The paper argues that this produces a denser and more task-aligned distribution of head importance for contextual QA (Fu et al., 2024).

The ablation structure in the paper isolates both components. HeadKV-ER keeps retrieval examples but adopts the improved whole-answer scoring, while HeadKV-R2 adds retrieval-reasoning examples on top of that scoring change. This separation is important because it shows that both the scoring form and the calibration example design contribute to the final gains (Fu et al., 2024).

4. Head-level budget allocation and token retention

After head scores are obtained, HeadKV-R2 performs a global head-level reallocation of KV budget. Each head starts with a fixed budget $l$0. A fraction $l$1 is removed from every head to form a shared pool: $l$2 The final budget for head $l$3 is then

$l$4

Under this rule, every head retains a base budget, while high-scoring heads receive a larger share of the pooled dynamic budget. The paper emphasizes that this is a global head-level allocation rather than a per-layer heuristic (Fu et al., 2024).

Token selection within each head is then handled by the SnapKV-style selector. The method preserves the last $l$5 instruction tokens as a local window, with $l$6. Attention from these local-window tokens to the earlier context is pooled, earlier tokens are ranked by that pooled score, and the top-scoring KV entries are retained. The retained cache for a head therefore consists of the local window, the base budget, and the dynamic budget share implied by $l$7 (Fu et al., 2024).

The paper evaluates KV sizes $l$8 and $l$9. In the reported setting, KV size denotes the initial per-head retention budget $h$0. Thus KV size $h$1 means each head begins from $h$2 before redistribution, and KV size $h$3 means each head begins from $h$4 (Fu et al., 2024). The reported improvements are largest in the low-resource regime, especially at KV sizes $h$5 and $h$6, which the paper interprets as evidence that precise head-aware allocation matters most when the budget is tight (Fu et al., 2024).

5. Empirical performance and analytical findings

The main experiments use Llama-3-8B-Instruct and Mistral-7B-Instruct, evaluated on LongBench, LooGLE, Needle-in-a-Haystack, and Reasoning-in-a-Haystack. For LongBench and LooGLE QA tasks, the appendix reports F1 as the evaluation metric. The paper’s headline claim is that retaining just $h$7 of the KV cache preserves about $h$8 of full-KV performance on contextual QA (Fu et al., 2024).

On the contextual QA average for Llama-3-8B-Instruct at KV size $h$9, the reported scores are FullKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$0, SnapKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$1, PyramidKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$2, Ada-SnapKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$3, HeadKV-R $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$4, and HeadKV-R2 $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$5. On Mistral-7B-Instruct at the same KV size, the scores are FullKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$6, SnapKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$7, PyramidKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$8, Ada-SnapKV $$X = \{x_1, x_2, \dots, x_n\} \in \mathbb{R}^{n \times d},$$9, HeadKV-R $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$0, and HeadKV-R2 $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$1 (Fu et al., 2024). These are the core quantitative results supporting the paper’s claim that head-level reasoning-aware allocation is more robust than token-level or layer-level baselines under severe compression.

The gains are stronger at KV size $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$2. For Llama-3-8B-Instruct, the contextual QA average is SnapKV $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$3, PyramidKV $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$4, Ada-SnapKV $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$5, HeadKV-R $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$6, and HeadKV-R2 $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$7. For Mistral-7B-Instruct, the corresponding scores are SnapKV $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$8, PyramidKV $$Q_h^l = XW_Q^l \in \mathbb{R}^{n \times d_h}, \quad K_h^l = XW_K^l \in \mathbb{R}^{n \times d_h}, \quad V_h^l = XW_V^l \in \mathbb{R}^{n \times d_h}.$$9, Ada-SnapKV $n$0, HeadKV-R $n$1, and HeadKV-R2 $n$2 (Fu et al., 2024). This is the regime where budget-allocation errors are most costly, and the method’s head-aware design yields its clearest advantage.

At larger budgets, the method approaches or slightly exceeds FullKV on some averages. On Llama-3-8B-Instruct contextual QA at KV size $n$3, HeadKV-R2 reports $n$4 against FullKV $n$5. On the LooGLE average at KV size $n$6, HeadKV-R2 reports $n$7 against FullKV $n$8 for Llama-3-8B-Instruct. The authors interpret these cases as evidence that reasoning-aware selection can improve focus (Fu et al., 2024).

The reasoning-specific evaluation also aligns with the method’s design. On Reasoning-in-a-Haystack at KV size $n$9, the average for Llama-3-8B-Instruct is FullKV $n$0, Ada-SnapKV $n$1, HeadKV-R $n$2, and HeadKV-R2 $n$3. For Mistral-7B-Instruct, the corresponding scores are FullKV $n$4, Ada-SnapKV $n$5, HeadKV-R $n$6, and HeadKV-R2 $n$7 (Fu et al., 2024). These results support the specific claim that R2 preserves long-context reasoning in addition to retrieval.

The ablations identify the source of improvement more precisely. On LongBench at KV size $n$8, Llama-3-8B-Instruct yields HeadKV-R $n$9, HeadKV-ER $2nd_h$0, and HeadKV-R2 $2nd_h$1; Mistral-7B-Instruct yields HeadKV-R $2nd_h$2, HeadKV-ER $2nd_h$3, and HeadKV-R2 $2nd_h$4 (Fu et al., 2024). The paper’s interpretation is that whole-answer scoring already improves over the sparse exact-match rule, and retrieval-reasoning examples add a further gain beyond the scoring refinement alone.

6. Related systems work, adjacent methods, and limitations

HeadKV-R2 is best understood as a compression policy, not a distributed-inference system. That distinction matters because head-aware budgeting induces nonuniform per-head KV retention, which is beneficial on a single device but can create deployment issues in tensor-parallel multi-GPU inference. FairKV studies exactly this failure mode: imbalanced per-head KV budgeting can produce skewed per-GPU memory usage and degraded throughput under tensor parallelism, and FairKV addresses it through best-effort assignment and Fair-Copying while leaving the underlying per-head compression policy unchanged. On LLaMA 70B and Mistral 24B, it reports up to $2nd_h$5 throughput relative to standard tensor-parallel inference, positioning itself as a systems-layer complement to HeadKV-style methods rather than a replacement for head scoring or cache selection (Zhao et al., 19 Feb 2025).

A different adjacent direction is hidden-dimension compression. ReCalKV reduces the hidden dimensionality of cached Keys and Values by low-rank factorization, uses head-wise similarity-aware reordering only for Keys, and applies offline calibration plus matrix fusion for Values. Its head-aware component is therefore head clustering for better grouped SVD, not head ranking, head pruning, or per-head budget allocation. This makes it complementary to HeadKV-R2 rather than directly comparable in objective or mechanism (Yan et al., 30 May 2025).

There is also a naming ambiguity in later literature. “Head-Aware Key-Value Compression for Efficient Autoregressive Image Generation” introduces another method called HeadKV for autoregressive image generation, but that work explicitly does not mention HeadKV-R2. Its scope is visual autoregressive generation rather than long-context textual QA, and any connection to HeadKV-R2 is only by name rather than by method lineage (Liang et al., 20 May 2026).

Several limitations of HeadKV-R2 are explicit or strongly implied in the original paper. The scoring method is designed for contextual QA requiring retrieval and reasoning, so generality to unrelated tasks is not fully established, even though appendix results on summarization, few-shot learning, synthetic, and code tasks are encouraging (Fu et al., 2024). The method is training-free but not cost-free, because it relies on an offline importance-estimation stage. The importance signal is attention-based rather than causal, gradient-based, or ablation-based, so it may not fully capture head necessity. The paper also does not introduce explicit per-task adaptation beyond the designed retrieval-reasoning examples, and it suggests that future work could use gradients or more general task-specific scoring (Fu et al., 2024).

Taken together, these connections place HeadKV-R2 in a specific part of the KV-compression landscape. It is neither a pure token selector nor a layer-budget heuristic, neither a low-rank hidden-dimension compressor nor a multi-GPU scheduling system. Its distinctive contribution is the claim that long-context contextual QA should allocate cache at the level of attention heads, and that the most useful head-importance signal for such tasks is a retrieval-and-reasoning-aware distribution rather than a retrieval-only one (Fu et al., 2024).