技术摘要：RDKV - 用于 KV 缓存联合淘汰与量化的率失真比特分配

1. 问题陈述

大语言模型（LLM）在长上下文推理期间面临内存带宽和容量的关键瓶颈。KV 缓存的大小随序列长度线性增长，导致在每一步解码时都需要从片外高带宽内存（HBM）重新读取数十 GB 的数据到片上内存。这使得推理过程受限于内存而非计算。

现有的压缩策略通常分为两个孤立的类别：

1. 淘汰（Eviction）： 对 token 或通道进行评分，并完全丢弃排名最低的项（二元的保留或丢弃）。这导致中等重要性的 token 缺乏精度层级。
2. 量化（Quantization）： 保留所有 token 但降低其比特宽度。这会保留甚至微不足道的 token，从而浪费内存。

尽管最近的混合方法试图结合这两者，但它们通常依赖启发式路由分数或通过离线搜索确定的预分配比例，未能将淘汰和量化视为一个统一的优化问题。

2. 方法论：RDKV

本文提出了RDKV（率失真 KV 缓存压缩），这是一个将 KV 缓存压缩形式化为率失真问题的框架。在此视角下，淘汰和量化并非独立的操作，而是单一比特分配曲线上的两个端点。

2.1 理论形式化

RDKV 根据每个缓存单元（V 缓存中的 token，K 缓存中的通道）的移除或量化对注意力计算产生的扰动，推导出其失真权重。

• V 缓存（Token）： token 的失真权重 $w_t$ 是淘汰后注意力分布的总变差距离。它计算为该 token 在所有查询位置接收到的注意力分数之和：$w_t = \sum_{\tau} a_{\tau,t}$。
• K 缓存（通道）： 通道的失真权重 $w_c$ 是将该通道置零引起的 logit 偏差的谱范数。它计算为对应查询和键列范数的乘积：$w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$。

这些权重作为率失真目标函数中的成本系数：
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{s.t.} \quad \sum_{u} b_u \leq B$$
其中 $b_u$ 是分配给单元 $u$ 的比特宽度，$B$ 是总比特预算。

2.2 比特分配流程

分配通过反向注水（reverse water-filling）（拉格朗日松弛）求解，并针对离散硬件支持的比特宽度 $\mathcal{B} = \{0, 2, 4, 8, 16\}$ 进行了适配。

1. 阶段 1（加权）： 在预填充阶段之后，使用完整的未压缩缓存计算一次 token 权重（$w_t$）和通道权重（$w_c$）。
2. 阶段 2（Token 分配）： 为 V 缓存求解多选择背包问题（MCKP）。Token 被分配来自 $\{0, 2, 4, 8, 16\}$ 的比特宽度。比特宽度为 0 对应于淘汰。
3. 阶段 3（通道分配）： 为 K 缓存求解 MCKP，但仅针对阶段 2 中保留的 token。通道的比特宽度分配独立于 token 的 V 比特宽度，允许单个 token 具有混合精度的 K 和 V 组件。

2.3 高效打包解码（TriZone）

为了实现实际的内存和延迟节省，RDKV 引入了TriZone缓存布局，将反量化融合到注意力核中，避免了在处理前将数据解包为 FP16 的需求。

• 区域 A： 打包的旧缓存流，包含保留的 K 行和 V 行，其比特宽度为 $\{2, 4, 8\}$。这些行按比特宽度分组，以允许统一的反量化路径。
• 区域 B： 分配了 16 比特精度的 token 的 FP16 V 行。
• 区域 C： 新生成 token 的 FP16 K 和 V 行。

注意力核通过组合这三个区域来计算 logit 和输出，反量化直接融合在矩阵乘法中，消除了额外的内存带宽开销。

3. 主要贡献

1. 率失真形式化： 本文将 KV 缓存压缩表述为一个统一的率失真问题，证明了最优解本质上混合了零速率（淘汰）和有限速率（量化）单元。这消除了对启发式比例或离线搜索的需求。
2. 基于反向注水的离散分配： 作者将分配实例化为一个可通过拉格朗日松弛求解的离散背包问题，根据源自注意力动态的失真权重，从 $\{0, 2, 4, 8, 16\}$ 中分配最优比特宽度。
3. TriZone 打包布局： 一种硬件感知的实现，将反量化融合到注意力核中，将理论上的比特宽度节省转化为实际的峰值内存和解码延迟降低。

4. 实验结果

实验在 LLaMA-3.1-8B-Instruct、Mistral-7B、Qwen3-4B、LLaMA-2-13B 和 Qwen2.5-72B 模型上进行，涵盖了 LongBench、RULER 和 InfiniteBench 基准测试。

• 准确率： RDKV 在平均表现上优于最佳评估基线（Ada-SnapKV）9.1%。在 LongBench 上，它仅用 2.48% 的缓存保留率就恢复了 97.81% 的全缓存准确率。
• 长上下文性能：
  • RULER： 在 128K 上下文下，RDKV 实现了 66.95% 的准确率，而最强基线（AdaKV）为 58.43%，差距为 8.5 个百分点。
  • InfiniteBench： RDKV 在压缩方法中排名第一，平均得分为 39.46，超过 AdaKV 1.40 分。
  • 大海捞针（Needle-in-a-Haystack）： RDKV 在不同深度下保持了近乎均匀的检索准确率（平均 0.99），而仅基于淘汰的基线在中间深度表现显著失败。
• 效率：
  • 延迟： 在 128K 上下文下，与全缓存 FlashAttention-2 解码相比，RDKV 实现了 4.5 倍 的解码加速。
  • 内存： 在 128K 上下文下，它实现了 1.9 倍 的峰值内存降低，并成功在单张 A100 64GB GPU 上运行 256K 上下文，而 FullKV 在此情况下内存溢出（OOM）。

5. 意义与主张

本文主张，在单一优化框架内联合处理淘汰和量化，优于分阶段或基于比例的方法。通过直接从注意力计算产生的失真中推导分配分数，RDKV 避免了仅淘汰方法的“二元陷阱”（丢弃中等重要性的 token）和仅量化方法的“浪费”（保留微不足道的 token）。

作者强调，他们的方法：

• 无需模型重训练。
• 与架构无关，在不同模型大小和系列中表现一致。
• 在延迟和准确率之间提供帕累托最优权衡，在长上下文长度下在两项指标上均严格优于基线。
• 证明了率失真理论的“连续最优解”在针对硬件离散化后，自然会产生一种混合比特策略，其性能优于当前最先进的混合方法。

本文结论指出，这种联合方法在不牺牲效率的情况下提高了准确率，使得在资源受限的硬件上进行长上下文推理成为可能。