通过联合优化缓存淘汰与量化,实现长序列推理的高效内存压缩。
RDKV: Rate-Distortion Bit Allocation for Joint Eviction and Quantization of the KV Cache

- 将缓存压缩建模为率失真问题,统一处理淘汰与量化。
- 在长文本任务中仅保留2.48%缓存即恢复97.81%精度。
- 支持128K上下文,解码速度提升4.5倍,内存减少1.9倍。
大语言模型在长序列推理中受限于内存容量与带宽,其键值(KV)缓存大小随序列长度线性增长,需每步从片外高带宽内存读取至片上内存,导致内存瓶颈。现有方法通常单独处理缓存淘汰或量化,未能协同优化。本文将KV缓存压缩建模为率失真问题,使淘汰与量化成为同一比特分配方案的两端。RDKV(Rate-Distortion KV缓存压缩)基于注意力计算受压缩影响的失真度,为每个标记或通道分配从全精度到零比特的位宽,采用反向水填法,在预填充阶段一次性完成。在LongBench、RULER和InfiniteBench上的实验表明,RDKV平均优于最优基线9.1%。在LongBench上,仅用2.48%缓存即可恢复97.81%全缓存精度。相比FlashAttention-2全缓存解码,128K上下文下实现4.5倍解码加速与1.9倍峰值内存降低,性能相近。
原文摘要 · Abstract (English)
Large language models (LLMs) have shown strong performance across diverse tasks, but their inference with long input contexts is bottlenecked by memory size and bandwidth. The Key-Value (KV) cache size grows linearly with sequence length and needs to be re-read from off-chip high-bandwidth memory (HBM) to on-chip memory at every decoding step, resulting in memory-bound inference. Existing methods reduce the cache by either eviction or quantization, but typically treat the two in isolation. In this paper, we cast KV cache compression as a rate-distortion problem, under which eviction and quantization are two end-points of the same bit allocation scheme. This exposes the need to optimize them jointly, motivating our method, RDKV (Rate-Distortion KV cache compression). RDKV derives the weight of each token or channel from the distortion that compression induces on the attention computation. Based on these weights, it assigns each token or channel a bit-width ranging from full precision down to zero bits guided by reverse water-filling, applied once after the prefilling stage. Experiments on LongBench, RULER, and InfiniteBench show that RDKV outperforms the best evaluated baseline by 9.1% on average. On LongBench it recovers 97.81% of full-cache accuracy with only 2.48% cache retention. Compared with full-cache FlashAttention-2 decoding, it achieves 4.5x decode speedup and 1.9x peak memory reduction with 128K context length, while maintaining comparable performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。