RDKV: Rate-Distortion Bit Allocation for Joint Eviction and Quantization of the KV Cache

TL;DR

This paper introduces RDKV, a joint rate-distortion approach for KV cache compression in large language models, optimizing eviction and quantization simultaneously to improve inference efficiency.

Contribution

It formulates KV cache compression as a rate-distortion problem and proposes RDKV, the first method to jointly optimize eviction and quantization for LLM inference.

Findings

RDKV outperforms baselines by 9.1% on average.

Recovers 97.81% of full-cache accuracy with only 2.48% cache retention.

Achieves 4.5x speedup and 1.9x memory reduction at 128K context length.

Abstract

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.