RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Authors: Payman Behnam
Yaosheng Fu
Ritchie Zhao
Po-An Tsai
Zhiding Yu
Alexey Tumanov
Venue: ICML 2025 poster (2025)
Keywords: LLM inference
KV cache compression
Efficiency
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: RocketKV achieves significant KV cache memory bandwidth and storage savings  during decode stage of LLM inference while maintains comparable accuracy to full KV cache attention.

Abstract: Transformer-based Large Language Models rely critically on the KV cache to efficiently handle extended contexts during the decode phase. Yet, the size of the KV cache grows proportionally with the input length, burdening both memory bandwidth and capacity as decoding progresses. To address this challenge, we present RocketKV, a training-free KV cache compression strategy containing two consecutive stages. In the first stage, it performs coarse-grain permanent KV cache eviction on the input sequence tokens. In the second stage, it adopts a hybrid sparse attention method to conduct fine-grain top-k sparse attention, approximating the attention scores by leveraging both head and sequence dimensionality reductions. We show that RocketKV provides a compression ratio of up to 400×, end-to-end speedup of up to 3.7× as well as peak memory reduction of up to 32.6% in the decode phase on an NVIDIA A100 GPU compared to the full KV cache baseline, while achieving negligible accuracy loss on a variety of long-context tasks. We also propose a variant of RocketKV for multi-turn scenarios, which consistently outperforms other existing methods and achieves accuracy nearly on par with an oracle top-k attention scheme. The source code is available here: https://github.com/NVlabs/RocketKV.

RocketK V: Accelerating Two-Stage KV Cache Compression
LLM Inference via
Payman Behnam 12 Yaosheng Fu "
Ritchie Zhao
Po-An Tsai
Zhiding Yu'
Alexey Tumanov
Abstract
tation becomes major bottleneck during the decode phase, as its size scales linearly with both the sequence length and batch size. For example, the Llama3. 1-70B- Instruct (AI@ Meta 2024) model with a batch size of 32, and a context length of 32K requires around 320GB of KV cache storage at FP16 precision, which even advanced hard-
Transformer- basseddeer---ddddd--e Large Language Models rely critically on the KV cache to efficiently handle extended contexts during the decode phase. Yet. the size of the KV cache grows proportionally with the input length, burdening both memory bandwidth and capacity as decoding progresses. To address this challenge, we present RocketK training-free KV cache compression strategy containing two consecutive stages. In the first stage, it performs coarse-grain permanent KV cache eviction on the input sequence tokens. In the second stage, it adopts a hybrid sparse atten- tion method to conduct fine-grain top-k sparse attention, approximating the attention scores by leveraging both head and sequence dimensional- ity reductions. We show that RocketKV provides compression ratio of up to 400x, end-to-end speedup of up to 3.7x as well as peak memory reduction of up to 32.6% in the decode phase on an NVIDIA A 100 GPU compared to the full KV cache baseline, while achieving negligible accu- racy loss on variety of long-context tasks. We also propose variant of RocketKV for multi-turn scenarios, which consistently outperforms other existing methods and achieves accuracy nearly on par with an oracle lop-k attention scheme. The source code is available here: https:// github. www..ccoo//w.a..ee coo.////lllll//ccekkkkkk RRcketkK/
ware (NVIDIA) 2024
Jouppi al. 2023
AMD 2024
hardly handle.
Fortunately, previous work has shown that only small subset of KV tokens is required at each decode step to maintain accuracy (Zhang et al. 2023 Tang et al. 2024 Ribar et al. 2024 ei et al. 2024 Singhania et al. 2024 Ge et al., 2024 Therefore, if those KV tokens can be accu- rately predicted in advance, dense attention operations can be replaced with sparse attention operations with significant memory bandwidth and capacity improvement. These meth- ods often fall into two categories: 1) permanent KV token eviction, and 2) dynamic KV token selection. The former results in both memory bandwidth and storage savings, but could lead to noticeable accuracy loss if KV tokens dropped earlier are needed by later decode steps. The latter avoids this shortcoming by keeping all KV tokens in memory and dynamically selecting subset each time. Hence, it only results in memory bandwidth savings, but often requires extra memory storage overhead for auxiliary data.
To understand the effectiveness of existing KV token drop- ping methods, Figure I presents the accuracy comparison of four different methods (two in each category) for qasper benchmark in LongBench (Bai et al.
2024 the Mistral- 7B-Instruct-v0.2 model (Jiang et al. 2023). We observe that the accuracy of all four methods drops significantly as the token budget becomes lower than 1024 while an oracle top-k attention scheme achieves negligible accuracy drop even with token budget of 256. This indicates that existing practical methods fail accurately predict top-k KV tokens under low token budgets.
1. Introduction
In Transformer-basee LLM inference (Vaswani al. 2017 Pawar et al.| (2024), the key-value cache (KV cache)-whiee which stores past attention keys and values to avoid recompu-
'Equal contribution: work done during internship/part-iin at NVIDIA 'NVIDIA. Santa Clara, USA GGeorgia Institute of Tech- nology, Atlanta, USA. Correspondence to: Payman Behnam <pay- man. behnam @ gatech.edu>>, Yaosheng Fu <-fu@@ mvidia.com>.
To improve the prediction accuracy, we propose RocketK a two-stage KV cache compression method that combines permanent KV token eviction with dynamic KV token selec- tion to accelerate the decode phase of LLM inference. The RocketKV framework enables flexible integration of wide range of KV cache compression techniques at each stage. To
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

RocketKV; Accelerating
LLM Inference via Two-Stage KV Cache Compression
SnapKV in the first stage. We design adaptive com- pression decomposition mechanism to balance the com- pression between the two stages.
30
QASPER Score
25
Full-KV
20
Exact-TopK
We conduct comprehensive evaluation of RocketKV and MocketKV-MT on wide variety of models and downstream tasks. RocketKV consistently demon- strates comparable accuracy with full KV attention at up to x00x compression ratio, while achieving up 3.7x end-to-end speedup and 32.6% peak memory saving the decode phase NVIDIA A A100 GPU, Meanwhile, RocketK is more suitable for multi- turn decoding and performs par with oracle top-k attention.
15
DuoAltention
10
SnapkV
Quest
5
SparQ
256
512
1024
2048
4096
Token Budget
Figure 1: Existing KV token dropping methods fail to match the accuracy scores of oracle top-k attention (Exact- TopK) on Mistral-7B-Ins-v0.22 in the qasper benchmark.
2. Related Work
A number of recent approaches have focused on improving the efficiency of attention mechanisms in LLMs, particu- larly when dealing with long contexts. One feasible solu- tion is KV cache sharing across multiple layers (Brandon) et al. 2024) or selectively dropping attention for some lay- ers (Ling et al. 2024). Others (Sun et al. 2024 Ho et al.. 2024) propose mixed attention designs where some layers use global attention while others use local attention. Multi- Query Attention (MQA) (Shazeer
achieve optimal performance, we directly adopt an existing method SnapKV (Li et al. (224) for coarse-grain permanent KV token eviction in the first stage. In the second stage, we propose hybrid sparse attention (HSA) to perform fine-grain dynamic KV token selection, which estimates KV token indices with top-k attention scores via two-dimensional re- ductions. To balance the compression between two stages, RocketKV introduces an adaptive compression decompo- sition mechanism that intelligently splits any given target compression ratio across two stages. Combining these two stages together, RocketKV achieves significant memory bandwidth and storage savings with negligible accuracy loss across a wide variety of models and downstream tasks.
2019). Grouped-Query 2023), and Multi-head La- 2024 widely adopted
Attention (GQA) (Ainslie et al.
tent Attention (MLL)) (L.iu et al. by many recent LLMs (Al@ @ Metal
2024 Jiang et al. 2023
Team et al. 2024 Liu et al. cache by sharing key-value pairs across multiple attention heads. All above techniques directly modify the attention architecture and require integration since pre-training.
2024b), which reduce the KV
Moreover, recent work (Li et al. 2025) demonstrates that permanent KV token eviction suffers in multi-turn decod- ing because important KV tokens vary significantly across queries. To address this challenge, we propose a variant of RocketKV called PocketKV-MT for multi-turn scenarios where we do not evict the unselected KV tokens in the first stage but keep them all for later turns. Yet, the decode phase is still restricted to perform dynamic selection on the filtered KV tokens from the first stage in each turn. By doing this, RocketKV-MF V-T achieves the same memory traffic savings as RocketKV V but does not introduce memory storage savings.
Another direction is to improve attention efficiency with training-free techniques. These techniques can be further categorized as prefill phase acceleration, decode phase ac- celeration, or accelerating both phases together. For exam- plc, StreamingLLM (Jiang et al.) 2023) combines initial and locdl-window attention reduce the KV cache into constant size regardless of the sequence length. DuoAt- tention (Xiao et al. 2025) and RazorAttention (Tang et al. 2025) improve upon StreamingLLM by applying global attention on retrieval heads and StreamingLLMM L.MiinggLLL---yleeeee M--style atten- tion on rest heads. While these techniques can be applied to both the prefill and decode phases, other techniques focus on accelerating only one of the two phases.
In summary, make the following contributions:
We analyze and identify the limitations of existing KV token dropping methods and then propose RocketKV as two-stage KV cache compression scheme. We further introduce a variant RocketK V-MT for multi- turn decoding scenarios.
For prefill phase acceleration, MInference (Jiang et al. 2024) identifies three distinct patterns in the attention ma- trix that can be harnessed for efficient sparse operations with customized GPU kernels. SeerAttention (Gao et al.) 2024) explores dynamic block-level sparsity in the attention module with a learnalle gate. Mcanwhile, XAttention (Xu et al. 2025) leverages the sum of antidiagonal values the
We propose a hybrid sparse attention (HSA) for dy- namic KV token selection with wwo-dimensional re- ductions in the second stage while directly adopting
2

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
attention matrix for estimating block importance.
CDF
A common approach for attention acceleration at the decode phase is through permanent KV cache eviction, which can save both memory bandwidth and storage requirements. H2O (Zhang et al. 2023) observes that a small subset of tokens, known as heavy-hitters, dominates the attention computation, thus only keeps recent and heavy-hitter tokens. SnapKV (Li al., 2024) employs an observation window at the end of the prompt to identify critical KV tokens of the input prompt. Then it uses a clustering algorithm via pooling to retain critical KV token clusters without losing information completeness. Ada-KV (Feng et al.] 2024) proposes an adaptive budget allocation strategy to provide better token budget utilization across individual heads.
0.8
0.6
Max seq. length
0.2
Unique topkindices
0
5,000
10,000
15,000 20,000 25,000
Figure 2: The CDF of maximum sequence length and num- ber of unique top-k indices (k=256) across all decoding steps for 200 questions in the gasper benchmark. Data was collected from random head (layer 31 head 0) in Mistral- 7B- 1B-Instruut-vO.2.
Meanwhile, Quest (Tang al. 2024) observes that perma- nent KV cache eviction could lead inevitable accuracy loss and proposes query-aware selection on top-k KV to- kens based on approximation attention with representative vectors of contiguous key cache pages. On the other hand, QparQ and Loki (Ribar et al., 2024; Singhania et al., (022)) conduct approximation attention by selecting only impor- tant indices on the head dimension instead. The above approaches can save KV cache compute and data fetching from memory, but not KV cache storage. InfiniGen (Lee et al. 2024 tackles this challenge by offloading the entire KV cache to CPU memory and only fetching the selected KV tokens to GPU memory when needed. MagicPIG et al., 2025) discovers that using importance sampling is more efficient than top-k estimation and proposes an ap- proximation attention solution leveraging locality sensitive hashing (LSH) and CPU offloading.
qasper benchmark. The reason behind this analysis is that in order to match the accuracy of Exact-TopK, we need to keep all important KV tokens that are selected by at least one top-k attention operation across all decoding steps. As shown in Figure 2 although the maximum sequence length can reach as high as 25000, the number of unique top-k indices only goes up to 1200. This implies that, ideally, permanent KV token eviction method should be able to close the accuracy gap under a token budget of 1200. To further reduce the token budget, we realize that dynamic KV token selection can be applied on the filtered KV token set after permanent KV token eviction. Since the filtered set much smaller than the original full KV cache, the difficulty for accurate top-k prediction is greatly reduced. Therefore, an ideal solution would be to perform permanent KV cache eviction with a larger token budget first and then conduct dynamic KV token selection on the remaining KV tokens. This fusion evicts unimportant tokens and also makes the dynamic selection more accurate, motivating our RocketKV design.
In contrast to prior work, RocketKV combines both perma- nent KV cache eviction and dynamic KV token selection as two consecutive stages for accelerating the decode phase and utilizes the advantages of both worlds. As result, it achieves remarkable memory bandwidth and storage savings without the need for sophisticated system-level optimiza- tions such as CPU offloading.
3.2. RocketKV Overview
Based on our observation, we propose RocketKV, a two- stage KV cache compression method for decode phase ac- celeration. As shown in Figure 3 RocketKV performs boarse-grain KV cache eviction in the first stage. The pur- pose of this stage is to remove KV tokens with low impor- tance while keeping the majority of important tokens. In the second stage, conducts fine-grain dynamic KV token selection on the remaining KV tokens, followed by top-k sparse attention. The RocketKV framework is generie, and many existing KV cache compression methods can fit into the corresponding stage. For example, SnapK' (Li et al. 2024) or Ada-KV (Feng et al., 2024) can be used for the first stage, and Quest (Tang et al. 2024) or SparQ (Ribar et al.) 2024] can be applied to the second stage. To achieve the best performance at each stage, we directly adopt SnapKV
3. Proposed Method: RocketKV
3.1. Observation
As shown previously in Figure [ existing token dropping methods, regardless of permanent KV token eviction or dynamic KV token selection, fail to match the accuracy of oracle top-k attention (Exact-Topk) under low token budgets. To further understand what causes the accuracy mismatch, we analyze a random attention head (layer 31 head 0) from -B-nlltttttntttt--o.... and present the cu- mulative distribution function (CDF) of both the maximum sequence length and number of unique KV indices selected by Exact-TopK (k=256) across all decoding steps in the
3

---

RocketKV; Accelerating Long-Contert LLM Inference via Two-Stage KV Cache Compression
Input prompt
Generated takens one by one
sky blue and clear today,
what
the color
the sky
blue
B
14 15
8
11
0
15
First Stage: Permanent partial KV cache eviction
2
3 8
11
14
15
Second Stage: Dynamic partial KV token selection at each decode step
Prefill phase
Decode phase
Figure 3: Overview of RocketKV with two consecutive stages.
in the first stage while proposing a hybrid sparse attention (HSA) method in the second stage.
one-dimensional sparsity can only achieve a certain degree of compression ratio, beyond which the accuracy could drop rapidly, as shown previously in Figure In contrast, we pro- pose hybrid sparse attention (HSA), which takes advantage of two-dimensional reduction, in both sequence and head dimensions together achieve better estimation accuracy on KV token indices with top-k attention scores.
3.3. First Stage: SnapKV
In the first stage, we directly adopt SnapKV for permanent KV token eviction on the input KV tokens. SnapKV's key idea to rely on the aggregated attention scores between the input context and observation window in the end select the most relevant tokens to keep within the input prompt. The original SnapKV method selects crucial KV tokens attention head basis. In case of grouped-query attention (GQA), each attention head within an attention group keeps a separate set of KV cache tokens, which could introduce re- dundant storage of the same KV token. To reduce KV token storage with GQA, we follow the V work to perform token selection on per-group basis according aggregated per-group attention scores that the selected KV tokens are shared across the entire attention group. SnapKV uses pooling along the sequence dimension to ensure critical KV tokens are selected along with their neighbor tokens. It demonstrates better accuracy with pooling because it retains the completeness of selected information. The employed kernel sizes for pooling are quite small (e.g., kernel size of 7 for LongBench (Bai al., 20244). Since in our case, SnapKV is only used for coarse-grain KV token eviction at the first stage, we discovered that the optimal kernel sizes for pooling are much larger. We empirically set the kernel size to 63 in all our experiments.
F the detailed implementation of our proposed algorithm, which inspired by Quest (Tang et al. 2024) and SparQ (Ribar et al. (2024). Our HSA algorithm can be decomposed into three steps:
Group tokens in key tensor into consecutive pages along the sequence dimension and store element- wise maximum (Kmax) and minimum (Kmin) values of each page auxiliary storage similar to Quest. Un- like Quest, they are stored with different layout by aligning along the head dimension to enable efficient gathering in Step 2. The auxiliary storage is updated accordingly each time a new key token is generated.
Step 1:
Step 2: For each query q, find ki largest absolute values along the head dimension. Then, fetch only the cor- responding indices in either element-wise maximum or minimum tensors, depending on the sign of q at those indices. The goal is to compute element-wing max(q x Kmax. @ x K for each page to approx- imate the highest possible attention scores within a page. To further reduce the approximation overhead, we only calculate on 1 partial positions along the head dimension with large magnitude of q and ignore oth- similar to SparQ. Once the approximation attention scores are calculated, ke indices with the largest atten- tion scores along the sequence dimension are selected.
3.4. Second Stage: Hybrid Sparse Attention
Previous methods on dynamic KV token selection often estimate top-k KV indices with reduced computation along a single dimension. For example, Quest (Tang et al., 20244 uses element-wis. minimum and maximum values to rep- resent continuous pages along the sequence dimension. Meanwhile, SparQ (Ribar et al. 2024] and Loki (Singha- nia et al.) 2024) leverage sparsity in the head dimension to conduct low-rank estimations. Unfortunately. relying on
Perform sparse attention by fetching the origi- nal key and value vectors from the predicted k indices.
Step 3:
Our HSA algorithm fully compatible with GQA. To achieve this, we perform all key tensor selections on per at- tention group basis. More details are shown in Algorithm ]

---

RocketKV; Accelerating Long-Contexl LLM Inference via Two-Stage KV Cache Compression
Sequence dimesion
Algorithm HSA Algorithm (only contains step and 3) Input:query vector key tensor K, value tensor V, element-wise max/min key tensor Kmaz/ Kmin
Head cimertioon
K
# get top-kı indices along head dim from sum of lgl in group dim
Page
Page
Page 3
Page 4
i1 t- argtopk(sung
dim group). kn)
Element-wise Max Element-wise Min
# get signs of top-k1 indices from sum of q group dim
Heed dima immrraao
g sign(sum( dim group))
ki
Approximate ettention score
# fetch corresponding indices from paged min max
P &- Kmanjigg 20] Kminlig <0)
k
# compute approximation attention scores
81 t score( P)
# get indices with top-k2 attention scores along seq. dim
b2 & argtopk(s;, k2)
# perform sparse attention
y f attn(g, K (x2,, V(2))
Page
Page 2
Page 3
Page
return
Figure 4: Illustration hybrid sparse attention with figure (left) and algorithm (right).
where we perform the sum of q or |g| in the group dimen- sion as needed to guarantee that all attention heads within group are making the same selection at each step.
on this whole input KV cache to select a new subset of important tokens for this turn's decode phase. By follow- ing this strategy, RocketK V-MT achieves similar decoding speedups to RocketKV in each turn, while retaining the full KV cache history across all turns. This approach effectively climinates the accuracy degradation caused by permanent KV token eviction in multi-ture scenarios at the cost of no memory storage savings.
3.5. RocketKV-MF
In multi multi-turn conversation setting, KV tokens that were pruned in earlier turns may become essential for answering queries in later turns because KV token importance could vary significantly between different turns (Li et al. 2025) Consequently, permanently removing these tokens could lead to a noticeable accuracy drop in subsequent turns.
3.6. Adaptive Compression Decomposition
The decode phase of LLM inference is typically memory bound (Ribar al., 2024); so, the time spent in the attention module is roughly proportional to the total memory traffic. In this work, we use the token budget t to estimate the amount of memory traffic for each attention operation in the decode phase (we mainly focus on KV cache traffic since it contributes to the majority of memory traffic in this scenario). For example, a token budget of 512 means each attention module needs to fetch an equivalent total amount of 512 key and value pairs from memory. Unlike prior work (Tang et al. 2024) where the token budget only reflects the memory traffic for top-k attention (Step 3 in HSA), we define the token budget to also include the memory traffic of top-k estimation (Step 2 in HSA). By doing this, the token budget can reflect the overall memory traffic in the attention module more precisely. For simplicity, we evenly split the token budget between these two steps for HSA and all other dynamic KV token selection methods in our later experiments. For models with GQA, this token budget is defined for the entire attention group rather than each attention head. For a given sequence length of S, the total
To mitigate this issue, introduce multi-turn variant of approach called RocketKV-MTI The core idea in RocketK is to avoid permanently evicting any KV tokens during the first stage; instead, all KV tokens are re- tained in memory across turns, ensuring that no potentially useful context is lost. Mcanwhile, to preserve computational efficiency, the second stage still performs dynamic selec- tion over the subset of KV tokens filtered by the first stage, similar to the original RocketKV. In other words, the model generates responses using reduced KV token set for speed, while reserving the full KV history for future turns. For example, suppose the first stage of RocketK V-MT retains only out of M total KV tokens from the input prompt in the first turn. KocketKV-MTT will still keep all M tokens in memory but restrict the second stage to dynamically select from these N input tokens (plus any new generated tokens) during the decode phase. In the next turn, the full set of previously stored KV tokens (all M input tokens plus all output tokens) is added to the new input KV cache. The filtering process (with SnapKV V in our case) is applied again
5

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
compression ratio of can be defined as S/L.
Table I: Normalized KV cache storage (including auxiliary data) and traffic comparison between RocketKV, RocketKV. Rockket MT and other methods.
Since RocketKV is two-stage KV cache compression framework where the filtered KV token set by the first stage serves as the imput of the second stage for dynamic selection, it is important to determine the intermediate token budget for the filtered KV token set. For an overall compression ratio of c, we define a split factor r so that is split into c" for the first stage and c(1-r) for the second stage, where 0 <= <= 1. We use the following formula to adaptively determine r:
Method
Compression Ratio
Storage
Traffic
Fu
DuoAttention
C
1/c
1/c
SnapKV
1/c
1/c
Quest
C
1 1/c
1/c
SparQ
C
2
1/c
RocketKV
1/c² + 22/2tt///2
1/c
RocketKV-MI
C
1+ 1+2///t///
1/c
r= (0.2 + 0.06 (c) 0.8)
The insight behind this formula is that when is small, we would like to minimize the number of permanent evicted KV tokens in the first stage to prevent information loss. As c increases, the accuracy drop caused by HSA in the second stage gets larger thus it is better to assign a higher compression ratio the first stage because SnapKV can estimate important KV tokens more precisely using exact attention scores. We further limit the range of r between 0.2 and 0.8 to balance the compression decomposition between these two stages.
both element- maximum and minimum tensors intro- duce a storage overhead of 1/c(l-r)/2 on top of it. There- fore, the total KV cache storage and traffic in Rocket V are 1/c" + 2/ (1++)/22 and 1/c of the full KV baseline, respec- tively. RocketKV-MT does not lead to storage saving in the first stage, so its relative KV cache storage is 1 1 2/c(++r//2 instead. Table compares the KV cache storage and traffic of RocketKV and RocketKV-MT against other methods at a given compression ratio We can see that while all methods lead to the same KV cache traffic savings, only RocketKV, DuoAttention and SnapKV provide additional KV cache storage savings, but RocKetKV-MT, Quest and SparQ require extra storage for auxiliary data.
For compression decomposition within HSA, we simply split the compression ratio c(1-r) evenly between the se- quence and head dimensions so that each dimension gets compression ratio of c!!-r)/2 Notice that the compression ratio along the sequence dimension is equivalent to the page size, so we need to round it up the nearest integer [c(1-r)/2 and the other dimension gets c(1-r)/ / (c(22)//21
3.7. System Implications
RocketKV is fully compatible with FlashAttention (Dao et al. 2022) because it does not modify attention in the pre- fill phase. Additionally, it seamlessly integrates with tensor parallelism (Shoeybi et al. 2019) because all operations are symmetric across attention heads/groups. It is worth noting that both RocketKV and RocketKV-MTT work well with the disaggregated serving system, where different GPUs are used for prefill and decode phases (NVIDIA) 2025 Qin et al. 2025). For RocketKV, the KV cache storage is re- duced on both the prefill and decode GPU, as well as the KV cache transfer traffic between them. While the full KV cache needs to be stored in the prefill GPU for RocketK MT. only the filtered set needs to be transferred and stored in the decode GPU, resulting in the same communication and decoding benefits as RocketKV.
Because RocketKV strategically decomposes KV cache compression across multiple stages and dimensions, it sig- nificantly enhances the potential for high compression ratios while maintaining strong securacy-surpasaee that rely on single-stage, single-dimension approaches. For example, given compression ratio of 64x, the split factor can be calculated as r = 0.2 + 0.06 * 1og2(64) = 0.56. Therefore, the compression ratio is split into 64 = 10.3x in the first stage and (4.))))6)) 6.2x in the second stage. HSA further splits its compression ratio into 3x in the sequence dimension (with a page size of 3) and 2.1x in the head dimension. We can see that each individual stage and dimension gets assigned with a much smaller compression ratio after decomposition.
methods
4. Experiments
In RocketKV, the first stage results in both KV cache storage and traffic reduction of c". In the second stage, we need to take the additional memory storage overhead introduced by the approximation attention into consideration. Since we evenly split the compression ratio c(1-r) into ((1-r)/2 (ignore the round up operation here for simplicity) between two dimensions in HSA, it introduces a memory storage overhead of (1/c") x (1/c()-r) x 2 = 2/c(1+r//2 where 1 1/c" is the relative KV cache storage after the first stage and
4.1. Experimental Settings
We conduct our experiments on three widely used long- context models: Ilama3. 1 --8t-unstt (AI @ Meta 2024 (Jiang et al. 2023), and LongChat- 7B-v1.5-32k (Li et al. 2023). We refer to these models as Llama3. 1amma.11-BB-o 1-8--Is,, Mistral- 7 Ins-v0.2, LongChat-7B--i 1.5, respectively, throughout the paper. The first two models are
6

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Full-KV
Exact-TopK
DuoAttention
SnapKV
Quest
SparQ
RocketKV
Average Score
Average Score
Average Score
30
20
10.
256
512
1024
2048
4096
256
S12
1024
2048
4096
256
512
1024
2048
4096
Token Budget
Taken Budget
Token Budget
(a) LongBench, Llama3.1-58 1-8B-Ins
(b) LongBench, Mistral-BB-. Mistal-7B-Ins-vO..
(c) LongBench, LongChat-77-v1.5 78-v1.5
100
100
Average Score
Average Score
Average Score
80
60
60
40
40
20
20
$12
1024
2048
4096
1024
2048
4096
256
512
1024
2048
4096
Taken Budget
Token Budget
Taken Budget
(d) NIAH, Llama3,1-88 InS
(c) NIAH, Mistral-7-IIns-v0...
NIAH. LongCha-7B--155
Average Score
Average Score
Average Score
60
8
40
Average
&
2008
2090
31
1024
2049
1024
2000
4006
1024
4090
Token Budget
Token Budget
Token Budget
Taken Budget
(g) Iaama, 1. I-SB 1-BB-In,, Ins, Seql 1en166
(h) Llama3.1-8P 1-8--Ins, Ins, Seql enn 32K
Jama3, 1-8--Ins, en
(j) Ilama3, 1-8B.In.. S6K
Average Score
Average Score
Average Score
Average Score
& 8
80
80
88
1924
2008
1054
2040
102
2048
Budget
Token Budget
Token Budget
Taken Budget
(k)
Mistral-7B-Ims--d 2..ta........LLe33lt
(I) Histral-7B-Ins--DD Stggs--0.....LLe--eeeeee.
(m) Mistral-BB -a- Nt
(n) Misttal-B--In-- 1.--ian-...2.Sqqleeee
Average Score
Average Score
8
Average Score
Average Score
8
8
1024
2048
4096
1024
2018
4086
Token Budget
Taken Budget
Taken Budget
Taken Budget
LongChat-7B..o S.gL
(p) LongChat- 1.5 Seglenn 16K
(q)
LongChat- B... Seglen
Figure 5: Comparing the accuracy of RocketKV with other methods on LongBench (a-c), Needle-in -aaale-i-a-Hayystckk (NIAH) (d-f), and RULER with various sequence lengths (g-r).
based GQA, while the last one is based MHA. For downstream tasks, we utilize LongBench (Bai al. 2024). Beedle-i---Haasttad (Kamradt 2023), RULER HHiih et al. 2024), and SCBench (Li et al. 2025). We compare RocketKV and RocketK V-MT with several other methods:
Full-KV, Exact-TopK, DuoAttention (iiao et al. 2025) SnapKV (Li et al. 2024), Quest (Tang et al. 2024). and
SparQ Ribbar et al. 2024 Exact-TopK serves as an ora- cle method for sparse attention with exact top-k KV token selection. Notice that RocketK acts the same as Rock-

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
e in single-turn scenarios, so we only evaluate it on SCBench under multi-turn mode. We evaluate all methods across various KV token budgets per attention group, except for Vull-KV.. RocketKV, PocketKV-MMT Quest, and SparQ involve additional memory traffic for top-k approximation, which is converted into an equivalent KV token budget such that the total token budget precisely reflects the total amount of memory traffic in the attention module. More details on our experimental settings are available in Appendix [A] We also conduct additional ablation studies in Appendix B.1
Full-Ky
Exact-TopK
Average Score
DuoAltention
40
SnapKy
Quest
30
Spart
Rockettty
RocketKV.M
20
1024
2048
4096
8192
18384
Taken Budget
Figure 6: Comparing the accuracy of RocketKV and
4.2. Accuracy Results
RocketKV-MT with other methods on Llama3.1 1-8B-Ins for SCBench.
In our accuracy evaluation, we vary the token budget of each method from 256 to 4096 for all single-turn benchmarks and from 1024 to 16384 for SCBench under multi-turn mode. We compare the average accuracy across all individual tasks for each benchmark. More detailed results with accuracy breakdown for individual tasks are available in Appendix
lengths as shown in Figure (rows 3-5). Again, RocketK V shows robust performance and a clear advantage over all other methods across various token budget and sequence length settings. Overall, we can see that the accuracy loss of RocketK negligible under short sequence lengths and gradually becomes larger as sequence length increases. We believe this is because these models are less robust sparse attention beyond their effective sequence lengths as defined in the RULER paper (Hsieh et al. (2024). Notice that the ac- curacy gaps between RocketK and other methods become even wider under longer sequence lengths.
LongBench Benchmark:
The first row in Figure shows the average score comparison of RocketKV with other meth- ods LongBench across all three models. Based the figure, we can see that RocketKV consistently outperforms all other methods, especially under lower token budgets. For Llama3. 1-8B--8Bnn, RocketKV achieves almost no accuracy loss with a token budget of 512 and above, and only 1.1% av- erage accuracy drop with a token budget of 256. RocketKV results in slightly higher accuracy losses for Mistral-7B- Ins-v0.2 and LongChat-7B-- BB-11., which might be because Llama3. 1-BB is better-trained than the other two models, making it more robust to sparse attention methods, as sim- ilar trend can be found with Exact-TopKe SparQ performs well on Llama3.1-88
Figure(opreeents the accuracy com- parison of various methods on SCBench under the multi- turn setting. While RocketKV still outperforms other meth- ods under low token budgets, it provides lower accuracy than SparQ for token budgets of 8192 and beyond. More- over, there is still a noticeable accuracy gap between Rock- etKV and Exact-TopK under all token budgets, showing room for further improvement. As we discussed earlier, this is primarily caused by the fact that KV token importance varies significantly across different turns, the unimportant KV tokens evicted by earlier turns could lead to signif- icant accuracy drop in later turns. After fixing this issue in RocketK V-MT, we can see that RocketKV-MT achieves a significant accuracy boost with comparable accuracy to Expct-TopK across all token budgets.
SCBench Benchmark:
and LongChat-7B-vl with to- ken budget of 1024 and above but underperforms all other methods on Mistral-7BBIIn--00.2 SnapKV alone achieves relatively good accuracy under low token budgets, but we still observe widening accuracy gaps between SnapKV and RocketKV as the token budget decreases.
Heedle-in-a-Haasttakk (NIAH) Benchmark: Presented in the second row of Figure 5 RocketKV achieves near the Full-KV accuracy across all models, even with token bud- get of 256. In fact, it reaches 100% accuracy on Llama3. 1- 8B-Ins under 256 token budget, which corresponds to compression ratio of over 400x with a maximum sequence length of 109K tokens. In contrast, all other methods suffer from substantial accuracy drops. While SnapKV leads to comparable accuracy to RocketKV on Llama3. --ama..--B-IIns. its accuracy for Mistral-7B- -i-tral-7B-Ins-v0..2 and LongChat-7B-v1.. decreases by more than 20% compared to Full-KV, even with token budget of 1024. We present the heatmaps of RocketKV under various token budgets in Appendix B.2
4.3. Efficiency Results
Our efficiency experiments are conducted with Llama3. 1- 8B-Ins (AI@ Metal 2024) under FP16 precision, running on NVIDIA H100 and A100 GPUs at a batch size of 1. Sim- ilar SparQ, we leverage gpt-fast (PyTorchLabs) 2023). low-latency python-native LLM framework for rumning the efficiency experiments. We found pbthon-based im- plementation of RocketKV under gpt-fast is sufficient to demonstrate its efficiency benefit, but it could be further im- proved with customized CUDA kernels and more advanced frameworks such FlashInfer (Ye et al. 2025). Figure|
RULER Benchmark: For the RULER benchmark, we eval- uate all methods across three models with varying sequence
8

---

RocketKV; Accelerating Long- Long-Context LLM Inference via Two-Stage KV Cache Compression
Speedup over Full-VV
Speedur over Full-KV
Memory Savings (%)
RacketKV
RocketKV
30
3
RocketKV-512
RocketKV-1024
BocketKV-1024
RacketKV 2048
RocketKV 2048
20
RocketKV-256
RacketKV.400
RocketKV.4096
Racket
2
10
RacketkV.1024
RocketiRV-2048
>
RacketKKV4499
20k
40k
60k
80k
100k
20k
40k
60k
80k
100k
20k
40k
60k
80k
100k
Sequence Length
Sequence Length
Sequence Length
(a) End-to-end speedup on A100
(a) End-to-end speedup on H100
(b) Peak memory saving
Figure 7: End-to-end speedup and peak memory savings of RocketKV with various token budgets compared to Full-KV.
demonstrates the end-to-end speedups and peak memory savings of RocketKV with varying token budgets at the de- code phase, with all values normalized to Full-KV, V. Notice that the peak memory usage is the same regardless of the underlying GPU, and it measured includes memory al- locations for weights, activations, KV cache, and all other metadata at the decode phase.
Acknowledgment
We thank the anonymous ICML reviewers for their valuable and constructive feedback. We are grateful to David Nellans for overseeing the internship process of Payman Behnam, during which this work was supported by the NVIDIA Grad- uate Fellowship. We also thank Song Han and Christos Kozyrakis for their insightful discussions and helpful feed- back.
As shown in the Figure 7 RocketKV V achieves up to 3.7x x and 3.3x end-to-end speedups on an A100 and H100 GPU, respectively, as well as up to 32.6% peak memory saving. The maximum speedup on A100 is 12% higher than H100 because A100 has lower memory bandwidth to compute ratio compared to H100. As result, L.M inference execu- tion is more memory- -ound on A100 and can benefit more from the memory traffic savings of the KV cache offered by RocketKV. We expect the speedup of RocketKV will be even higher on cheaper GPUs such as RTX 4090/5090 since they are not equipped with High Bandwidth Memory (HBM).
Impact Statement
This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
References
Al@ Llama Model Card. 2024.
Ainslie, J., Lee-Thorp. J., de Jong, M., Zemlyanskiy, Y.,
5. Conclusion
Lebrón, F., and Sanghai, S. GQA: Training generalized multi-query transformer models from multi-head check- points. Conference on Empirical Methods in Natural Language Processing (EMNLP), 2023.
RocketKV presents novel, training-free approach to KV cache compression, addressing the challenges of memory bandwidth and capacity demands during the decode phase of
LLM inference. RocketKV contains two consecutive stages: SnapKV for coarse coarse-grain grain permanent KV cache eviction and hybrid sparse attention (HSA) for fine-grain dynamic KV token selection. Our evaluations various models and long-context benchmarks demonstrate that RocketKV main- tains comparable accuracy to full KV cache attention while significantly lowering memory bandwidth and capacity us- age with a compression ratio of up to 400x, as well as up to 3.7x end to end speedup and 32.6% peak memory reduction at the decode phase, highlighting its efficiency and potential for widespread application in optimizing LLM performance. We also propose variant of RocketKV called RocketK V-MT that is optimized for multi-tun turiii-tunn scenarios.
AMD. InstinctUM MI300X Accelerators. https://
www. amd. http://////oooooccee/eeeeeeeee/e lerators/
2024. Accessed:
2025-01-07.
Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., et al. LongBench: A bilin- gual, multitask benchmark for long context understanding. Annual Meeting of the Association for Computational Lin- guistics (ACL), 2024.
Brandon, W. Mishra, M.. Nrusimha, A., Panda, R., and Ragan Kelly, J. Reducing transformer key-value cache size with cross- attention. Conference on Neural Information Processing Systems (Neurips), 2024.

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Chen, Z., Sadhukhan, R., Ye, Z., Zhou, Y., Zhang, J., Nolte, N., Tian, Y., Douze, M., Bottou, L.,, Jia, Z.., and Chen, B. MagicPIG: LSH Sampling for Efficient LLM Generation. International Conference on Learning Representations (ICLR), 2025.
Lee, W., Lee, Seo, J., and Sim, J. InfiniGen: Efficient gen- erative inference of large language models with dynamic KV cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI), pp. 155-172, 2024.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory-efficient exact attention with io-awareness. Conference on Neural Information Pro- cessing Systems (Neurips), 35:16344-16359 2022.
Li, D., Shao, R., Xie, A., Sheng, Y., Zheng, L.,, Gonzalez, J., Stoica, I., Ma, X., and Zhang, H. How Long Can Context Length of Open-Source LLMs truly Promise? NeurlPS 2023 Workshop on Instruction Tuning and Instruction Following. 2023.
Feng, Y., Lv, J., Cao, Y., Xie, X., and Zhou, S. K. Ada- kv: Optimizing kv cache eviction by adaptive budget allocation for efficient Ilm inference. arXiv preprint arXiv:2407.11550, 2024.
Li, Y., Huang. Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., Cai, T., Lewis, P., and Chen, D. SnapKV: LLM knows what you are looking for before generation. Conference on Neural Information Processing Systems (Neurips), 2024.
Gao, Y., Zeng. Z., Du, D., Cao, S., So, H. K.-H., Cao, T., Yang, F., and Yang, M. SecrAttention: Learn- ing Intrinsic Sparse Attention in Your L.MMs. arXiv preprint arXiv:2410.13270 2024. URL https:// arxiv.
Li, Y., Jiang, H., Wu, Q.. Luo, X., Ahn, S., Zhang, C., Abdi, A. H., Li, D., Gao, J., Yang, Y. et al. Schench: A kv cache-centric analysis of long-context methods. Interna- tional Conference on Learning Representations (ICLR), 2025.
Ge, S., Zhang. Y., Liu, L., Zhang, M., Han, J., and Gao, J. Model tells you what to discard: Adaptive kv cache com- pression for llms. International Conference on Learning Representations (ICLR), 2024.
Ling, G., Wang, Z., Yan, Y., and Liu, Q. SlimGPT: Layer- wise Structured Pruning for Large Language Models. Conference Neural Information Processing Systems (Neurips), 2024.
Ho, N., Bae, S., Kim, T., Jo. H., Kim, Y., Schuster, T., Fisch, A., Thorne, J., and Yun, S.-Y. Block Transformer: Global-to- Language Modeling for Fast Inference. Conference on Neural Information P Processing Systems (Neurips), 2024.
Liu, A., Feng, B., Wang, B., Wang, B., Liu, B., Zhao, C., Dengr, C., Ruan, C., Dai, D., Guo, D., et al. Deepseck-v2: A strong, economical, and efficient mixture-of-expers language model. arXiv preprint arXiv:2405.0444 0444444 2024a
Hsieh, C.-P. Sun, S., Kriman, S., Acharya, S., Rekesh, D., Jia, F., and Ginsburg, B. RULER: What's the real context size of your long-context language models? Conference on Language Modeling (COLM), 2024.
Liu, A., Feng, B., Xue, B., Wang. B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek- v3 technical report. arXiv preprint arXiv:2412.19437, 2024b.
Jiang, A. Q.. Sablayrolles, A., Mensch, A., Bamford, C.,, Chaplot, D. S., Casas, D. d. I., Bressand, F., Lengyel, G., Lample, G., Saulnier, L., et al. Mistral 7b. arXiv preprint arXiv:2310.06825. 2023.
NVIDIA. H100 Tensor Core GPU. https:/ www nvidia.com en- httm///euus/aatt-ceeerrr/// /ddddaaatteeeeettaaaae
Jinng H., Li, Y., Zhang, C., Wu, Q., Luo, X., Ahn, S., Han, 7., Abdi, A. H., Li, D., Lin, C.-Y., et al. MInference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. Conference on Neural Information Pro- cexsing Systems (Neurips), 2024.
2024. Accessed: 2025-01-06.
NVIDIA. Dynamo: A datacenter scale distributed
inference serving framework. https:
com/a /ai dynamo/dyaaaaaaooooooooo 2025. URL https://
github. com/ ai ai-dynamo/ /wy.daammmm//mmmmo.oo.. dynamo
Version 0.
Jouppi, N., Kurian, G., Li, S., Ma, P., Nagarajan, R., Nai, L.., Patil, N., Subramanian, S., Swing, A., Towles, B., al. TPUv4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings. In Annual International Symposium on Computer Archi- fecture (ISCA), pp. 1-14, 2023.
Pawar, S., Tonmoy, S., Zaman, S., Jain, V., Chadha, A., and
Das, A. The What, Why, and How of Context Length Extension Techniques in Large Language Models-A De- tailed Survey. arXiv preprint arXiv:2401.07872, 2024.
py Porr gpt-fast: Simple and efficient PyTorch-ntive transformer text generation. htttp:://githtt com/ pytorch- Sachhhlabs/aee / fast 2023. Accessed: 2024-
Kamradt, G.
LLMTest Needle In A Haystack. com/ haam/g/ke http///kaaaaaaLL///a.a /LLMTest SheelleIAAHaystacl 2023. Accessed: 2024-12-14,
12-12.
10

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Qin, R., Li, Z., He, W. Zhang, M., Wu, Y., Zheng, W., and Xu, X. Mooncake: A kveache-centric disaggregated architecture for IIm serving. USENIX Conference File and Storage Technologies (FAST 25), 2025.
Ye, Z., Chen, L., Lai, R., Lin, W., Zhang, Y., Wang, S., Chen, T., Kasikei, B., Grover, V., Krishnamurthy, A., and Ceze, L. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. Conference on Machine Learning and Systems (MLSys)) 2025. URL. /ttp::///aaxi...or/aa//////////
Ribar, L., Chelombiev, I,, Hudlass- Hudlass-Galley, L., Blake, C., Luschi, C., and Orr, D. SparQ Attention: Bandwidth- Efficient LLM Inference. In International Conference Machine Learning (ICML), 2024.
Yuan, J., Liu, H., Zhong, S., Chuang, Y.-N,, S., Wang. G., Le, D., Jin, H., Chaudhary, V., Xu, Z., Zirui, L., and Hu, X. KV Cache Compression, But What Must We Give in Return? Comprehensive Benchmark Long Context Capable Approaches. Conference on Empirical Methods in Natural Language Processing (EMNLP), 2024.
Shazeer, N. Fast transformer decoding: One write-head is all you need. arXiv preprini arXiv: 191 2019.
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LMM: Training multi- billion parameter language models using model paral- lelism. arXiv preprint arXiv:1909.08000 2019.
Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H2O: Heavy- hitter oracle for efficient generative inference large language models. Conference on Neural Information Processing Systems (Neurips), 36:34661-34711 2023.
Singhania, P., Singh, S., He, S., Feizi, S., and Bhatele, A. Loki: Low-Rank Keys for Efficient Sparse Attention. Conference on Neural Information Processing Systems (Neurips), 2024.
Sun, Y., Dong, I.,, Zhu, Y., Huang, S., Wang, W., Ma, S., Zhang, Q., Wang, J., and Wei, F. You only cache once: Decoder-decoder architectures for language models. Conference on Neural Information Processing Systems (Neurips), 2024.
Tang. H., Lin, Y., Lin, J., Han, Q., Hong, S., Yao, Y., and Wang, G. RazorAttention: Efficient kv cache compres- sion through retrieval heads. International Conference Learning Representations (ICLR), 2025.
Tang, J., Zhao, Y., Zhu, K., Xiao, G., Kasikci, B., and Han, S. Quest: Query-Awar Sparsity for Efficient Long-Conter LLM Inference. In International Conference on Machine Learning (ICML), 2024.
Team, G., Mesnard, T., Hardin, C., Dadashi, R., Bhupatiraju, S., Pathak, S., Sifre, L., Rivière, M., Kalc, M. S., Love, J., et al. Gemma: Open models based on gemini research and technology. arXiv preprint arXiv:2403.0829. 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention Is All You Need. Conference on Neural Information Processing Systems (Neurips), 2017.
Xiao, G., Tang, J., Zuo, J., Guo, J., Yang, S., Tang, H., Fu, Y., and Han, S. DuoAttention: Efficient Long- ConteCtt LLM Inference with Retrieval and Streaming Heads. Interna- tional Conference on Learning Representations (ICLR), 2025.
Xu, R., Xiao, G., Huang, H., Guo, J., and Han, S. Xattention: Block sparse attention with antidiagonal scoring. arXiv preprint 2rv3: 16428, 2025.

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
A. Detailed Experiment Settings
A.I. Models
We evaluate our methods on the three widely used long-context models: Llama3.1-8B- (AI@ 2024). Mistral-7B-Ins-
v0.2 Diang et al., 2023). and LongChat-7B-ve
(Bai et al. 2024). They are all decoder-only Transformer models [Vaswani al. 2017). Llama3.1-8B2 188BIIns and Mistral-7B-Ins--00.. use GQA, while LongChat-7B-vl., uses MHA in the attention module. In terms of sequence length, Llama3. 1-8B-Ins supports maximum sequence length of 128K while the other two support up to 32K sequence length. Prior work (Tang et al. 2024 Chen et al., 2025) usually skips the first two layers during KV cache compression to maintain accuracy. We found that this is only necessary for LongChat-77 7B-v1.5, so, we conduct KV cache compression on all attention layers for Llama3. I- and Distral-7B-Ins-OO. but skip the first two layers for LongChat-7B--v/ 3..
A.2. Benchmarks
Our LongBench and Boadle-in-a-Haystadee settings mostly follow the prior work (Yuan et al. 2024). For LongBench, we set the maximum prompt length to 127,500 for Llama3.1-8B-Ins and 31,500 for Mistral-7B-Ins-voe and LongChat-7B-vl..e Prompts beyond the maximum prompt length will be middle-truncated by keeping the first and last half of the input tokens. For Needle-in Lartlle-in----aaae Haystack, we evaluate with 10 different input sequence lengths uniformly spanning from 2048 81,920 words and 10 different depths for each sequence length in 1-ama..1-B---l.. The maximum word count of 81,920 roughly converts to 109K tokens by the tokenizer in Llama3. I-8B-Ins. For the other two models, the input sequence lengths range from 512 to 20480 words since they support up to 32K sequence length, which convert up to about 30K and 31K tokens for In---l-----In-O-... and LongChat-7B-v
respectively. For RULER, we mostly follow the configurations in the original benchmark, except for reducing the number of examples per task from 500 to 50 to speed up the evaluation. For SCBench, we follow the multi-turn setting and set the maximum prompt length to 127,500 for Llama3.1-8B-In -88-
A.3. Baselines
Since RocketKV primarily focuses on accelerating the decode phase of LLM inference, we do not perform any KV cache compression or sparse attention mechanism in the prefill phase across all comparing baselines, including DuoAltention, for a fair comparison.
Exact-TopK serves as oracle method to demonstrate the effectiveness of top-k sparse attention. We assume the top-k KV tokens can be directly identified with no cost a given token budget will directly correspond to its top-k value.
Exact-TopK:
DuoAttention: We set the initial and recent tokens to 128 and 512. respectively, for streaming heads if the token budget is larger than 640, and vary the ratio of retrieval heads to match the given token budget on average. If the token budget is lower than 640, all attention heads become streaming heads. this case, we set the number of initial tokens for the attention sink to 128 if the token budget is larger than 256, otherwise, it is set to 20% of the token budget. The rest of the token budget assigned recent tokens within the sliding window.
SnapKV: We set the observation window size 32 for all single-turn benchmarks and 128 for multi-turn SCBench. The kernel size for pooling is set to 7, following the SnapKV paper. The token budget is divided by the number of heads per attention group in case of GQA as the SnapKV performs on each attention head separately and the pruned KV cache is not shared within the attention group.
We evenly split the token budget between approximate attention for identifying top-k KV token indices and top-k sparse attention so that the token budget can accurately reflect the total memory fetching bandwidth required in the attention module. The original Quest method is not compatible with GQA. We modify it select indices of top-k accumulated attention scores per attention group rather than per attention head.
Quest:
SparQ: Similar to Quest, we evenly split the token budget between approximate attention for identifying top-k KV token indices and lop-k sparse attention.
RocketKV/Rockette V-MT: For SnapKV in the first stage, We set the observation window size to 32 for all single-turn benchmarks and 128 for multi-tum SCBench. For HSA in the second stage, we evenly split the token budget between approximate attention for identifying top-k indices and top-k sparse attention.
12

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
B. Additional Results
B.1. Ablation Studies
In this subsection, we present a series of ablation studies to further demonstrate the effectiveness of RocketKV.
100
Average Score
Average Score
Average Score
80
40
60
Full-KV
Full-KV
30
HSA
40
HSA
40
HSA
Quest
Quest
Quest
Spart
20
SparQ
20
SparQ
25
512
1024
2048
4096
S12
1024
2048
4096
256
512
1024
2048
4096
Token Budget
Token Budget
Token Budget
(a) LongBench
(b) NIAH
(c) RULER, SeqLen 16K
Average Score
80
Average Score
80
Average Score
60
60
60
40
Full-KV
-40
Full-KV
40
Full-KV
HSA
HSA
HSA
20
Quest
20
Quest
Quest
Spart
SparQ
Spart
0
256
2
1024
2044
4096
2552
1024
2048
4096
1024
2048
4096
Token Budget
Taken Budget
Token Budget
(d) RULER, LeqLen= 32K
(e) RULER, SeqLen=64K
(f) RULER, SeqLen=96K
Figure 8: Accuracy comparison among HSA, Quest, and SparQ on Llama3. 1 In8B-Ins.
Average Score
100
92
Average Score
Average Score
90
Full-KV RocketKV
98
Full-V
Full-KV
51
RocketKV
88
RocketKV
RocketKV r0.3
96
RocketKV
ro.3
RocketKV r0.5
RocketKV
RocketKV
50
RocketKV r0.7
RocketKV
RocketKV r0.7
512
1024
2048
4096
23
1024
2048
4096
258
512
1024
2048
4096
Token Budget
Token Budget
Token Budget
(a) LengBench
(b) NIAH
(c) RULER, SeqLen=16K
Average Score
85
Average Score
Average Score
Full-VV
Full-KV
Full-Ky
RocketKV
RocketKV
RocketKV
RocketKV r0.3
RocketKV r0.3
RocketKV r0.3
RocketKV r0.5
70
RocketKV ro.5
RocketKV r0.5
75
RocketKV r0.7
RocketKV ro.7
RocketKV r0.7
256
512
1024
2048
4096
256
512
1024
2048
4096
256
512
1024
2048
4096
Token Budget
Token Budget
Taken Budget
(d) RULER, SeqLen=32K
(e) RULER, SeqLen=64K
(f) RULER, SeqLen=96K
Figure 9: Comparing adaptive against static split factors on 1lama3.1 1-8B-Insss Ins.
13

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
B. 1. COMPARING HSA, QUEST, AND SPARQ
To illustrate the effectiivness of hybrid sparse attention (HSA), we compare the accuracy of the standalone HSA mechanism against Quest and SparQ. Figure & demonstrates the results 1-8B Ins across multiple different benchmarks. In all cases, HSA consistently outperforms Quest and SparQ, especially at low token budgets. This clearly demonstrates the advantage of HSA, which intelligently leverages approximations in both sequence and head dimensions compared to single dimension approximation methods such as Quest and SparQ.
1.1.2. SPLIT FACTOR
In this study, we compare the adaptive compression decomposition method against statically determined split factors r ranging from 0.3 to 0.7. A split factor of 0.5 indicates an even split between the first and second stage of RocketKV. As shown in Figure 9 the best static split factor varies with different sequence lengths and token budgets, while adaptive compression decomposition provides comparable accuracy to the best static split factor in most cases.
B.2. Needle-in-a-Haystadl Visualization
Needle-in-a-haystad (NIAH) is type of synthetic challenge designed to test how effectively LLM can retrieve specific information in a large volume of text (Kamradi 2023). In Figures 10 12 the x-axis shows the document length (i.e., "haystack"), " while the y-axis marks the relative position of the "needle" (i.e., a short sentence) within the text. As shown in the results, VocketKV can retrieve the needle with almost the same accuracy as Full-KVV throughout the whole text across all three models with token budgets as low as 256.
100
100
100
0
11
11
11
22
22
22
33
33
33
60
44
44
44
56
56
56
40
40
40
67
67
67
78
78
78
89
89
89
100
100
100
10K
81k
13k
81K
Word Count
Word Count
Word Count
(a) Token Budget = 256
(b) Token Budget = 512
(c) Token Budget= 1024
100
0
0
11
11
11
80
80
80
22
22
22
33
33
33
Depth
44
44
44
56
40
40
67
67
67
78
78
78
20
20
20
89
89
89
100
100
100
10K
81K
10k
13K 81K
Word Count
Word Count
Word Count (I) Full-KV
(d) Token Budget= = 2048
(e) Token Budget= = 4096
Figure 10: Seedle-in-a-Haystare visualization results of RocketKV on Llama3.1-8B-In
14

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
100
100
100
0
11
11
11
22
22
22
33
Deoth
44
60
44
60
44
60
56
56
56
67
67
78
78
78
89
89
89
100
100
100
16k
18k
20K
18
20k
16k
18k
20K
Word Count
Word Count (b) Token Budget 512
Word Count
(a) Token Budget 256
(c) Token Budget 1024
100
100
0
11
11
11
22
22
22
33
33
33
Depth
50
Depth
44
60
44
60
56
56
56
67
67
67
78
78
78
20
89
89
89
100
100
100
18K
20k
18k
20k
16K
18K
20K
Word Count
Word Count
Word Count (f) Full-KV
(d) Token Budget= = 2048
(e) Token Budget= 4099
Figure 11: Haale--an-aHHyssackk visualization results of RocketKV on Mistral-7B-IIn-v0...
100
100
0
11
11
11
80
80
80
22
22
22
33
33
33
Depth
44
44
44
56
56
56
67
67
67
78
78
78
20
20
20
89
89
89
100
100
100
16
18V
20k
18k
20k
36
18K
20k
Word Count
Word Count
Word Count Taken Budget= 1024
(a) Token Budget = 256
(b) Token Budget = 512
(c)
100
100
100
0
11
11
11
80
80
80
22
22
22
33
33
33
Depth
44
Depth
44
44
56
40
40
40
67
67
67
78
78
78
89
89
89
100
100
100
20k
16
1,8
20k
18k
20K
Word Count
Word Count
Word Count (f) Full-KV
(d) Token Budget = 2048
(c) Token Budget= 4096
Figure 12: -eedle-in-a-Haueee
visualization results of RocketKV on Lonv
15

---

RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
B.3. Detailed Accuracy Results
In this section, we present 20 detailed tables that display the accuracy results discussed in previous sections. These tables offer additional information and insights into the results we captured.
Table 2: Average results of LongBench and NIAH for Llama3.1-8B-11
Taken Budgel
Method
Single QA
Malti. Q4
Summ.
Few Few-shol
Syatheric
Code
L. Avg.
NIAH
N/A
Full-KV
43.5
44.5
29.2
69.3
53.9
62.0
52.2
100.0
Exact- Exact-TopK
43.6
44.4
28.9
53.6
65.3
100.0
DuoAttention
24.4
33.5
18.2
39.0
51.2
60.1
SnapKV
32.5
39.1
20.3
50.7
53.0
56.8
42.6
98.3
256
Quest
16.0
14.0
14.3
14.4
4.0
41.4
17.8
10.7
SparQ
26.3
18.7
18.2
24.5
15.5
47.0
25.7
10.3
HSA
41.3
41.9
26.8
64.2
51.7
61.1
49.5
44.3
RocketKV
43.7
43.5
25.2
67.3
53.8
64.0
51.1
100.0
Exact- TopK
43.7
44.5
28.8
53.7
63.1
52.2
100.0
DuoAttention
27.7
34.9
51.5
61.7
41.2
21.0
SnapKV
36.0
41.9
53.4
512
Quest
27.9
27.7
19.8
52.6
35.1
SparQ
37.1
35.1
24.5
48.0
57.3
43.8
14.7
HSA
43.8
43.2
68.6
53.1
65.2
52.1
73.7
RocketKV
43.7
43,8
26.8
69.5
53.4
63.5
51.9
100.0
Exact- TopK
43.8
44.6
29.0
69.3
53.5
61.9
52.2
100.0
DuoAttention
31.8
36.4
23.7
54.8
53.3
64.2
44.5
26.3
SnapKV
38.2
42.7
23.5
59.0
53.7
59.6
47.3
99.7
1024
Quest
36.5
37.5
24.8
54.9
50.9
57.3
45.0
37.3
SparQ
41.5
42.9
52.0
62.2
51.1
26.0
HSA
43.9
44.0
28.8
54.0
64.6
52.5
94.3
RocketKV
43,9
43.8
53.4
63.2
52.1
100.0
Exact-Topk
43.6
44.5
29.2
69.5
53.8
61.1
52.1
100.0
AttAnttttion
39.3
40.9
25.4
64.5
52.2
64.8
49.2
28.3
SnapKV
40.6
43.4
25.5
65.0
52.8
61.7
49.8
99.7
2048
Quest
42.5
43.0
28.4
65.5
52.3
59.8
50.5
64.3
SparQ
43.5
43.5
29.3
69.4
53.6
62.9
52.1
56.3
HSA
43.7
43.8
29.2
69.4
53.6
62.3
52.2
100.0
RocketKV
43.6
53,9
62.2
52.1
100.0
Exact- Exact-TopK
43.5
53.3
60.4
$2.0
100.0
AttAnttttion
42.4
41.7
27.5
68.9
53.2
63.9
51.2
36.7
SnapKV
42.3
43.7
26.9
67.5
53.4
61.5
50.9
100.0
4096
Quest
43.9
44.3
28.9
69.0
53.5
60.0
51.9
94.3
SparQ
43.7
44.0
29.1
69.2
53.1
63.2
52.3
89.3
HSA
43.9
44.1
29.1
69.4
53.5
61.1
52.1
100.0
RocketKV
43.3
43.9
29.0
69.6
53.6
61.1
$1.9
100.0
16

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 3: Individual results of LongBench for Llama3.1 1-B------IIm-
Tokor Budder
Method
parratreaua
gorper
halpotea
Jwlkingge
musique
gov-report
mulid
trieiage
pass-
repebench-
A7A
Fall-VV
30.2
45.5
31.5
45.8
8.4
Exact Topk
30.7
44.7
15.4
55.0
46.5
31.7
34,8
25.1
26.8
71.
92.2
48.8
99.3
67.1
63.6
7.8
DunAltention
19.4
22.8
31.1
4233
35.9
22.4
15.9
20.6
17.2
200
57.4
39.7
98.0
6122
99.1
4.5
SnapKV
19.2
39,0
4.1
20.5
7.5
49.5
0.5
SpurQ
5.0
38.0
35.8
19.7
28.8
7.5
13.7
15.0
25.8
16.5
32.1
24,8
28.8
60.1
34.0
2.2
HSA
26.4
43.5
663
26.6
8.7
Exteet Topk
30.9
45.0
55.1
55.3
47.0
31.0
34.5
25.0
20.9
71.5
92.0
43.9
99.5
03.0
61.1
7.9
DuoAmention
SnapKV
21.5
59.7
512
Quest
9.0
36.5
38.2
32.2
36.4
144
16.1
17.7
25.6
35.2
51.2
27.5
82.0
59.7
45.4
5.0
SparQ
13.6
45.4
52.1
39.4
25.6
26.7
49.5
72.2
37.9
61.2
4.1
HSA
7.1
RockerKV
30.6
44.8
55.7
54.7
46.5
30.2
30.2
24.7
25.5
7300
91.8
43.6
99.5
64.8
62.
7.3
Extet-Topk
50.8
DinAttenionn
30.4
24.5
21.2
25.3
45.5
8.0
SnapKV
28.1
36.6
50.0
$6.0
44.2
28.0
23.8
23.8
23.1
44.5
91.3
41.3
99.5
61.8
57.4
8.0
1924
Quest
15.4
62.6
Sparl)
HSA
31.1
45.3
55.3
54.2
46.5
31.2
34,6
24.9
26.9
72.5
92.1
44.5
99.5
66.3
62.9
8.5
RocletVV
30.7
44.9
56.0
$4.6
46.7
30,0
33.0
25.2
26.6
72.0
43.5
99.5
64.7
61.8
7.3
Exuet-Topk
DunAlteniion
25.6
42.3
50.0
$2.3
42.1
28.1
27.5
22.1
26.6
66.0
$4.9
42.8
98.5
66.1
63.6
6.0
SnapKY
28.8
39.9
46.5
60.0
92.4
42.6
63.8
33.1
SparQ
30.3
45.2
54.9
53.7
46.5
30/4
35,3
25.4
27.1
73.0
91.2
44,0
99.5
65.3
60.5
7.7
USA
30.7
45.2
46.2
35.1
27.4
73.0
44.2
61.3
7.7
EescI-TopX
30.2
45.5
54.8
55.5
47.1
31.8
35.2
25.1
27.2
72.5
91.7
43.7
99.5
63.4
57.5
7.2
DunAttenionn
26.5
45.1
43.3
31,9
27.0
735
13.3
65.5
62.3
7.0
SnapKY
29.0
4096
Quest
30.7
45.5
55.4
35.6
46.0
31.3
34,9
348
27.1
720
00.9
44.0
99.5
63.3
56.8
74
SparQ
30.8
45.3
54.9
46.4
31.0
36,9
25.2
27.2
72.5
91.5
13.7
65.6
6.6
72.
RockerKV
30.3
45.3
54.3
54.
46.3
30.4
34.7
25.0
27.1
730
91.3
44.3
99.5
63.7
58.5
766
17

---

RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 4: Average SCBench results for 1lamaa.1-8B---.
Token Budget
Retr. Semantiiattic
Retr.String
Multi-task
Method
Global
Avg.
N/A
Full-KV
49.5
40.7
36.2
64.6
49.9
Exact- Exact-Topk
32.3
41.0
36.8
61.8
45.0
DuoAttention
0.1
26.6
33.5
11.6
20.9
SnapKV
0.2
17.0
34.7
10.7
17.6
Quest
0.0
22.4
25.2
4.0
15.7
1024
SparQ
0.0
19.2
23.6
0.9
13.6
HSA
4.1
27.2
27.2
20.0
22.1
RocketKV
14.3
22.7
35.3
26.2
26.6
RocketK V-MT
40.5
35.7
36.8
56.2
44.3
Exact-Toph Topk
38.0
41.2
36.6
62.9
46.7
DuoAttention
0.1
26.6
33.9
12.2
21.2
SnapKV
1.1
17.1
35.9
12.1
18.5
Quest
2.0
24.4
26.6
11.8
18.8
2048
SparQ
0.1
21.1
25.8
4.1
15.5
HSA
8.4
35.3
35.7
53.4
34.6
RocketKV
16.5
25.0
35.4
30.5
28.8
RocketK V-MT
42.2
38.4
37.0
59.5
46.4
Exact- TopK
43.4
42.1
36.4
63.8
48.7
DuoAttention
0.1
26.7
33,9
12.7
21.3
SnapKV
3.5
19.1
35.1
13.2
19.9
Quest
8.6
26.8
28.5
23.0
24.0
4096
SparQ
3.0
29.9
29.2
30.8
25.2
HSA
13.1
38.7
36.3
60.9
38.6
RocketKV
18.9
28.5
35.9
33,9
31.5
RocketK V-MT
44.8
39.9
36.8
60.3
47.8
Exact- Topk
48.0
42.4
36.4
64.7
50.2
DuoAttention
27.3
34.0
16.5
22.2
SnapKV
5.0
20.3
35.7
19.0
21.9
Quest
10.1
31.6
32.0
31.9
28.8
8192
SparQ
6.8
41.5
35.6
64.1
38.3
HSA
34.1
41.8
36.5
63.8
46.0
RocketKV
22.9
33.2
35.8
39.1
35.2
RocketK V-MT
42.0
36.2
62.9
49.6
Exact- Exact-TopK
TopK
49.5
42.8
36.4
65.0
50.8
DuoAttention
0.4
28.9
34.0
30.5
25.4
SnapKV
6.9
22.0
34.5
24.7
23.8
Quest
12.5
40.6
36.2
61.8
39.3
16384
SparQ
40.3
42,4
36.1
64.5
48.0
HSA
43.1
36.6
64.4
50.0
RocketKV
29.4
37.2
36.1
44.6
39.5
RocketK V-MT
49.9
42.1
36.3
64.3
50.5
18

---

RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 5: Individual SCBench results for 1 1-8B-Ins
summarr aecelles
Token Budget
reporaa Ladd Jry
Method
preft suffis
repogg
quaug
qa.chn
cholce.eng
many-stot
summary
N/A
Full KV
16.2
34.6
37.6
48.0
21.5
20,9
72.6
59.3
37.5
60.4
68.8
Topk
43.0
35.4
72.5
59.3
DunAttettion
0.2
0.0
0.5
18.3
15.8
71.7
60.0
28.3
12.2
22.9
0.3
SnapKV
0.0
0.0
0.6
0.9
8.7
12.2
46.2
69.6
27.2
7,2
21.2
0.1
Quest
0.0
0.1
0.0
10.1
72.1
50.7
11.7
0.0
1024
SparQ
0.0
0.0
3.7
68.6
47.8
HSA
0.0
0.0
12.4
5.7
22.9
9.0
71.2
59.6
9.7
12.2
35.3
4,7
RocketKY
10.2
2.4
30.4
5.0
16.5
15.2
54.1
65.6
30.0
10.3
49.1
3.3
RocketKV MT
Exact TopK
53.4
23.8
36.7
43.9
29.6
18.7
72.5
60.0
37.2
12.5
59.7
66.1
DucAtiention
0.2
0.0
U.1
0.5
18.8
16.0
71.2
61.1
28,9
11.8
24.0
0.3
SnapKV
0.2
1.4
12.0
44.9
70.0
27.9
9.7
23.6
Quest
0.0
5.9
0.0
12.8
12.8
72.1
50.0
14.1
15.5
23.2
2048
SparQ
0.0
0.0
04
0.0
9.3
6.9
68.1
61.1
6.0
10.2
8.2
0.0
ISS
0.0
25.3
23.6
29.6
16.4
71.7
61.1
31.9
14.0
57.2
49.6
RocketKV
RockerKV-MT
61.2
30.6
34,7
31.8
30.4
19.5
71.7
61.5
36,9
12.5
61.0
$8.1
Exact-lopk
66.0
26.8
37.4
46.6
29.6
18.3
73.9
39.6
37.3
12.2
39.7
67.9
DueAttention
0.0
U.0
0.5
19.7
15.4
71.2
60.7
28.6
12.5
25.1
0.3
SnapKV
0.2
10.3
2.7
12.2
48.9
67.8
28.6
8.8
25.6
Quest
0.0
0.0
25.7
0.5
21.4
13.5
71.7
55.9
16.5
13.0
44.8
1.3
SparQ)
0.0
9.0
17.0
22.3
8.8
71.2
60.0
16.2
11.3
42.6
18.9
29.3
RocketKV
17.6
4.8
34.2
13.2
20.3
17.8
62.9
62.2
34.6
10.8
58.2
9.7
RocKerKV-MT
68.8
32.6
36.1
37.5
30.7
18.5
73.0
60.7
37.4
12.3
60.4
60.2
Exact-TopK
72.6
33.0
38.3
45.9
30.7
18.9
73.9
59.6
37.4
60.4
68.9
DunAttention
0.0
0.5
21.2
71.7
60.7
29.3
32.4
SnapKV
0.0
0.2
14.9
4.1
13.0
14.4
49.7
69.6
28.5
9.0
36.5
1.4
8192
Ques
0.2
0.0
30.0
10.7
27.5
17.1
71.2
$9.6
22.3
44.3
19.6
SparQ)
67.6
HSA
47.4
20.0
34.8
44.8
29.9
19.7
73.0
59.6
36.8
13.0
60.1
67,5
RocketKV
25.0
7.4
36.3
23.2
23.5
17.8
68.2
60.4
36.2
10.7
60.6
17.6
RockerKV-MT
70.4
37.7
31.0
18.7
73,9
60.0
37.2
60.5
Exact-TopK
74.8
37.4
DuoAttention
0.2
0.0
1.1
3.6
24.1
16.7
71.2
59.6
29.9
12.5
46.4
14,5
SnapKV
0.4
20.4
5.9
14.2
14.5
53.2
64.1
29.9
9.5
47.8
1.7
16384
28.7
36.1
SparQ
60.8
26.0
34.1
48.4
29.9
18.0
73.4
59.3
36.6
12.3
60.5
68.5
HSA
70.8
30.8
38.0
50.0
29.9
19.0
73.4
39.6
37.6
12.7
60.7
68.2
RocketKV
15.5
60.9
RockerKV-MT
74.6
37.2
38.0
30.1
73.4
59.6
37.6
60.8
67.8
19

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 6: Results of RULER-16K for Llama3.1-8B-Ml
Tokon Budge
Method
single1
single2
shngle?
multikeyl
multikey2
multikey
multivalue
mulliquery
CWC
twe
ge
N/A
Full-KV
100.0
100.0
100.0
100.0
100.0
100.0
100.0
100.0
99.2
60.6
92.7
84.0
50.0
91.3
Exact- Topk
100,0
100.0
100.0
100.0
100.0
98.0
99.0
100.0
94.8
65.8
87.3
80.0
50,0
90.4
DunAttention
0.0
2.0
2.0
2.0
1.6
2.8
67.3
42.0
28.0
12.2
SnapKV
86.0
0.0
0.0
32.7
78.0
26.3
256
Quest
2.0
4.0
0.0
12.0
0.0
0.6
33.3
6.0
8.0
5.3
SparQ
4.0
0.0
0.0
0.0
2.0
0.0
1.0
0.0
5.2
5.6
53.3
26.0
16.0
8.7
HSA
66.0
80.5
85.0
82.8
80.7
80.0
44.0
69.0
RockerKV
100.0
100.0
100.0
100,0
100.0
98.0
92.0
9.6
82.0
48.0
86.3
Exact-TopK
100.0
100.0
100.0
100.0
100.0
100.0
99.0
100.0
96.8
69.8
86.7
80.0
50.0
90.9
DonAttention
2.0
2.0
4.0
2.0
4,0
0.0
6.0
5.0
3.2
0.4
90.0
42.0
20.0
13.9
SnapKV
512
Quest
98.0
74.0
0.0
100.0
26.0
22.0
16.0
60.4
4.0
52.7
34.0
22.0
39.2
SparQ
60.0
54.0
0.0
28.0
2.0
0.0
23.0
11.0
21.2
13.0
73.3
36.0
30.0
27.0
HSA
100.0
98.0
100.0
90.0
66.0
95.0
99.5
90.4
24.0
88.7
82.0
52.0
83.5
RockerKV
100.0
100.0
100.0
100,0
100.0
99,5
100.0
92.8
92.7
82.0
54.0
88.6
Exact-Topk
100.0
100.0
100.0
100.0
100.0
100.0
V9,3
100.0
98.8
71.4
89.3
80.0
50.0
91.3
DunAtenttin
2.0
2.0
4.0
2.0
6.0
2.8
0.0
90.0
34.0
22.0
13.6
SnapKV
92.0
0.0
64.0
41.0
68.0
78.8
3.2
84.0
80.0
46.0
56.8
1024
Quest
94.0
100.0
38.0
100.0
84.0
66.5
79.5
80.4
9.2
70.0
68.0
50.0
64.6
SparQ)
100.0
100.0
100.0
100.0
98.0
54.0
97.0
99.0
85.2
66.0
86.7
84.0
54.0
86.5
HSA
100.0
100.0
100.0
98.0
99,5
100.0
94.4
51.8
88.0
82.0
50.0
89.5
RocketKV
100.0
$2.0
Exact-TopK
100.0
100.0
100.0
100.0
100.0
100.0
99.5
100.0
99.2
64.2
90.0
82.0
50.0
91.1
DonAttention
4.0
20.0
6.0
20.0
24.0
4,0
6.0
8.0
2.4
2.6
89.3
52.0
34.0
20.9
SnapKV
2048
Quest
94.0
96.0
60.0
98.0
92.0
28.0
91.5
90.5
80.8
19.8
86.7
80.0
50,0
74.4
SparQ
100.0
100.0
100.0
100.0
100.0
100.0
100.0
100.0
99.2
77.2
88.0
82.0
50.0
92.0
HSA
100.0
100.0
100.0
100.0
100.0
100.0
99.5
100.0
97.6
66.8
86.7
80.0
50.0
90.8
RockerKV
100.0
100.0
100.0
100,0
100.0
99.5
95.6
95.3
80.0
50.0
91.4
Exact-Topk
100,0
100.0
100.0
100.0
100.0
100.0
98.5
100.0
99.2
57.4
900.7
84.00
50.0
90.8
DueAtrention
20.0
88.0
86.0
98.0
100.0
46.0
71.5
77.0
$8.0
14.2
88.0
82.0
52.0
67.7
SnapKV
100.0
96.0
96.0
76.5
94.5
94.4
92.0
78.0
54.0
70.7
4096
Quest
100.9
94.0
96.0
100.0
96.0
98.0
98.0
98.0
98.8
54.6
88.0
82.0
50.0
88.7
SparQ
100.0
100.0
100.0
100.0
100.0
100.0
100.0
100.0
99.2
75.2
92.0
82.0
50.0
92.2
HSA
100.0
100.0
100.0
100.0
100.0
100.0
100.0
99.2
77.0
91.3
82.0
50.0
92.3
RocketKV
20

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 7: Results of RULER-32K for Llama3.1-8B-Ml
Tokon Budge
Method
single1
single2
shngle?
multikeyl
multikey2
multikey
multivaluc
multiquery
CWC
twe
@
N/A
Full-KV
100.0
100.0
100.0
100.0
100.0
100.0
99.0
99.5
96.4
6.0
94.0
74.0
54.0
86.4
Exact- Topk
100,0
100.0
100.0
100.0
100.0
90.0
99.0
99.3
91.6
24.6
81.3
76.0
54,0
85.8
DunAttention
0.0
4.0
4.0
4.0
2.0
4.5
2,4
1.0
52.0
38.0
32.0
11.4
SnapKV
96.0
0.0
0.0
26.2
256
Quest
0.0
0.0
0.0
D.0
1.2
22.0
0.0
0.0
SparQ
44.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
2.8
0.4
18.7
4.0
8.0
6.0
HSA
98.0
$2.0
57.5
70.4
0.4
48.0
62.0
38.0
$1.6
RockerKV
100.0
100.0
100.0
100,0
100.0
96.5
98.0
87.6
0.8
77.3
78.0
54.0
84.0
Exact-TopK
100.0
100.0
100.0
100.0
100.0
100.0
99.0
99.0
90.8
17.2
82.0
78.0
52.0
86.0
DonAttention
0.0
4.0
8.0
4.0
2.0
0.0
4,5
4,5
3.6
0.0
70.7
40.0
30.0
13.2
SnapKV
512
Quest
100.0
26.0
16.0
4,0
1.2
45.3
10.0
16.0
17.7
SparQ
62.0
0.0
0.0
D.0
0.0
0.0
0.5
0.5
3.2
1.6
26.7
24.0
26.0
11.1
HSA
96.0
100.0
72.0
89.0
96.0
83.2
1.2
60.7
80.0
52.0
71.2
RocketKV
100.0
100.0
100.0
100,0
100.0
98.0
100.0
89.2
1.6
78.0
76.0
56.0
84.5
Exact-Topk
100.0
100.0
100.0
100.0
100.0
100.0
59.00
98.3
91.6
10.6
87.3
74.0
52.0
83.6
DunAtentiin
0.0
4.0
8.0
4.0
2.0
4.5
4.0
4.0
0.0
71.3
42.0
34.0
13.7
SnapKV
92.0
0.0
60.0
48.5
68.5
72.0
1.4
56.7
72.0
52.0
54.4
1024
Quest
90.0
78.0
18.0
76.0
46.0
34,0
31.0
38.8
1.0
50.0
20.0
22.0
38.8
SparQ)
84.0
54.0
0.0
46.U
8.0
0.0
27.0
23.0
21.6
4.2
53.3
44.0
38.0
31.0
HSA
100.0
100.0
100.0
97.5
99.5
88.0
6.0
72.7
80.0
58.0
83.4
RocketKV
$2.0
84.9
Exact-TopK
100.0
100.0
100.0
100.0
100.0
100.0
99.0
98.5
95.6
7.0
90.7
74.0
54.0
86.1
DonAttention
0.0
4.0
8.0
6.0
8.0
0.0
4.5
3.5
3.2
0.0
72.0
40.0
36.0
14,2
SnapKV
50.0
2048
Quest
86.0
100.0
6.0
86.0
76.0
43.5
66.0
54.4
2.2
68.0
60.0
42.0
53.1
SparQ
100.0
100.0
98.0
98.0
96.0
36.0
98.0
98.0
82.0
34.2
74.7
76.0
58.0
80.7
HSA
100.0
100.0
100.0
100.0
100.0
100.0
99.0
99.0
92.0
14.6
76.7
76.0
56.0
85.6
RockerKV
100.0
100.0
100.0
100,0
100.0
100.0
93.2
80.0
76.0
52.0
86.0
Exact-Topk
100.0
100.0
100.0
100.0
100.0
100.0
99.0
98.5
96.8
7.0
93.3
74.0
54.0
86.4
DueAtrention
4.0
26.0
14.0
36.0
16.0
2.0
5.5
10.5
3.2
0.0
76.7
50.0
48.0
22.5
SnapKV
96.0
98.0
80.0
78.5
94.5
88.8
8.8
69.3
74.0
54.0
66.0
4096
Quest
82.0
100.0
70.0
90.0
92.0
38.0
87.0
87.5
88.4
5.6
78.7
78.0
52.0
73.0
SparQ
100.0
100.0
100.0
100.0
100.0
100.0
99.5
99.0
95.2
20.6
86.7
74.0
52.0
86.7
HSA
100.0
100.0
100.0
100.0
100.0
99.5
99.0
93.2
19.8
79.3
74.0
52.0
85.9
RocketKV
21

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 8: Results of RULER-64K for Llama3.1-8B-ml
Token Budge
Methed
single1
single1
single3
multkeel
mullikey2
multikey
multivalue
multiquery
CWP
fwe
98
90
ge
N/A
Full-KV
100.0
100.0
100.0
100.0
100.0
100.0
95.5
97.5
94.8
0.0
83.3
76.0
56.0
84.9
Exact-Tope
100.0
100.0
100.0
100.0
100.0
58.0
94.5
97.5
93.6
0.4
63.3
72.0
56.0
79.6
DunAtiention
0.0
8.0
0.0
0.0
5.0
3.0
0.4
84.0
44.0
32.0
13.7
SnapKV
23.0
256
Quesi
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.4
25.3
0.0
2.0
2.1
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
100.0
24.0
2.0
0.0
4,5
6.0
48.8
0.2
22.7
22.8
RocketKV
100.0
100.0
84,8
70.0
81.2
Exact-TopK
100.0
100.0
100.0
100.0
100.0
V8.0
92.8
0.4
67.3
76.0
$4.0
82.0
DuoAttention
0.0
2.0
0.0
8.0
0.0
0.0
5.0
3,5
0.4
0.8
96.0
38.0
28.0
14.0
SnapKV
512
Quest
0.0
0.0
0.0
0.0
0.0
0.5
0.4
34.7
4.0
6.0
3.5
SparQ
20.0
0.0
0.0
0.0
0.0
0.0
0,0
0,0
2.0
1.0
15.3
4.0
14.0
4.3
HSA
100.0
92.0
26.0
82.0
0.0
$2.5
71.5
74.8
0.2
56.7
62.0
46.0
52.7
RocketKV
100.0
100.0
100.0
100.0
88.8
76.0
74.0
83.3
Exact-Topk
100.0
100.0
100.0
100.0
100.0
96.0
V8.3
98.5
93.6
0.2
74.7
74.0
$4.0
83.8
DunAftention
0.0
2.0
0.0
8.0
0.0
2.0
5.0
3.5
0.4
0.8
96.0
40.0
30.0
14.4
SnapKV
100.0
82.0
0.0
28.0
56.0
46.6
1024
Quest
100.0
32.0
10.0
24.0
2.0
0.0
1.5
6.0
35.6
0.4
49.3
12.0
12.0
21.9
SparQ)
48.0
0.0
0.0
2.0
0.0
0.0
U.5
0.0
12.4
0.8
18.7
20.0
24.0
9.7
100.0
100.0
2.0
92.0
84.4
0.2
70.0
72.0
72.5
RocketKV
100.0
Exact-TopK
100.0
100.0
100.0
100.0
100.0
98500
97.0
94.0
0.0
78.7
74.0
52.0
84.1
DuoAticution
0.0
2.0
0.0
8.0
0.0
2.0
5.0
3,5
0.4
0.8
94.0
44.0
28.0
14.4
SnapKV
2048
Quesi
100.0
78.0
70.0
88.0
48.0
0.0
19.5
76.8
0.0
63.3
38.0
38.0
50.7
SparQ
98.0
46.0
0.0
$4.0
10.0
0.0
31.0
27.0
41.2
5.6
$2.0
36.0
36.0
33.6
HSA
100.0
100.0
98.0
100.0
96.0
54.0
96.0
99.0
85.6
0.8
68.7
78.0
56.0
79.4
RocketKV
100.0
100.0
100.0
91.6
72.0
83.6
Exact-Topk
100.0
100.0
100.0
100.0
100.0
100.0
96.5
99.0
95.6
0.0
80.0
76.0
52.0
84.5
DuoAttention
0.0
4.0
0.0
8.0
0.0
2.0
4.5
3.5
0.4
0.8
91.3
44.0
32.0
14.7
SnapKV
100.0
90.0
20.0
2.0
91.2
0.8
60.0
61.3
4096
Quest
100.0
76.0
20.0
100.0
76.0
0.0
48.0
77.5
78.8
0.0
72.0
76.0
52.0
59.7
SparQ
100.0
100.0
100.0
100.0
94.0
32.0
94.0
96.5
86.4
2.4
73.3
70.0
52.0
77.0
HSA
100.0
100.0
100.0
100.0
98.0
98.0
90.8
0.2
70.0
76.0
56.0
83.2
RocketKV
22

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 9: Results of RULER-96K for Llama3.1-8B-mn
Token Budge
Methed
single1
single1
single3
multkeel
mullikey2
multikey
multivalue
multiquery
CWe
90
N/A
Full-KV
100.0
100.0
100.0
100.0
94.0
88.0
98.5
100.0
66.8
0.0
74.7
78.0
50.0
80.8
Exact-Tope
100.0
100.0
100.0
100.0
84.0
14.0
95.3
99.0
67.2
0.4
56.0
82.0
50.0
72.9
DunAtiention
4.0
6.0
0.0
0.0
3.5
5.0
1.6
0.0
38.7
40.0
32.0
10.2
SnapKV
90.0
0.0
17.3
256
Quesi
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.6
19.3
0.0
2.0
1.7
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
11.3
4.0
4.0
RocketKV
98.0
53,6
$9.3
73.2
Exact-TopK
100.0
100.0
100.0
100.0
84.0
32,0
V6.5
100.0
67.6
0.4
59.3
80.0
50.0
74.6
DuoArtention
0.0
2.0
6.0
6.0
0.0
0.0
3.5
5,5
1.6
0.0
70.0
38.0
32.0
12.7
SnapKV
512
Quest
0.0
0.0
0.0
0.0
0.0
U.D
0,0
0.0
26.0
0.0
4.0
2.3
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
HSA
94.0
46.0
0.0
18.0
0.0
0.0
4.0
7.5
42.0
0.6
28.0
20.0
16.0
21.2
RocketKV
98.0
100.0
98.0
100.0
64.0
61.3
44.0
76.4
Exact-Topk
100.0
100.0
100.0
100.0
90.0
$8.0
V6.0
100.0
68.4
0.0
61.3
80.0
48.0
77.1
DunAftention
2.0
6.0
6.0
0.0
0.0
3.5
5.5
1.6
0.0
70.7
40.0
32.0
12.9
SnapKV
100.0
76.0
0.0
36.0
36.0
40.0
36.2
1024
Quest
94.0
0.0
0.0
10.0
0.0
0.0
2.5
2.5
0.2
58.0
2.0
8.0
13.7
SparQ)
36,0
0.0
0.0
0.0
0.0
0.0
0.U
0.0
1,2
0.6
15.3
2.0
0.0
4.2
100.0
88.0
12.0
0.0
66.0
43.5
42.0
48.5
RocketKV
100.0
Exact-TopK
100.0
100.0
100.0
100.0
92.0
82.0
95.5
100,0
64.4
0.0
66.7
80.0
48.0
79.1
DuoAticution
0.0
2.0
6.0
6.0
0.0
0.0
3,5
5.0
1.6
0.0
70.7
44.0
28.0
12.8
SnapKV
2048
Quesi
100.0
56.0
0.0
62.0
4.0
0.0
9.5
$4.0
0.2
$9.3
32.0
22.0
31.6
SparQ
12.0
0.0
0.0
0.0
0.0
0.0
0.0
2.0
3.6
0.2
32.7
8.0
14.0
5.6
HSA
100.0
96.0
74.0
98.0
54.0
0.0
96.5
77.2
0.4
$2.0
78.0
48.0
66.4
RocketKV
100.0
100.0
100.0
62.0
78.1
Exact-Topk
100.0
100.0
100.0
100.0
96.0
$8.0
96.3
100,0
65.6
0.0
71.3
78.0
46.0
80,1
DuoAttention
0.0
2.0
6.0
8.0
0.0
0.0
3.5
5.0
1.2
0.0
72.0
44.0
32.0
13.4
SnapKV
100.0
100.0
100.0
0.0
62.8
0.2
50.7
72.0
44.0
53.9
4096
Quest
98.0
100.0
0.0
94.0
36.0
2.0
40.5
63.0
45.2
0.0
58.7
36.0
36.0
46.9
SparQ
100.0
62.0
2.0
60.0
12.0
0.0
68.5
62.5
46.0
0.4
46.7
48.0
42.0
42.3
HSA
100.0
100.0
98.0
100.0
99.0
70.4
0.2
54.7
76.0
48.0
72.8
RocketKV
100.0
23

---

RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 10: Average results of LongBench and NIAH for Mistral-7B-Ina -B-1-----.--
Taken Dudgel
Method
Single QA
Malti Q4
Summ.
Few-shol
Synthetic
Code
LB Avg.
NIAH
NIA
Full-KV
36.4
29.6
28.1
66.7
44,9
55.6
45.4
99.0
Exact-Topk
35.5
42.6
58.2
98.7
DunAttention
17.7
19.6
18.5
52.3
28.3
22.0
SnapKV
23.9
21.1
19.3
50.6
31.6
50.4
33.6
34.0
256
Quest
8.7
5.3
15.3
12.6
3.2
32.6
13.1
14.0
SparQ
1.6
0.9
8.6
2.5
6.4
26.0
7.0
10.0
IISA
29.7
21.9
25.5
58.5
25.0
50.9
37.0
56.7
RocketKV
34.0
27.0
24.7
65.7
38,0
56,2
42.7
97.3
Exact-TopK
35.8
29.2
28.2
67.1
42.9
57.4
45.3
99.3
AuoAttention
19.8
19.2
53.9
7.1
53.8
30.8
24.0
SnapKV
27.3
23.7
37.3
53.3
37.6
45.7
512
Quest
17.4
23.0
44.8
39.7
SparQ
1.7
13.0
8.2
35.3
11.7
IISA
34.2
27.4
27.2
65.9
38.3
$7.2
43.5
85.7
RocketKV
35.2
28.1
26.4
66.6
42.3
56.2
44.2
98.3
Exact-TopK
35.8
29.5
28.2
67.0
42.8
56.7
45.2
99.3
DuoAttention
21.2
19.4
22.4
59.1
8.1
55.0
32.6
30.0
SnapKV
30.4
25.6
22.7
61.2
42.3
54.9
40.7
$4.7
1024
Quest
25.4
18.6
25.5
56.2
36.3
53.3
36.9
79.3
SparQ
13.1
4.1
16.9
43.5
16.3
20.0
HSA
35.3
28.0
42.1
44.8
RocketKV
36.0
55.9
44.7
Exact-TopK
36.2
29.7
28.2
42.9
56.8
45.2
99.3
DuoAttention
27.1
21.5
23.5
62.4
12.1
55.6
35.7
41.0
SnapKV
33.3
26.8
24.4
64.1
42.6
55.7
42.6
65.7
2048
Quest
32.2
26.9
27.0
62.5
39.3
56.0
42.3
90.3
SparQ
22.0
10.3
20.1
34.1
4.0
50.5
24.5
27.0
HSA
36.2
29,0
28.2
67.2
42.2
57.2
45.2
99.0
RocketKV
35.9
28.9
27.9
67.0
42.8
55.7
44,9
99.0
Exact- TopK
36.2
29.6
43.0
56.5
45.2
99.3
DuoAttention
32.1
26.7
56.4
SnapKV
34.2
28.0
42.8
55.7
43.5
76.7
4096
Quest
35.1
29.1
27.9
66.3
42.7
56.3
44.7
96.0
SparQ
30.6
20.3
24.8
59.4
37.7
53.3
38.9
50.3
HSA
36.4
29.5
28.1
66.8
42.8
57.0
45.3
99.3
RocketKV
36.2
29.5
28.3
66.8
42.8
55.0
45.0
99.3
24

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 11: Individual results of LongBench for Mistral-7B-Ins-v0..2
Tokor Budder
Method
parratveaja
gorper
halpotga
Jwlkingg
musique
gov-report
multi-nee
repebench-p
A7A
Fall-KV
26.8
33.1
18.8
27.1
710
87.0
569
Exact Topk
26.2
31.3
48.9
40.0
26.2
19.6
33.3
24.2
27.1
73.0
36.6
43.3
82.9
59.0
51.3
2.3
DunAltention
13.0
11.7
28.5
29.5
20.1
9.3
17.6
19.4
18.6
26.0
818
37.7
7.1
5044
54.1
3.2
SnapKV
18.5
21.4
320
4,0
6.0
7.8
2.1
13.
39.6
04
SpurQ
0.3
1.7
2.9
0.9
1.5
0.1
5.3
6.7
13.8
0.0
7.1
0,4
12.7
30.0
22.0
0.2
HSA
17.1
28.1
27.3
22.4
63.0
33.6
25.0
24.7
2.4
Exteet Hopk
27.2
31.6
48.6
50.5
27.5
19.7
33.2
24.2
27.1
720
86.6
42.8
83.1
58.3
50.5
26
DuaAmentio
SnapKV
19.6
21.9
19.6
2.5
512
Quest
6.0
17.0
29.1
18.1
13.8
5.1
22.8
18.9
25.5
35.0
50.7
21.8
43.1
48.3
41.2
299
SparQ
1.8
2.8
0.5
11.9
20.3
4.4
45.2
25.4
0.2
HSA
RoekerKV
25.9
30.9
48.9
38.0
27.0
10.2
29.0
24.2
25.9
70.5
86.9
42.3
61.6
56.4
$6.0
3.0
Extet-Topk
27.2
720
DinAttetion
15.7
20
22.0
24.8
53.0
40.5
55.1
54.9
3.1
SnapkV
21.7
23.1
46.3
35.7
23.4
17.7
21.8
22.8
23.5
57.5
85.2
40.8
81.4
53.8
55.9
32
1924
Quest
66.0
Sparl)
5.8
HSA
26.1
31.6
48.2
38.4
26.8
18.8
32.1
23.5
27.3
7300
86.5
42.5
81.8
58.8
56.7
2.5
RocletVV
27.1
32.7
48.4
39.7
26.2
19.6
31.2
24.0
27.2
71.3
86.2
42.7
81.7
56.3
55.5
2.8
Exuet-Topk
27.9
DunAlteniion
17.2
24.5
39.5
32.9
21.8
9.9
23.5
20.6
26.6
61.5
$4.3
41.3
20.7
$6.1
55.0
344
SnapKY
24.8
23.8
25.9
66.0
40.6
56.0
32
SparQ
2.9
29.3
33.8
9.0
21
0.6
16.7
16.9
26.6
35.5
47.6
19.2
6.9
S7.2
43.8
12
USA
27.2
32.2
27.0
33.2
27.4
720
42.9
$7.7
56.8
3.1
RocketKV
EescI-TopX
27.7
32.77
48.2
41.5
27.9
19.4
33.1
24.1
27.2
710
86.2
42.4
83.3
57.8
55
288
DunAttenionn
19.1
31.6
36.6
26.8
12.4
27.7
21.8
27.2
700
85.0
11.2
$1.2
56.8
56.1
22
SnapKY
28
4096
Quest
24.6
32.1
48.7
4044
27.1
19.7
32.8
23.5
27.2
72.5
86.3
40.3
82.7
57.7
549
28
SparQ
12.3
31.7
26.1
25.8
9.0
26.5
27.1
71.0
37.5
71.5
574
39
RockerKV
27.8
32.1
48.7
41.8
27.5
19.2
33.4
24.5
27.1
71.5
86.2
42.5
82.8
%600
54.0
29
25

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 12: Results of RULER-SK for Mistral-7B-Ins-wol
Tokon Budge
Method
single1
single2
shngle?
multikeyl
multikey2
multikey
multivalic
mulliquery 94.0
Swc
q0
N/A
Full-VV
100.0
100.0
100.0
98.0
100.0
82.0
79.5
97.6
90.6
71.3
82.0
64.0
89.2
Exact- Topk
100,0
100.0
100.0
98.0
100.0
84.0
80.3
98.0
96.0
88.0
08.7
84.0
66.0
89,3
DunAttention
4.0
0.0
2.0
2.0
0.0
2.0
2,4
26.8
48.7
32.0
40.0
12.5
SnapKV
16.0
0.0
5.8
28.0
52.0
256
Quest
92.0
32.0
0.0
24.0
15.2
19.6
38.0
22.0
28.0
22.0
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
2.2
9.3
2.0
10.0
1.8
HSA
98.0
74.0
42.5
60.5
73.3
84.0
66.0
65.0
RockerKV
100.0
100,0
78.0
83,5
95.5
94.4
72.7
64.0
85.7
Exact-TopK
100.0
100.0
100.0
V8.0
100.0
82.0
80.5
98.0
96.0
89.6
72.7
82.0
66.0
89.6
DonAttention
8.0
12.0
2.0
4.0
4,0
8.0
9.5
8.0
6.0
2.4
69.3
26.0
42.0
15.5
SnapKV
512
Quest
94.0
76.0
2.0
52.0
$3.5
59.5
76.4
29.8
65.3
64.0
54.0
55.3
SparQ
0.0
0.0
0.0
D.0
0.0
6.0
0.0
0.0
0.0
2.6
67.3
0.0
32.0
8.3
HSA
98.0
98.0
96.0
80,0
75.5
93.0
95.2
58.0
71.3
84.0
66.0
85.5
RocketKV
100.0
98.0
100.0
84,0
87.5
95.5
96.4
74.0
84.0
64.0
87.7
Exact-Topk
100.0
100.0
100.0
98.0
100.0
82.0
80.0
96.0
96.8
90.8
73.3
80.0
66.0
89.3
DunAtentiin
12.0
12.0
2.0
4.0
14.0
10.0
9.5
8.0
9.2
14.2
72.0
42.0
48.0
19.8
SnapKV
96.0
0.0
78.0
75.3
82.0
58.0
48.9
1024
Quest
92.0
84.0
26.0
94.0
76.0
22.0
74.5
78.0
89.6
51.8
70.0
84.0
64.0
69.7
SparQ)
0.0
62.0
6.0
72.0
0.0
72.0
70.0
61.0
0.0
8.6
68.0
30.0
44.0
38.0
HSA
100.0
98.0
100.0
84,0
79.5
98.0
96.4
68.7
82.0
66.0
88.7
RocketKV
89.5
Sxact-TopK
100.0
100.0
100.0
98.0
100.0
82.0
79.5
95.5
97.6
91.2
74.0
80.0
66.0
89.5
DonAttention
38.0
40.0
30.0
34.0
60.0
24.0
15.0
22.0
24.0
48.0
70.0
62.0
62.0
40.7
SnapKV
2048
Quest
100.0
96.0
98.0
94.0
98.0
80.0
80.5
95.5
97.2
74.2
70.7
82.0
66.0
87.1
SparQ
96.0
90.0
100.0
96.0
100.0
$4.0
78.5
95.0
81.6
90.0
70.7
82.0
66.0
86.9
HSA
100.0
100.0
100.0
98.0
100.0
82.0
80.5
96.0
96.4
88.4
69.3
80.0
66.0
89.0
RockerKV
100.0
100,0
82.0
91,5
98.0
98.8
73.3
84.0
66.0
90.4
Exact-Topk
100.0
100.0
100.0
98.0
100.0
82.0
79.5
94..5
97.6
90.8
73.3
82.0
64.0
89.4
DueAtrention
100.0
98.0
100.0
96.0
100.0
$4.0
71.0
90.5
96.0
85.4
70.0
86.0
66.0
87.9
SnapKV
96.0
2.0
94.0
92.0
27.5
$2.5
98.0
57.8
73.3
80.0
66.0
65.5
4096
Quest
100.9
100.0
100.0
98.0
100.0
$4.0
80.0
96.0
96.8
87.8
71.3
82.0
66.0
89.4
SparQ
100.0
100.0
100.0
95.0
100.0
82.0
79.0
95.5
97.6
89.8
72.7
82.0
64.0
89.3
HSA
100.0
100.0
98.0
100.0
82.0
79.5
96.0
97.6
90.6
74.0
80.0
64.0
89.4
RocketKV
26

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 13: Results of RULER-16K for
Tokon Budge
Method
single1
single2
shngle?
multikeyl
multikey2
multikeys
multivalic
mulliquery 85.5
Swe
@
N/A
Full-VV
100.0
100.0
100.0
86.0
90.0
74.0
84.0
87.6
74.2
86.7
78.0
54.0
84.6
Exact- Topk
100,0
100.0
100.0
84,0
92.0
82.0
87.0
80.0
85.2
70.6
76.7
80.0
58.0
84.3
DunAttention
2.0
0.0
2.0
2.0
0.0
0.8
12.6
48.7
30.0
36.0
10.8
SnapKV
16.0
0.0
20.0
56.0
256
Quest
0.0
2,0
0.0
D.0
3.6
18.7
8.0
12.0
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.4
0.0
2.0
0.0
0.2
HSA
96.0
2.0
24.0
16.0
72.0
70.7
74.0
54.0
39.7
RockerKV
100.0
82.5
84.4
82.0
76.0
$2.0
78.0
Exact-TopK
100.0
100.0
100.0
84.0
92.0
80.0
85,0
84.0
86.0
72.6
79.3
80.0
58.0
84.7
DonAttention
2.0
2.0
4.0
2.0
0.0
0.0
6,5
4,5
2.8
0.4
80.7
32.0
32.0
13.0
SnapKV
27.9
512
Quest
100.0
72.0
4.0
56.0
12.0
18.5
16.5
41.2
13.4
62.7
22.0
36.0
34,9
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
1.2
2.7
0.0
6.0
0.8
HSA
100.0
94.0
80.0
76.0
18.0
47.5
59.0
81.6
34.4
75.3
80.0
58.0
68.4
RocketKV
100.0
86.0
42.0
88.5
95.5
85.6
81.3
80.0
52.0
79.8
Exact-Topk
100.0
100.0
100.0
84.0
92.0
78.0
83.0
85.3
88.4
74.4
84.0
78.0
58.0
85.0
DunAmentiin
6.0
8.0
4.0
2.0
0.0
4.4
0.2
80.7
38.0
46.0
15.66
SnapKV
94.0
62.0
0.0
40.0
78.0
78.0
72.0
48
41.4
1024
Quest
84.0
94.0
10.0
82.0
44.0
61.5
$4.5
73.6
28.8
70.0
68.0
56.0
55.9
SparQ)
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
6.6
60.0
0.0
2.0
5.3
HSA
90.0
72.0
76.5
77.5
82.4
57.8
78.0
78.0
56.0
80.9
RocketKV
84.0
56.0
83.4
Exact-TopK
100.0
100.0
100.0
84.0
92.0
76.0
83.0
85.5
88.8
74,2
86.7
80.0
58.0
85.2
DonAttention
10.0
8.0
6.0
6.0
6.0
0.0
8.0
6,0
4,8
2.2
79.3
40.0
48.0
17.3
SnapKV
98.0
2048
Quest
92.0
88.0
44.0
76.0
72.0
10.0
68.0
73.0
84.0
51.6
74.7
80.0
58.0
67.0
SparQ
0.0
6.0
0.0
16.0
0.0
4.0
22.0
6.0
0.4
12.8
79.3
4.0
20.0
13.1
HSA
100.0
100.0
100.0
84.0
92.0
76.0
82.5
82.0
85.2
70.8
78.0
76.0
58,0
83.4
RockerKV
100.0
91.5
96.5
87.2
84.7
74.0
56.0
84.1
Exact-Topk
100.0
100.0
100.0
84.0
92.0
74.0
83.0
86.0
88.0
75.2
86.7
80.0
56.0
85.0
DueAtrention
48.0
46.0
40.0
32.0
48.0
14.0
24.0
20.0
22.8
16.6
80.7
62.0
58.0
39.4
SnapKV
92.0
2.0
60.0
20.5
42.0
92.0
37.6
81.3
72.0
54.0
56.0
4096
Quest
100.9
100.0
96.0
86.0
92.0
$2.0
82.5
83.0
86.0
65.8
80.7
80.0
58.0
84.0
SparQ
82.0
86.0
84.0
82.0
80.0
78.0
67.5
64.5
31.2
73.2
83.3
76.0
52.0
72.3
HSA
100.0
100.0
84.0
92.0
78.0
84.0
86.5
84.4
74.6
81.3
78.0
58.0
84.7
RocketKV
27

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 14: Results of RULER-24K for Mistral-7B-Ins-OO..
Tokon Budge
Method
single1
single2
single?
multikeyl
multikey2
multikey
multivalic
multiquery 83.5
twe
@*
ge
N/A
Full-VV
100.0
100.0
92.0
86.0
94.0
78.0
77.5
93.2
40.6
86.7
68.0
58.0
81.3
Exact- Topk
100,0
98.0
88.0
90.0
92.0
78.0
78.3
79.0
92.8
62.4
80.0
70.0
58.0
82.7
DunAttention
2.0
0.0
2.0
6.0
2.0
0.8
1.2
56.7
24.0
34.0
10.2
SnapKV
32.0
0.0
2.6
22.0
256
Quest
0.0
2,0
0.0
D.0
1.8
15.3
2.0
2.0
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
6.0
0.0
0.5
HSA
90.0
0.0
73.6
6.4
56.0
52.0
44.0
30.4
RockerKV
100.0
44,0
78.0
87.0
69.3
70.0
54.0
75.7
Exact-TopK
100.0
100.0
83.0
92.0
80.0
77.0
81,0
93.2
59.0
86.0
70.0
58.0
82.6
DonAttention
4.0
4.0
4.0
6.0
6.0
2.0
4,5
4,5
1.6
0.4
68.7
26.0
28.0
12.3
SnapKV
0.0
512
Quest
100.0
12.0
32.0
8.0
7.2
3.0
56.7
8.0
16.0
19.8
SparQ
0.0
0.0
0.0
D.0
0.0
0.0
0.0
0.0
0.0
0.6
0.0
2.0
0.0
0.2
HSA
100.0
74.0
48.0
$4.0
42.0
33,0
33.5
86.0
18.0
66.7
68.0
52.0
$1.9
RockerKV
100.0
94,0
50.0
80,5
91.5
94.4
26.2
72.7
72.0
52.0
77.8
Exact-Topk
100.0
100.0
90.0
86.0
92.0
80.0
78.0
84.0
93.6
35.4
86.7
70.0
58.0
82.6
DunAtentiin
4.0
4.0
4.0
6.0
8.0
4.5
2.0
0.6
70.0
28.0
32.0
13.0
SnapKV
96.0
0.0
32.0
76.4
7.8
62.0
46.0
37.9
1024
Quest
100.9
82.0
2.0
78.0
20.0
32.0
34.5
45.2
9.6
66.0
36.0
32.0
41.3
SparQ)
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
2.2
1.3
0.0
4,0
0.6
HSA
98.0
72.0
36.0
62.0
71.5
93.2
36.4
72.7
72.0
56.0
72.0
RocketKV
84.0
79.6
Exact-TopK
100.0
100.0
90.0
86.0
94.00
80.0
78.0
83.0
93.6
50.2
88.7
70.0
58.0
82.4
DonAttention
4.0
4.0
4.0
6.0
10.0
2.0
4.5
4,5
2.4
0.4
69.3
26.0
30.0
12.9
SnapKV
0.0
70.7
2048
Quest
96.0
88.0
2.0
64.0
65.0
66.0
88.0
18.0
75.3
62.0
50.0
58.0
SparQ
0.0
0.0
0.0
D.0
0.0
34.0
0,0
0.0
0.0
27.2
80.0
2.0
4,0
11.3
HSA
100.0
98.0
88.0
88.0
90.0
68.0
73.5
80.5
94.8
53.6
74.7
68.0
60.0
79.8
RockerKV
100.0
87.0
94,0
98.4
81.3
54.0
82.2
Exact-Topk
100,0
100.0
90.0
86.0
94.0
78.0
78.0
84.0
93.6
41.2
88.7
70.0
58.0
81.7
DueAttention
24.0
18.0
10.0
18.0
20.0
2.0
6.0
6.0
6.4
0.8
71.3
52.0
44.0
21.4
SnapKV
84.0
0.0
$2.0
19.5
40.5
96.8
75.3
68.0
48.0
51.3
4096
Quest
90.0
86.0
72.0
72.0
72.0
38.0
71.5
75.0
90.4
43.4
80.0
70.0
56.0
70.5
SparQ
0.0
6.0
0.0
4.0
0.0
8.0
2.0
10.5
0.0
6.2
86.0
0.0
18.0
10.8
HSA
98.0
88.0
92.0
78.0
73.5
83.0
94.8
62.0
82.7
70.0
58.0
82.3
RocketKV
28

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 15: Results of RULER-32K for Mistral-7B-Ins-vO..
Tokon Budge
Method
single1
single2
shngle?
multikeyl
multikey2
multikey
multivalic
mulliquery 81.0
twe
@
qe
N/A
Full-KV
100.0
98.0
86.0
82.0
88.0
68.0
83,0
90.0
29.6
94.7
70.0
48.0
78.3
Exact- Topk
100,0
92.0
68.0
78.0
80.0
00.0
78.0
58.5
91.2
40.4
86.0
66.0
48.0
72.8
DunAttention
0.0
2.0
4.0
4.0
0.0
0.0
2.0
34.0
24.0
30.0
SnapKV
36.0
0.0
0.2
10,9
256
Quest
0.0
0.0
0.0
D.0
0.6
22.0
0.0
2.0
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
2.0
0.0
0.2
HSA
76.0
0.0
55.6
2.2
57.3
26.0
32.0
21.4
RockerKV
100.0
20.0
78.5
83,0
81.6
7.2
80.7
62.0
48.0
71.0
Exact-TopK
100.0
94.0
78.0
86.0
64.0
83.0
60.5
93.2
37.4
87.3
66.0
46.0
74.1
DonAttention
2.0
6.0
6.0
4.0
0.0
2.0
3,5
4,0
2.8
0.8
75.3
24.0
30.0
12.3
SnapKV
512
Quest
88.0
2.0
4.0
2.0
4.0
64.0
8.0
16.0
14.7
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.4
0.0
0.0
2.0
0.2
HSA
96.0
28.0
4.0
28.0
26.0
16.5
12.5
77.2
9.2
77.3
72.0
44.0
37.7
RockerKV
100.0
32.0
84,0
87.5
93.6
12.4
84.7
70.0
46.0
74.5
Exact-Topk
100.0
96.0
72.0
80.0
90.0
64.0
81,3
66.0
92.4
37.0
91.3
68.0
50.0
76.0
DunAtentiin
2.0
6.0
6.0
4.0
0.0
4.5
3.6
0.6
80.7
20.0
32.0
12.7
SnapKV
58.0
0.0
22.0
56.0
2.6
72.0
58.0
40.0
33.0
1024
Quest
94.0
74.0
2.0
44.0
18.0
15.5
19.0
69.6
5.6
77.3
28.0
32.0
36.8
SparQ)
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.6
0.0
0.0
8.0
0.7
HSA
98.0
50,0
10.0
46.5
38,0
85.2
24.8
78.7
74.0
48.0
56.6
RocketKV
100.0
48.0
75.3
Exact-TopK
100.0
98.0
76.0
82.0
90.0
64.0
84,0
73.5
90.8
34.0
92.7
68.0
48.0
77.0
DonAttention
2.0
6.0
6.0
4.0
0.0
2.0
3,5
5.0
3.6
0.8
81.3
28.0
30.0
13.2
SnapKV
2048
Quest
98.0
94.0
2.0
70.0
32.0
49.5
46.0
91.2
7.0
84.7
$2.0
40.0
51.3
SparQ
0.0
0.0
0.0
D.0
0.0
0.0
0,0
0.0
0.0
4.4
67.3
0.0
4,0
5.8
HSA
100.0
96.0
72.0
72.0
40.0
71.5
59.5
90.8
44.2
82.0
68.0
50.0
70.9
RockerKV
100.0
85,0
96.4
64.0
77.8
Exact-Topk
100,0
98.0
76.0
82.0
92.0
64.0
82.5
77.0
90.4
34,8
93.3
68.0
48.0
77.4
DueAttention
10.0
8.0
8.0
6.0
6.0
2.0
4,0
4,5
4.8
1.8
84.0
34,0
32.0
15.8
SnapKV
74.0
0.0
32,0
12.0
22.0
92.8
16.2
80.0
62.0
46.0
46.4
4096
Quest
92.0
78.0
32.0
68.0
64.0
67.0
$9.0
90.0
24.8
89.3
68.0
46.0
60.5
SparQ
0.0
2.0
0.0
0.0
0.0
0.0
0.5
0.0
0.0
9.4
88.0
0.0
4,0
8.0
HSA
98.0
78.0
84.0
60.0
74.0
68.5
94.4
43.4
85.3
66.0
48,0
75.0
RocketKV
29

---

RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 16: Average results of LongBench and NIAH for LongChat-7B--v/
Taken Dudgel
Method
Single QA
Malti Q4
Summ.
Few-shol
Synthetic
Code
LB Avg.
NIAH
NIA
Full-KV
31.1
24.1
26.7
57.7
15.2
57.0
37.5
96.3
-xpct-TopK
30.5
23.5
14.8
96.7
DunAttention
14.7
15.8
14.4
2.4
48.5
21.7
20.7
SnapKV
25.9
23.0
20.8
45.1
15.2
51.1
31.7
43.3
256
Quest
9.2
7.2
7.8
11.4
0.9
34.4
11.8
15.3
SparQ
18.9
11.2
15.2
16.8
1.5
32.4
16.9
21.0
IISA
28.2
23.7
25.1
49.4
13.0
54.4
34.1
58.0
RocketKV
30.1
24,5
23.0
51.8
13.8
56,5
35.1
96.3
Exact-TopK
30.8
23.9
26.7
54.4
14.5
57.6
36.8
96.3
AuoAttention
18.1
20.4
19.5
49.4
26.1
23.7
SnapKV
28.1
23.8
14.5
53.3
34.1
66.7
512
Quest
20.7
15.6
46.4
48.0
SparQ
26.5
21.4
23.6
41.5
13.5
49.1
30.9
IISA
30.6
24.8
26.5
52.3
14.2
$7.2
36.3
90.7
RocketKV
30.9
24.1
24.5
$2.9
14.5
56.4
35.9
96.7
Exact-TopK
30.7
23.9
26.8
54.3
14.2
57.4
36.7
96.3
DuoAttention
20.8
20.7
21.0
42.0
4.0
53.6
28.5
28.0
SnapKV
28.9
24.0
23.6
$3.0
14.5
54.1
35.0
80.0
1024
Quest
28.2
23.0
23.6
47.9
17.2
53.3
33.8
77.0
SparQ
29.7
22.9
26.4
17.2
55.9
36.1
76.7
HSA
30.7
24.2
27.0
14,0
57.9
36.7
95.7
RocketKV
14.8
Exact-TopK
30.7
24.0
15.0
57.3
96.3
DuoAttention
26.0
20.2
23.1
46.3
3.8
55.0
30.9
42.3
SnapKV
30.2
23.8
24.9
54.0
14.5
54.4
35.8
85.0
2048
Quest
30.0
24.2
26.5
53.3
15.5
56.3
36.3
89.7
SparQ
30.6
24.1
26.9
54.6
14.5
56.8
36.7
96.0
HSA
30.7
24.1
26.8
54.6
14.5
57.3
36.8
96.3
RocketKV
31.1
24.3
26.3
54.6
14.5
55.0
36.5
96.3
Exact- TopK
31.0
24.0
26.6
15.5
57.2
36.9
DuoAttention
28.6
22.4
10.5
55.2
34.1
SnapKV
31.1
24.4
26.1
14.8
55.1
36.6
89.3
4096
Quest
30.7
24.2
27.0
54.3
14.2
57.3
36.8
96.3
SparQ
30.8
23.9
26.7
54.6
15.0
57.2
36.8
94.3
HSA
31.0
23.9
26.8
54.4
14.2
57.3
36.8
96.3
RocketKV
30.7
24.3
26.8
54.3
14.2
55.2
36.5
96.3
30

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 17: Individual results of LongBench for LongChat-7B-vL.
Tokor Budder
Method
parratreaja
gorper
halpoteaa
Jwlkingg
musique
gov-report
mulli-nre
trieiage
pass-
repebench-
A7A
Fall-VV
20.8
29.4
43.1
26.7
22.5
0.0
Exact Topk
19.6
30.2
41.6
33:7
23.5
13.5
30.6
22.5
26.3
65.5
83.5
12.7
29.5
57.7
58.9
0.0
DunAltention
12.4
10.1
21.6
18.6
20.5
8.1
12.6
18.7
11.8
13.5
65.9
9.8
4.8
48.6
48.3
0.0
SnapKV
17.5
23.1
44.0
0.9
08
SpurQ
2,5
27.3
26.8
12.9
19.4
L.2
10.4
9.8
25.3
22.5
25.4
2.5
2.5
44.8
20.0
0.5
HSA
14.1
27.3
57.8
200
18.6
24.5
1.5
Exteet Topk
20.1
30.5
41.9
33.9
25.5
14.0
30.6
22.9
26.6
06.0
84.4
12.9
29.0
50.4
58.8
0.0
DuoAmentio
SnapKV
17.5
21.5
23.1
59.0
52.4
512
Quest
5.0
26.8
10.2
21.3
19.8
5.8
114
13.9
25.2
42.5
$1.4
3.5
13.8
$4.0
38.8
1.5
SparQ
10.3
29.6
39.5
24.1
23.9
26.5
62.0
51.6
43.5
1.0
HSA
30.5
RockerKV
20.2
29.8
42.6
34.8
23.6
13.8
26.9
22.1
24
64.5
84.1
10.1
28.5
$5.9
$6.9
0.5
Extet-Topk
DinAttetionn
13.5
18.1
28.2
24.8
9.2
20.2
22.9
43.5
6.5
56.6
SnapKV
18.9
27.1
40.5
35,5
23.1
13.4
23.9
22.1
24.9
63.0
$4.2
11.9
29.0
$2.0
56.2
0.0
1924
Quest
63.0
2.5
Sparl)
HSA
20.4
S0LS
41.2
34.2
23.7
14,9
31.4
23,0
26.5
65.5
84.6
11.4
28.0
57.2
58.6
0.0
RocletVV
20.2
29.4
43.3
34.6
23.8
14.3
28.9
22.5
25.5
67.5
84.1
11.8
29.5
54.1
56.4
0.0
Exuet-TopK
24.7
31.1
DunAlteniion
13.9
26.6
37.4
27.8
25.4
7.4
23.3
19.9
26.1
55.5
74.4
9.1
6.5
56.5
53.5
1.0
SnapKY
20.7
23.4
26.5
25.6
12.8
52.6
0.0
30.2
56.2
SparQ
20.3
20.5
42.0
33.3
24.7
14.4
31.3
22,9
26.5
66.0
84.5
13,3
29.0
55.1
58.
0.0
USA
20.3
29.9
31,2
26.4
67.0
12.4
55.5
0.0
EescI-TopX
21.0
20.4
42.6
33.0
24.3
14.6
30.7
22.6
26.5
66.5
84.0
12.8
31.0
55.1
50.4
0.0
DunAttenionn
14.2
29.6
322
23.9
28.6
26.5
61.0
78.3
20.5
55.7
51.6
0.5
SnapKY
4096
Quest
20.5
29.8
41.9
33.9
24.8
14.0
31.4
22.8
26.7
67.0
84.3
11.6
28.5
55.4
50.2
0.0
SparQ
21.0
29.0
33.2
24.2
14.3
30.7
22.8
26.6
66.5
$4.2
13,1
30,0
55.1
99.4
0.0
55.3
RockerKV
20.3
29.2
42.8
33.5
24.9
14.5
30.5
23.3
26.5
660
842
12.8
28.5
53.5
56.9
0.0
31

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 18: Results of RULER-8K for LongChat-7B-vl..
Tokon Budge
Method
single1
single2
single?
multikeyl
multikey2
multikeys
multivalue
mulliquery
Swe
qe
N/A
Full-KV
100.0
100.0
96.0
98.0
100.0
74.0
96.0
69.0
94.0
33.8
80.7
70.0
58.0
82.3
Exact- Topk
100,0
96.0
96.0
94,0
100.0
84.0
88.0
00.0
90.8
39.4
76.7
66.0
58.0
80.7
DunAttention
4.0
0.0
0.0
2.0
2.0
1.2
1.6
3.3
28.0
24.0
SnapKV
0.0
54.0
46.0
256
Quest
52.0
56.0
0.0
28.0
30.8
12.4
46.0
6.0
18.0
20.1
SparQ
4.0
0.0
0.0
2.0
0.0
0.0
0.5
0.0
20.0
30.4
35.3
46.0
40.0
13.7
HSA
98.0
78.0
92.0
80.0
27.5
71.3
72.0
56.0
RockerKV
100,0
56.0
87.2
33.4
69.3
66.0
56.0
74.5
Exact-TopK
100.0
98.0
966.0
96.0
100.0
82.0
92.5
64.0
91.2
36.6
78.0
64.0
58.0
81.3
DonAttention
8.0
12.0
2.0
4.0
4,0
0.0
8.5
3.0
0.0
0.4
52.7
30.0
26.0
11.6
SnapKV
512
Quest
98.0
74.0
2.0
36.0
68.0
45.5
21.0
83.6
30.4
71.3
38.0
40.0
46.8
SparQ
98.0
94.0
$6.0
76.0
92.0
16.0
39.0
23.0
78.0
38.4
70.0
72.0
$6.0
62.2
HSA
96.0
84.0
94.0
100.0
74,0
85.0
$1.5
88.8
37.0
71.3
68.0
58.0
77.5
RockerKV
98.0
100,0
44.0
92.0
66.0
90.0
33.6
72.0
60.0
78.0
Exact-Topk
10000
100.0
96.0
98.0
100.0
80.0
97.0
68.0
93.2
34,4
82.0
62.0
56.0
82.0
DunAtentiin
12.0
12.0
0.0
4.0
8.5
1.2
0.8
62.0
32.0
30.0
13.2
SnapKV
92.0
92.0
40.0
70.5
24.5
90.8
60.0
58.0
56.8
1024
Quest
100.9
90.0
30.0
50.0
78.0
24.0
80.0
49.5
92.4
32.8
73.3
66.0
54.0
65.4
SparQ)
100.0
98.0
94.0
98.0
100.0
66.0
93.0
67.0
91.2
35.2
79.3
66.0
56.0
80.3
HSA
96.0
100.0
86,0
91.0
62.0
92.4
37.0
74.7
66.0
58.0
81.3
RocketKV
Exact-TopK
100.0
100.0
96.0
98.0
100.0
80.0
96.5
69.0
92.8
34.4
82.7
66.0
58.0
82.6
DonAttention
48.0
76.0
18.0
78.0
66.0
26.0
61.5
28.5
14.8
1.6
76.7
60.0
48.0
46.4
SnapKV
64.0
66.7
2048
Quest
100.0
98.0
96.0
98.0
100.0
84,0
96.0
70.0
92.8
34.2
74.0
68.0
58.0
82.2
SparQ
100.0
100.0
96.0
98.0
100.0
80.0
97.0
68.5
92.8
34.0
82.0
66.0
58.0
82.5
HSA
100.0
100.0
96.0
98.0
100.0
80.0
98.0
67.0
91.2
36.6
79.3
62.0
58.0
82.0
RockerKV
100,0
76.0
94.5
66.5
79.3
58.0
81.3
Exact-Topk
100,0
100.0
96.0
98.0
100.0
74.0
96.5
70.0
93.6
33.8
82.7
70.0
58.0
82.5
DueAtrention
98.0
100.0
88.0
96.0
100.0
64.0
97.0
74.5
85.6
35.0
82.0
66.0
58.
80.3
SnapKV
100.0
98.0
90.0
34.0
93.0
66.5
94.0
74.0
66.0
60.0
77.2
4096
Quest
100.9
100.0
98.0
100.0
78.0
96.5
69.5
93.6
34.0
80.0
66.0
58.0
82.3
SparQ
100.0
100.0
96.0
9%.0
100.0
76.0
96.5
69.5
93.2
33.8
82.0
70.0
58.0
82.5
HSA
100.0
98.0
100.0
80.0
96.5
69.5
92.8
34.4
82.7
66.0
58.0
82.6
RocketKV
32

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 19: Results of RULER-I6K for LongChat-7B-v1.5
Tokon Budge
Method
single1
single2
single?
multikeyl
multikey2
multikey
multivaluc
mulliquery
twc
@
qe
N/A
Full-KV
100.0
100.0
96.0
94.0
98.0
42.0
84.5
50.0
82.4
18.0
72.0
64.0
44.0
72.7
Exact- Topk
100,0
98.0
98.0
86.0
100.0
54.0
74,0
44.3
75.6
38.4
70.0
58.0
44,0
72.3
DunAttention
2.0
0.0
0.0
2.0
0.0
0.8
0.4
0.7
26.0
30.0
SnapKV
0.0
48.8
2.0
40.7
50.0
256
Quest
0.0
0.0
0.0
D.0
0.0
0.2
12.0
0.0
6.0
SparQ
6.0
2.0
0.0
0.0
0.0
0.0
0.0
0.0
2.8
0.0
12.0
2.0
14.0
3.0
HSA
96.0
22.0
36.5
76.8
39.3
50.0
41.2
RockerKV
78.0
49.5
73.2
58.0
46.0
65.3
Exact-TopK
100.0
100.0
98.0
84.0
100.0
54.0
78.5
50.0
76.8
29.8
73.3
58.0
44.0
72.8
DonAttention
2.0
2.0
2.0
2.0
2.0
0.0
5.5
4,5
0.4
0.6
68.0
20.0
30.0
10.7
SnapKV
512
Quest
58.0
46.0
0.0
20.0
16.0
61.2
8.4
48.0
4.0
16.0
21.5
SparQ
46.0
54.0
0.0
6.0
10,0
0.0
5.5
0.5
64.8
32.6
22.0
30,0
36.0
23.6
HSA
96.0
62.0
72.0
47.5
78.8
19.6
$2.0
60.0
46.0
61.5
RockerKV
100.0
98.0
24,0
83.0
61.5
76.8
18.4
66.0
66.0
46.0
70.4
Exact-Topk
100.0
100.0
96.0
90.0
100.0
$2.0
82.0
49.3
77.6
26.4
69.3
64.0
44.0
73.1
DunAtentiin
6.0
8.0
2.0
2.0
2.0
0.6
65.3
28.0
34.0
12.4
SnapKV
0.0
27.5
77.6
60.7
54.0
44.0
43.9
1024
Quest
94.0
98.0
8.0
74.0
42,0
55.5
19.0
80.0
26.0
58.7
36.0
34.0
48.1
SparQ)
100.0
96.0
88.0
86.0
90.0
28,0
66.0
42.0
79.2
27.6
68.0
$8.0
38.0
66.7
HSA
98.0
88.0
44.0
76.0
46.5
77.6
29.2
60.0
56.0
46.0
69.3
RocketKV
100,0
81.5
46.0
72.1
Exact-TopK
100.0
100.0
96.0
98.0
50.0
82.5
$1.0
81.6
25.0
71.3
62.0
44,0
73.3
DonAttention
6.0
14.0
2.0
4.0
2.0
0.0
8.5
5.0
1.2
0.4
70.0
38.0
28.0
13.8
SnapKV
78.0
2048
Quest
100.0
92.0
64.0
82.0
80.0
16.0
80.5
48.0
77.6
26.8
69.3
$2.0
48.0
64.3
SparQ
100.0
100.0
94.0
90.0
100.0
54.0
82.0
53.5
77.6
21.2
74.7
64.00
44.0
73.5
HSA
100.0
100.0
98.0
88.0
98.0
56.0
81.5
$2.0
78.8
29.6
69.3
62.0
46.0
73.8
RockerKV
100.0
100,0
80.5
78.4
69.3
64.0
46.0
73.0
Exact-Topk
100,0
100.0
96.0
92.0
98.0
44.0
84.5
$2.5
82.0
20.6
70.7
62.0
44.0
72.8
DueAtrention
82.0
90.0
38.0
90.0
58.0
4.0
68.5
46.0
34.8
0.4
70.7
58.0
46.0
$2.8
SnapKV
96.0
92.0
42.0
69.0
37.0
78.4
21.4
65.3
62.0
44.0
58.7
4096
Quest
100.9
100.0
92.0
96.0
62.0
81.5
53.5
78.4
26.6
68.7
62.0
44.0
74.1
SparQ
100.0
100.0
96.0
92.0
98.0
48.0
81.5
$2.0
81.6
22.4
72.0
62.0
44.0
73.0
HSA
100.0
92.0
98.0
$2.0
83,0
$3.0
79.2
26.4
70.7
62.0
44.0
73.6
RocketKV
33

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 20: Results of RULER-24K for LongChat-7B-v1.5
Tokon Budge
Method
single1
single2
shngle?
multikeyl
multikey2
multikeys
multivalue
mulliquery
A
twe
qe-3
N/A
Full-KV
100.0
100.0
98.0
90.0
84,0
18.0
60.5
38.5
79.6
5.8
70.0
50.0
36.0
63.9
Exact- Topk
100,0
98.0
98.0
88.0
78.0
16.0
46.3
39.0
60.0
24.8
64.7
46.0
38.0
61.3
DunAttention
0.0
0.0
2.0
4.0
0.0
0.4
20.0
24.0
SnapKV
0.0
59.2
1.0
28.7
40.0
256
Quest
0.0
0.0
0.0
0.0
0.0
16.0
0.0
2.0
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
2.0
0.2
HSA
78.0
0.0
7.8
47.3
28.0
28 0
24.7
RockerKV
100.0
51,5
37,5
70.4
3.8
53.3
48.0
34.0
56.2
Exact-TopK
100.0
100.0
98.0
88.0
78.0
22.0
49.5
38.5
63.6
19.6
68.7
46.0
36.0
62.1
DonAttention
2.0
4.0
4.0
4.0
2.0
0.0
4,5
4,0
0,4
0.6
55.3
22.0
22.0
9.6
SnapKV
28.0
512
Quest
96.0
0.0
0.0
0.0
32.8
0.6
50.7
2.0
12.0
15.1
SparQ
10.0
0.0
0.0
2.0
0.0
0.0
0.0
0.0
6.4
12.2
34.0
20.0
18.0
7.9
HSA
90.0
86.0
44.0
34,0
28.0
64.0
17.6
$1.3
50.0
42.0
49.3
RockerKV
100.0
68.0
78.0
$3.0
38,0
71.6
7.6
53.3
50.0
38.0
59.8
Exact-Topk
100.0
100.0
98.0
86.0
82.0
20.0
$8.0
35.3
06.8
12.3
70.7
46.0
34.0
62.2
DunAtentiin
4.0
4.0
4.0
6.0
2.0
4.5
0.4
0.6
58.7
16.0
24.0
9.9
SnapKV
2.0
71.6
4.0
38.7
50.0
30.0
33.1
1024
Quest
98.0
44.0
2.0
72.0
40.0
17.0
73.6
13.2
56.0
20.0
34.0
36.6
SparQ)
86.0
62.0
0.0
20.0
12.0
0.0
4.5
5.5
76.8
27.6
26.0
48.0
34.0
31.0
HSA
98.0
76.0
10.0
40.0
36.5
72.0
26.2
$2.0
50.0
42.0
60.1
RocketKV
62.3
Exact-TopK
100.0
100.0
98.0
88.0
84.0
22.0
60.0
36.5
66.8
9.4
69.3
48.0
34,0
62.8
DonAttention
4.0
4.0
4.0
6.0
2.0
0.0
4.5
4,5
0.4
0.4
$7.3
22.0
24.0
10.2
SnapKV
77.6
39.1
2048
Quest
90.0
74.0
42.0
68.0
50.0
57.5
24.5
79.2
34.2
56.7
44.0
36.0
50.5
SparQ
100.0
100.0
98.0
90.0
82.0
10.0
$7.0
38.5
63.6
11.0
65.3
48.0
38.0
61.6
HSA
100.0
100.0
92.0
86.0
78.0
28.0
$3.0
39.5
76.8
21.4
62.7
46.0
36.0
63.0
RockerKV
100.0
60.0
75.2
36.0
62.9
Exact-Topk
100,0
100.0
98.0
90.0
84.0
20.0
63.5
36.5
68.4
7.6
72.0
48.0
36.0
63.4
DueAtrention
16.0
14.0
4.0
10.0
2.0
0.0
6.5
5.0
0.4
0.6
65.3
28.0
22.0
13.4
SnapKV
22.0
33,5
23.5
78.0
8.6
56.7
48.0
34.0
45.1
4096
Quest
100.9
98.0
90.0
78.0
10.0
65.0
48.0
80.0
25.2
64.0
50.0
36.0
64.6
SparQ
100.0
100.0
98.0
90.0
82.0
18.0
61.0
38.0
65.2
7.0
69.3
48.0
36,0
62.5
HSA
100.0
85.0
82.0
22.0
58.5
38.5
67.2
12.2
66.0
46.0
34.0
62.5
RocketKV
8.2
34

---

RocketKV; Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Table 21: Results of RULER-32K for LongChat-7B-v1.5
Tokon Budge
Method
single1
single2
single?
multikeyl
multikey2
multikeyd
multivalic
mulliquery
twc
@
qe
g
N/A
Full-KV
100.0
94.0
80.0
78.0
68.0
4,0
58.5
31.0
47.6
38.6
70.7
22.0
34.0
55.9
Exact- Topk
100,0
88.0
82.0
78.0
70.0
8.0
49.0
25.0
39.2
53.4
57.3
18.0
30.0
53.7
DunAttention
0.0
2.0
4.0
4.0
0.0
0.4
0.4
0.7
20.0
26.0
SnapKV
0.0
19.2
0.2
24.0
256
Quest
0.0
0.0
0.0
D.0
0.0
0.0
0.0
SparQ
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
HSA
60.0
0.0
38.0
2.0
22.0
12.1
RockerKV
51,5
27.5
33.6
8.8
32.7
26.0
30.0
47.9
Exact-TopK
100.0
90.0
84.0
78.0
72.0
51,0
27.5
43.2
52.4
65.3
20.0
30.0
55.5
DonAttention
2.0
6.0
8.0
4.0
0.0
0.0
2,5
1.5
0.0
0.4
33.3
20.0
28.0
8.1
SnapKV
512
Quest
24.0
0.0
0.0
0.0
0.4
18.7
0.0
10.0
SparQ
0.0
0.0
0.0
D.0
0.0
0.0
0.0
0.0
6.8
1.8
10.0
2.0
14.0
2.7
HSA
76.0
28.0
2.0
34,0
13,5
35.2
14.0
34.7
20.0
32.0
22.3
RockerKV
98.0
54,0
32,0
37.6
14.8
36.7
22.0
34.0
48.9
Exact-Topk
100.0
90.0
84.0
78.0
72.0
50.0
26.0
44.0
49.8
68.0
20.0
32.0
55,4
DunAtenttin
2.0
6.0
8.0
4.0
2.0
2.0
0.4
0.4
33.3
20.0
24.0
SnapKV
0.0
32.0
38.8
3.4
24.0
32.0
23.0
1024
Quest
96.0
20.0
0.0
8.0
35.6
3.4
42.0
2.0
18.0
18.2
SparQ)
$4.0
8.0
0.0
0.0
0.0
0.0
0.0
0.0
20.8
33.0
22.0
18.0
34,0
14.6
HSA
94.0
70.0
$4.0
17.5
29.2
42.7
22.0
30.0
42.1
RocketKV
50.5
Exact-TopK
100.0
90.0
82.0
78.0
68.0
51,5
27.0
47.2
47.6
73.3
20.0
30.0
55.4
DonAttention
2.0
6.0
8.0
4.0
2.0
0.0
2,5
2.0
0.8
0.2
35.3
18.0
24.0
8.1
SnapKV
7.8
2048
Quest
88.0
90.0
0.0
46.0
14.0
50.5
17.5
50.0
26.6
52.0
20.0
26.0
37.0
SparQ
96.0
58.0
62.0
78.0
34.0
0.0
34.0
20.5
41.2
46.2
60.0
22.0
28.0
44.6
HSA
86.0
86.0
76.0
44.0
$2.0
30.5
40.4
46.8
58.0
20.0
28.0
$1.7
RockerKV
50.5
39.2
62.0
34,0
53.8
Exact-Topk
100.0
94.0
82.0
78.0
70.0
$3.0
30.0
46.8
41.2
71.3
20.0
32.0
55.7
DueAttention
6.0
8.0
8.0
6.0
2.0
2.0
2,5
2,5
1.6
0.0
49.3
18.0
26.0
10.1
SnapKV
6.0
70.0
10.0
22.0
17.5
40.8
12.6
38.0
18.0
32.0
34.1
4096
Quest
100.9
94.0
66.0
78.0
54.0
$8.0
27.0
45.2
44.8
57.3
22.0
26.0
51.7
SparQ
100.0
94.0
84.0
78.0
58.0
2.0
49.5
26.5
45.6
47.4
72.0
20.0
32.0
54.5
HSA
90.0
76.0
70.0
49.5
30.5
43.2
44.2
66.0
20.0
32.0
55.2
RocketKV
35


Reviews:
[{"review_id": "UScpM1HwrA", "reviewer": "Reviewer_utBU", "summary": "This paper presents RocketKV, a method that leverages observation made upon existing permanent and dynamic token eviction. Specifically, RocketKV aims to conduct a permanent eviction with a large budget first and refine it to target a budget with fine-grained dynamic evictions.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer’s insightful feedback and recognition of our work in several aspects.\n\n**Needle Dataset:**\nOur needle dataset already follows reference [2] mentioned by the reviewer which adopts PGraham Essay as background and a passkey-like needle as explained in Appendix A.2. \n\n\n**SCBench Results:**\nThis is a great suggestion. We have conducted additional experiments with SCBench under multi-turn mode and below are the preliminary results. Due to time and resource constraints, we only present results with 4K token budget on Llama3.1-8B-Instruct, but do plan to show more comprehensive SCBench results in the final version of the paper. Notice that some results for SnapKV are missing because it only compresses the KV cache of the input prompt but not the generated tokens, therefore it cannot meet the 4K token budget requirement for tasks with more than 4K generated tokens.\n\nOur results below demonstrate that RocketKV still greatly outperforms other baseline methods in multi-turn scenarios, especially for string retrieval tasks. However, there is still a noticeable gap between RocketKV and Exact-TopK, showing room for further improvement. As you already pointed out, we believe this is due to the limitation of SnapKV which relies on the query at the end of the input prompt to filter out unimportant KV tokens. In multi-turn scenarios, the unimportant KV tokens evicted by the previous turns might be essential for queries in later turns, and this could cause a significant accuracy drop in later turns. To address this challenge, we propose a variant of RocketKV called RocketKV-MT where we do not perform permanent KV token eviction in SnapKV++ but keep them all for later turns. However, the decode phase is still restricted to perform dynamic selection on the filtered KV tokens by SnapKV++. For instance, if SnapKV++ identifies 4K important tokens out of 16K input tokens in the prefill phase, all 16K KV tokens are kept in memory but the hybrid attention method only selects among the 4K important tokens in the decode phase. In the next turn, all 16K input tokens are added as prefixes in the prefill phase and SnapKV++ will perform another round of filtering based on unfiltered history. By doing this, RocketKV-MT still achieves the same performance benefit as RocketKV but does not introduce memory storage savings. Our results below demonstrate that RocketKV-MT performs much better than RocketKV in SCBench and achieves almost the same accuracy as Exact-TopK. We would like to also point it out that RocketKV-MT would be a great fit for disaggregated serving where different GPUs are used for prefill and decode (see NVIDIA’s Dynamo[4] and Moonshot AI’’s Mooncake[5]), as it can keep the full KV cache on the prefill node and send only the filter KV cache to the decode node in each turn. This would lead to not only memory storage savings on the decode node but also significant communication traffic savings between the prefill and decode node.\n\n| Method           | Retr.String | Retr.Semantic | Global | Multi-task | AVG.  |\n|------------------|-------------|----------------|--------|-------------|--------|\n| Full-KV          | 49.3        | 40.9           | 36.3   | 64.7        | 50.0   |\n| Exact-TopK       | 43.4        | 40.0           | 36.4   | 63.8        | 47.9   |\n| DuoAttention     | 0.1         | 25.2           | 34.0   | 12.7        | 20.8   |\n| SnapKV           | 2.9         | N/A            | 35.4   | N/A         | N/A    |\n| Quest            | 8.6         | 26.0           | 28.0   | 23.2        | 23.7   |\n| SparQ            | 3.0         | 27.9           | 29.1   | 28.7        | 24.1   |\n| RocketKV         | 36.9        | 26.2           | 35.9   | 31.8        | 35.2   |\n| RocketKV-MT      | **47.8**    | **37.3**       | **37.0** | **61.2**    | **47.8** |\n\n[4] NVIDIA Dynamo, https://github.com/ai-dynamo/dynamo\n\n[5] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving\n\n**LongBench Input Truncation:**\nWe follow the original setting from LongBench to truncate the input if the sequence length is longer than the maximum sequence length enabled by the model. However, the maximum sequence length across all tasks in LongBench is within 30K tokens (as shown in Table 3 of Appendix D in the SnapKV paper), while all three models we evaluated have a maximum sequence length of at least 32K tokens. Therefore, no truncation occurs in practice during our evaluation. \n\n**Proper Quotation Marks:**\nThank you for pointing this out, we will correct those quotation marks in the final version of the paper.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The needle dataset seems to be basic if it strictly follows the GKamradt setup (with the background filler being the reputation of texts). This needle setup is known to be weak per findings like [1] and [2] and it is recommended for the authors to opt for a more comprehensive needle setup. A common practice is to adopt PGraham Essay as background and a passkey-like needle, as done in [2]. This is possibly not much of a concern due to the adaptation of RULER, which is a much more standardized needle task, but it should still be considered.\n\nFurther, while LongBench and RULER are popular datasets for long context evaluation, these alone might be a bit outdated by today's standard. As the authors are familiar with SnapKV and Razor/DuoAttention, one major challenge of SnapKV is it is query-position sensitive, which is showcased in literature like SCBench. I am interested in seeing if the proposed method can perform well on such multi-round datasets. \n\nLast, per A.2, the LongBench input is truncated even for long context-capable models like Llama 3.1. It is recommended to have non-truncated input for such models or adopt a longer dataset like $\\infty$Bench.\n\n\n[1] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory \n[2] KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches\n[3] SCBench: A KV Cache-Centric Analysis of Long-Context Methods", "theoretical_claims": "-", "experimental_designs_or_analyses": "Yes. All three of them.", "relation_to_broader_scientific_literature": "-", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "See other sections for weaknesses. Generally I am very fond of this work as it is:\n- Nicely written\n- Makes some clear and correct taxonomy of existing works, and pay proper tribute when the improvement is incremental in nature (e.g., SnapKV++).\n- Goes for a simple but effective approach when such a design works.\n- Solid evaluation on executed ones — I might have some reservations regarding its dataset comprehensiveness, but for the conducted experiments, the model and setting coverage seem nicely done.\n\nI am open to improve the rating upon a satisfactory rebuttal.", "other_comments_or_suggestions": "Nothing major but it might be worth noting the proper quotation marks for LaTeX are `` and ''. It seems like sometimes only the right quotation marks are used.\n\n\n---\n\n**Post Rebuttal Update**\n\n\n\nI thank the authors for the added results and clarifications. Without memory efficiency gains, RocketKV-MT remains a rather pilot study akin to methods like DeepSeek NSA (though one may argue NSA actually generates *more* KV cache), and while I agree there are channels for optimization in a disaggregated service scenario, significant consideration would incur under that setting. I hope to see a faithful discussion of RocketKV-MT in this regard in the updated manuscript. Further, on the presentation side, the current writing of SnapKV++ can also use a bit more background for unfamiliar readers. I understand it is a small module of the presented method, but Quest/SparQ/(and even AdaKV)'s contribution should also be explicitly stated.\n\nBack to the evaluation note: I encourage authors to construct more complete tests on newer, more challenging benchmarks like SCBench, LongBench v2, and HELMET with more compression settings and baseline features. KV cache studies have been long limited under simple needle tasks and LongBench, which is no longer a proper standard for modern KV cache evaluation. While that might reveal gaps between full precision and the compressed model, making many authors reluctant to feature such evaluations, I firmly believe such gaps would guide future development and should be highlighted more than perfect results. **SnapKV should also be fully featured, by aligning with SCBench's report setting it should be doable.**\n\nFurthermore, I wonder why DuoAttention underperforms. It relies on the mechanistic proper of attention heads, and it seems unintuitive to perform badly under a multi-round scenario if the setting is right. Some investigations in this regard would be appreciated.\n\nLast, regarding the LongBench truncation, apologize for reading A.2. wrong, the current setting sounds proper. The needle test is also proper if it is following [2] (and as a kindly reminder, this paper is wrongly cited).\n\nI now improved the rating to 4 as promised.", "questions_for_authors": "-", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PzErqzXCPq", "reviewer": "Reviewer_eef9", "summary": "This paper introduces RocketKV, a two-stage KV cache compression approach. The first stage applies permanent KV eviction through adaptive pooling and GQA-compatible SnapKV methodology, while the second stage efficiently retrieves necessary KV components dynamically based on queries via a hybrid attention mechanism. This approach effectively retrieves essential key-value pairs during each decoding step, maintaining high accuracy. The proposed method demonstrates superior accuracy compared to other KV cache compression techniques (SnapKV, Quest, SparQ, etc.) while also achieving memory savings and end-to-end acceleration.\n\n## Update After Rebuttal\n\nAfter reading other reviewers' opinions and the authors' rebuttal, I have gained a better understanding of RocketKV's contribution and methods. However, I still have remaining concerns regarding the weaknesses I initially pointed out. First, regarding presentation, while the authors' additional explanations resolved many of my questions, considering the state of the initial submission, I believe this paper still requires significant improvements in structure and writing.\n\nSecond, regarding contribution, as reviewer eFmj mentioned, I am not fully convinced about the novelty of RocketKV compared to SnapKV and QUEST. This concern is related to the aforementioned presentation issues. For example, looking back at the paper, it is difficult to consider the heuristic search for adaptive pooling size (a differentiating point from SnapKV++) as a fundamental contribution. Additionally, it is challenging to understand why GQA compatibility improves performance in Figure 6 (it is difficult to find even minimal insight on this). Furthermore, rather than briefly mentioning the concept of QUEST in section 3.1, the paper should provide more detailed explanations and insights about the foundational methodology. Such comprehensive understanding would help readers better appreciate the distinguishing points of RocketKV compared to QUEST.\n\nOverall, while I now better understand the core of RocketKV—ensuring accuracy and achieving better efficiency through two-stage KV cache control, as the authors claim—I find it difficult to consider this paper suitable for ICML publication in its current state, even taking into account the additional rebuttal results. This is mainly due to insufficient presentation and structure. Considering all of my points, I will update my score from 1 to 2.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comprehensive review and providing constructive feedback.\n\n**Lack of Methodology Explanation:**\nDue to Space constraints, we decided to prioritize RocketKV’s performance results, resulting in a briefer method explanation. In the final version, we will move certain ablation studies to the appendix to provide a more comprehensive, self-contained description of our approach as you suggested.\n\n**Why RocketKV Is Necessary:**\nFigure 1 illustrates that existing KV cache compression methods struggle to match the accuracy of Exact-TopK on Qasper under low token budgets. \nFor further motivation, we analyzed a random attention head (layer 31 head 0 from Llama 3.1-8B) and calculated the maximum number of unique KV tokens used by Exact-TopK attention across all decoding steps in the Qasper benchmark. As shown in the table below, to keep all important KV tokens (used by at least one TopK attention across all decode steps), a permanent KV cache eviction method needs to keep 2197 out of 6110 tokens (where 6110 is the total seq. length) when K=256. Moreover, a dynamic token selection algorithm needs to accurately select a small set of TopK tokens out of a large set (e.g. 256 out of 6110 when K=256), which is a challenging task. An ideal solution is to keep only important KV tokens and then conduct dynamic token selection among them. Thus, we propose a two-stage approach that first retains only important tokens and then performs dynamic selection on this smaller set. This fusion evicts unimportant tokens and makes the dynamic selection more accurate, motivating our RocketKV design.\n\n| TopK value                    | 256  | 512  | 1024 | 2048 | 4096  |\n|------------------------------|------|------|------|------|-------|\n| Max num of unique TopK tokens  | 2197 | 3223 | 5229 | 8272 | 11993 |\n| Total num of tokens            | 6110 | 6110 | 17789| 17789| 17789 |\n\n**Explanation of Calculating Approximate Attention Scores:**\nRocketKV’s hybrid attention focuses on accurately pointing KV tokens with the TopK attention scores. We split the K tensor into pages of consecutive tokens. Within each page, we record element-wise min *(Kmin)* and max *(Kmax)* to estimate the upper bound of attention scores for a query *q*. Specifically, we compute *max⁡(q x Kmin⁡, q x Kmax⁡)* for each page to approximate highest possible attention scores within a page, then select pages with higher scores. To further reduce the approximation overhead, we only calculate on partial positions along the head dimensions where the magnitude of *q* is large and ignore other positions, and only fetching either from *Kmin* or *Kmax* at a given position based on the corresponding sign of q as shown in Figure 3. Hence, we approximate attention scores via both sequence and head dimension reductions. To be fully compatible with GQA, we select based on the sum of *q* or *|q|* at group dimension as needed to guarantee that all attention heads within a group are making the same selection at each step.\n\n**Measurement of Exact-TopK:**\nAs explained in Appendix A-3, Exact-TopK is an oracle-based method that assumes prior knowledge of token importance and dynamically chooses TopK KV tokens for each attention head and decoding step. By contrast, SnapKV permanently prunes tokens of the input prompt via one-time TopK filtering.\n\n**Are Different KV Tokens Retrieved at Each Decoding Step?**\nYes. Our second-stage hybrid attention dynamically selects KV tokens at each decoding step, aiming to pick those most relevant to the current query vector.\n\n**Is Page Always Divided Into 4?**\nNo. Page size depends on the overall compression ratio. For a token budget of 512 and a total length of 128K, the compression ratio is 256x. We then split this ratio evenly between the first and second stages (16x each), and again between sequence and head dimensions (4x each) in hybrid attention as mentioned in Section 3.5. Consequently, the page size is four tokens in that scenario, but it can vary for other compression ratios.\n\n**Novelty of SnapKV++:**\nThank you for pointing out that a similar GQA-compatible enhancement for SnapKV has already been proposed in AdaKV. This feature was added to GitHub in Nov.'24 and the arXiv paper in Jan.'25. We proposed GQA enhancement for SnapKV independently and concurrently. We are happy to give AdaKV credit for this concurrent contribution in the final version. However, per reviewer instructions outlined in ICML website, authors should not be held responsible for papers that were made public within 4 months of the submission deadline. We are discussing with AC to see whether we should give up GQA enhancement as a claimed contribution. Importantly, this does not diminish our core novelty: a two-stage KV cache compression framework uniquely combining permanent KV eviction with dynamic token selection, which is broadly generalizable to various compression methods at each stage. \n\nRegarding adaptive pooling size, please see our response to reviewer eFmJ.", "claims_and_evidence": "The paper lacks sufficient explanation of the proposed methodology, making it difficult to thoroughly examine the evidence supporting its claims. For instance, the Stage 2 Hybrid Attention appears to be written with the assumption that readers are already familiar with Quest and SparQ, as detailed explanations are notably absent. It's challenging to understand from the main text alone why a sign function is applied to the sum of q, or how approximate attention scores are calculated by retrieving from K_max and K_min based on these values.\n\nRegarding the proposed SnapKV++, the GQA-compatible aggregation method has already been widely used in recent approaches (e.g., AdaKV - https://github.com/FFY0/AdaKV). Similarly, the adaptive pooling method relies on rule-based, heuristic pooling sizes determined by sequence length without clear justification. Despite potentially contributing to higher accuracy, these two contributions are difficult to recognize as novel and central to the paper's contribution.", "methods_and_evaluation_criteria": "The paper presents performance comparisons across various long context benchmarks at different compression ratios.", "theoretical_claims": "There appear to be no theoretical claims in the paper. The adaptive pooling size is presented heuristically based on empirical experimental results, offering both size candidates and determination methods.", "experimental_designs_or_analyses": "The experimental design appears sound, though the paper lacks detailed analysis.", "relation_to_broader_scientific_literature": "The proposed KV cache compression methodology seems closely related to H2O and SnapKV, while the dynamic query-based KV retrieval approach appears closely connected to Quest.", "essential_references_not_discussed": "No issue", "other_strengths_and_weaknesses": "A strength of the paper is that the proposed method demonstrates high performance across various compression ratios. However, there are significant weaknesses in presentation that cannot be overlooked. Without detailed background on closely related works like Quest and SparQ, readers would struggle to properly understand Stages 1 and 2. Furthermore, there appears to be no analysis in the main text explaining why this method is necessary.\n\nWhile the paper makes valid contributions regarding performance improvements, considering this is a submission to a top-tier conference, insights and detailed explanations supporting the proposed methodology are essential for readers. Despite the performance results, the current state of the paper makes it difficult to recommend for acceptance. Below are aspects that should be addressed in more detail:\n\n- How exactly is Exact-TopK measured? Does it involve computing Full KV during prefill and retaining only the highest attention scores? Then what would be the difference between SnapKV and Exact-TopK?\n- Does Figure 2's illustration of the second stage indicate that different KV tokens are retrieved at each decoding step?\n- The process of calculating approximate attention scores using element-wise Max, Min, and the sign result of sum of q is difficult to understand from this paper alone.\n- Is the Page always divided into 4 in the dimension? Is there a specific rationale for this approach?", "other_comments_or_suggestions": "Experimental results (especially ablation studies) seem to occupy excessive space. This could be adjusted using tables or other methods to allow for more comprehensive background explanations. If my overall perspective is somewhat misguided or differs from other reviewers' opinions, I am open to reconsidering my assessment.", "questions_for_authors": "See weakness", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pTvW7jK992", "reviewer": "Reviewer_6pQ8", "summary": "This paper combines the advantages of permanent KV token eviction and dynamic KV token selection. It uses a two-staged kv cache compression method to give strong results and shows that it reduces GPU memory usage.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for finding our work interesting and providing valuable suggestions.\n\n**Add More Benchmarks to Prove the Losslessness:**\nWe would like to clarify that RocketKV is not a lossless approach, as evident from the provided accuracy results. And none of the other methods we compared are 100% lossless including the Exact-TopK method because they all replace dense attention operations with sparse attention. The primary goal is to achieve comparable (but may not exactly the same) accuracy to Full KV cache with much smaller KV token budgets. We have evaluated RocketKV across three models and three datasets under various token budgets to demonstrate the superiority of our method against existing work on achieving minimal accuracy loss with an extremely high KV cache compression ratio (up to ~500x in our evaluation). Compared to other existing works, we believe we have provided a sufficiently comprehensive evaluation. Moreover, we have added some additional evaluation on the recent SCBench in multi-turn scenarios and demonstrated that RocketKV outperforms other methods by a significant margin. Please refer to our response to reviewer utBU for more details.\n\n**Other Architectural Models Like Deepseek MLA or the Reasoning Models:**\nThank you for your suggestion. We are also interested to see how RocketKV and other KV cache compression methods perform on top of DeepSeek MLA models. Note that none of the other works on KV cache compression have done it before. Since RocketKV is fully compatible with GQA, we believe it can be applied to MLA with simple modifications (MLA can be considered as a variant of GQA where all attention heads in a layer share the same KV cache tensors). Considering RockeKV outperforms current methods over various models and datasets, we expect to see similar trends when we evaluate various methods on MLA. Unfortunately, due to time and resource constraints during the rebuttal period, we could not add either DeepSeek V3 or R1 model into our evaluation and have decided to leave it for future work.  \n\n**The Real Inference Cost:**\nWe have provided end-to-end inference results of RocketKV on an NVIDIA H100 GPU showing significant latency speed-up and memory saving against Full KV cache that already includes the time overhead of RocketKV (Section 4.3).", "claims_and_evidence": "I mostly agree with it. However, given that \"permanent KV token eviction\" could lose important information in the tokens, I do not quite understand why combining two kinds of kv cache compression could be lossless.\n\nIn this sense, I think two things can be done.\n\n1. add more benchmarks to prove the losslessness.\n2. use other architectural models? like deepseek MLA models? this may be too expensive. But I really ponder about the performance of RocketKV on the reasoning models, given they have extremely long reasoning tokens.", "methods_and_evaluation_criteria": "As above, I think we should add more model types and widely used benchmarks.", "theoretical_claims": "No proof.", "experimental_designs_or_analyses": "See above coments, thanks.", "relation_to_broader_scientific_literature": "Lossless KV cache compression is what a lot of people need. Have the authors tried to implement it on widely used inference engines, say sglang, vllm?", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "see above, I am pondering about the real inference cost, like the time overhead introduced by RocketKV, and whether the overhead can be resolved. If it works quickly and losslessly, it would be perfect.", "other_comments_or_suggestions": "Nice work! Hope to see further improvements.", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "No"}}, {"review_id": "9ayd4Wuuzr", "reviewer": "Reviewer_eFmJ", "summary": "RocketKV is a training-free KV cache compression strategy designed to optimize the inference efficiency of long-context LLMs during the decode phase. The main challenge it addresses is the exponential memory overhead due to KV cache storage, which scales with sequence length. The method is empirically validated on Mistral-7B, LLaMA-3.1-8B, and LongChat-7B, showing up to 3× speedup and 31% peak memory reduction on NVIDIA H100 GPUs, while preserving accuracy across long-context benchmarks (LongBench, Needle-in-a-Haystack, RULER). RocketKV outperforms SnapKV, Quest, SparQ, and DuoAttention, achieving near full-KV accuracy at significantly lower memory footprints.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and finding RocketKV results promising.  \n\n\n**Novelty: RocketKV is SnapKV + QUEST:**\nAs discussed in the paper, existing methods for KV cache compression typically fall into two categories: permanent KV token eviction and dynamic KV token selection. We would like to clarify that the primary novelty of our work is that we introduce a two-stage approach that effectively combines the strengths of both paradigms into a single framework. While we propose SnapKV++ for the first stage and hybrid attention for the second stage, they can be directly replaced with various other methods of the same category. \n\nIn terms of SnapKV++, we admit that our improvement over SnapKV is not significant which is reflected in the name (this is also recognized by Reviewer utBU). However, it is still much more effective than the original SnapKV method as demonstrated in our ablation study (Section 4.4 and Appendix B.1).\n\nWe would like to point out that our hybrid attention method is NOT the same as QUEST. QUEST relies on K tensor reduction along the sequence dimension to conduct approximate attention, while our hybrid attention method relies on K tensor reduction along both sequence and head dimensions. This two-dimensional reduction scheme can achieve much higher accuracy than one-dimensional reduction at a given compression ratio. For example, with a compression ratio of 16, hybrid attention can evenly split it into a compression ratio of 4 at each dimension, introducing much lower accuracy loss compared to directly compressing the sequence dimension by 16x in QUEST. This is further confirmed in the ablation study (Section 4.4.3 and Appendix B.1) where we demonstrate the standalone hybrid attention method consistently outperforms QUEST and SparQ. \n\nOverall, we believe our work introduces sufficient novelty and is qualified as a top-tier conference publication.\n\n**Potential Sensitivity to Kernel Size in SnapKV++:**\nWe thoroughly examined the impact of pooling kernel size in the ablation study (Section 4.4.2 and Appendix B.1). While the adaptive pooling size is indeed empirically determined based on accuracy results on the RULER benchmark as shown in Figure 7, we found this simple method is quite effective and generalizes well on other benchmarks as shown in Figure 10 in Appendix B.1. The insight we observe from this study is that tasks with longer sequence lengths usually perform better with larger kernel sizes, which can provide practical guidelines for better pooling kernel size selection in future work.\n\n**Test on Cheaper GPUs:**\nBelow are the end-to-end speedup numbers of RocketKV over Full KV cache running Llama3.1-8B on A100 with 256 token budget.\n| Sequence length | 16K  | 32K  | 64K  | 96K  |\n|-----------------|------|------|------|------|\n| Speedup         | 1.3x | 1.5x | 2.7x | 3.6x |\n\nCompared to efficiency data on H100 shown in Figure 5(a), we can see that the maximum speedup on A100 is 20% higher (3.6x versus 3x). This is because A100 has a lower memory bandwidth to compute ratio compared to H100. As a result, LLM inference execution is more memory-bound on A100 and can benefit more from memory traffic savings of KV cache offered by RocketKV. We believe the speedup of RocketKV will be even higher on cheaper GPUs such as RTX 4090/5090 since they are not equipped with High Bandwidth Memory (HBM). Unfortunately, these GPUs also have much smaller memory capacity which prevents us from conducting long-context experiments on them. Notice that the memory savings are the same between H100 and A100 so we didn’t show them here.\n\n\n**Lack of Direct Latency Comparison:**\nSince different sparse attention methods are usually implemented under different frameworks with different levels of code optimizations, it is difficult to provide an apple-to-apple comparison between them. In this work we use the token budget to estimate memory traffic, including attention score approximation. For example, with RocketKV at a token budget of 256, half is used for attention approximation (Steps 1 and 2 in Section 3.4) and half for sparse attention (Step 3). Thus, the token budget itself can directly reflect attention latency in highly-optimized implementations since attention operations are mostly memory-bound.\n\n**Is Needle Among Evicted Tokens:**\nRocketKV evicts KV cache tokens independently across attention heads and layers, making it unclear if the needle's tokens are consistently retained. Additionally, later layers can spread the needle's information to other positions through attention, reducing the effectiveness of using needle positions to track information loss.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No theory.", "experimental_designs_or_analyses": "The experiments are sound but could benefit from:\n- Latency breakdown: token retrieval and sparse attention computation.\n- Ablation studies on kernel size impact.", "relation_to_broader_scientific_literature": "- Advancing Sparse Attention Methods for LLMs", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths\n- Hardware-Friendly Optimization.\n- Significant Performance Gains in Speed and Memory Efficiency.\n\nWeaknesses:\n- Novelty is not enough. The two main stages come from SnapKV + QUEST.\n- Potential Sensitivity to Kernel Size Selection in SnapKV++. The adaptive pooling kernel size in SnapKV++ is empirically determined, raising concerns about its generalizability to unseen domains.\n- Evaluation on Only One Hardware Setup (NVIDIA H100). Test on cheaper GPUs such as A100.\n- Lack of Direct Latency Comparison with Alternative Sparse Attention Methods.", "other_comments_or_suggestions": "Regarding the needle-in-the-haystack benchmark, since RocketKV evicts previous tokens, it is unclear whether it might also evict the needle—despite still producing the correct output. Could the authors provide an analysis of which tokens RocketKV evicts and whether the needle is among them?", "questions_for_authors": "See Strengths And Weaknesses", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: RocketKV introduces a practical and effective two-stage KV cache compression framework for long-context LLM inference that yields strong efficiency gains while maintaining accuracy. The paper presents a comprehensive evaluation across multiple models and tasks, and features thoughtful ablation studies that support the design choices. Despite limited novelty, the authors demonstrate clear improvements through a synergistic integration and introduce unique elements like hybrid attention over both sequence and head dimensions. While some reviewers noted concerns about presentation clarity and comparisons with combined baselines, the rebuttal offers detailed clarifications and new results, partially mitigating these concerns. Given the practical impact, clarity of exposition, and responsiveness during the discussion, I recommend acceptance.
Decision: Accept (poster)