RateQuant：基于率失真理论的优化混合精度KV Cache量化

# 基于率失真理论的混合精度 KV Cache 量化最优解

来源: https://arxiv.org/html/2605.06675

Fei Zuo${}^{1,\*}$, Zikang Zhou${}^{2,\*}$, Hao Cong${}^{3,\*}$, Xiaoyan Xi${}^{1,\*}$, Ho Fai Leung${}^{1,\dagger}$

1. BA TechWorks (BMW Group)
2. 新加坡国立大学
3. 清华大学

###### 摘要

大型语言模型在生成过程中缓存所有先前计算得到的键-值（KV）对，且 KV Cache 的大小随序列长度线性增长，使其成为服务部署中的主要内存瓶颈。将 KV Cache 量化为更低位数可以降低成本，然而，当前所有的量化器都为每个注意力头分配相同的位宽，忽视了头重要性存在巨大差异这一事实。一个自然的思路是为重要的头分配更多位宽，为其余头分配较少位宽。然而，我们表明这种混合精度分配存在一个隐藏的陷阱：每个量化器遵循不同的失真曲线 $D(b)=\alpha\beta^{-b}$，且不同量化器设计的衰减速率 $\beta$ 在 3.6 到 5.3 之间变化。将一个量化器的失真模型应用于另一个量化器会颠倒分配顺序，导致性能*劣于*均匀量化。我们将这种失效模式称为*失真模型失配*，并提出 **RateQuant** 来解决该问题。RateQuant 从小校准集中拟合每个量化器的失真模型，然后通过率失真理论中的逆水填法（reverse waterfilling）以闭式解求解由此产生的位宽分配问题。在 Qwen3-8B 模型上，使用 2.5 平均位宽时，经过校准的 RateQuant 将 KIVI 的困惑度（PPL）从 49.3 降至 **14.9**（减少 70%），并将 QuaRot 的性能提高了 6.6 PPL。整个校准过程在单张 GPU 上仅需 1.6 秒，且在推理时零额外开销。

${}^{11}$脚注文本: 同等贡献。 ${}^{22}$脚注文本: 通讯作者。

## 1 引言

大规模部署大型语言模型（LLM）需要缓存所有先前计算的键-值（KV）对，以便每个新令牌都能关注完整上下文（Pope et al., 2023 (https://arxiv.org/html/2605.06675#bib.bib20)）。KV Cache 的内存占用随序列长度、批量大小和模型深度线性增长。对于处理 4k 令牌序列的 32B 参数模型，仅在 FP16 精度下，KV Cache 每次请求就占用超过 1 GB 的内存，这通常超过了模型权重本身消耗的内存。KV Cache 量化通过以较低精度存储缓存状态来降低这一成本，近期大量研究工作产生了越来越有效的量化器（Liu et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib16); Ashkboos et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib1); Hooper et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib10); Zandieh et al., 2026 (https://arxiv.org/html/2605.06675#bib.bib26)）。然而，所有这些量化器都为每个注意力头应用*均匀*位宽，忽视了文献中充分记录的头部重要性异质性（Voita et al., 2019 (https://arxiv.org/html/2605.06675#bib.bib22)）。近期的混合精度方法在层级别（Chen et al., 2025a (https://arxiv.org/html/2605.06675#bib.bib2); Li et al., 2025b (https://arxiv.org/html/2605.06675#bib.bib13)）或通道级别（Liao and Wen, 2026 (https://arxiv.org/html/2605.06675#bib.bib14); Wang et al., 2025 (https://arxiv.org/html/2605.06675#bib.bib23)）放宽了这一限制，但每种方法都依赖于与特定量化器绑定的启发式分配规则，限制了可转移性且缺乏理论保证。

![Figure 1](https://arxiv.org/html/2605.06675#S1.F1)
**图 1：失真模型失配**（Qwen3-8B, 2.5 平均位宽）。左：不同量化器之间的失真曲线 $D(b)=\alpha\beta^{-b}$ 发生发散（$\beta$ 变化达 1.5 倍）。右：带有失配 $\beta$ 的朴素混合精度*恶化*了 PPL（KIVI: 49.3 $\to$ 87）；带有 K/V 分离的校准 RateQuant 达到了 14.9（70% $\downarrow$）。

我们识别出一个更深层的障碍：*失真模型失配*。不同的量化器具有根本不同的失真-率曲线 $D(b)=\alpha\cdot\beta^{-b}$，其中衰减速率 $\beta$ 从 3.6（TurboQuant）变化到 5.3（QuaRot）。如图 1 所示， naive 地将一个量化器的失真模型应用于另一个量化器，使得混合精度分配*劣于均匀量化*，因为失配的 $\beta$ 颠倒了各头之间的边际增益顺序。解决这一失配是构建真正通用分配框架的关键。

我们提出了 **RateQuant**，一个将每头 KV Cache 位宽分配形式化为率失真优化的框架。三个实证发现推动了其设计。首先，*敏感度代理至关重要*：在 3.5 位时，基于梯度的头部重要性比基于激活的重要性提高了 1.07 PPL。其次，*失真模型必须与量化器匹配*：应用错误的模型会使 KIVI 在 2.5 位时的 PPL 从 49.3 恶化到 87.0。第三，*键和值需要单独的位宽预算*：对于 2.5 位的 KIVI，向键分配 2.85 位，向值分配 2.15 位可将 PPL 降低 70%。

我们的贡献包括：

- **率失真框架**。我们将混合精度 KV 量化形式化为加权失真最小化问题，并通过逆水填法推导了最优连续分配（定理 2 (https://arxiv.org/html/2605.06675#Thmtheorem2)）。可实现失真减少量等于头部敏感度算术平均值与几何平均值之比（AM/GM）（定理 3 (https://arxiv.org/html/2605.06675#Thmtheorem3)），该比值无需任何量化即可计算，并可作为混合精度是否有益的廉价预测指标。
- **失真校准与 K/V 分离**。我们识别出失真模型失配这一先前未被认识的失效模式，并通过从校准数据中拟合每个量化器的失真模型来解决它。我们进一步将分配泛化，将键和值视为 $2N$ 个独立组件，使 RateQuant 适用于任意基础量化器。
- **一致的实证增益**。校准后的 RateQuant 将 KIVI 在 2.5 位时的困惑度从 49.3 降至 **14.9**（70% 减少），并将 QuaRot 提升了 6.6 PPL。在 4.0 位的 TurboQuant 上，RateQuant 在三个 Qwen3 模型中恢复了 66% 的量化引起的退化。所有增益均无运行时开销，且一次性校准时间少于 2 秒。

![Figure 2](https://arxiv.org/html/2605.06675#S2.F2)
**图 2：RateQuant 流程**。阶段 1–3 是一次性离线成本（8B 模型 << 2 秒）；阶段 4 零运行时开销。

## 2 相关工作

#### KV Cache 量化

减少 KV Cache 内存的方法包括淘汰（Zhang et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib28); Ge et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib9)）、令牌合并（Nawrot et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib18)）和量化（Hooper et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib10); Liu et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib16); Yue et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib25)）。在量化方法中，KIVI（Liu et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib16)）应用每通道对称键和每令牌非对称值；QuaRot（Ashkboos et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib1)）通过 Hadamard 旋转抑制异常值；KVQuant（Hooper et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib10)）处理异常通道采用非均匀量化；TurboQuant（Zandieh et al., 2026 (https://arxiv.org/html/2605.06675#bib.bib26)）引入了基于旋转的矢量量化。所有这些方法都为每个头分配相同的位宽。

几种近期方法转向混合精度分配。在层级别，KVmix（Chen et al., 2025a (https://arxiv.org/html/2605.06675#bib.bib2)）通过梯度范数分配精度，KVTuner（Li et al., 2025b (https://arxiv.org/html/2605.06675#bib.bib13)）使用帕累托搜索，PM-KVQ（Chen et al., 2025b (https://arxiv.org/html/2605.06675#bib.bib3)）应用整数规划。在通道级别，ChanMix（Liao and Wen, 2026 (https://arxiv.org/html/2605.06675#bib.bib14)）按动态范围聚类，KITTY（Wang et al., 2025 (https://arxiv.org/html/2605.06675#bib.bib23)）增强异常键，MixKVQ（Zhang et al., 2025 (https://arxiv.org/html/2605.06675#bib.bib27)）通过查询-激活幅度对通道评分。KV-AdaQuant（Kim et al., 2025 (https://arxiv.org/html/2605.06675#bib.bib11)）提供全局 K/V 预算分割，CoKV（Li et al., 2025a (https://arxiv.org/html/2605.06675#bib.bib12)）通过沙普利值优化每组令牌预算。这些方法中的每一种都是围绕特定的基础量化器设计的，这使得难以将一种方法的分配规则应用于不同的量化器。RateQuant 以每头粒度进行操作，通过闭式解进行分配，并通过校准支持任意基础量化器；表 1 (https://arxiv.org/html/2605.06675#S2.T1) 总结了定位。

#### 量化中的率失真理论

逆水填算法（Cover and Thomas, 2006 (https://arxiv.org/html/2605.06675#bib.bib4)）是高斯率失真问题的经典解决方案，在并行信道间最优地分配位宽预算。在 LLM 权重量化中，Radio（Tseng et al., 2025 (https://arxiv.org/html/2605.06675#bib.bib21)）通过随机对偶上升应用率失真优化，BAQ（Zheng et al., 2025 (https://arxiv.org/html/2605.06675#bib.bib29)）在 Hessian 加权目标下推导闭式水填解。对于混合精度权重，HAWQ（Dong et al., 2019 (https://arxiv.org/html/2605.06675#bib.bib6)）使用顶部 Hessian 特征值，HAWQ-V2（Dong et al., 2020 (https://arxiv.org/html/2605.06675#bib.bib7)）使用平均迹来分配每层位宽。所有这些方法都针对模型*权重*。KV Cache 在两个关键方面有所不同：头形成具有不同敏感度的自然量化组，且键和值具有不对称的错误特征。RateQuant 是第一个将率失真分配应用于 KV Cache 的方法，解决了这两个结构差异。

#### 敏感度估计

基于 Hessian 的敏感度驱动 HAWQ（Dong et al., 2019 (https://arxiv.org/html/2605.06675#bib.bib6), 2020 (https://arxiv.org/html/2605.06675#bib.bib7)）进行权重量化，而基于激活的指标（如通道幅度）在训练后量化中很常见（Dettmers et al., 2022 (https://arxiv.org/html/2605.06675#bib.bib5); Xiao et al., 2023 (https://arxiv.org/html/2605.06675#bib.bib24); Lin et al., 2024 (https://arxiv.org/html/2605.06675#bib.bib15)）。我们表明，对于 KV Cache 位宽分配，基于梯度的敏感度在质量上优于基于激活的替代方案，且代理的选择比分配算法的选择更重要（第 4.4 节 (https://arxiv.org/html/2605.06675#S4.SS4)）。

**表 1：混合精度 KV Cache 量化全景。** Cal. = 校准失真模型; Q-Agn. = 量化器无关; Closed = 闭式解。RateQuant 是唯一结合每头粒度、率失真理论、校准、K/V 分离和闭式分配的方法。

## 3 RateQuant

我们从四个方面介绍 RateQuant：问题形式化和最优分配（第 3.1 节 (https://arxiv.org/html/2605.06675#S3.SS1)）、敏感度估计（第 3.2 节 (https://arxiv.org/html/2605.06675#S3.SS2)）、整数分配算法（第 3.3 节 (https://arxiv.org/html/2605.06675#S3.SS3)）以及量化器无关扩展（第 3.4 节 (https://arxiv.org/html/2605.06675#S3.SS4)）。

### 3.1 问题形式化和最优分配

考虑一个具有 $L$ 层且每层有 $H$ 个 KV 头的 LLM，产生 $N=L\times H$ 个量化组。每个组 $i$ 有一个敏感度权重 $w_i > 0$，反映其对模型输出的重要性（定义见第 3.2 节 (https://arxiv.org/html/2605.06675#S3.SS2)）。

###### 假设 1（指数失真-率）。每头量化均方误差遵循 $D(b)=\alpha\cdot\beta^{-b}$，其中常数 $\alpha>0, \beta>1$ 取决于量化器设计和头维度 $d$。我们通过实证验证这一点：拟合 $d=128$ 时 TurboQuant 的 Lloyd-Max MSE 得到 $\alpha\approx 1.36, \beta\approx 3.48$，且 $R^2 > 0.99$（附录 E (https://arxiv.org/html/2605.06675#A5)）。

优化问题将总位宽预算 $B=\lfloor \bar{b}\cdot N \rceil$ 分配以最小化加权失真：

$$
\min_{\mathbf{b}\in\mathbb{R}^N} \mathcal{J}(\mathbf{b}) \triangleq \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{s.t.} \quad \sum_{i=1}^{N} b_i = B, \quad b_{\min} \le b_i \le b_{\max} \tag{1}
$$

###### 定理 2（逆水填法）。在假设 1 (https://arxiv.org/html/2605.06675#Thmtheorem1) 下，方程 (1) (https://arxiv.org/html/2605.06675#S3.E1) 在连续 $b_i$ 且边界约束非活跃时的解为：

$$
b_i^* = \bar{b} + \frac{\ln w_i - \overline{\ln w}}{\ln \beta} \tag{2}
$$

其中 $\bar{b}=B/N$ 且 $\overline{\ln w} = \frac{1}{N}\sum_j \ln w_j$。

###### 证明概要。拉格朗日平稳性给出 $w_i \alpha (\ln \beta) \beta^{-b_i} = \lambda$ 对所有 $i$ 成立，得出 $b_i \propto \ln w_i / \ln \beta$。常数由 $\sum_i b_i = B$ 确定。包含边界处理的完整证明见第 A.1 节 (https://arxiv.org/html/2605.06675#A1.SS1)。 $\blacksquare$

**解释。** 敏感度更高的头获得更多位宽。权衡由 $\beta$ 控制：对于 TurboQuant（$\beta=3.48$），敏感度比平均值高 $e$ 倍的头将获得 $1/\ln 3.48 \approx 0.80$ 个额外位宽。

###### 定理 3（增益比）。设 $\mathcal{J}^*$ 和 $\mathcal{J}_u$ 分别表示最优和均匀加权失真（无活跃边界）。则：

$$
\frac{\mathcal{J}_u}{\mathcal{J}^*} = \frac{\bar{w}}{\widetilde{w}} \ge 1 \tag{3}
$$

其中 $\bar{w}=\frac{1}{N}\sum_i w_i$ 是头部敏感度的算术平均值，$\widetilde{w}=(\prod_i w_i)^{1/N}$ 是几何平均值。比值 $\bar{w}/\widetilde{w}$ 仅从敏感度即可计算，无需量化，作为潜在增益的廉价*先验*预测器。对于 Qwen3 模型，经验 AM/GM $\approx 2.0$，表明有巨大的改进空间。

###### 推论 4。若 $\ln w_i \sim \mathcal{N}(\mu, \sigma^2)$，则 $\mathcal{J}_u/\mathcal{J}^* = \exp(\sigma^2/2)$。

### 3.2 敏感度估计

我们通过 KV 投影输出的平方梯度范数来估计每头的重要性：

$$
w_{l,h}^K = \mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial \mathbf{K}_{l,h,t}} \right\|^2 \right], \quad w_{l,h}^V = \mathbb{E}_{\mathbf{x}\sim\mathcal{D}}\left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial \mathbf{V}_{l,h,t}} \right\|^2 \right] \tag{4}
$$

其中 $\mathcal{L}$ 是自回归语言模型损失，$\mathcal{D}$ 是小校准集（16 个长度为 512 的序列）。

###### 命题 5（损失-失真联系）。在带有对角 Fisher 近似的二阶泰勒展开下：

$$
\mathbb{E}[\mathcal{L}(\hat{\theta}) - \mathcal{L}(\theta)] \approx \sum_{l,h} \left[ w_{l,h}^K \cdot D(b_{l,h}^K) + w_{l,h}^V \cdot D(b_{l,h}^V) \right]
$$