PaperlayerView exact revision on arXiv

Source-linked AI summary

ChampionKV: Head-Wise Adaptive Budget Allocation for KV Cache Compression

Yuechuan Chen, Yanda Guo, Xiaoxi Liu, Weilin Ruan, Huanbo Xu, Ruyi Wang, Wei Zhu

arXiv:2610.18026v1cs.LGcs.AI

TL;DR

Long-context KV cache compressors use prompt attention to select important tokens but may leave input-dependent head-wise budget demand unaddressed. CHAMPIONKV uses a maximum-attention concentration proxy to adapt budgets without changing token selection. On LongBench and RULER, its SnapKV and PyramidKV variants consistently improve over their base compressors, with significant gains on the tested 70B model under tight budgets.

  • Problem

    Observation-window compressors use prompt attention to identify important tokens but may not adapt budgets to input-dependent differences in head-wise demand.

    Source: §1 · p. 1

  • Method

    CHAMPIONKV uses the maximum aggregated observation-window attention score to adapt head-wise budgets, with abnormal-head masking and floor-budget protection, while leaving token selection unchanged.

    Source: §1 · pp. 1–2

  • Results

    On LongBench and RULER, Champion variants consistently improve over their SnapKV and PyramidKV bases, with statistically significant gains on the 70B model under tight budgets; tested 7B/8B results are competitive but mixed.

    Source: §6 · p. 8

  • Takeaways & Limitations

    CHAMPIONKV is a plug-in for upgrading observation-window compressors in scarce-budget, heterogeneous-head regimes, rather than replacing the compressor or global cross-head selection.

    Source: §5.3.4 · p. 8

    The evaluation covers limited models and tasks, does not include repeated-trial sampling-based significance tests, and inherits the underlying compressor’s token-selection limitations.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Some observation-window-based KV cache compressors use prompt attention to identify important tokens but allocate uniform or predefined budgets across heads, overlooking input-dependent differences in budget demand. We propose ChampionKV, an input-adaptive head-wise budget allocator for long-context KV cache compression. It uses the maximum aggregated observation-window attention score as a lightweight empirical proxy for attention concentration, assigning larger budgets to heads with dispersed attention and smaller budgets to concentrated heads. An abnormal-head mask and a floor-budget mechanism improve robustness under weak attention responses. ChampionKV plugs into compressors such as SnapKV and PyramidKV without changing their token-selection procedures. On LongBench and RULER, both Champion variants consistently improve over their SnapKV and PyramidKV base compressors; the improvements are largest and statistically significant on Llama3.1-70B-Instruct, especially under tight budgets, whereas results on the tested 7B/8B models remain competitive but mixed. ChampionKV does not uniformly outperform Ada-KV, whose global cross-head Top-K selection retains an advantage in several settings. Controlled analysis further shows improved retained attention mass over the corresponding base compressors at negligible overhead.

1 Introduction

KV caches become a major memory bottleneck as long-context inference grows, while existing compressors do not fully address input-dependent differences in head-wise budget demand. CHAMPIONKV adapts budgets using observation-window attention and improves on its SnapKV and PyramidKV bases, with the clearest gains on the tested 70B model under tight budgets.

1 Introduction

  • KV cache size grows with context length, batch size, model depth, and KV-head dimensionality, making it a major memory bottleneck for long-context inference.Compression addresses this by retaining KV entries most useful for future generation.Source: §1 · p. 1
  • The paper investigates whether a lightweight per-head statistic available during prefilling can support input-adaptive budget allocation without changing the underlying token-selection procedure.Its contributions include an inverse-max allocation rule, abnormal-head masking, floor-budget protection, and evaluation of downstream performance, retained attention mass, and hardware efficiency.Source: §1 · pp. 1–2
  • CHAMPIONKV uses the maximum aggregated observation-window attention score as a concentration proxy, assigning larger budgets to dispersed heads and smaller budgets to concentrated heads.An abnormal-head mask and floor-budget mechanism improve robustness under weak or noisy attention responses, without changing token selection.Source: §1 · p. 1, Figure 1 · p. 2
  • Champion variants consistently outperform their SnapKV and PyramidKV base compressors, with statistically significant gains on the 70B model under tight budgets; tested 7B/8B results are competitive but mixed.Neither variant uniformly surpasses Ada-KV.Source: §1 · p. 2

2 Related Work

KV cache compression research spans token selection, layer-wise allocation, and head-wise allocation. CHAMPIONKV belongs to the head-wise line, adapting budgets from current-input attention while leaving the compressor’s token-selection procedure unchanged.

2.1 Taxonomy of KV Cache Compression

  • Observation-window methods estimate historical-token importance from recent prompt attention, while layer-wise methods allocate different cache budgets across Transformer layers.SnapKV uses end-of-prompt observation-window attention for per-head Top-K selection; PyramidKV allocates non-uniform budgets across layers.Source: §2.1 · p. 2

2.2 Head-Wise KV Cache Compression

  • Head-wise methods target intra-layer allocation, reflecting heterogeneity in head importance and function.Studies describe specialized induction and retrieval heads, and FastGen selects among discrete compression strategies for different head patterns.Source: §2.1 · p. 2, §2.2 · pp. 2–3
  • Ada-KV uses input-dependent attention for global cross-head Top-K selection, whereas HeadKV allocates budgets using offline head-importance estimates.FastGen profiles head patterns but primarily selects among discrete strategies rather than continuously allocating budgets.Source: §2.2 · p. 3
  • CHAMPIONKV adapts budgets during prefilling with a scalar concentration proxy and inverse-max weighting, rather than global cross-head token ranking or offline importance estimates.It follows observation-window token selection and can be combined with layer-wise allocation.Source: §2.2 · p. 3

3 Pilot Experiment

Observation-window attention varies across heads, motivating head-wise budgets that reflect attention concentration. The proposed maximum-score proxy provides a lightweight basis for this allocation, but can be unreliable for low-response outliers.

3.1 Head-Wise Distribution Heterogeneity

  • The last W prompt tokens form the observation window; their attention to historical tokens is aggregated into a score vector used to estimate historical KV-entry importance.The historical-token attention matrix excludes the observation-window positions.Source: §3.1 · p. 3
  • Aggregated observation-window attention scores vary in standard deviation across layers and heads, indicating heterogeneous attention concentration.Figure 2 shows this variation for Llama3.1-8B-Instruct on NarrativeQA.Source: §3.1 · p. 3, Figure 2 · p. 3
  • Concentrated heads place most attention mass on a few historical tokens, while dispersed heads spread scores across many tokens and are harder to cover with small budgets.This motivates assigning larger budgets to dispersed heads than to concentrated heads.Source: §3.1 · p. 3

3.2 Motivation for Champion-Based Budget Allocation

  • Within a head, selecting its highest-scoring entries maximizes retained attention mass for that head, making allocation of the limited layer budget across heads the central problem.The stated motivation draws on an upper bound for eviction loss based on attention-output perturbation.Source: §3.2 · p. 4
  • The maximum aggregated attention score serves as a lightweight ranking-level proxy for concentration, and smaller values lead to larger adaptive budgets among unmasked heads.The proxy is motivated by cases where second-order statistics are dominated by the maximum score; the cited analysis also identifies low-response outliers where it can be unreliable.Source: §3.2 · p. 4

4 Methodology

The method computes a champion score from each head’s aggregated observation-window attention and distributes historical-cache budget adaptively. Floor budgets and an abnormal-head mask constrain that allocation before per-head token selection.

4 Methodology

  • The champion score is the maximum aggregated observation-window attention score, and remaining budget is allocated among non-abnormal heads using inverse-max weights.The allocation applies only within the non-abnormal head set.Source: §4 · pp. 4–5
  • Each head receives a fixed floor budget controlled by ρ, which trades guaranteed historical-token retention against budget available for adaptive redistribution.The paper uses ρ = 0.2 as a practical default; the preferred balance can vary across models and workloads.Source: §4 · p. 4
  • Abnormal heads retain only the floor budget and are excluded from adaptive competition; the threshold is an empirical default whose suitability may depend on model and context length.This prevents very small champion scores from receiving disproportionately large inverse-max weights.Source: §4 · pp. 4–5
  • Unused budget from rounding or per-head allocation limits enters a shared residual pool and is reassigned by non-abnormal-head weight order.Reassignment continues until the budget is exhausted or eligible heads reach their upper bounds.Source: §4 · p. 5
  • Within each head, local pooling smooths aggregated attention scores before Top-K selection using that head’s allocated budget.The selected historical-token indices are obtained from the smoothed score vector.Source: §4 · p. 5

11 end

The procedure selects historical tokens within each head, retains recent tokens, and uses the resulting indices to form the compressed KV cache.

11 end

  • For each head, the procedure selects Top-bl,h historical indices using the allocated budget.The algorithm lists this selection after assigning leftover units to eligible heads in descending weight order.Source: §11 · p. 5
  • The final retained index set combines the selected historical indices with the recent-token window.The recent tokens are always retained.Source: §11 · p. 5
  • Indexing each head’s key and value tensors with the retained set yields the compressed KV cache.The complete procedure is shown in Algorithm 1.Source: §11 · p. 5

5 Experiments

Across LongBench and RULER, Champion variants improve over their SnapKV and PyramidKV bases most clearly on the tested 70B model, while results on smaller models are mixed. Controlled analyses find higher retained attention mass with head-wise reallocation, alongside model- and budget-dependent robustness and efficiency findings.

5.1 Experimental Setup

  • The evaluation uses LongBench and RULER with three 7B/8B/70B models, comparing Champion-adjusted SnapKV and PyramidKV against their bases and Ada-KV.LongBench budgets are 64, 128, 512, and 2048, with W = 32; comparisons preserve each base compressor’s token-selection procedure.Source: §5.1 · pp. 5–6, §5.2 · p. 6

5.2 Main results

  • On Llama3.1-8B-Instruct with a 1024-token budget, Champion-Snap and Champion-Pyramid use 0.375 GB and 0.365 GB less peak GPU memory than Ada-KV, with comparable generation time.SnapKV retains the lowest latency and total time; this profile used a 32,100-token input on an H200.Source: §5.2 · p. 6

5.3.1 Retained Attention Mass Coverage

  • With token selection fixed, Champion allocation raises retained attention mass across all six model–compressor pairs and moves closer to the upper-bound reference.The controlled comparison uses LongBench prompts, a 512-token cache budget, and W = 32.Source: §5.3.1 · pp. 6, 8, Figure 4 · p. 7

5.3.2 Robustness Analysis of Hyper-parameters

  • The default (ρ, τ) = (0.2, 0.01) is robustness-oriented rather than universal; abnormal heads comprise 0.57–1.66% on Llama3.1-8B and 0.81–1.50% on Mistral-7B.Preferred hyperparameter settings vary across models and budgets.Source: §5.3.2 · p. 8

5.2 Main results

  • On the tested 70B model, Champion variants improve over their base compressors with significant LongBench gains and Champion-Snap RULER gains, while 7B/8B results remain mixed.Ada-KV still leads several settings; Champion-Snap’s RULER advantage is clearest at lower retention ratios.Source: §5.2 · p. 6, §5.3.3 · p. 8

5.3.3 Paired Consistency Analysis

  • Paired tests find significant LongBench gains over both baselines on 70B (p < 0.001), but insignificant LongBench deltas on 7B/8B; 70B gains peak at tight budgets.Champion-Snap also has a significant RULER gain with zero losing pairs; gains diminish as retention approaches 80%.Source: §5.3.3 · p. 8

5.3.4 Discussion

  • The champion-score proxy needs no offline estimation and adds negligible overhead, positioning ChampionKV as a plug-in upgrade for observation-window compressors under scarce budgets and heterogeneous heads.It does not replace the underlying compressor or global cross-head selection.Source: §5.3.4 · p. 8

6 Conclusion

ChampionKV reallocates cache budgets across heads using an attention-concentration proxy while leaving token selection unchanged. It improves on its base observation-window compressors most clearly on the tested 70B model, and is presented as a plug-in upgrade rather than a replacement for those compressors.

6 Conclusion

  • ChampionKV combines an attention-concentration proxy with inverse-max allocation, abnormal-head masking, and floor-budget protection without changing token selection.The proxy uses the maximum aggregated observation-window attention score during prefilling, with no offline estimation or retraining.Source: §6 · p. 8
  • On LongBench and RULER, Champion-Snap and Champion-Pyramid improve over their base compressors, with significant gains on the tested 70B model and mixed results on tested 7B/8B models.Gains are strongest under tight budgets and moderate retention; the controlled analysis also finds improved retained attention mass over matched baselines.Source: §6 · p. 8
  • ChampionKV is a practical plug-in for observation-window compressors that complements, rather than replaces, global cross-head selection.Its allocation is decoupled from token scoring.Source: §6 · p. 8

Limitations

The conclusion is limited by proxy reliability, transfer of fixed calibration settings, evaluation scope, and inherited token-selection constraints.

Limitations

  • The maximum-attention score is an empirical proxy, and abnormal-head masking cannot prevent every unreliable case.Source: Limitations · p. 9
  • Fixed ρ and τ may not transfer across models or workloads, motivating adaptive calibration.Source: Limitations · p. 9
  • Evaluation covers limited models and tasks, lacks repeated-trial sampling-based significance tests, and does not assess generation-heavy, long-chain-of-thought, or multi-turn settings.Source: Limitations · p. 9
  • Because ChampionKV leaves token scoring unchanged, it inherits the token-selection limitations of its underlying compressor.Source: Limitations · p. 9

Declaration for Using AI Assistants

The authors describe AI-assisted language polishing and retain responsibility for the paper. The appendix characterizes concentrated and dispersed attention patterns as relevant to head-wise cache allocation.

Declaration for Using AI Assistants

  • AI assistants supported language polishing, grammar correction, and clarity; the authors state they did not generate research ideas, analyze data, or draw conclusions.The authors say all suggestions were reviewed and verified and that they take responsibility for the paper.Source: Declaration for Using AI Assistants · p. 9

A.1 Head-Wise Distribution Heterogeneity

  • The appendix analyzes prefilling-stage aggregated observation-window attention on Llama3.1-8B-Instruct using NarrativeQA as a representative dataset.Source: §A.1 · p. 11
  • Concentrated heads focus attention on a few historical tokens and may need relatively few KV entries, whereas dispersed heads spread attention more broadly and motivate larger budgets.The appendix relates the two patterns to attention that is quick or difficult to accumulate.Source: §A.1 · p. 11

A.2 Validation of the Max-Attention Proxy

The max aggregated observation-window attention score closely tracks attention-distribution standard deviation across tested models, window sizes, and sequence lengths. The proxy is useful for head-wise allocation, but weak-response outliers constrain its reliability.

Limitations of the proxy

  • The max-score approximation is unreliable for some heads, especially low-response outliers, so those abnormal heads receive only the floor budget.The proxy assumes sparse or approximately unimodal attention whose second-order statistics are dominated by the maximum score.Source: §A.2 · p. 12

Computational comparison

  • Unlike Ada-KV’s global cross-head Top-K selection, ChampionKV uses a lighter head-wise proxy that avoids large-scale cross-head concatenation and selection.Both finding the maximum and computing standard deviation require a linear scan, but the maximum avoids second-moment computation and integrates more easily with attention-score accumulation or Top-K procedures.Source: §A.2 · pp. 12–13

Correlation between maximum attention and standard deviation

  • On NarrativeQA with Llama3.1-8B-Instruct at W = 8, the maximum score and standard deviation correlate at r = 0.9993.Their heatmaps also show consistent patterns, and the scatter fit is approximately linear.Source: §A.2 · p. 13, Figure 7 · p. 14

Robustness across model architectures and scales

  • Across Mistral-7B-Instruct-v0.3 and Llama3.1-70B-Instruct, the measures remain highly consistent and near-linearly related, with Pearson correlations close to one.This supports the proxy across the tested architectures and model scales.Source: §A.2 · p. 13, Figure 8 · p. 14, Figure 9 · p. 14

Robustness across observation window sizes

  • Across Llama3.1-8B-Instruct analyses at window sizes W = 4, 16, 32, and 64, maximum attention score and standard deviation remain highly linearly correlated, indicating the proxy does not depend on a fixed observation window size.Source: §A.2 · p. 13, Figure 10 · p. 14, Figure 11 · p. 14

Scaling relation with sequence length

  • Across tested sequence lengths, the fitted slope remains close to its theoretical value; across LongBench datasets and model scales, it remains close to 1/√L.The sequence-length validation uses NarrativeQA with Llama3.1-8B-Instruct; the dataset validation covers three model scales.Source: Table 2 · p. 13, §A.2 · p. 15

B Appendix: Additional Experimental Results

Additional experiments show mixed results on the tested 7B/8B models, alongside configuration-specific stability and profiling findings. Hyperparameter defaults are practical rather than universal, and the long-context slice does not establish performance in untested generation settings.

Hyperparameter analysis

Hardware profiling across input lengths

  • Champion variants use slightly less peak GPU memory than Ada-KV across the profiled input lengths, but timing differences vary by method and length.The profiling uses Llama3.1-8B-Instruct, a 1024 cache budget, 15 generated tokens, and NVIDIA H200; it is configuration-specific, not a general speedup claim.Source: B Appendix: Additional Experimental… · p. 16, Table 10 · p. 18

Abnormal-head prevalence

Long-context slice

LongBench results on 7B/8B-scale models

  • On LongBench, Ada-KV leads the Llama3.1-8B averages at budgets 64, 128, and 2048 and all four Mistral budgets; Champion variants win selected tasks but not the strongest average.Champion-Pyramid scores 43.08 versus Ada-KV’s 43.42 on Llama3.1-8B at budget 2048; Champion-Snap scores 47.04 versus Ada-KV’s 47.27 on Mistral at that budget.Source: B Appendix: Additional Experimental… · p. 17, Table 14 · p. 20, Table 15 · p. 20

RULER results on 7B/8B-scale models

C Appendix: Detailed Information of Benchmarks

The appendix describes two benchmarks used to evaluate long-context performance retention. LongBench spans six task domains, while RULER comprises 11 subtasks focused on retrieval and aggregation.

C Appendix: Detailed Information of Benchmarks

C.1 LongBench Benchmark

  • LongBench covers 16 English datasets across single-document QA, multi-document QA, summarization, few-shot learning, synthetic tasks, and code completion.The task descriptions include long-document fact retrieval, cross-document reasoning, global summarization, in-context learning, localization, and repository-level code completion.Source: §C.1 · p. 23

C.2 RULER Benchmark

  • RULER includes 11 subtasks covering single- and multi-key retrieval, multi-value and multi-query aggregation, word extraction, and variable tracking.It is presented as a more challenging and flexible evaluation than the foundational Needle-in-a-Haystack test.Source: §C.2 · pp. 23–24
Loading 2610.18026v1…