Technical Summary: COMPILERKV

1. Problem Statement

Large Language Models (LLMs) face a memory bottleneck due to the linear growth of the Key-Value (KV) cache with context length. Prefill-only KV compression has emerged as a deployment-friendly solution, where a subset of tokens is selected at the end of the prefill stage and frozen for the entire decoding pass. This approach converts runtime memory usage into a fixed cost with zero decoding-time intervention.

However, the paper identifies a fundamental statistical mismatch in existing prefill-only methods (e.g., SnapKV, PyramidKV, HeadKV). These methods attempt to estimate corrective signals—specifically per-head reliability (which attention heads are critical for retrieval) and prompt-level compression sensitivity (how much a specific prompt can be compressed)—online, using statistics from a single, noisy prompt. The authors argue this is the wrong statistical unit:

• E1 (Single-prompt noise): Per-prompt attention scores are heavy-tailed and scale-biased; large magnitudes do not guarantee semantic importance, and small estimation errors cause irreversible selection errors.
• E2 (Head heterogeneity): Attention heads are functionally specialized. Some are "retrieval-critical," while others are consistently noisy. Naïve aggregation allows noisy heads to dominate scores.
• E3 (Prompt-level risk variance): Different prompts have different compression limits. A fixed retention rule fails on high-entropy or high-perplexity prompts, leading to catastrophic information loss.

Existing methods re-estimate these signals online for every prompt, failing to leverage the fact that these signals exhibit high cross-prompt regularity rather than within-prompt stability.

2. Methodology: COMPILERKV

The authors propose COMPILERKV, a KV-retention policy that compiles corrective tables offline from a calibration corpus, rather than estimating them online. The core insight is that head reliability rankings and the mapping from risk features to safe thresholds are stable architectural priors that can be learned once and reused.

The framework consists of three coupled components operating during the prefill stage:

A. Stabilized Token Utility (Addressing E1)

To remove transient attention spikes and inter-head scale bias without learning, the method computes a parameter-free utility score:

• Window-aggregated attention: Smooths transient spikes by aggregating attention over an observation window $\Omega$ at the prompt tail.
• Head-relative normalization: Normalizes token scores by the mean value norm of the head, removing scale bias.
• Formula: $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$, where $\alpha_t$ is the window mass and $\rho$ is the normalized value norm.

B. Head Heterogeneity Table (Addressing E2)

Instead of treating all heads equally or re-estimating importance online, COMPILERKV compiles a frozen table $W_{head}[l, h]$ of continuous reliability weights.

• Offline Compilation: Learned via Conservative Q-Learning (CQL) on a calibration corpus. The state is $(l, h, B_l)$, and the action is a discretized reliability weight.
• Mechanism: During inference, the utility of a token is rescaled by the reliability of its head. Crucially, the aggregation uses weighted max-pooling ($\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$). This gives "veto power" to retrieval-critical heads: if a single critical head attends strongly to a token, it is retained, regardless of consensus from noisy heads.

C. Risk-Adaptive Threshold Gate (Addressing E3)

To handle prompts with varying compression limits, a Risk-Adaptive Threshold Gate maps prompt-level risk features to a retention threshold.

• Risk Signals: Uses Attention Entropy (structural risk) and Local Perplexity (semantic risk) computed over the observation window.
• Offline Compilation: A lookup table (LUT) maps discretized risk bins $(b_{ent}, b_{ppl})$ to a retention threshold $\tau^{(l)}$. This is also compiled via CQL.
• Conservative Bias: The CQL objective penalizes aggressive thresholds on rare, high-risk states (high entropy/perplexity), biasing the table toward safer retention where online heuristics typically fail.
• Elastic Selection: The final selection retains all tokens with utility above the threshold $\tau^{(l)}$. If the candidate set exceeds the budget $B_l$, it is clamped to the top-$B_l$ tokens. If it is below, all candidates are kept (elastic under-retention).

D. Theoretical Grounding

The paper provides a formal analysis showing that the approximation error is controlled by the discarded tail mass (Proposition 4.1). It argues that finite-table calibration error is controlled at the prompt level (across independent prompts) rather than assuming independence within a single prompt's attention rows, thereby reducing variance.

3. Key Contributions

1. Reframing the Problem: Identifies that the statistical unit for correction in prefill-only compression should be the cross-prompt scale, not the single prompt.
2. Offline Compilation: Introduces a method to compile head reliability and risk-adaptive thresholds offline, turning them into portable architectural priors.
3. Coupled Policy: Proposes a joint optimization where head reliability weights and risk thresholds are calibrated against the same downstream reward, rather than being independent modules.
4. Portability: Demonstrates that compiled tables transfer across disjoint corpora and even across different model families (e.g., LLaMA to Mistral) with minimal performance degradation.

4. Experimental Results

The method was evaluated on four 7B–8B backbones (LLaMA-3, Mistral, Qwen2, InternLM) across LongBench, RULER, and Needle-in-a-Haystack (NIAH) benchmarks.

• LongBench (512-token budget): COMPILERKV achieves compressed-SOTA on all four backbones, outperforming the strongest baseline (e.g., ChunkKV, HeadKV) by an average of +1.67 points (95% CI [+1.08, +2.37]). On some backbones, it slightly exceeds FullKV performance, which the authors attribute to denoising rather than information gain.
• Pressure Regimes:
  • RULER (128k context): Under a fixed 512/32k cache ratio, COMPILERKV maintains a score of 73, significantly outperforming SnapKV (38) and KVZip (66.5), while staying below FullKV (79).
  • NIAH (32k): Achieves 0.89 accuracy vs. SnapKV's 0.42.
  • Batching: At 32k input with a batch size of 16, COMPILERKV remains feasible, whereas FullKV runs out of memory (OOM).
• Transferability:
  • Cross-Corpus: Head reliability rankings show high Spearman correlation ($\bar{\rho} \approx 0.90$) across disjoint calibration corpora.
  • Cross-Model: Direct table transfer between models (e.g., LLaMA-3 to Mistral) incurs a loss of only 0.4–0.8 LongBench points compared to target-specific compilation.
• Efficiency: The online overhead is minimal: $O(1)$ table lookups per layer and a single sort over the candidate pool. It achieves up to 2.45× decoding throughput compared to FullKV at 32k context.

5. Significance and Claims

The paper claims that COMPILERKV closes the gap to FullKV at fixed budgets by correcting the statistical mismatch in how retention policies are computed. By shifting the computation of corrective signals from an online, noisy, single-prompt estimation to an offline, compiled, cross-prompt prior, the method stabilizes the selection process against noise and heterogeneity.

The authors emphasize that the novelty is not merely "head awareness" (which prior methods like HeadKV possess) but the joint offline compilation of continuous head reliability and prompt-risk thresholds. This transforms head reliability from a tuning trick into a reusable scientific object that can be compiled once and deployed across models and contexts with $O(1)$ online overhead. The work suggests that for prefill-only compression, the "right" signals are architectural properties of the model, not transient features of the input prompt.