Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference

Yifei Gao; Lei Wang; Rong-Cheng Tu; Qixin Zhang; Jun Cheng; Dacheng Tao

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.

$\mathbf{x}$	Input token embedding	$\mathbf{q},\,\mathbf{k},\,\mathbf{v}$	Query, key, value
$L$	Input length	$d$	Per-head dimension
$H$	Number of heads	$A(\cdot)$	Attention weight
$\mathbf{K},\mathbf{V}$	KV cache matrices	$t$	Step count
$\mathcal{S}$	Selector returning a subset of indices	$S_{t}$	Selected KV index set at step $t$
$N$	KV budget (sparsity: $N/t$ )	$\tau(\cdot)$	Retained attention mass
$\delta(\cdot)$	Dropped attention mass	$I_{\mathrm{full}},I_{\mathcal{S}}$	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

\frac{\text{Number of True Positive}}{\text{Number of Critical Tokens}}

. Tokens in violet are also critical for query 405.

Fig. 5: Visual illustration of PSAW and ETF in prefill stage. When

\ell<\ell_{s}

, attention is unchanged. For

\ell\geq\ell_{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	$\hat{\rho}$	GSM8K $\uparrow$	COQA $\uparrow$	Comp* $\downarrow$
LLaMA2 -7b-chat	Original	-	0.2297/0.2297	0.5997	-
	H2O [ 52 ]	-	0.0986/0.0144	0.4952	$\mathcal{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

\mathcal{O}(1)

denotes constant time complexity, and

T

represents the theoretical computation time for a dense attention mechanism.

Model	Eval	Method
Model	Eval	Original	Quest [ 31 ]	HShare [ 44 ]		CIS (ours)		CIS ∗ (ours)		CPE (ours)
Qwen2.5 -7B -Instruct	$\hat{\rho}$			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

\downarrow

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

\uparrow

) 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	$\hat{\rho}\downarrow$	Avg. KV	$\overline{\beta}_{\mathrm{emp}}\downarrow$	GSM8K $\uparrow$	COQA $\uparrow$	LongBench v2 $\uparrow$
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.

\overline{\beta}_{\mathrm{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

\overline{\Delta I}_{\mathrm{emp}}

versus the averaged theoretical certificate

\overline{B}_{\mathrm{MI}}

; 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.

Figure 14

Figure 15

Figure 16

Figure 17

Figure 18

Figure 19

Fig. 8: Query-norm statistics of CIS-triggered sharing pairs on Qwen3-4B and GLM4.5-Air. Statistics are computed only over query pairs

(\mathbf{q},\mathbf{q}^{\prime})

for which CIS sharing is actually triggered, i.e.,

\mathsf{sim}(\mathbf{q},\mathbf{q}^{\prime})\geq\tau

. Left: the query-norm ratio

\|\mathbf{q}^{\prime}\|/\|\mathbf{q}\|

, where the dashed line indicates the norm-preserving reference value

1

. Right: the absolute query-norm difference

|\|\mathbf{q}^{\prime}\|-\|\mathbf{q}\||

. Results are reported on COQA, GSM8K, and LongBench v2.

Variant	Parameter setting	$\hat{\rho}\downarrow$	Avg. KV	$\overline{\beta}_{\mathrm{emp}}\downarrow$	GSM8K Ret. $\uparrow$	COQA Ret. $\uparrow$	LongBench v2 Ret. $\uparrow$
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.

\overline{\beta}_{\mathrm{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=\lfloor k/n_{\mathrm{den}}\rfloor

and a smaller

n_{\mathrm{den}}

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

\tau

. (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$	$\tau$	$r$	$\phi$	$\psi$	$\alpha$	$\gamma$	$\hat{\rho}$	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
PSAW	-	-	-	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
Key	16	0.1405 / 0.1334	0.5608
Hidden	8	0.1303 / 0.1289	0.5416
Hidden	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.

Model	Eval	Method
Model	Eval	Original	Quest [ 31 ]	HShare [ 44 ]		CIS (ours)		CIS ∗ (ours)		CPE (ours)
Qwen2.5 -7B -Instruct	$\hat{\rho}$			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

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference

Abstract

Figures & tables

Explore similar work

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

RaBitQCache: Rotated Binary Quantization for KVCache in Long Context LLM Inference

Sparse Attention as a Range Searching Problem: Towards an Inference-Efficient Index for KV Cache

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

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

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