RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression

We thank the reviewer for the valuable feedback and finding RocketKV results promising.

Novelty: RocketKV is SnapKV + QUEST: As 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.

In 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).

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

Overall, we believe our work introduces sufficient novelty and is qualified as a top-tier conference publication.

Potential Sensitivity to Kernel Size in SnapKV++: We 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.

Test on Cheaper GPUs: Below are the end-to-end speedup numbers of RocketKV over Full KV cache running Llama3.1-8B on A100 with 256 token budget.

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

Lack of Direct Latency Comparison: Since 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.

Is Needle Among Evicted Tokens: RocketKV 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.

We thank the reviewer for finding our work interesting and providing valuable suggestions.

Add More Benchmarks to Prove the Losslessness: We 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.

Other Architectural Models Like Deepseek MLA or the Reasoning Models: Thank 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.

The Real Inference Cost: We 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).

Thank you for your comprehensive review and providing constructive feedback.

Lack of Methodology Explanation: Due 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.

Why RocketKV Is Necessary: Figure 1 illustrates that existing KV cache compression methods struggle to match the accuracy of Exact-TopK on Qasper under low token budgets. For 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.

Explanation of Calculating Approximate Attention Scores: RocketKV’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.

Measurement of Exact-TopK: As 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.

Are Different KV Tokens Retrieved at Each Decoding Step? Yes. Our second-stage hybrid attention dynamically selects KV tokens at each decoding step, aiming to pick those most relevant to the current query vector.

Is Page Always Divided Into 4? No. 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.

Novelty of SnapKV++: Thank 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.

Regarding adaptive pooling size, please see our response to reviewer eFmJ.

We sincerely appreciate the reviewer’s insightful feedback and recognition of our work in several aspects.

Needle Dataset: Our 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.

SCBench Results: This 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.

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

[4] NVIDIA Dynamo, https://github.com/ai-dynamo/dynamo

[5] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving

LongBench Input Truncation: We 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.

Proper Quotation Marks: Thank you for pointing this out, we will correct those quotation marks in the final version of the paper.