Technical Summary: PolyKV: Heterogeneous Retention and Allocation for KV Cache Compression

Problem Statement

Long-context inference in Large Language Models (LLMs) faces a significant memory bottleneck due to the growth of the Key-Value (KV) cache. While existing KV cache eviction methods effectively reduce memory usage by retaining specific subsets of tokens, they typically apply a uniform compression policy and a uniform cache budget across all transformer layers.

This uniform design overlooks the inherent heterogeneity of transformer layers. Prior research indicates that different layers vary in their attention patterns, represented abstractions, contextual roles, and sensitivity to pruning. Consequently, a single global eviction rule is a coarse abstraction that fails to account for the fact that:

1. Different layers may prefer different eviction patterns (e.g., heavy hitters vs. attention sinks).
2. Different layers exhibit varying sensitivity to KV cache reduction, implying they require different cache capacities.

Current adaptive methods have begun to explore non-uniform budgets but often fail to treat the eviction pattern and the cache budget as coupled, layer-wise design choices.

Methodology: PolyKV

The authors propose PolyKV, a framework that reformulates KV cache eviction as a layer-wise design space. PolyKV explicitly couples two decisions for each layer and inference phase (prefill and decoding):

1. What to evict: Selecting a specific eviction pattern (e.g., H2O, TOVA, PyramidKV) suitable for that layer.
2. How much to retain: Assigning a specific cache budget to that layer.

Core Components

• Design Space Formulation: PolyKV defines a strategy $\Pi$ where each layer $l$ is assigned an eviction pattern $r_{l,\phi}$ and a budget $b_{l,\phi}$ for each phase $\phi$. The total budget is constrained by a fixed global limit $B$.
• Offline Calibration: To avoid exhaustive online search, PolyKV uses a small calibration set to generate a fixed, static strategy. It collects two categories of layer-specific signals:
  • Eviction Preference Signals: Measures which eviction pattern best preserves a layer's FullKV computation. Metrics include cosine similarity between perturbed and FullKV activations, and Perplexity (PPL) sensitivity (degradation caused by a specific pattern).
  • Cache Budget Sensitivity Signals: Estimates how strongly a layer is affected by KV reduction. Metrics include attention entropy (diffuse attention requires larger budgets) and PPL sensitivity (aggregated worst-case degradation).
• Strategy Construction:
  • Pattern Selection: For each layer-phase pair, the eviction pattern with the highest preference score (e.g., highest cosine similarity or lowest PPL degradation) is selected.
  • Budget Allocation: Budgets are allocated proportionally based on sensitivity signals. Layers with higher sensitivity (or entropy) receive larger budgets, subject to a minimum budget constraint ($b_{min}$).
• Implementation: PolyKV is implemented as a hook-based framework. It injects logic at common points in the generation pipeline (pre-kernel, post-kernel, cache update) to support modular integration of various eviction methods without modifying the underlying model code.

Key Contributions

1. Reformulation of KV Eviction: The paper reframes KV cache eviction as a layer-wise design space that couples the choice of eviction pattern with the allocation of cache capacity under a shared memory constraint.
2. The PolyKV Framework: A system that operationalizes this design space by composing eviction patterns and budgets at layer granularity. It integrates multiple eviction primitives via a unified interface and uses offline calibration signals to guide strategy construction.
3. Empirical Validation: The authors demonstrate that heterogeneous, layer-wise strategies outperform uniform, single-policy baselines under fixed memory budgets.

Experimental Results

Experiments were conducted on LLaMA-3.1-8B and Qwen3-8B using the LongBench and RULER benchmarks.

• Performance on LongBench (512-token average budget):
  • LLaMA-3.1-8B: PolyKV recovered 54.5% of the performance gap between the strongest single-policy baseline and FullKV.
  • Qwen3-8B: PolyKV recovered 25.7% of the gap.
  • PolyKV consistently outperformed the strongest baseline across a budget sweep (128–1024 tokens), with improvements ranging from 1.7% to 6.4%.
• Performance on RULER:
  • Results were mixed. PolyKV slightly exceeded the baseline on LLaMA but underperformed on Qwen (by 3.98 points). The authors attribute this to the nature of RULER tasks, which rely on sparse evidence tokens that may be discarded by heterogeneous policies if calibration signals do not perfectly identify them.
• Ablation Studies:
  • Eviction Pattern vs. Budget: The primary performance gain stems from assigning different eviction patterns to different layers, particularly during the prefill phase.
  • Budget Allocation: Non-uniform budget allocation provides additional gains but is more model-dependent and less robust than pattern selection. It is most effective when the total budget is sufficient to avoid starving sensitive layers.
• Task Analysis: Gains were most pronounced in tasks requiring broad contextual understanding (e.g., QA, few-shot learning). Performance degraded on tasks requiring exact retrieval of sparse tokens (e.g., NIAH in RULER) or strict adherence to local form (e.g., code generation).

Significance and Claims

The paper claims that treating KV cache eviction as a compositional, layer-aware optimization problem rather than a single global retention policy yields superior results.

• Heterogeneity is Key: The study validates that transformer layers are not homogeneous; a "one-size-fits-all" eviction strategy is suboptimal.
• Static Efficiency: By using offline calibration, PolyKV achieves these heterogeneous benefits without the runtime overhead of online policy search, making it practical for deployment.
• Limitations: The authors modestly note that current calibration signals do not always preserve "sparse but decisive" evidence tokens required for exact retrieval tasks, indicating a direction for future refinement in signal design.

In summary, PolyKV demonstrates that by aligning eviction strategies and cache capacities with the specific characteristics of each transformer layer, it is possible to significantly reduce the memory-quality trade-off in long-context LLM inference.