RDKV：面向联合驱逐与量化的 KV 缓存速率失真 bit 分配
机器学习
2026-05-12 v1 人工智能
摘要
大型语言模型（LLMs）在各类任务上表现突出，但因内存大小和带宽限制，在长输入上下文上的推理成为瓶颈。键值（KV）缓存的大小随序列长度线性增长，需在每一步解码时将其从离芯片高带宽内存（HBM）重新读取到片上内存，导致内存受限的推理。现有方法通过驱逐或量化来压缩缓存，但通常二者处理独立。本文将 KV 缓存压缩建模为速率-失真问题，在此框架下，驱逐与量化成为相同的 bit 分配方案的两个端点。这暴露了需联合优化它们的需求，催生了我们的方法——RDKV（Rate-Distortion KV cache compression）。RDKV 根据压缩对注意力计算造成的失真，推导每个 token 或通道的权重。基于这些权重，它为每个 token 或通道分配 bit 宽度，从全精度到零位，依据逆水塔填充法则在预填充阶段后仅执行一次。实验在 LongBench、RULER 和 InfiniteBench 上表明，RDKV 在平均性能上超过最佳基线方法 9.1%。在 LongBench 上，仅保留 2.48% 的缓存，即可恢复 97.81% 的全缓存准确率。相较于全缓存 FlashAttention-2 解码，在 128K 上下文长度下，实现 4.5 倍解码加速和 1.9 倍峰值内存降低，同时保持可比性能。
引用
@article{arxiv.2605.08317,
title = {RDKV: Rate-Distortion Bit Allocation for Joint Eviction and Quantization of the KV Cache},
author = {Junkai Zhang and Hang Guo and Luca Benini and Yawei Li},
journal= {arXiv preprint arXiv:2605.08317},
year = {2026}
}