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

📄 arXiv: 2502.14051v3 📥 PDF

作者: Payman Behnam, Yaosheng Fu, Ritchie Zhao, Po-An Tsai, Zhiding Yu, Alexey Tumanov

分类: cs.CL, cs.LG

发布日期: 2025-02-19 (更新: 2025-08-13)

备注: ICML 2025

🔗 代码/项目: GITHUB

💡 一句话要点

RocketKV：通过两阶段KV缓存压缩加速长文本LLM推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: KV缓存压缩 长文本推理 稀疏注意力 大语言模型 模型加速

📋 核心要点

  1. 长文本LLM推理面临KV缓存随上下文长度线性增长的挑战，导致内存带宽和容量瓶颈。
  2. RocketKV通过粗粒度KV缓存淘汰和细粒度混合稀疏注意力两阶段压缩，有效降低KV缓存大小。
  3. 实验表明，RocketKV在长文本任务上实现了显著的压缩率、加速和内存减少，同时保持了精度。

📝 摘要（中文）

基于Transformer的大语言模型在解码阶段严重依赖KV缓存来高效处理长文本。然而，KV缓存的大小与输入长度成正比增长，增加了内存带宽和容量的负担。为了解决这个问题，我们提出了RocketKV，一种无需训练的KV缓存压缩策略，包含两个连续的阶段。第一阶段对输入序列token执行粗粒度的永久KV缓存淘汰。第二阶段采用混合稀疏注意力方法，通过利用头部和序列维度缩减进行细粒度的top-k稀疏注意力，从而近似注意力分数。实验表明，与完整KV缓存基线相比，RocketKV在NVIDIA A100 GPU上的解码阶段提供了高达400倍的压缩率、高达3.7倍的端到端加速以及高达32.6%的峰值内存减少，同时在各种长文本任务上实现了可忽略不计的精度损失。我们还提出了一种用于多轮场景的RocketKV变体，该变体始终优于其他现有方法，并实现了几乎与oracle top-k注意力方案相当的精度。源代码已开源。

🔬 方法详解

问题定义：长文本LLM推理过程中，KV缓存的大小会随着输入序列长度线性增长，这导致了巨大的内存占用和带宽需求，成为推理速度的瓶颈。现有的方法要么需要重新训练模型，要么在压缩率和精度之间做出妥协。

核心思路：RocketKV的核心思路是通过两阶段的KV缓存压缩，在不重新训练模型的前提下，大幅降低KV缓存的大小，从而加速推理过程。第一阶段通过粗粒度的淘汰不重要的token对应的KV值，第二阶段通过细粒度的稀疏注意力进一步压缩KV缓存。

技术框架：RocketKV包含两个主要阶段：1) 粗粒度KV缓存淘汰：根据某种策略（例如，token的重要性得分）永久性地移除一部分token对应的KV值。2) 细粒度混合稀疏注意力：使用一种混合稀疏注意力机制，通过头部维度和序列维度的缩减，近似计算注意力分数，从而减少计算量和内存占用。

关键创新：RocketKV的关键创新在于其两阶段的KV缓存压缩策略，以及混合稀疏注意力机制。两阶段压缩能够更有效地降低KV缓存的大小，而混合稀疏注意力机制能够在保证精度的前提下，进一步减少计算量和内存占用。与现有方法相比，RocketKV无需重新训练模型，并且能够在压缩率、精度和推理速度之间取得更好的平衡。

关键设计：在粗粒度KV缓存淘汰阶段，需要选择合适的token重要性评估策略，例如基于注意力权重或梯度的方法。在细粒度混合稀疏注意力阶段，需要选择合适的头部维度和序列维度缩减比例，以及稀疏模式。此外，还需要设计合适的损失函数来训练稀疏注意力机制，以保证精度。

📊 实验亮点

实验结果表明，RocketKV在NVIDIA A100 GPU上实现了高达400倍的KV缓存压缩率，3.7倍的端到端推理加速，以及32.6%的峰值内存减少。同时，在各种长文本任务上，RocketKV的精度损失可忽略不计。在多轮对话场景中，RocketKV变体优于其他现有方法，并接近oracle top-k注意力方案的精度。

🎯 应用场景

RocketKV可应用于各种需要处理长文本的LLM推理场景，例如长文档摘要、机器翻译、对话系统等。通过降低内存占用和提高推理速度，RocketKV可以使LLM在资源受限的设备上运行，并支持更大规模的应用。

📄 摘要（原文）

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. The source code is available here: https://github.com/NVlabs/RocketKV.