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
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
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 Qin and Qmid with respect to Qout . Right: attention mass retained by tokens selected using each query, evaluated under the Qout 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
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
Δ
Qwen2.5-32B
Quest
256
47.62
47.61
−0.01
RocketKV
256
47.10
47.14
+0.04
RocketKV
1024
48.16
48.15
−0.01
Qwen2.5-7B
Loki
128
38.36
38.44
+0.08
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. Δ denotes ES minus the corresponding base selector.
Figure 10: Auxiliary proxy-query diagnostics on the LongAlpaca-12k calibration set. Left: cosine similarity to Qout . 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
Δ
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; Δ denotes ES-RocketKV minus RocketKV.
Task
Sparse
Real- Q
Exact- WQ
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- WQ 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 Wqℓ+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 ± s.d.
Appendix
Table 7: Key reproducibility settings for EarlySelectKV.
Budget
Method
NarrativeQA
QASPER
MultiFieldQA
HotpotQA
2WikiMQA
MuSiQue
GovReport
QMSum
MultiNews
TREC
TriviaQA
SamSum
PassageRetrieval
PassageCount
LCC
RepoBench
Llama-3-8B-Instruct
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.
Budget
Method
NarrativeQA
QASPER
MultiFieldQA
HotpotQA
2WikiMQA
MuSiQue
GovReport
QMSum
MultiNews
TREC
TriviaQA
SamSum
PassageRetrieval
PassageCount
LCC
RepoBench
Llama-3.1-8B-Instruct
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.
Budget
Method
NarrativeQA
QASPER
MultiFieldQA
HotpotQA
2WikiMQA
MuSiQue
GovReport
QMSum
MultiNews
TREC
TriviaQA
SamSum
PassageRetrieval
PassageCount
LCC
RepoBench
Mistral-7B-Instruct-v0.2
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.
Budget
Method
NarrativeQA
QASPER
MultiFieldQA
HotpotQA
2WikiMQA
MuSiQue
GovReport
QMSum
MultiNews
TREC
TriviaQA
SamSum
PassageRetrieval
PassageCount
LCC
RepoBench
Qwen2.5-7B-Instruct
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.
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.
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.
Yifei Gao, Lei Wang, Rong-Cheng Tu +3
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
Efficient long-context inference faces two coupled bottlenecks: KV-cache memory grows linearly with context length, while attention computation grows quadratically. Existing approaches typically address one at the expense of irreversible token eviction, full-cache retention, or full-history reconstruction, limiting their effectiveness for multi-turn interaction and long-form reasoning. Motivated by two empirical properties, Long-Range Inter-Token Similarity and Smooth Residual Distribution, we propose ResidualKV, which factorizes the KV cache into a sparse set of globally retrieved references and compact, quantized residual codes for the remaining tokens. This representation preserves token-specific information without permanent eviction and, when combined with sparse attention, reconstructs only the selected states on demand. Dynamic-stride scheduling further reduces reference growth from linear to approximately logarithmic at ultra-long contexts. Across Llama, Qwen, LLaVA-OV, and Qwen3-VL backbones, ResidualKV maintains near-full-cache performance using only 13%-16% KV storage and 30% attention computation on LongBench, and 8%-10% storage and 10% computation in matched-budget multimodal evaluation. It also accelerates decoding by up to 1.5× with KV-cache quantization and 3.4× without it. These results show that global cross-token redundancy supports accurate, memory-efficient, and computation-efficient long-context inference. The source code is available at https://github.com/CURRENTF/ResidualKV.
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