EarlySelectKV: Anticipatory KV Selection via Cross-Layer Query Proxies

Yuhao Liu; Yuehao Li; Linghe Zhang; Gaobin Huang; Junran Li; Hongchang Wang; Haifang Jian

cs.CLNov 1, 2026

EarlySelectKV: Anticipatory KV Selection via Cross-Layer Query Proxies

Authors: Yuhao Liu, Yuehao Li, Linghe Zhang, Gaobin Huang, Junran Li, Hongchang Wang, Haifang Jian

Organizations: Laboratory of Solid State Optoelectronics Information Technology, Institute of Semiconductors, Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

Abstract

Query-conditioned sparse KV selection reduces long-context inference cost by attending to a subset of historical KV entries at each decoding step. Existing selectors must wait for the target-layer query before scoring candidates and generating indices, placing routing on the immediate pre-attention path. We introduce EarlySelectKV, a training-free transformation that uses proxy queries from earlier hidden states to precompute selected KV indices while retaining the real target-layer query for final attention. Tests with RocketKV, Quest, and Loki show no consistent pattern of quality degradation across the evaluated settings. In the primary RocketKV evaluation, ES-RocketKV differs by +0.018 points on LongBench, approximately -0.01 points on RULER, and +0.25 percentage points on representative NIAH/Passkey settings. In a TensorRT-LLM deployment with CPU-offloaded routing, ES-RocketKV improves end-to-end output throughput by up to 12.4% and reduces GPU-resident memory by 11.1% of the combined KV/K-summary pool. This makes routing timing an explicit optimization target alongside KV subset selection in sparse long-context decoding. Code is available at https://github.com/sea-with-sakura/EarlySelectKV.

Figures & tables

Figure 1: Conventional query-conditioned KV selection and EarlySelectKV . Conventional selectors wait for the target-layer query, so index generation remains on the immediate pre-attention path. EarlySelectKV derives a proxy query from preceding hidden states and precomputes KV indices in an earlier computation window.

Figure 2: Theoretical routing overhead under sparse attention budgets. We analyze a single Llama-3.1-8B layer with an 8K KV cache and vary the selected-token budget. As the budget decreases, sparse attention becomes cheaper, and routing accounts for a larger fraction of the remaining attention-side computation.

Figure 3: Overview of EarlySelectKV. The main flow gathers selected KV using precomputed page indices and performs sparse attention with the real query. The side flow constructs a proxy query from an earlier hidden state and precomputes next-layer page indices. Early and Mid taps represent two timing–fidelity trade-offs: the former offers a longer preparation window, while the latter provides higher selection fidelity.

Figure 4: Task score comparison between ES-RocketKV and RocketKV on LongBench and RULER. The first row reports LongBench results, while the second and third rows report RULER results under different context lengths.

Figure 5: Additional selector results on LongBench using Mistral-7B-Instruct-v0.2 with a per-query KV-access budget of 1024 tokens. The upper panel reports task-level scores, and the lower panel shows the score difference between each ES variant and its corresponding base selector.

Budget	Method
Budget	TopK	RocketKV	ES-RocketKV
Mistral-7B-Instruct-v0.2
64	100.0	100.0	100.0
128	100.0	100.0	100.0
Llama-3-8B-Instruct
64	100.0	98.0	98.0

Table 1: Matching accuracy (%) on NIAH/Passkey under low per-query KV-access budgets. Across four instruction-tuned backbones, ES-RocketKV closely matches RocketKV in passkey retrieval accuracy.

Figure 6: Proxy-query fidelity across the decoder layers of Mistral-7B-Instruct-v0.2. Left: signed top-

r

channel recall of

Q_{\mathrm{in}}

and

Q_{\mathrm{mid}}

with respect to

Q_{\mathrm{out}}

. Right: attention mass retained by tokens selected using each query, evaluated under the

Q_{\mathrm{out}}

attention distribution. The dotted line marks the transition layer used by ES-Mix.

Figure 7: Fixed-trajectory routing fidelity over decode position on MultiNews. Points connect five interval means; shaded regions and vertical bars denote 95% bootstrap confidence intervals, and dashed lines show OLS position trends. Retained mass and page overlap are ES-RocketKV-minus-RocketKV gaps, whereas signed top-

r

recall is measured between the proxy and target-layer queries.

Figure 8: End-to-end output throughput on a 40K-token input with a 4096-token Stage-1 prompt budget, a 128-token Stage-2 selected-token budget, and batch size 1. Solid lines report the output throughput of RocketKV and ES-RocketKV, while the dashed line shows the relative throughput gain of ES-RocketKV over RocketKV.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Effect of the low-rank proxy-query projection on LongBench using Mistral-7B-Instruct-v0.2 with a per-query KV-access budget of 512. Scores are compared against RocketKV with the full target-layer query projection.

Budget	Method
Budget	FullCache	TopK	RocketKV	ES-RocketKV
Llama-3-8B-Instruct
128	43.23	43.72	41.68	41.78
256	43.23	43.58	42.43	42.72
512	43.23	43.50	43.07	43.03
1024	43.23	43.31	43.18	43.21

Appendix

Table 2: Average LongBench scores of different methods under different per-query KV-access budgets.

Model	Selector	Budget	Base	+ES	$\Delta$
Qwen2.5-32B	Quest	256	47.62	47.61	$-0.01$
	RocketKV	256	47.10	47.14	$\mathbf{+0.04}$
	RocketKV	1024	48.16	48.15	$-0.01$
Qwen2.5-7B	Loki	128	38.36	38.44	$\mathbf{+0.08}$
Qwen2.5-7B	Quest	128	43.30	42.93	$-0.37$
Mistral-7B	Loki	128	29.51	28.65	$-0.86$

Appendix

Table 3: Additional LongBench averages across model scales, selectors, and KV-access budgets.

\Delta

denotes ES minus the corresponding base selector.

Figure 10: Auxiliary proxy-query diagnostics on the LongAlpaca-12k calibration set. Left: cosine similarity to

Q_{\mathrm{out}}

. Right: page recall against Exact-TopK selection.

Model	Routed	Pearson $r$	Agreement
Mistral-7B	31	0.9988–0.9995	100.0%
Llama-3-8B	31	0.9786–0.9975	96.8%–100.0%
Llama-3.1-8B	31	0.9724–0.9942	100.0%

Appendix

Table 4: Cross-domain stability of ES-Mix profiles. Pearson correlations and assignment-agreement ranges are computed over the three pairwise comparisons among LongAlpaca-12k, Qasper, and QMSum.

Model	Full	RocketKV	+ES	$\Delta$
Llama-3.1-8B	45.6	44.4	45.4	$+1.0$
Qwen2.5-7B	72.4	74.2	73.4	$-0.8$

Appendix

Table 5: Independent greedy-generation accuracy (%) on MATH-500. RocketKV and ES-RocketKV use a per-query KV-access budget of 1024;

\Delta

denotes ES-RocketKV minus RocketKV.

Task	Sparse	Real- $Q$	Exact- $W_{Q}$	Rank-256
GovReport	13/16	25.21	25.91	24.68
QMSum	15/16	24.05	24.47	23.80
MuSiQue	16/16	6.96	6.89	6.78
HotpotQA	15/16	15.11	15.11	14.96

Appendix

Table 6: Post-hoc EarlySelectKV results on InfLLM-V2-Long-Sparse-Base. Real-

Q

is native routing; Exact-

W_{Q}

and Rank-256 use exact and rank-256 proxy projections. Scores cover all 16 examples per task; “Sparse” counts those activating sparse routing. GovReport/QMSum use ROUGE-L; MuSiQue/HotpotQA use QA F1.

Figure 11: Steady-state decode throughput under different output lengths. The input length is 40K tokens, the Stage-1 prompt budget is 4096 tokens, and the Stage-2 selected-token budget is 128. Solid lines report post-context decode throughput of RocketKV-HSA and ES-RocketKV, while the dashed line reports the relative gain of ES-RocketKV.

Figure 12: End-to-end throughput gain of ES-RocketKV over RocketKV on 8K-token inputs across batch sizes. Shaded regions show sample standard deviations over three runs.

Item	Setting
Method defaults	ES-Mid; layer 0 dense; local tail outside routing
Proxy projection	Offline SVD of $W_{q}^{\ell+1}$ ; rank-128 BF16 for CPU routing
Runtime configuration	40K main and 8K batch-scaling inputs; Stage-1: 4096; $p=4$ , $k=32$ , $pk=128$ ; outputs 128–8192
Backend / hardware	TensorRT-LLM 1.3.0rc15 with CUDA Graphs; RTX 6000D; Xeon Platinum 8481C
Routing offload	Routing and FP8 K-summary on CPU; KV cache on GPU; no prefetch
Measurement	40K: one run; 8K: three runs, mean $\pm$ s.d.

Appendix

Table 7: Key reproducibility settings for EarlySelectKV.

Llama-3-8B-Instruct
Budget	Method	NarrativeQA	QASPER	MultiFieldQA	HotpotQA	2WikiMQA	MuSiQue	GovReport	QMSum	MultiNews	TREC	TriviaQA	SamSum	PassageRetrieval	PassageCount	LCC	RepoBench
FULL	FullCache	21.17	43.17	46.54	46.81	39.93	23.05	30.94	22.70	27.97	73.00	90.00	42.16	66.50	8.00	60.63	49.07
64	TopK	19.97	42.53	46.29	47.32	39.90	22.75	29.20	22.37	27.10	72.50	90.36	42.16	67.00	7.50	64.88	57.64
	RocketKV	18.44	30.47	46.57	46.31	34.83	21.34	20.72	21.14	23.82	55.00	89.68	38.20	65.50	6.00	62.36	57.92
	ES-RocketKV	19.19	28.77	46.42	44.88	34.78	21.47	20.62	21.09	23.72	57.00	89.49	38.66	65.50	5.50	64.22	58.16
128	TopK	20.55	43.15	45.16	46.81	41.52	22.12	29.19	22.42	27.45	72.50	90.86	42.17	66.50	8.50	64.76	55.91

Appendix

Table 8: Task-level LongBench scores for Llama-3-8B-Instruct.

Llama-3.1-8B-Instruct
Budget	Method	NarrativeQA	QASPER	MultiFieldQA	HotpotQA	2WikiMQA	MuSiQue	GovReport	QMSum	MultiNews	TREC	TriviaQA	SamSum	PassageRetrieval	PassageCount	LCC	RepoBench
FULL	FullCache	32.12	46.58	55.73	57.84	48.61	30.82	34.78	25.42	27.07	73.00	92.14	43.76	99.50	7.33	65.36	55.42
64	TopK	31.35	45.44	55.80	58.40	48.92	31.56	33.73	24.41	26.28	73.00	92.25	43.06	99.50	7.92	65.12	53.72
	RocketKV	27.17	29.99	57.32	56.65	44.88	28.77	22.85	23.24	23.30	56.50	91.07	40.10	99.00	9.00	63.50	52.23
	ES-RocketKV	27.35	31.46	56.42	56.87	44.19	27.29	23.17	23.15	23.37	58.00	91.70	39.98	98.50	7.83	63.20	51.39
128	TopK	30.99	45.27	55.84	58.07	49.03	31.84	34.25	25.57	26.56	73.00	92.25	44.23	99.50	7.58	65.83	54.78

Appendix

Table 9: Task-level LongBench scores for Llama-3.1-8B-Instruct.

Mistral-7B-Instruct-v0.2
Budget	Method	NarrativeQA	QASPER	MultiFieldQA	HotpotQA	2WikiMQA	MuSiQue	GovReport	QMSum	MultiNews	TREC	TriviaQA	SamSum	PassageRetrieval	PassageCount	LCC	RepoBench
FULL	FullCache	28.00	32.61	48.43	42.26	26.44	19.01	33.16	24.29	27.16	71.00	86.92	43.61	83.80	2.95	59.97	61.52
64	TopK	25.46	24.03	44.39	37.87	24.23	17.75	30.42	19.76	27.05	71.00	87.11	40.80	75.60	2.43	58.55	57.99
	RocketKV	22.24	19.12	48.10	37.73	23.80	16.08	23.04	22.61	24.26	59.00	86.43	39.42	50.01	4.01	55.47	55.84
	ES-RocketKV	22.58	19.71	47.56	36.95	22.90	16.29	23.41	22.70	24.72	59.00	86.41	39.94	43.52	3.99	55.03	55.30
128	TopK	25.75	30.06	48.21	39.28	25.44	17.88	31.68	21.86	27.13	71.00	87.27	42.95	82.37	2.53	59.21	59.62

Appendix

Table 10: Task-level LongBench scores for Mistral-7B-Instruct-v0.2.

Qwen2.5-7B-Instruct
Budget	Method	NarrativeQA	QASPER	MultiFieldQA	HotpotQA	2WikiMQA	MuSiQue	GovReport	QMSum	MultiNews	TREC	TriviaQA	SamSum	PassageRetrieval	PassageCount	LCC	RepoBench
FULL	FullCache	27.77	45.33	51.63	58.00	46.13	29.88	33.51	24.28	25.81	73.00	88.41	47.02	100.00	9.50	59.86	64.92
64	TopK	28.90	43.05	48.84	54.16	43.16	27.74	29.42	23.25	24.73	70.50	87.62	46.04	96.09	9.00	59.26	61.96
	RocketKV	20.88	29.82	50.17	50.90	39.19	24.96	21.37	21.41	20.77	50.50	83.07	40.87	86.58	8.00	50.80	50.79
	ES-RocketKV	21.56	27.97	50.13	49.92	37.50	25.09	20.38	21.09	20.61	49.00	82.91	39.56	83.83	7.50	49.01	50.40
128	TopK	28.30	43.73	49.13	54.26	45.12	28.62	28.52	23.37	24.41	69.50	87.58	46.44	97.25	9.50	59.13	62.50

Appendix

Table 11: Task-level LongBench scores for Qwen2.5-7B-Instruct.

Model	Method	Acc.@64	First 100% Budget
Mistral-7B-v0.2	TopK	100.0	64
	RocketKV	100.0	64
	ES-In	100.0	64
	ES-Mid	100.0	64
	ES-Mix	100.0	64
Llama-3-8B	TopK	100.0	64

Appendix

Table 12: NIAH/Passkey exact matching accuracy. All settings use 10x10x3_7digits ; context lengths are 4096 words for Llama-3-8B-Instruct, 20480 words for Mistral-7B-Instruct-v0.2, and 81920 words for Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct. For compactness, models reaching perfect retrieval are summarized by the lowest budget achieving 100% accuracy, while Qwen2.5-7B-Instruct is shown separately.

Panel of Figure 13 — Figure 13: NIAH/Passkey exact-match accuracy heatmaps on Qwen2.5-7B-Instruct across context lengths and needle depths. All methods use a per-query KV-access budget of 256 tokens.

Context	Model	Budget	FullCache	TopK	RocketKV	ES-RocketKV
RULER-8K	Mistral-7B-Instruct-v0.2	128	90.65	89.52	81.33	81.16
		256	90.65	90.34	85.09	84.95
		512	90.65	90.45	88.53	88.50
		1024	90.65	90.80	89.65	89.63
		2048	90.65	90.61	89.70	90.07
	Llama-3.1-8B-Instruct	128	95.31	92.71	85.62	85.13

Appendix

Table 13: Average matching accuracy (%) on RULER across context lengths.

Model	Budget	RocketKV	ES-In	ES-Mid	ES-Mix
Llama-3-8B-Instruct	64	39.89	39.54	39.97	40.02
	128	41.68	41.62	41.78	41.92
	256	42.43	42.45	42.72	42.65
	512	43.07	43.02	43.03	43.15
	1024	43.18	43.16	43.21	43.14
	2048	43.32	43.36	43.29	43.38

Appendix

Table 14: Average LongBench scores of RocketKV and EarlySelectKV routing variants.

Context	Model	Budget	RocketKV	ES-In	ES-Mid	ES-Mix
RULER-8K	Mistral-7B-Instruct-v0.2	128	81.33	80.72	81.16	80.73
		256	85.09	84.76	84.95	85.18
		512	88.53	87.88	88.50	88.31
		1024	89.65	89.49	89.63	89.58
		2048	89.70	89.96	90.07	89.77
	Llama-3.1-8B-Instruct	128	85.62	84.54	85.13	85.38

Appendix

Table 15: Average matching accuracy (%) of RocketKV and EarlySelectKV routing variants on RULER.

KV Cache Selection

See all 9 papers

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference
Feb 9, 2026Yifei Gao, Lei Wang, Rong-Cheng Tu +3KV Cache Selection Self-Attention
University College London, London, UK · Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences (CAS), Shenzhen 518055, China · The Chinese University of Hong Kong +1
ResidualKV: Residual-Based KV Cache Compression for Efficient Long-Context Inference
Feb 8, 2026Jitai Hao, Qiang Huang, Yaowei Wang +2Self-Attention KV Caching
Harbin Institute of Technology (Shenzhen), China
CompKV: Compensation-Aware KV Selection for Long-Context LLM Inference
Sep 22, 2026Zhen Huang, Ruizhe Yao, Danyi Liu +8Long-Context Language Model Inference KV-Cache Management
Tsinghua University · Harbin Institute of Technology, Shenzhen · Northeastern University

Feb 9, 2026J/K move · Enter open · S save

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

Yifei Gao, Lei Wang, Rong-Cheng Tu, Qixin Zhang+2

University College London, London, UK · Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences (CAS), Shenzhen 518055, China · The Chinese University of Hong Kong+1

KV Cache Selection Self-Attention KV Caching Long-Context Language Model Inference KV-Cache Management LLM Inference Acceleration+2

Open paper AskPDF arXiv

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.

Abstract

KV Cache Selection

Explore similar work

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

ResidualKV: Residual-Based KV Cache Compression for Efficient Long-Context Inference

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