PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Junkai Zhang, Hang Guo, Luca Benini, Yawei Li

arXiv:2605.08317v1cs.LGcs.AI

TL;DR

Long-context inference is bottlenecked by the growing KV cache and the memory traffic of rereading it during decoding, while existing compression often treats eviction and quantization separately. RDKV jointly allocates bits from attention-derived distortion weights using reverse water-filling. Across the evaluated benchmarks and models, it outperforms evaluated baselines in accuracy, decoding speed, and peak memory, and reaches a context length where Full KV does not.

  • Problem

    Long-context decoding is memory-bound because the KV cache grows with sequence length and is reread at every step, while existing methods often handle eviction and quantization separately.

    Source: §1 · p. 1

  • Method

    RDKV derives token and channel weights from attention-induced distortion and uses reverse water-filling to jointly assign bit-widths, including zero-bit eviction.

    Source: §1 · p. 2

  • Results

    Across LongBench, RULER, and InfiniteBench on five open-source LLMs, RDKV consistently outperforms evaluated baselines in accuracy, decoding speed, and peak memory, and reaches 256K on one A100 64 GB where Full KV does not.

    Source: §5 · p. 9

  • Takeaways & Limitations

    RDKV’s joint allocation and packed-decode implementation turns mixed-bit cache assignments into actual HBM savings while supporting long-context decoding.

    Source: §1 · p. 2, §5 · p. 9

    RDKV allocates once after prefill, so its cache assignment is frozen and does not adapt to attention-pattern shifts during generation.

    Source: K Limitations and Future Work · p. 28

Abstract

from arXiv · show

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.

1 Introduction

Long-context decoding is memory-bound because the KV cache grows with sequence length and must be reread for every generated token. RDKV treats eviction and quantization as one rate-distortion allocation problem, assigning bit-widths according to attention-induced distortion.

1 Introduction

  • The KV cache grows linearly with context length, and rereading it at every decoding step makes long-context inference memory-bound.The paper frames cache-footprint reduction without retraining while preserving generation quality as a central efficiency problem.Source: §1 · p. 1
  • Eviction and quantization leave a gap: binary eviction has no middle precision tier, while mixed-precision quantization retains even negligible tokens.Existing hybrids combine the actions, but their routing scores remain heuristic and their budget ratios are preset or searched before assignment.Source: §1 · pp. 1–2
  • RDKV derives token or channel weights from attention-distribution or logit deviations and uses reverse water-filling to assign bit-widths from full precision to zero bits.This makes eviction and quantization endpoints of a single allocation scheme rather than staged actions.Source: §1 · p. 2
  • RDKV formulates cache compression as rate-distortion allocation and realizes mixed-bit caches with TriZone packed decode, fusing dequantization into attention.The introduction reports that RDKV approaches FullKV performance across six categories while the strong baseline lags consistently.Source: §1 · p. 2

2 Related Work

Prior KV-cache compression methods can be grouped by their target: a prescribed compression ratio or tier schedule, or a fixed active-cache budget. RDKV extends budget-driven methods by jointly selecting eviction and quantization through one allocation problem.

2 Related Work

  • Compression-ratio-driven methods target a specified compression ratio, using quantization and, in some hybrids, eviction as the lowest precision tier.HqeKV searches tier ratios across 200 trials per setting, while ARKV searches scoring thresholds offline.Source: §2 · p. 2
  • Budget-driven methods retain selected cache units under a fixed active-cache budget, with scores based on attention, output perturbation, or other signals.Some methods distribute budgets across heads or layers; ThinK applies eviction to key channels.Source: §2 · p. 3
  • RDKV generalizes budget-driven methods by deriving distortion weights from attention and jointly allocating bit-widths without offline ratio or threshold optimization.Unlike binary keep-or-drop methods, it provides intermediate precision tiers for tokens or channels.Source: §2 · p. 3

3 Methodology

RDKV formulates KV-cache compression as a rate-distortion allocation that minimizes attention-computation distortion under a bit budget. Its per-head pipeline computes token and channel weights, allocates V before K, and stores mixed-bit entries in packed zones for fused decode.

3.1 Rate-Distortion Formulation

  • Token weights sum attention received across query positions, while channel weights reflect the effect of removing a K channel on the logit matrix.For K, the weight is the product of the corresponding Q and K column norms; a channel large in both is costly to remove.Source: §3.1 · pp. 3–4
  • Reverse water-filling allocates bits using each unit’s distortion weight and quantization hardness, with the budget setting the boundary between zero-bit eviction and finite-bit retention.The discrete implementation uses {0, 2, 4, 8, 16}-bit choices and a multiple-choice knapsack formulation.Source: §3.1 · p. 4

3.2 Per-Head Allocation Pipeline

  • The per-head pipeline computes token and channel weights once after prefill, allocates V first, then allocates K across only the retained tokens.V-side eviction determines which tokens remain for K allocation; the channel weights are reused without recomputation.Source: §3.2 · p. 5

3.3 Efficient Packed Decode

  • TriZone separates packed quantized entries, full-precision retained entries, and newly generated tokens so the attention kernel can use uniform dequantization paths.The compressed cache stays packed in HBM, with dequantization fused into attention rather than unpacked to FP16 beforehand.Source: §3.3 · p. 6

4 Experiments

Experiments compare methods under matched per-layer cache budgets across long-context benchmarks, then measure latency, memory, and ablations. RDKV leads on benchmark accuracy and reduces long-context decoding cost.

4.1 Main Results

  • 49.47 at 1024L is within 0.35 points of FullKV, and RDKV has the highest LongBench average at every tested budget.At 64L, its lead over the strongest baseline widens to 2.7 points.Source: Table 1 · p. 7, §4.1 · p. 7
  • 0.99 average retrieval accuracy keeps RDKV near FullKV’s 1.00, while SnapKV and AdaKV lose accuracy in mid-depth bands.The comparison is at 32K context and Btotal = 64L; a 2- or 4-bit copy can recover the answer where eviction would discard the token.Source: §4.1 · pp. 7–8, Figure 4 · p. 8
  • 1.40 points over AdaKV and 1.55 over SnapKV are RDKV’s InfiniteBench leads on the 10-task average.The benchmark averages ∼200K-token sequences, with its longest task reaching ∼2M tokens.Source: Table 3 · p. 8, §4.1 · p. 8
  • 4.1–14.2 points is RDKV’s lead over the strongest baseline across RULER lengths; it stays within 1 point of FullKV up to 8K.Eviction-only baselines fall over 6 points below FullKV at 4K and 17 points at 128K.Source: Table 2 · p. 8, §4.1 · p. 8

4.2 Memory and Latency

  • 4.5× speedup at 128K: RDKV stays near 18 ms per token while FullKV latency rises to 82 ms.On a single A100 64 GB, RDKV remains near 18 ms from 8K to 128K; AdaKV is ∼25 ms at the same token budget.Source: §4.2 · p. 8, Figure 5 · p. 9
  • 30.5 GB at 128K is a 1.9× reduction vs. FullKV; at comparable latency, RDKV leads AdaKV by 8.5 points (67.0 vs. 58.4).FullKV reaches 58.1 GB at 128K and OOMs at 256K, while RDKV uses 44.5 GB at 256K.Source: Figure 5 · p. 9, §4.2 · p. 9

4.3 Ablation Studies

  • 8.4 / 8.6 points are the LongBench costs of removing eviction; removing quantization costs 4.9 / 2 points, while tri-state closes most of the gap.Values are for Btotal = 128L and 512L; the tri-state {0, 4, 16} gaps are 0.5 / 0.2 points.Source: Table 4 · p. 9, §4.3 · p. 9
  • ∼18 ms per token with TriZone versus ∼36 ms without it shows a 2× latency reduction with bit-identical output.Without TriZone, the mixed-bit cache is unpacked to FP16 before FA2, reading the same HBM as full precision.Source: §4.3 · p. 9

5 Conclusion

The conclusion frames eviction and quantization as one rate-distortion allocation, with packed decoding turning mixed-bit assignments into HBM savings. RDKV performs well across long-context benchmarks and motivates streaming re-budgeting.

5 Conclusion

  • RDKV assigns cache units bit-widths from {0, 2, 4, 8, 16}, with eviction as the zero-rate boundary of quantization.Packed decoding keeps the mixed-bit cache in HBM and fuses dequantization into attention.Source: §5 · p. 9
  • Across three benchmarks and five open-source LLMs, RDKV outperforms evaluated baselines in accuracy, decoding speed, and peak memory at long context, reaching 256K where FullKV does not.The conclusion states that joint eviction and quantization can improve accuracy at no cost to efficiency.Source: §5 · p. 9
  • Streaming re-budgeting is proposed as an extension to adapt bit assignments during decoding as the cache grows, rather than freezing them after prefill.Source: §5 · p. 9

A Proofs

The proofs derive token and channel importance weights from the attention changes caused by removal, then establish continuous and discrete allocation bounds. The continuous optimum follows from convexity and KKT conditions, while the discrete Lagrangian expression lower-bounds the optimum.

A Proofs

  • Removing a V-cache token changes the attention distribution by a total amount proportional to its attention mass, so its weight is the sum of that mass across queries.The proof sums the evicted entry's attention mass and the shifts to surviving entries.Source: A Proofs · p. 14
  • Removing a K-cache channel changes the attention logits by a rank-one matrix whose spectral norm is determined by the norms of the query and key channel vectors.The proof notes that the same weight follows under unitarily invariant norms.Source: A Proofs · p. 14
  • The continuous allocation minimizes weighted distortion under an active bit budget; convexity makes KKT conditions sufficient, and zero-bit allocation occurs below the stated weight threshold.For the high-rate model, per-unit RMSE scales as σ_u 2^-b_u, yielding a convex weighted objective.Source: A Proofs · p. 14
  • For discrete bit-width choices, the Lagrangian expression g(λ) is a lower bound on the optimal cost, while a feasible allocation supplies an upper bound.The per-unit minimization and feasibility of the candidate allocation establish the respective bounds.Source: A Proofs · p. 15

B Detailed Implementation

RDKV computes scores, allocates bits, and packs the cache once at the end of prefill; its implementation uses a discrete knapsack solver and a packed layout designed for decoding. The implementation also fuses K-cache dequantization into the query–cache dot product.

B Detailed Implementation

  • RDKV performs probe-based scoring, bit allocation, and TriZone packing once after prefill, using bit-widths {0, 2, 4, 8, 16} and an equal K/V budget split.The probe shares a 32-token observation window and pooling kernel 5 across eviction-style baselines.Source: B Detailed Implementation · p. 15
  • Calibration estimates per-coordinate distortion for each bit-width from 32 LongBench-prefix sequences, separately at token-level V and channel-level K granularity.Sequences are truncated to 4k tokens, and estimates are cached per layer and KV-head slice.Source: B Detailed Implementation · p. 15
  • TriZone sorts V rows and K channels into 2-, 4-, and 8-bit segments, enabling shift-and-mask dequantization without per-element branching.The query is permuted to match the sorted K-channel order.Source: B Detailed Implementation · pp. 15–16
  • The K-cache kernel folds per-channel scaling into the query, multiplies by raw integer codes, and subtracts a query-specific bias after accumulation.This avoids writing and rereading a dequantized FP16 tile for retained tokens at each decode step.Source: B Detailed Implementation · p. 16
  • Padding and token/channel masks zero out aligned or out-of-range positions so they contribute nothing to the computed result.The 2-bit and 4-bit K segments are padded to channel counts divisible by four and two, respectively.Source: B Detailed Implementation · p. 16

C Algorithm

  • The MCKP solver uses Lagrangian bisection: each unit selects the bit-width minimizing weighted distortion plus a λ-weighted bit cost until the average budget is met.The algorithm takes the prefill cache and per-head budget as input and produces the TriZone cache for decoding.Source: C Algorithm · p. 16

D Additional Experiments on LongBench

Across the evaluated models and cache budgets, RDKV achieves the highest LongBench average without architecture-specific tuning. Its advantage is largest under aggressive compression, and its results remain close to FullKV at larger budgets.

D Additional Experiments on LongBench

  • The evaluation spans 16 LongBench tasks, four additional models, and budgets from 64L to 1024L, with shared probe settings across methods.The additional models test cross-architecture and cross-scale generality.Source: D Additional Experiments on LongBench · p. 16
  • At Btotal = 64L, RDKV leads the strongest baseline by 1.5, 5.5, and 2.7 points on Mistral-7B, Qwen3-4B, and LLaMA-3.1-8B, respectively.At Btotal = 1024L it recovers 98.6–99.4% of FullKV across models; LLaMA-2-13B and Qwen2.5-72B reach 99.5–99.6% at Btotal = 512L.Source: D Additional Experiments on LongBench · p. 16
  • At aggressive compression, water-filling retains moderate-importance tokens at low bit-width rather than evicting them entirely.The reported consistency across model architectures and scales suggests the distortion weights transfer without modification.Source: D Additional Experiments on LongBench · p. 16

E Additional Experiments on RULER

On RULER, RDKV attains the highest reported average accuracy across context lengths and leads HqeKV at each length in the comparison. Its gains include difficult retrieval and aggregation tasks, though HqeKV wins on some individual tasks.

E Additional Experiments on RULER

  • RULER covers needle retrieval variants, variable tracking, and word-extraction aggregation tasks that test recall, multi-hop reasoning, and synthesis of distributed signal.The evaluation also compares methods from 4k to 128k tokens on LLaMA-3.1-8B-Instruct.Source: E Additional Experiments on RULER · p. 17
  • At 64k tokens, RDKV averages 83.20 versus 78.73 for AdaKV, a 4.5-point lead; at 4k it matches FullKV at 98.59 versus 98.58.Across the evaluated context lengths, RDKV has the highest average accuracy.Source: E Additional Experiments on RULER · p. 18
  • On S-NIAH-3, RDKV scores 99.80 at 64k and 93.80 at 128k; on CWE at 16k, it scores 77.26, 22 points above the next-best method.The reported CWE comparison attributes baseline scores of 46–55 to eviction discarding low-attention tokens carrying frequency signal.Source: E Additional Experiments on RULER · p. 18
  • Against HqeKV at Btotal = 1024L, RDKV leads by 5.3 points at 16k, 1.8 at 32k, and 5.0 at 64k, while HqeKV wins on FWE at 32k and 64k and MK-NIAH-3.RDKV's task-level leads include CWE, MV-NIAH, and MQ-NIAH.Source: E Additional Experiments on RULER · p. 18

F Additional Experiments on Needle-in-a-Haystack

At Btotal = 128L, RDKV retains near-uniform Needle-in-a-Haystack retrieval success comparable to FullKV, while SnapKV and AdaKV show failure bands at intermediate depths for longer contexts.

F Additional Experiments on Needle-in-a-Haystack

G Additional Experiments on InfiniteBench

On InfiniteBench, RDKV leads the evaluated compression methods on average, with advantages concentrated in retrieval and several language-understanding tasks. The accompanying overhead and budget-split experiments report a modest one-time prefill cost and favor an equal K/V budget split.

G Additional Experiments on InfiniteBench

  • 39.46 is RDKV’s highest InfiniteBench average among compression methods, 1.40 points above AdaKV’s 38.06.On KV Retrieval, RDKV scores 7.20 versus at most 2.20 for baselines, a 3.3× improvement.Source: G Additional Experiments on… · p. 20, Table 11 · p. 23
  • RDKV leads on En.Dia (13.00 vs. ≤11.50), En.Sum (25.13 vs. ≤23.66), and En.QA (14.49 vs. ≤13.43), and matches FullKV on En.MC (68.56).On Math.Find and Debug, compression methods cluster near each other.Source: G Additional Experiments on… · pp. 20–21

H.1 Prefill Overhead

  • At 128K prefill on one A100 64 GB, RDKV’s 1,740 ms net overhead is 6.0% of FullKV prefill, with break-even within the first ∼50 generated tokens.The measured TTFT is 30,583 ms; the one-time cost is amortised over decoding.Source: §H.1 · p. 21, Table 12 · p. 23

H.2 K/V Budget Split Ratio

  • Across cache budgets 128L, 256L, and 512L, the equal split rK = 0.5 scores highest, with margins no larger than 0.63 points.The reported margins indicate robustness to moderate K/V ratio shifts.Source: §H.2 · pp. 21–22
  • Allocating more budget to K (rK = 0.6) consistently hurts, while rK = 0.4 is largely absorbed because K distortion concentrates on a small subset of outlier channels.The passage attributes the greater penalty to starving V of per-token information that directly scales attention output.Source: §H.2 · pp. 21–22

I Visualization of Bit Allocation

The visualized allocations assign more precision to high-weight V tokens and K channels while compressing the broader remainder. The accompanying calibration, limitations, and impact discussion describes model-dependent K distortion, frozen allocations during decoding, and potential resource and energy savings.

I Visualization of Bit Allocation

J Calibrated Distortion Tables

  • V-cache distortion is consistent across architectures, whereas K-cache distortion varies more, chiefly due to elevated per-channel 8-bit error in Qwen2.5-72B embedding-adjacent layers.The calibrated values are averaged across layers, heads, and 32 calibration sequences and are consumed by the MCKP solver.Source: J Calibrated Distortion Tables · p. 28, Table 14 · p. 28

K Limitations and Future Work

  • RDKV allocates once after prefill, so its cache assignment is frozen if attention patterns shift during generation.This is the prefill-dominated regime targeted by RDKV and the compared methods.Source: K Limitations and Future Work · p. 28
  • Future work could periodically apply the same MCKP allocation to a recent decode-token buffer to keep decode-phase memory sub-linear in generated tokens relative to the uncompressed baseline.The proposal accumulates attention weights over the buffer before reallocating.Source: K Limitations and Future Work · p. 28

L Impact Statement

  • RDKV may lower hardware, financial, and energy costs for long-context serving, while the impact statement says it introduces no new model capabilities.The statement also notes the possibility that lower serving costs could lower the barrier to large-scale deployment.Source: L Impact Statement · p. 28

Related papers

Loading 2605.08317v1…