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

• URL: http://arxiv.org/abs/2502.14051v2
• Date: Mon, 30 Jun 2025 19:01:18 GMT
• Title: 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,
• Abstract summary: We present RocketKV, a training-free KV cache compression strategy containing two consecutive stages.<n>In the first stage, it performs coarse-grain permanent KV cache eviction on the input sequence tokens.<n>In the second stage, it adopts a hybrid sparse attention method to conduct fine-grain top-k sparse attention.
• Score: 25.190765258589707
• License: http://creativecommons.org/licenses/by/4.0/
• 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$\times$, end-to-end speedup of up to 3.7$\times$ 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.

Related papers

• DBudgetKV: Dynamic Budget in KV Cache Compression for Ensuring Optimal Performance [125.81664663201282]
  We introduce a new KV cache compression method dubbed DBudgetKV. It features an attention-based metric to signal when the remaining KV cache is unlikely to match the full-cache performance, then halting the pruning process. Our method is easy to integrate within LLM inference, not only optimizing memory space, but also showing reduced inference time compared to existing methods.
  arXiv  Detail & Related papers  (2025-02-24T06:33:39Z)
• QuantSpec: Self-Speculative Decoding with Hierarchical Quantized KV Cache [67.84112700032007]
  Large Language Models (LLMs) are increasingly being deployed on edge devices for long-context settings. In these scenarios, the Key-Value ( KV) cache is the primary bottleneck in terms of both GPU memory and latency. We propose a novel self-speculative decoding framework, QuantSpec, where the draft model shares the architecture of the target model but employs a hierarchical 4-bit quantized KV cache and 4-bit quantized weights for acceleration.
  arXiv  Detail & Related papers  (2025-02-05T20:43:48Z)
• FastKV: KV Cache Compression for Fast Long-Context Processing with Token-Selective Propagation [4.856070170902535]
  Large language models (LLMs) excel at handling long-context sequences.<n>They require substantial key-value ( KV) caches to store contextual information.<n>FastKV is a KV cache compression method designed to enhance latency for long-context sequences.
  arXiv  Detail & Related papers  (2025-02-03T05:25:09Z)
• ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression [10.003118268356017]
  Long context poses significant challenges for inference efficiency.<n>We introduce ClusterKV, which recalls tokens at the granularity of semantic clusters.<n>Experiment results show that ClusterKV attains negligible accuracy loss across various tasks with 32k context lengths.
  arXiv  Detail & Related papers  (2024-12-04T10:58:27Z)
• VL-Cache: Sparsity and Modality-Aware KV Cache Compression for Vision-Language Model Inference Acceleration [7.463830743649754]
  Vision-Language Models (VLMs) have demonstrated impressive performance across a versatile set of tasks. Key-Value (KV) cache encodes long visual contexts, such as images or videos. Existing KV cache compression methods are effective for Large Language Models (LLMs) We propose a novel KV cache compression recipe tailored for accelerating VLM inference.
  arXiv  Detail & Related papers  (2024-10-29T20:04:34Z)
• KVSharer: Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing [58.29726147780976]
  We propose a plug-and-play method called textit KVSharer, which shares the KV cache between layers to achieve layer-wise compression. Experiments show that textit KVSharer can reduce KV cache computation by 30%, thereby lowering memory consumption. We verify that textit KVSharer is compatible with existing intra-layer KV cache compression methods, and combining both can further save memory.
  arXiv  Detail & Related papers  (2024-10-24T08:06:41Z)
• Lossless KV Cache Compression to 2% [22.98828332096935]
  This work introduces a novel architecture, Cross-Layer Latent Attention (CLLA), aimed at compressing the KV cache to less than 2% of its original size. CLLA integrates attention head/dimension reduction, layer sharing, and quantization techniques, into a cohesive framework.
  arXiv  Detail & Related papers  (2024-10-20T02:17:35Z)
• LoRC: Low-Rank Compression for LLMs KV Cache with a Progressive Compression Strategy [59.1298692559785]
  Key-Value ( KV) cache is crucial component in serving transformer-based autoregressive large language models (LLMs) Existing approaches to mitigate this issue include: (1) efficient attention variants integrated in upcycling stages; (2) KV cache compression at test time; and (3) KV cache compression at test time. We propose a low-rank approximation of KV weight matrices, allowing plug-in integration with existing transformer-based LLMs without model retraining. Our method is designed to function without model tuning in upcycling stages or task-specific profiling in test stages.
  arXiv  Detail & Related papers  (2024-10-04T03:10:53Z)
• UNComp: Uncertainty-Aware Long-Context Compressor for Efficient Large Language Model Inference [38.11539884622708]
  UNComp is an uncertainty-aware compression scheme that adaptively compresses both the hidden states and the KV cache. Our method achieves a 1.6x speedup in the prefilling stage and reduces the KV cache to 4.74% of its original size. Remarkably, in needle-in-a-haystack tasks, UNComp outperforms the full-size KV cache even when compressed to 9.38% of its original size.
  arXiv  Detail & Related papers  (2024-10-04T02:32:36Z)
• CSKV: Training-Efficient Channel Shrinking for KV Cache in Long-Context Scenarios [13.144156413032896]
  We introduce CSKV, a training-efficient Channel Shrinking technique for KV cache compression. We show that CSKV can reduce the memory overhead of the KV cache by 80% while maintaining the model's long-context capability. Our method can be seamlessly combined with quantization to further reduce the memory overhead, achieving a compression ratio of up to 95%.
  arXiv  Detail & Related papers  (2024-09-16T17:36:50Z)
• ThinK: Thinner Key Cache by Query-Driven Pruning [63.13363917871414]
  Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. We propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels.
  arXiv  Detail & Related papers  (2024-07-30T17:59:08Z)
• LoCoCo: Dropping In Convolutions for Long Context Compression [77.26610232994508]
  This paper presents a novel approach, Dropping In Convolutions for Long Context Compression (LoCoCo) LoCoCo employs only a fixed-size Key-Value ( KV) cache, and can enhance efficiency in both inference and fine-tuning stages.
  arXiv  Detail & Related papers  (2024-06-08T01:35:11Z)
• PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling [53.08975547824068]
  We investigate whether attention-based information flow inside large language models (LLMs) is aggregated through noticeable patterns for long context processing. Our observations reveal that LLMs aggregate information through Pyramidal Information Funneling where attention is scattering widely in lower layers. Motivated by these insights, we developed Pyramid KV, a novel and effective KV cache compression method.
  arXiv  Detail & Related papers  (2024-06-04T07:51:30Z)
• Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs [82.08922896531618]
  We introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs) We conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens.
  arXiv  Detail & Related papers  (2023-10-03T05:17:08Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.