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

1. Problem Statement

Large Language Models (LLMs) face a critical bottleneck during long-context inference: memory bandwidth and capacity. The Key-Value (KV) cache size grows linearly with sequence length, requiring the re-reading of tens of gigabytes of data from off-chip High-Bandwidth Memory (HBM) to on-chip memory at every decoding step. This makes inference memory-bound rather than compute-bound.

Existing compression strategies generally fall into two isolated categories:

1. Eviction: Scoring tokens or channels and discarding the lowest-ranked ones entirely (binary keep-or-drop). This leaves no precision tier for tokens of moderate importance.
2. Quantization: Retaining all tokens but reducing their bit-width. This retains even negligible tokens, wasting memory.

While recent hybrids attempt to combine these, they often rely on heuristic routing scores or pre-allocated ratios determined via offline search, failing to treat eviction and quantization as a unified optimization problem.

2. Methodology: RDKV

The paper proposes RDKV (Rate-Distortion KV cache compression), a framework that formulates KV cache compression as a rate-distortion problem. In this view, eviction and quantization are not separate actions but two endpoints of a single bit-allocation curve.

2.1 Theoretical Formulation

RDKV derives a distortion weight for each cache unit (token in the V cache, channel in the K cache) based on the perturbation induced by its removal or quantization on the attention computation.

• V Cache (Tokens): The distortion weight $w_t$ for a token is the total variation distance of the attention distribution after eviction. It is calculated as the sum of attention scores received by the token across all query positions: $w_t = \sum_{\tau} a_{\tau,t}$.
• K Cache (Channels): The distortion weight $w_c$ for a channel is the spectral norm of the logit deviation caused by zeroing that channel. It is calculated as the product of the norms of the corresponding query and key columns: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

These weights serve as cost coefficients in a rate-distortion objective function:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{s.t.} \quad \sum_{u} b_u \leq B$$
where $b_u$ is the bit-width assigned to unit $u$, and $B$ is the total bit budget.

2.2 Bit Allocation Pipeline

The allocation is solved via reverse water-filling (Lagrangian relaxation) adapted for discrete hardware-supported bit-widths $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. Stage 1 (Weighting): Compute token weights ($w_t$) and channel weights ($w_c$) once after the prefill stage using the full uncompressed cache.
2. Stage 2 (Token Allocation): Solve a Multiple-Choice Knapsack Problem (MCKP) for the V cache. Tokens are assigned bit-widths from $\{0, 2, 4, 8, 16\}$. A bit-width of 0 corresponds to eviction.
3. Stage 3 (Channel Allocation): Solve the MCKP for the K cache, but only for the tokens retained in Stage 2. Channels are assigned bit-widths independently of the token's V-bit-width, allowing a single token to have mixed-precision K and V components.

2.3 Efficient Packed Decode (TriZone)

To realize actual memory and latency savings, RDKV introduces a TriZone cache layout that fuses dequantization into the attention kernel, avoiding the need to unpack data to FP16 before processing.

• Zone A: Packed old-cache stream containing retained K rows and V rows with bit-widths $\{2, 4, 8\}$. These are grouped by bit-width to allow uniform dequantization paths.
• Zone B: FP16 V rows for tokens assigned 16-bit precision.
• Zone C: FP16 K and V rows for newly generated tokens.

The attention kernel computes logits and outputs by combining these three zones, with dequantization fused directly into the matrix multiplication, eliminating extra memory bandwidth overhead.

3. Key Contributions

1. Rate-Distortion Formulation: The paper casts KV cache compression as a unified rate-distortion problem, proving that the optimal solution inherently mixes zero-rate (eviction) and finite-rate (quantization) units. This eliminates the need for heuristic ratios or offline searches.
2. Discrete Allocation via Reverse Water-Filling: The authors instantiate the allocation as a discrete knapsack problem solvable via Lagrangian relaxation, assigning optimal bit-widths from $\{0, 2, 4, 8, 16\}$ based on distortion weights derived from attention dynamics.
3. TriZone Packed Layout: A hardware-aware implementation that fuses dequantization into the attention kernel, converting the theoretical bit-width savings into actual reductions in peak memory and decode latency.

4. Experimental Results

Experiments were conducted on LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B, and Qwen2.5-72B across LongBench, RULER, and InfiniteBench.

• Accuracy: RDKV outperforms the best evaluated baseline (Ada-SnapKV) by 9.1% on average. On LongBench, it recovers 97.81% of full-cache accuracy with only 2.48% cache retention.
• Long-Context Performance:
  • RULER: At 128K context, RDKV achieves 66.95% accuracy compared to 58.43% for the strongest baseline (AdaKV), a gap of 8.5 points.
  • InfiniteBench: RDKV ranks first among compression methods with an average score of 39.46, surpassing AdaKV by 1.40 points.
  • Needle-in-a-Haystack: RDKV maintains near-uniform retrieval accuracy (avg. 0.99) across depths, whereas eviction-only baselines fail significantly at intermediate depths.
• Efficiency:
  • Latency: With a 128K context, RDKV achieves a 4.5× decode speedup compared to full-cache FlashAttention-2 decoding.
  • Memory: It achieves a 1.9× peak memory reduction at 128K context and successfully runs 256K context on a single A100 64GB GPU, where FullKV runs out of memory (OOM).

5. Significance and Claims

The paper claims that treating eviction and quantization jointly within a single optimization framework is superior to staged or ratio-based approaches. By deriving allocation scores directly from the distortion induced on the attention computation, RDKV avoids the "binary trap" of eviction-only methods (which discard moderate tokens) and the "waste" of quantization-only methods (which retain negligible tokens).

The authors emphasize that their method:

• Requires no model retraining.
• Is architecture-agnostic, performing consistently across different model sizes and families.
• Provides a Pareto-optimal trade-off between latency and accuracy, strictly outperforming baselines in both metrics at long context lengths.
• Demonstrates that the "continuous optimum" of rate-distortion theory, when discretized for hardware, naturally yields a mixed-bit strategy that outperforms current state-of-the-art hybrids.

The paper concludes that this joint approach improves accuracy without sacrificing efficiency, enabling long-context inference on resource-constrained hardware.