Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference
Authors: Yifei Gao, Lei Wang, Rong-Cheng Tu, Qixin Zhang, Jun Cheng, Dacheng Tao
Organizations: University College London, London, UK · Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences (CAS), Shenzhen 518055, China · The Chinese University of Hong Kong · College of Computing & Data Science, Nanyang Technological University, Singapore
A core bottleneck in large language model (LLM) inference is the cost of attending over the ever-growing key-value (KV) cache. Although near-oracle top-k KV selection can preserve the quality of dense attention while sharply reducing computation and bandwidth, existing sparse methods generally rely on posterior heuristics, i.e., selectors conditioned on observed attention or proxy scores. Such conditioning introduces posterior bias: it tends to distort true token importance and miss salient tokens, thereby impairing long-range reasoning. To tackle this problem, we propose Pre-hoc Sparsity (PrHS), which selects KV entries before attention scoring and provides explicit accuracy control. Let the attention mass of discarded entries be delta (the dropped mass). Through a marginal-to-mutual-information analysis, we derive an upper bound on the mutual-information loss that depends only on the dropped mass. This relation explains failure modes of posterior heuristics and enables verifiable guarantees by controlling the dropped mass in advance. Within PrHS, we instantiate three orthogonal pre-hoc selectors along the axes of time, depth, and layer. Extensive experiments on LLaMA and Mistral families validate PrHS. Across GSM8K and CoQA, PrHS reduces retrieval overhead by over 90%, achieving 3x higher retrieval sparsity than HShare at matched or better accuracy. It incurs under 1% average degradation on LongBench, lowers attention FLOPs by about 15% versus prior sparse baselines, and yields a 9.9x speedup in attention-operator latency and 2.8x higher throughput on NVIDIA A100-80GB GPUs than the dense baseline.
Figures & tables
Fig. 1: Overall analysis of performance and efficiency . Across induced attention-approximation error and accuracy-efficiency trade-offs, our method consistently surpasses prior SOTA approaches and closely tracks the top- k oracle for optimal KV compression accuracy.
x
Input token embedding
q,k,v
Query, key, value
L
Input length
d
Per-head dimension
H
Number of heads
A(⋅)
Attention weight
K,V
KV cache matrices
t
Step count
S
Selector returning a subset of indices
St
Selected KV index set at step t
N
KV budget (sparsity: N/t )
τ(⋅)
Retained attention mass
δ(⋅)
Dropped attention mass
Ifull,IS
Routing MI under full and sparse attention
TABLE I: Summary of key variables and notation.
Fig. 2: Distribution of critical indices in LLaMA2-7B-Chat (Layer 10, Head 2). For each query, we retrieve 64 critical indices using the top- k oracle on WikiText2. Columns show five temporally adjacent queries with cosine similarity >0.8 . Left : overall distribution across keys 400–710; clusters are outlined in blue boxes . Right : zoom-ins of three clusters (annotated by green arrows in the left) at keys 426–436, 591–700, and 691–704. Critical indices are in red . Similar patterns across other models are shown in Appendix G .
Fig. 3: Attention distribution on LLaMA2-7B-Chat. Heatmaps across representative layer–head pairs, with darker color indicating stronger attention.
Fig. 4: Critical-index set (CIS) dilation. For three adjacent queries (405–407), we examine two key-index clusters [120,125] and [220,225] . Critical tokens from query 405 are shared to queries 406 and 407. Overlap is measured as Number of Critical TokensNumber of True Positive . Tokens in violet are also critical for query 405.
Fig. 5: Visual illustration of PSAW and ETF in prefill stage. When ℓ<ℓs , attention is unchanged. For ℓ≥ℓs , both methods prune redundant computation: PSAW computes a per-step sliding window, so the set of masked tokens can vary across steps (columns); ETF applies a fixed prune range and freezes earlier tokens so they no longer update.
Model
Method
ρ^
GSM8K ↑
COQA ↑
Comp* ↓
LLaMA2 -7b-chat
Original
-
0.2297/0.2297
0.5997
-
H2O [ 52 ]
-
0.0986/0.0144
0.4952
O(1)
Quest [ 31 ]
-
0.0478/0.0462
0.5713
0.1250 T
Hshare-1 [ 44 ]
28.1
0.1803/0.1703
0.5898
0.0180 T
Hshare-2 [ 44 ]
12.5
0.1524/0.1145
0.5672
0.0080 T
CIS (ours)
16.4
0.1842 / 0.1827
0.6085
0.0104 T
TABLE II: Evaluation of different methods on GSM8K and COQA. The best result (excluding origin) in each column is highlighted in bold. Comp* refers to the theoretical time complexity for each method to select critical KV cache tokens, where O(1) denotes constant time complexity, and T represents the theoretical computation time for a dense attention mechanism.
Model
Eval
Method
Original
Quest [ 31 ]
HShare [ 44 ]
CIS (ours)
CIS ∗ (ours)
CPE (ours)
Qwen2.5 -7B -Instruct
ρ^
28.1
12.5
19.4 ∣ 16.9
13.1 ∣ 9.3
19.2 ∣ 17.2
12.8 ∣ 9.1
18.8 ∣ 16.2
13.4 ∣ 9.5
Easy
30.2 ∣ 37.5
15.8 ∣ 21.3
19.8 ∣ 28.6
16.4 ∣ 25.2
29.7 ∣ 35.4
29.2 ∣ 31.8
29.2 ∣ 33.2
29.7 ∣ 34.9
26.6 ∣ 31.8
26.0 ∣ 30.7
Hard
32.5 ∣ 28.3
22.9 ∣ 26.5
23.2 ∣ 28.0
20.7 ∣ 24.5
30.9 ∣ 29.6
30.5 ∣ 31.5
31.2 ∣ 27.0
30.5 ∣ 26.4
33.4 ∣ 32.8
33.4 ∣ 33.1
Short
41.1 ∣ 35.5
29.8 ∣ 30.2
35.5 ∣ 34.4
29.2 ∣ 31.5
41.6 ∣ 37.8
41.1 ∣ 37.2
42.2 ∣ 34.3
42.2 ∣ 35.0
40.5 ∣ 38.9
40.5 ∣ 36.6
Medium
27.9 ∣ 32.5
13.1 ∣ 24.5
15.8 ∣ 25.1
11.1 ∣ 23.6
24.6 ∣ 30.7
24.6 ∣ 29.8
25.6 ∣ 27.0
24.2 ∣ 26.5
25.1 ∣ 28.4
24.6 ∣ 30.7
TABLE III: Evaluation of different methods on LongBench v2. Results under CoT prompting are highlighted with a gray background. The best result in each row (excluding the original) is highlighted in bold.
BS
Seqlen
Flash
Flex-BS
H2O [ 52 ]
Quest [ 31 ]
DS [ 48 ]
HShare-0 [ 44 ]
CIS-8
CIS-16
CPE-8
CPE-16
8
2k
0.830
0.119
0.093
0.460
0.241
0.160
0.147
0.106
0.144
0.095
4k
1.630
0.560
0.470
0.850
0.733
0.570
0.534
0.501
0.506
0.469
16k
7.140
2.500
2.080
3.420
3.160
2.430
2.280
2.250
2.220
2.110
16
2k
1.630
0.200
0.110
0.770
0.422
0.230
0.214
0.178
0.196
0.165
4k
3.230
1.000
0.850
2.210
1.350
1.041
0.953
0.869
0.904
0.821
16k
14.220
4.550
4.040
10.030
5.940
4.850
4.430
4.130
4.270
3.910
TABLE IV: Attention operator latency (ms ↓ ) across batch sizes and sequence lengths. Flex-BS denotes FlexAttention with BlockMask using the same logical sparse pattern as CPE; BlockMask construction is excluded from operator timing. HShare configuration follows Tab. II . The suffix of CIS/CPE denotes the CIS block size s .
BS
Seqlen
GPT-Fast
Flex-BS
H2O [ 52 ]
Quest [ 31 ]
DS [ 48 ]
HShare-0 [ 44 ]
CIS-8
CIS-16
CPE-8
CPE-16
8
2k
188
228
234
206
213
222
216
224
225
233
4k
118
210
228
152
201
214
216
220
219
227
16k
33
80
92
59
68
83
86
87
88
91
16
2k
233
408
416
287
360
398
402
415
412
421
4k
136
338
396
175
286
350
352
364
362
384
16k
34
99
112
47
79
95
102
108
105
114
TABLE V: End-to-end decoding throughput ( ↑ ) across batch sizes and sequence lengths. Flex-BS replaces our fused sparse-attention kernel with FlexAttention BlockMask while retaining the same CPE selection pipeline. The suffix of CIS/CPE denotes the CIS block size s .
Variant
Parameter setting
ρ^↓
Avg. KV
βemp↓
GSM8K ↑
COQA ↑
LongBench v2 ↑
HShare-matched
Matched
10.0
1060
0.061
0.7002
0.5550
24.25 ∣ 27.35
DirectReuse
Global-frozen
10.3
1060
0.022
0.7907
0.6068
31.30 ∣ 34.40
RandomExpansion
Global-frozen
10.1
1060
0.041
0.7187
0.5454
29.15 ∣ 31.40
CIS
Global-frozen
10.1
1064
0.019
0.8007
0.6050
31.85 ∣ 35.40
CIS
Per-model tuned
10.2
1065
0.018
0.8051
0.6082
31.95 ∣ 35.85
TABLE VI: Controlled attribution and hyperparameter-transfer analysis. Panel A compares CIS with budget-, retrieval-, and implementation-matched controls. Panel B compares PSAW with static and linear windows under matched query-key FLOPs. βemp denotes the mean of the empirical retained-mass gap. Accuracy retention is averaged over Qwen2.5-7B-Instruct, Qwen3-4B.
Fig. 6: Accuracy and KV-selection quality across retrieval ratios on GSM8K with Qwen3-4B. CIS consistently maintains higher accuracy and a smaller retained-mass gap than HShare, together with higher overlap with the top- k oracle critical indices.
Fig. 7: Empirical grounding of the information certificate. Left: averaged routing-MI perturbation ΔIemp versus the averaged theoretical certificate BMI ; the dashed line denotes y=x . Right: routing-MI perturbation versus downstream degradation (dense-minus-sparse GSM8K accuracy on Qwen3-4B). The certificate is conservative but order-preserving, while larger routing-information perturbations are associated with larger task-level degradation.
Fig. 8: Query-norm statistics of CIS-triggered sharing pairs on Qwen3-4B and GLM4.5-Air. Statistics are computed only over query pairs (q,q′) for which CIS sharing is actually triggered, i.e., sim(q,q′)≥τ . Left: the query-norm ratio ∥q′∥/∥q∥ , where the dashed line indicates the norm-preserving reference value 1 . Right: the absolute query-norm difference ∣∥q′∥−∥q∥∣ . Results are reported on COQA, GSM8K, and LongBench v2.
Variant
Parameter setting
ρ^↓
Avg. KV
βemp↓
GSM8K Ret. ↑
COQA Ret. ↑
LongBench v2 Ret. ↑
HShare-matched
Matched
10.0
1060
0.069
0.6217
0.5092
20.1 ∣ 25.5
DirectReuse
Global-frozen
10.2
1060
0.021
0.7655
0.5384
29.7 ∣ 29.3
RandomExpansion
Global-frozen
10.3
1060
0.036
0.6836
0.5084
28.8 ∣ 26.5
CIS
Global-frozen
9.8
1064
0.018
0.7787
0.5336
30.2 ∣ 29.6
CIS
Per-model tuned
10.2
1062
0.017
0.7876
0.5389
30.6 ∣ 30.2
TABLE VII: Controlled attribution and hyperparameter-transfer analysis. Panel A compares CIS with budget-, retrieval-, and implementation-matched controls. Panel B compares CIS variants under matched settings. βemp denotes the mean empirical retained-mass gap. Accuracy retention is averaged over Qwen2.5-7B-Instruct and Qwen3-4B.
Fig. 9: Failure-mode and sensitivity analysis of CIS on Qwen3-4B. (a) Downstream accuracy under different dilation strengths, where m=⌊k/nden⌋ and a smaller nden corresponds to more aggressive dilation. (b) Average processed-KV composition on LongBench v2 when varying the dilation strength. The orange segment denotes selected indices that also belong to the reference top- k oracle ( k=4096 , following Sec. V-C ), whereas the green segment denotes additional non-oracle ("minimal") indices introduced by dilation. (c) Downstream accuracy as a function of the CIS similarity threshold τ . (d) Overlap with the top-50% highest-attention-mass indices of the top- k oracle. Dilation causes only limited accuracy fluctuations even when it substantially enlarges the processed KV set, whereas an overly permissive similarity threshold sharply reduces the preservation of high-mass critical indices and leads to considerably larger downstream degradation. These results identify mismatched cross-query sharing, rather than dilation itself, as the primary failure mode of CIS.
Methods
s
τ
r
ϕ
ψ
α
γ
ρ^
Avg. Token
Wiki PPL
GSM8K(flexible)
COQA(EM/F1)
CIS
4
0.8
1
-
-
-
-
0.301
124
-
0.1834
0.6023 / 0.7632
8
0.7
1
-
-
-
-
0.152
131
-
0.1801
0.6090 / 0.7648
8
0.8
2
-
-
-
-
0.184
167
-
0.1653
0.6083 / 0.7650
32
0.8
1
-
-
-
-
0.0046
136
-
0.1693
0.6080 / 0.7641
PSAW
-
-
-
0.5
-
1
-
-
-
6.120
0.1896
0.5925 / 0.7568
-
-
-
0.7
-
1.5
-
-
-
6.113
0.1803
0.5938 / 0.7573
TABLE VIII: Hyperparameter tuning on LLaMA2-Chat-7B for CPE. The COQA setting for CIS follows the CIS ∗ configuration, and HyperKV follows the HyperKV configuration, as described in Sec. V-B . WikiText perplexity (PPL) is measured only during the prefilling stage. The Avg.Token measures the average processed tokens per head only in the decoding stage.
Method
s
GSM8K (flexible/strict)
COQA
Key
8
0.1547 / 0.1482
0.5613
16
0.1405 / 0.1334
0.5608
Hidden
8
0.1303 / 0.1289
0.5416
16
0.1121 / 0.1108
0.5214
TABLE IX: Variants of CIS on the similarity designs. The model is LlaMA2-7b-chat.
Fig. 10: Query-similarity distributions across LongBench v2 workloads. Cosine similarities are measured between consecutive per-head decoding queries. Boxes denote the 25th–75th percentiles with the median, and whiskers denote the 10th–90th percentiles. Similarity remains concentrated at high values, while harder and longer workloads exhibit lower and broader distributions.
Fig. 11: Distribution of critical indices on modern models. Each panel shows the top- 32 oracle critical indices of five consecutive decoding queries after excluding deterministic sink and local regions. Critical indices form spatially localized groups whose positions exhibit only small shifts across adjacent queries, consistently on Qwen3-4B and GLM-4.5-Air.
Fig. 12: Performance across total KV budgets on Qwen3-4B. We sweep the retained KV budget from 36 to 256 for HShare-1 and CIS. CIS remains stable under aggressive KV compression, whereas HShare degrades substantially as the budget decreases, particularly on GSM8K.
Despite their strong performance, large language models (LLMs) are bottlenecked by KV cache memory traffic during long-context inference. Sparse attention is widely used to accelerate LLM inference by computing exact attention over a selected subset of tokens. To recover the contribution of tokens excluded from exact attention, recent methods apply coarse-grained compensation to the omitted attention tail. However, existing methods typically select tokens based on attention mass and only then compensate for the unselected tokens. This decoupled design overlooks their interaction: selection should prioritize tokens that would leave the largest compensation error if omitted. To address this limitation, we introduce CompKV, the first compensation-aware sparse attention framework that divides tokens into blocks and explicitly optimizes selection for the downstream compensation mechanism. Our theoretical analysis shows that the residual left by block-level mean compensation is governed by both block attention mass and within-block logit variation. We approximate this residual using compact block-level statistics, yielding a deployable selection criterion. We further develop an efficient asynchronous implementation. Experiments on RULER and LongBench-Pro show that CompKV performs best among the evaluated sparse baselines while delivering up to a 6.85× self-attention speedup over full attention.
Zhen Huang, Ruizhe Yao, Danyi Liu +8
Tsinghua University · Harbin Institute of Technology, Shenzhen · Northeastern University
Long-context Large Language Model inference is severely bottlenecked by the massive Key-Value (KV) cache, yet existing sparse attention methods often suffer from static fixed-budget (Top-k) retrieval or rely on proxy scores that are computationally expensive and biased. To address these limitations, we propose RaBitQCache, a novel sparse attention framework that utilizes randomized rotated binary quantization and high-throughput binary-INT4 arithmetic to efficiently estimate attention weights. Our proxy score serves as an unbiased estimator with a proven error bound, enabling adaptive Top-p retrieval that dynamically adjusts the token budget based on actual attention sparsity. We further implement a hardware-aware system with asynchronous pipelining and lazy updates to mask overhead. Evaluations demonstrate that RaBitQCache significantly accelerates inference and reduces memory I/O while preserving generation quality compared to state-of-the-art baselines. Code is available at https://github.com/Sakuraaa0/RaBitQCache.git.
Wenhao Li, Jinhao Dong, Hailin Zhang +3
School of Information, Renmin University of China, Beijing, China · Peking University, Beijing, China
Sparse attention improves LLM inference efficiency by selecting a subset of key-value entries, but at the cost of potential accuracy degradation. In particular, omitting critical KV entries can induce substantial errors in model outputs. Existing methods typically operate under fixed or adaptive token budgets and provide empirical robustness or partial theoretical guarantees, yet they do not ensure zero false negatives in decoding steps, particularly since the set of relevant tokens is both query- and step-dependent. Our empirical observations confirm that missing even one critical key can lead to sharp error spikes, especially in long reasoning tasks where the set of important tokens varies throughout decoding. This observation motivates the need for indexing methods that dynamically adapt to these variations across decoding steps while guaranteeing a full recall of the relevant keys above a certain threshold. We address this challenge by reformulating sparse attention as the halfspace range searching problem. However, existing range searching indices are not suitable for modern LLM inference due to their computational and implementation overheads. To overcome this, we introduce Louver, a novel index structure tailored for efficient KV cache retrieval. Louver (i) guarantees zero false negatives with respect to a specified threshold in both theory and practice, (ii) is lightweight to integrate into existing LLM pipelines, and (iii) incorporates hardware-aware optimizations for both CPU and GPU executions. Our experiments demonstrate that Louver outperforms prior sparse attention methods in both accuracy and runtime, and is faster than highly optimized dense attentions such as FlashAttention. These results highlight that recall guarantees are a critical and overlooked dimension of sparse attention, and open a new direction for building theoretically grounded, efficient KV cache indices.
Mohsen Dehghankar, Abolfazl Asudeh
Department of Computer Science University of Illinois Chicago