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):

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

LongBench-Pro Scores

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 %)

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