通过两级压缩提升长文本生成速度,内存占用降低三成以上。
RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
- 先粗粒度剔除冗余键值,再用稀疏注意力精细优化。
- 最高压缩400倍,推理加速3.7倍,内存峰值减少32.6%。
- 无需训练,适合长对话和大上下文场景的部署应用。
基于Transformer的大语言模型在解码阶段依赖键值(KV)缓存处理长上下文,但其大小随输入长度线性增长,严重消耗内存带宽与容量。为此,我们提出RocketKV,一种无需训练的两阶段KV缓存压缩策略:第一阶段对输入序列进行粗粒度永久性键值剔除;第二阶段采用混合稀疏注意力机制,通过同时降低头维度和序列维度来近似注意力得分。实验表明,相较于完整KV缓存基线,在NVIDIA A100 GPU上,RocketKV实现最高400倍压缩比、3.7倍端到端加速及32.6%的峰值内存降低,且在多种长上下文任务中保持接近零误差。我们还提出了适用于多轮对话的变体,性能持续优于现有方法,精度接近理想top-k注意力方案。代码已开源:https://github.com/NVlabs/RocketKV。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。