# CompKV: Compensation-Aware KV Selection for Long-Context LLM Inference

> CompKV is the first sparse attention framework that selects KV blocks based on downstream compensation error, achieving near-full-attention accuracy with up to 6.85x speedup.

- **Source:** [arXiv](https://arxiv.org/abs/2609.26300)
- **Published:** 2026-09-26
- **Permalink:** https://picx.dev/p/bjdCcz
- **Whiteboard:** https://picx.dev/p/bjdCcz/image

## Summary

## Summary (Overview)

- **CompKV** is the first compensation-aware sparse attention framework for long-context LLM inference that explicitly optimizes KV block selection for the downstream compensation mechanism, rather than treating selection and compensation as separate stages.
- The paper provides a theoretical analysis showing that the post-compensation residual under block-mean compensation is jointly governed by a block's **attention mass** ($p_{g,b}$) and **within-block logit variance** ($\sigma_{g,b}^2$), leading to a selection score of $p_{g,b} \cdot \sigma_{g,b}^2$.
- CompKV estimates these quantities from compact block statistics (mean keys, grouped variances, mean values), requiring only $2d + r + 1$ scalars per block, and uses an efficient asynchronous CPU-offload implementation.
- On RULER and LongBench-Pro benchmarks across Llama-3.1-8B-Instruct, Qwen3-8B, and Qwen3-32B, CompKV achieves the **highest average scores among all evaluated sparse baselines**, with accuracy drops as low as 2.7 points from full attention (vs. 18.1 points for the best baseline on Qwen3-32B).
- The asynchronous implementation delivers **up to 6.85× self-attention speedup** over full attention, with the lowest mean decoding-step latency across all nine combinations of context lengths (32K–128K) and token budgets (512–2,048).

---

## Introduction and Theoretical Foundation

### Background and Motivation

During autoregressive decoding, the KV cache grows with context length, and dense attention reads the entire cache at every step. This makes long-context inference increasingly bottlenecked by **KV cache memory traffic**. Query-aware sparse attention reduces this traffic by reading only a small set of tokens selected for the current query.

Recent methods divide tokens into blocks and compensate for omitted blocks' attention contributions using compact summaries (e.g., mean values). However, **selection and reconstruction remain largely decoupled**: blocks are ranked by predicted relevance or attention mass *before* reconstruction is considered.

> **Key Insight:** A high-mass block may be easy to reconstruct, whereas a lower-mass block with greater within-block logit variation may leave a larger error. Therefore, selection should prioritize blocks that would leave the largest reconstruction error if omitted.

### Theoretical Foundation

The paper defines **Mean compensation** as replacing every logit in an unselected block by its block mean. The compensated logit is:

$$\hat{z}_{g,j} = \begin{cases} z_{g,j}, & b \in \mathcal{S}, \\ \bar{z}_{g,b}, & b \notin \mathcal{S}, \end{cases} \tag{1}$$

where $\bar{z}_{g,b} = |b|^{-1} \sum_{j \in b} z_{g,j}$ is the block-mean logit.

The selection objective minimizes the KL divergence between the compensated and full attention distributions:

$$\min_{\mathcal{S} \subseteq \mathcal{B}, |\mathcal{S}| = K} \mathcal{L}(\mathcal{S}) = \sum_{g=1}^{G} D_{\mathrm{KL}}\left(\widehat{P}_g(\mathcal{S}) \lVert P_g\right) \tag{2}$$

---

## Methodology

### 3.1 Deriving the Compensation-Aware Selection Criterion

For each query head $g$, define the attention mass and logit variance of block $b$:

$$p_{g,b} = \frac{Z_{g,b}}{\sum_{c \in \mathcal{B}} Z_{g,c}}, \quad \sigma_{g,b}^2 = \frac{1}{|b|} \sum_{j \in b} (z_{g,j} - \bar{z}_{g,b})^2 \tag{3}$$

The KL divergence expansion yields:

$$\mathcal{L}(\mathcal{S}) = -\sum_{g=1}^{G} \log\left[1 - \sum_{b \notin \mathcal{S}} p_{g,b}\left(1 - \frac{|b| e^{\bar{z}_{g,b}}}{Z_{g,b}}\right)\right] \tag{4}$$

A second-order Taylor expansion around the block-mean logit gives:

$$p_{g,b}\left(1 - \frac{|b| e^{\bar{z}_{g,b}}}{Z_{g,b}}\right) = p_{g,b}\left(\frac{1}{2}\sigma_{g,b}^2 + O(\sigma_{g,b}^3)\right) \tag{5}$$

This leads to the key result: **the optimal selection criterion is to maximize**

$$\max_{\mathcal{S} \subseteq \mathcal{B}, |\mathcal{S}| = K} \sum_{b \in \mathcal{S}} \sum_{g=1}^{G} p_{g,b} \sigma_{g,b}^2 \tag{7}$$

### 4.1 Block Scoring from Compact Statistics

CompKV estimates the score using:

1. **Grouped key-variance estimation** — partition key coordinates into $r$ groups, storing average variance per group:
   
$$\widehat{\sigma}_{g,b}^2 = \frac{1}{d} \sum_{t=0}^{r-1} \left(\sum_{i \in \mathcal{G}_t} q_{g,i}^2\right)(\bar{\mathbf{D}}_b)_t \tag{8}$$

2. **Attention-mass estimation** using the log-partition expansion:
   
$$\widehat{p}_{g,b} = \frac{|b| \exp\left(\bar{z}_{g,b} + \frac{1}{2}\widehat{\sigma}_{g,b}^2\right)}{\sum_{c \in \mathcal{B}} |c| \exp\left(\bar{z}_{g,c} + \frac{1}{2}\widehat{\sigma}_{g,c}^2\right)} \tag{9}$$

3. **Compensation-aware block scoring**:
   
$$S_b = \sum_{g=1}^{G} \widehat{p}_{g,b} \widehat{\sigma}_{g,b}^2 \tag{10}$$

### 4.2 Block Selection and Compensated Decoding

The selected set includes mandatory blocks (sink + local window) plus Top-K by score:

$$\mathcal{S} = \mathcal{F} \cup \operatorname{TopK}_{K-|\mathcal{F}|}\{S_b : b \in \mathcal{B} \setminus \mathcal{F}\} \tag{11}$$

Output is computed via joint normalization:

$$\widehat{\mathbf{o}}_g = \frac{\sum_{b \in \mathcal{S}} \sum_{j \in b} e^{z_{g,j}} \boldsymbol{v}_j + \sum_{b \notin \mathcal{S}} |b| e^{\bar{z}_{g,b}} \bar{\boldsymbol{v}}_b}{\sum_{b \in \mathcal{S}} Z_{g,b} + \sum_{b \notin \mathcal{S}} |b| e^{\bar{z}_{g,b}}} \tag{12}$$

### 4.3 Asynchronous Implementation

The implementation keeps KV cache in pinned CPU memory, using:
- Main CUDA stream: transfers selection region, runs block statistics + Top-K scoring
- Auxiliary stream: computes Mean compensation, overlapping with CPU gathering and exact attention
- The summary buffer stores $\{|b|, \bar{k}_b, \bar{\mathbf{D}}_b, \bar{v}_b\}$ — $2d + r + 1$ scalars per block

---

## Empirical Validation / Results

### RULER Accuracy (32K context, 512-token budget)

**Table 1 highlights (best sparse result per model):**

| Model | Full | Quest | InfLLM | Quest+RESA | CompKV |
|-------|------|-------|--------|------------|--------|
| Llama-3.1-8B | 87.0 | 77.0 (↓10.0) | 73.1 (↓13.9) | 78.2 (↓8.8) | **83.2 (↓3.8)** |
| Qwen3-8B | 91.5 | 76.5 (↓15.0) | 76.6 (↓14.9) | 75.8 (↓15.7) | **86.0 (↓5.5)** |
| Qwen3-32B | 93.7 | 75.7 (↓18.0) | 82.5 (↓11.2) | 75.6 (↓18.1) | **91.0 (↓2.7)** |

CompKV shows particularly large gains on **MK3** (multi-key retrieval), outperforming baselines by 35–90 points on this difficult task.

### LongBench-Pro Scores

| Model | Full | Quest | InfLLM | Quest+RESA | CompKV |
|-------|------|-------|--------|------------|--------|
| Llama-3.1-8B | 25.26 | 24.28 (↓0.98) | 23.60 (↓1.66) | 24.30 (↓0.96) | **24.44 (↓0.82)** |
| Qwen3-8B | 34.23 | 30.58 (↓3.65) | 31.75 (↓2.48) | 30.60 (↓3.63) | **32.74 (↓1.49)** |
| Qwen3-32B | 42.00 | 37.62 (↓4.38) | 39.62 (↓2.38) | 36.16 (↓5.84) | **40.30 (↓1.70)** |

### Efficiency Results

- CompKV ($r=4$) achieves the **lowest mean latency in all nine settings** (3 context lengths × 3 budgets)
- **Up to 6.85× speedup** over full attention
- At 512-token budget: CompKV outperforms Quest and InfLLM across 32K, 64K, and 128K contexts

### Ablation Studies

**Table 3: Factor ablations on RULER (13-task AVG %)**

| Variant | Llama 256 | Llama 512 | Qwen 256 | Qwen 512 |
|---------|-----------|-----------|----------|----------|
| **CompKV (full)** | **78.6** | **83.2** | **80.9** | **86.0** |
| $\widehat{p}^{(2)}$ only (no Mean) | 75.6 | 81.8 | 79.2 | 83.5 |
| $\widehat{p}^{(2)} + \text{Mean}$ (no outer var) | 78.1 | 82.8 | 79.7 | 85.1 |
| $\widehat{p}^{(1)}\widehat{\sigma}^2 + \text{Mean}$ (first-order mass) | 76.1 | 81.2 | 77.5 | 83.8 |

Key findings:
- Mean compensation adds 0.48–2.47 points
- Outer variance factor adds 0.35–1.19 points
- Second-order mass correction adds 1.96–3.39 points
- Finer variance groups ($r=128$) improve accuracy but at 32× metadata cost; $r=4$ is the practical default

---

## Theoretical and Practical Implications

### Theoretical Contributions

1. **First formalization of compensation-aware selection**: The paper identifies and formalizes the mismatch between block selection and downstream compensation, showing that optimal selection must consider reconstruction error, not just attention mass.

2. **Analytically tractable residual analysis**: By using Mean compensation, the KL divergence reduces to a log-partition gap, enabling closed-form analysis. The key theoretical result—that residual scales as $p_{g,b} \cdot \sigma_{g,b}^2$—provides a principled, deployable selection criterion.

3. **Unified framework**: The paper shows that removing Mean compensation from the KL objective recovers standard Top-K attention-mass selection, demonstrating that CompKV is a strict generalization.

### Practical Implications

- **Training-free**: CompKV requires no model fine-tuning, making it immediately deployable with existing LLMs.
- **Compact metadata**: Only $2d + r + 1$ scalars per block (e.g., ~1 KB per 16-token block for $d=128$, $r=4$), enabling efficient CPU offloading.
- **Generalizable**: Works across model families (Llama, Qwen) and scales (8B–32B), with consistent improvements over baselines.
- **Efficient**: The asynchronous CPU-offload design overlaps compensation with KV retrieval, achieving sub-millisecond attention latencies at 128K context.

---

## Conclusion

CompKV addresses a fundamental limitation of existing sparse attention methods by making block selection explicitly aware of the downstream compensation error. The paper's key contributions are:

1. **Theoretical**: Proves that post-compensation residual is jointly governed by attention mass and within-block logit variance, yielding the selection score $p_{g,b} \cdot \sigma_{g,b}^2$.

2. **Methodological**: Introduces a practical, training-free framework using compact block statistics (mean keys, grouped variances) to estimate this score without reading all keys.

3. **Empirical**: Achieves state-of-the-art accuracy among sparse baselines on RULER and LongBench-Pro across three models, while delivering up to 6.85× speedup over full attention.

**Future directions** suggested by this work include:
- Extending compensation-aware selection to other compensation mechanisms beyond Mean (e.g., low-rank reconstructions)
- Exploring finer-grained variance statistics for accuracy-critical applications
- Investigating adaptive block sizes and variance group allocations based on observed attention patterns

---

_Markdown view of https://picx.dev/p/bjdCcz, served by PicX — AI-generated visual whiteboard summaries of research papers._