Research map / Architectures and efficiency / Attention and transformers
RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes RocketKV, a training-free two-stage KV cache compression method (SnapKV eviction followed by hybrid sparse attention over head and sequence dimensions) for faster LLM decoding, plus a multi-turn variant RocketKV-MT.
Problem
KV cache size grows with context length during decoding, straining memory bandwidth and capacity; existing permanent eviction loses accuracy at low budgets and dynamic top-k selection keeps full storage and only saves traffic.
Approach
Chain a coarse permanent eviction stage (SnapKV with large pooling kernel, GQA-group-level selection) with a fine-grain dynamic top-k stage (hybrid sparse attention that approximates scores via reductions along both head and sequence dimensions), with an adaptive formula splitting the compression ratio between stages; RocketKV-MT retains evicted tokens for later turns.
Key results
- RocketKV reaches 100% NIAH accuracy on Llama3.1-8B-Instruct at a 256 token budget, a compression ratio above 400x at up to 109K tokens.
- RocketKV achieves up to 3.7x end-to-end decode speedup on A100 and 3.3x on H100 with up to 32.6% peak memory saving.
- On LongBench with Llama3.1-8B-Instruct, RocketKV loses only 1.1% average accuracy at a 256 token budget and almost none at 512 and above.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| combine first-stage SnapKV eviction with second-stage HSA instead of using HSA alone (ablation) | RocketKV on LongBench · LongBench (Llama3.1-8B-Ins, token budget 256) | average accuracy: 49.5 -> 51.1 | Helped |
| use HSA's reductions in both sequence and head dimensions rather than single-dimension Quest or SparQ approximations (ablation) | Llama3.1-8B-Ins · LongBench and other evaluated benchmarks | accuracy | Helped |
| use an adaptive compression split factor instead of a static split factor (ablation) | RocketKV · Long-context benchmarks across sequence lengths and token budgets | accuracy | No clear effect |
| retain the full KV history across turns instead of permanently evicting unselected tokens (ablation) | RocketKV in multi-turn SCBench · SCBench multi-turn (Llama3.1-8B-Ins, token budget 1024) | average accuracy: 26.6 -> 44.3 | Helped |
| use RocketKV-MT rather than RocketKV (ablation) | Llama3.1-8B-Ins · SCBench multi-turn | average accuracy | Mixed |
| use RocketKV rather than the Full-KV baseline during decode | Llama3.1-8B-Ins on NVIDIA A100 · Decode phase | end-to-end speedup: up to 3.7× | Helped |
| use RocketKV rather than the Full-KV baseline during decode | Llama3.1-8B-Ins · Decode phase | peak memory saving: up to 32.6% | Helped |
Surprising findings
- On SCBench, RocketKV trails SparQ at token budgets of 8192 and above, and RocketKV-MT is needed to approach Exact-TopK in multi-turn settings.
Limitations
- Permanent eviction in the first stage hurts multi-turn accuracy, leaving a gap to Exact-TopK on SCBench for RocketKV.
- Efficiency results rely on a Python-based gpt-fast implementation without custom CUDA kernels.
- RocketKV-MT gives no KV storage savings, only memory traffic savings.
Experimental setup
Datasets: LongBench Needle-in-a-Haystack RULER SCBench
Benchmarks: LongBench Needle-in-a-Haystack RULER SCBench
Metrics: accuracy end-to-end speedup peak memory reduction compression ratio
Baselines: Exact-TopK DuoAttention SnapKV Quest SparQ
Builds on: SnapKV Quest SparQ Ada-KV
Closest papers by meaning
- Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (ICLR 2024 · similarity 0.94)
- ChunkKV: Semantic-Preserving KV Cache Compression for Efficient Long-Context LLM Inference (NeurIPS 2025 · similarity 0.93)
- ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference (ICML 2025 · similarity 0.93)
- CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.93)
- Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.92)
- SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget (ICLR 2025 · similarity 0.92)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.