National University of Singapore
TL;DR: RefineKV adapts KV cache access to each decoding query. An attention-error bound selects which clusters to expand, delivering up to 1.76× faster decoding while retaining near-full-attention quality.
Repeatedly accessing the full key–value (KV) cache is a major bottleneck in long-context language model decoding. We find that small query-specific subsets can accurately approximate full attention, yet their union covers nearly the entire cache, motivating query-adaptive access for the KV cache. A natural approach is to split the KV cache into clusters, using compact representatives for clusters that incur negligible error while expanding the rest to their original KV entries.
We derive a principled selection criterion, an upper bound on attention-output error, which reveals two complementary factors governing the error: importance measures the impact of a cluster, while distortion measures how faithfully the cluster can be represented. Based on this criterion, we propose RefineKV, a decoding-time query-adaptive KV compression method. Experiments on two models and three benchmarks demonstrate that RefineKV retains near-full-attention quality while achieving up to 11.7× speedup in the attention phase and 1.76× faster decoding, establishing a new quality–latency Pareto frontier.
Full Attention and RefineKV generate 256 tokens from the same 128K context on the same GPU.
Selective access is useful when it reduces real decoding time. This replay uses measured token-arrival times: RefineKV reaches 256 tokens first, then waits while Full Attention finishes.
This replay uses real token-arrival timestamps from a LongBench-v2 literary-QA prompt: 130,801 input tokens and 256 output tokens. Both methods run sequentially on the same NVIDIA RTX PRO 6000 Blackwell GPU, using BF16 and vLLM 0.31.0.
Each method has one warmup and three measured runs; the video shows the median-TPOT run for each, aligned at the first output token. Playback is at 0.5× speed; prefill is excluded. RefineKV uses ε = 0.01 and GQA-mean. At 3.20 seconds of decode time, RefineKV has produced 256 tokens and Full Attention 140.
This is a separate vLLM demonstration. The paper’s 1.76× result below uses its benchmark protocol and GQA-any. This demo measures speed, not answer correctness or output equivalence.
RefineKV reduces attention time substantially; GEMV becomes the main remaining cost. These bars show profiled GPU kernel time for one representative request, separate from the wall-clock decoding demo above.
Compare scores across three benchmarks, then see how close RefineKV stays to Full Attention.
Qwen2.5-7B-Instruct15 ms/token budget
+3.86score points
vs. SnapKV 10% within this budget
1.09× the baseline score
0.23 points above Full Attention
1.27× faster decoding
+18.49score points
vs. SnapKV 10% within this budget
1.40× the baseline score
0.55 points below Full Attention
1.32× faster decoding
+4.14score points
vs. SnapKV 25% within this budget
1.08× the baseline score
1.79 points below Full Attention
1.29× faster decoding
Llama-3.1-8B-Instruct17 ms/token budget
+2.78score points
vs. PyramidKV 10% within this budget
1.06× the baseline score
0.05 points below Full Attention
1.76× faster decoding
+4.42score points
vs. SnapKV 10% within this budget
1.07× the baseline score
2.12 points below Full Attention
1.58× faster decoding
+0.51score points
vs. PyramidKV 25% within this budget
1.01× the baseline score
0.38 points below Full Attention
1.44× faster decoding
RefineKV and the best efficient-attention baseline fit within the shown time budget. Full Attention is the quality reference and takes longer. Higher task scores are better.
These examples use a fixed budget of 15 ms per output token for Qwen and 17 ms for Llama. Each card shows the highest-scoring measured RefineKV configuration (τ = 0.01 or 0.1) and the highest-scoring complete baseline configuration within that budget. Baselines include SnapKV, StreamingLLM, PyramidKV, Quest, and Double-P. Actual times are shown beside each method.
Scores are task averages on a shared 0–100 scale. Differences and ratios use unrounded scores. Incomplete evaluations and preliminary single-input timing results are excluded. Timings follow the paper’s 128-output-token decode protocol, excluding prefill; they are separate from the vLLM video. Figure 4 and Table 1 below show the wider quality–latency trade-offs.
Long-context decoding repeatedly reads a large KV cache, even though each query needs only a small part of it. The difficulty is that the useful subset changes with the query. In the paper’s WikiText-2 study, retaining 90% of attention mass takes 17.6% of KV entries per query, while the union of those subsets covers 96.4% of the cache.
This motivates choosing what to access at each decoding step. RefineKV keeps the full cache available, uses compact representatives where they are sufficient, and retrieves exact entries where the current query needs more detail. The next question is how to make that choice without first reading every key.
Summarize the cache once, estimate each cluster’s risk, and refine the clusters that need exact KV entries.
The selection step reads cached cluster summaries. Only clusters that need refinement fetch their original KV entries; the others contribute through size-weighted representatives. The animation follows one query through this process.
1. Summarize. Partition the KV cache into contiguous clusters and cache their mean keys, mean values, and metadata.
2. Select. Score each cluster using the query-dependent attention-error bound: global importance × local distortion.
3. Refine. Expand clusters with RC > τ and attend to the resulting mixed representation.
The full original KV cache remains available. RefineKV reduces decoding reads and computation; τ controls local approximation risk.
Attention mass alone does not tell us whether a cluster can be replaced by its mean. We also need to know how differently its keys respond to the current query. Computing the exact approximation error would require the full KV entries—the reads we are trying to save.
RefineKV instead derives an upper bound on attention-output error and uses its cluster-wise terms as risk scores. Each score combines two factors that can be estimated from cached summaries:
The cluster’s approximate share of attention mass. An error in a more influential cluster can have a larger effect on the output.
A bound on how far individual QK logits can deviate from the representative’s logit under the current query. Similar responses are easier to approximate.
RC > τ → expand to exact KV entries RC ≤ τ → use the representative
The trade-off creates the curved decision boundary below: more important clusters tolerate less distortion. τ is a per-cluster selection threshold, rather than a direct bound on the final generated answer’s error.
Figure 2 / Interactive
Adjust τ to move the decision boundary. Click a cluster to inspect its QK scores, or switch between two real 512-token samples.
Try it: lower τ to expand more clusters, or raise it to use more representatives. “Animate τ” sweeps the threshold; clicking a point reveals its original QK scores. Scores and point positions stay fixed while the decision boundary moves.
The two axes explain why a cluster is selected. The correlation analysis below checks whether higher risk also corresponds to larger measured approximation error.
Select a point to inspect its 10 keys. Arrow keys move between clusters.
Original Figure 2 data: Llama-3.1-8B, WikiText-2, layer 15, query head 31, KV head 7, query position 511 (zero-based). The plot contains 51 complete clusters of 10 keys. The remaining two visible keys contribute to the fixed normalization but are not plotted.
Both samples attend to the same 512-token prefix at query position 511, using query heads 31 and 8 respectively. Scatter axes fit each sample by default; select Shared to compare samples on identical axes. Color scales and the QK score axis remain shared in both views. Switching samples preserves τ. Within each sample, Q, K, the cluster summaries, and each risk score stay fixed as τ changes. The selected boundary is the level set RC = τ. As in the Figure 2 plotting code, this slice uses Vmax for both value-bound terms:
RC = 2 p̂C Vmax [(cosh δC − 1) / D + tanh(δC / 2)]
D = ΣC′ p̂C′ cosh δC′
δC on this plot is a bound on logit deviation, not the empirical variance of the 10 displayed scores. Point color records the original single-cluster replacement error: ‖o − o′C‖₂ with all other clusters exact. It is not recomputed as a total error for the selected τ.
All clusters were reproduced and matched the saved geometry.
Figure 2’s right panels compare RC with the measured error in a cluster’s attention contribution. Across 4,124 observations per model, Spearman rank correlation is 0.945 for Llama and 0.947 for Qwen. Values close to 1 mean that clusters ranked as riskier also tend to have larger approximation errors.
Spearman 0.945 · log Pearson 0.947
Spearman 0.947 · log Pearson 0.947
Figure 2 (right), replotted from the original data: WikiText-2, layer 15, four query heads and 16 query positions per model; clusters with at least two keys. Both panels share log scales. The error here is ‖oC − ôC‖₂ for a cluster’s attention contribution; the interactive view above reports single-cluster replacement error with other clusters exact.
This agreement supports using the score to prioritize which clusters to expand. Correlation measures the quality of that ranking; the decoding demo and benchmark results show how the resulting access decisions translate into speed and task quality.
Quality–latency trade-offs on LongBench, RULER, and HELMET, using Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct.
Read Figure 4 toward the upper left: higher task score with less time per output token is better. Varying τ traces RefineKV’s quality–latency trade-off; Table 1 breaks down representative operating points by task group.
At τ = 0.01, RefineKV retains near-full-attention quality while achieving up to 11.7× attention-phase acceleration and 1.76× faster decoding. The access pattern adapts to each query instead of following a fixed KV budget.
A more aggressive threshold, τ = 0.1, reaches up to 1.82× decode speedup in the paper, with quality effects that vary by task.
Table 1. Each cell reports accuracy / acceleration. Acceleration is relative to Full Attention; higher is better for both. Results from the paper, grouped by benchmark and task.
Scroll horizontally to see all task groups →
| Method | LongBench | RULER | HELMET | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Doc QA | Summ. | Few-shot | Synth. | Code | NIAH | A&T | QA | RAG | LC-QA | ICL | Rec./Rank. | |
| Full Attention | 43.8 / 1.0 | 29.1 / 1.0 | 69.2 / 1.0 | 53.7 / 1.0 | 59.8 / 1.0 | 88.4 / 1.0 | 45.3 / 1.0 | 55.0 / 1.0 | 60.4 / 1.0 | 48.6 / 1.0 | 77.0 / 1.0 | 76.6 / 1.0 |
| SnapKV 10% | 41.2 / 1.7 | 23.4 / 1.7 | 65.9 / 1.7 | 53.4 / 1.7 | 58.5 / 1.7 | 79.5 / 1.7 | 40.7 / 1.7 | 55.0 / 1.7 | 61.0 / 1.6 | 48.4 / 1.6 | 77.0 / 1.6 | 60.7 / 1.5 |
| SnapKV 25% | 42.8 / 1.5 | 26.3 / 1.5 | 67.5 / 1.5 | 53.2 / 1.5 | 59.8 / 1.5 | 84.2 / 1.5 | 41.3 / 1.5 | 55.0 / 1.5 | 60.6 / 1.5 | 49.4 / 1.5 | 78.0 / 1.4 | 71.8 / 1.4 |
| StreamingLLM 10% | 32.7 / 1.7 | 22.2 / 1.7 | 63.3 / 1.7 | 47.5 / 1.7 | 56.1 / 1.7 | 11.2 / 1.7 | 29.8 / 1.7 | 37.5 / 1.7 | 53.5 / 1.6 | 42.4 / 1.6 | 75.0 / 1.6 | 8.6 / 1.5 |
| StreamingLLM 25% | 35.3 / 1.5 | 24.7 / 1.5 | 66.1 / 1.5 | 46.2 / 1.5 | 58.1 / 1.5 | 30.6 / 1.5 | 34.2 / 1.5 | 50.0 / 1.5 | 57.9 / 1.5 | 41.5 / 1.5 | 73.0 / 1.4 | 17.3 / 1.4 |
| PyramidKV 10% | 41.3 / 1.7 | 23.6 / 1.7 | 66.0 / 1.7 | 53.3 / 1.7 | 58.4 / 1.7 | 79.7 / 1.7 | 39.2 / 1.7 | 55.0 / 1.7 | 61.0 / 1.6 | 49.0 / 1.6 | 75.0 / 1.6 | 63.9 / 1.5 |
| PyramidKV 25% | 43.0 / 1.5 | 26.2 / 1.5 | 67.5 / 1.5 | 53.0 / 1.5 | 59.8 / 1.5 | 84.8 / 1.5 | 41.8 / 1.5 | 55.0 / 1.5 | 60.4 / 1.5 | 49.4 / 1.5 | 76.0 / 1.4 | 74.4 / 1.4 |
| RefineKV τ = 0.01 | 43.6 / 1.8 | 29.2 / 1.8 | 69.5 / 1.7 | 53.4 / 1.8 | 59.9 / 1.8 | 84.5 / 1.6 | 46.6 / 1.6 | 55.0 / 1.5 | 60.4 / 1.5 | 48.9 / 1.5 | 76.0 / 1.4 | 76.3 / 1.4 |
| RefineKV τ = 0.1 | 43.3 / 1.8 | 28.4 / 1.8 | 67.6 / 1.8 | 52.4 / 1.8 | 59.4 / 1.8 | 70.5 / 1.8 | 40.0 / 1.8 | 55.0 / 1.8 | 58.5 / 1.7 | 47.9 / 1.7 | 79.0 / 1.6 | 54.2 / 1.6 |
| Method | LongBench | RULER | HELMET | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Doc QA | Summ. | Few-shot | Synth. | Code | NIAH | A&T | QA | RAG | LC-QA | ICL | Rec./Rank. | |
| Full Attention | 43.5 / 1.0 | 26.4 / 1.0 | 69.2 / 1.0 | 54.5 / 1.0 | 63.6 / 1.0 | 73.8 / 1.0 | 56.9 / 1.0 | 42.5 / 1.0 | 49.0 / 1.0 | 51.7 / 1.0 | 76.0 / 1.0 | 39.3 / 1.0 |
| SnapKV 10% | 41.3 / 1.3 | 21.3 / 1.3 | 62.2 / 1.3 | 54.5 / 1.3 | 59.5 / 1.3 | 45.3 / 1.3 | 53.6 / 1.3 | 37.5 / 1.3 | 46.7 / 1.3 | 49.9 / 1.3 | 74.0 / 1.3 | 12.7 / 1.3 |
| SnapKV 25% | 42.8 / 1.2 | 23.9 / 1.2 | 66.2 / 1.2 | 54.5 / 1.2 | 62.0 / 1.3 | 51.2 / 1.2 | 57.9 / 1.2 | 40.0 / 1.2 | 49.2 / 1.2 | 49.1 / 1.2 | 75.0 / 1.2 | 22.2 / 1.2 |
| StreamingLLM 10% | 30.1 / 1.3 | 19.7 / 1.3 | 62.8 / 1.3 | 23.2 / 1.3 | 53.8 / 1.3 | 6.9 / 1.3 | 27.1 / 1.3 | 22.5 / 1.3 | 38.1 / 1.3 | 46.3 / 1.3 | 70.0 / 1.3 | 6.3 / 1.3 |
| StreamingLLM 25% | 33.0 / 1.2 | 22.5 / 1.2 | 65.9 / 1.2 | 33.1 / 1.2 | 57.8 / 1.2 | 22.0 / 1.2 | 30.3 / 1.2 | 40.0 / 1.2 | 39.8 / 1.2 | 47.3 / 1.2 | 75.0 / 1.2 | 10.2 / 1.2 |
| PyramidKV 10% | 40.5 / 1.3 | 22.2 / 1.3 | 61.9 / 1.3 | 54.5 / 1.3 | 60.0 / 1.3 | 38.1 / 1.3 | 52.0 / 1.3 | 35.0 / 1.3 | 47.3 / 1.3 | 49.0 / 1.3 | 71.0 / 1.3 | 10.2 / 1.3 |
| PyramidKV 25% | 42.1 / 1.3 | 24.2 / 1.3 | 66.9 / 1.3 | 54.5 / 1.3 | 62.1 / 1.3 | 48.8 / 1.3 | 54.7 / 1.3 | 40.0 / 1.3 | 45.6 / 1.2 | 48.7 / 1.2 | 75.0 / 1.2 | 13.5 / 1.2 |
| RefineKV τ = 0.01 | 43.0 / 1.3 | 27.7 / 1.3 | 69.8 / 1.2 | 54.5 / 1.3 | 64.2 / 1.3 | 74.4 / 1.2 | 56.7 / 1.2 | 42.5 / 1.1 | 48.5 / 1.1 | 50.8 / 1.1 | 76.0 / 1.1 | 40.5 / 1.2 |
| RefineKV τ = 0.1 | 43.0 / 1.4 | 28.0 / 1.4 | 69.9 / 1.4 | 54.0 / 1.4 | 63.6 / 1.4 | 72.2 / 1.3 | 57.1 / 1.3 | 45.0 / 1.3 | 48.5 / 1.3 | 51.0 / 1.3 | 76.0 / 1.3 | 34.2 / 1.3 |
Accuracy is the task-macro benchmark score; acceleration is Full Attention’s task-macro time per output token divided by the method’s. Baseline percentages are KV budgets. Shading and underlines reproduce the paper’s table.
@misc{lian_refinekv,
title = {RefineKV: Query-Adaptive KV Cache Compression
via Attention-Error Bounds},
author = {Lian, Kewei and Zhai, Jiayi and Liu, Anji},
note = {Manuscript}
}