KV缓存压缩的消融、统计推断与验证

# 消融研究、统计推断与 KV-Cache 压缩验证
来源: https://arxiv.org/html/2607.09683
Ashish SirasaoElliott DelayeRajeev Patwari 超威半导体公司 (Advanced Micro Devices, Inc.) \{paolo.dalberto, ashish.sirasao, elliott.delaye, rajeev.patwari\}@amd.com

###### 摘要

本文对两类 KV-cache 压缩方案进行了系统的对比研究：TurboQuant (TQ)，它采用随机化的 Walsh-Hadamard 旋转与数据无关的 Beta Lloyd-Max 码本；以及 SpectralQuant (SQ)，它校准每个注意力头的特征基并通过注水法分配比特。两类方案均可选择在键路径、值路径或两者上附加 1 比特的 Johnson-Lindenstrauss (QJL) 残差草图。

我们做出三项贡献。首先，我们对多种 QJL 变体和嵌入维度进行了完整的消融实验，结果表明只有三种方案是非支配的：不使用旋转的标量量化、WHT 旋转结合 Beta Lloyd-Max 码本、以及后者在键上增加 QJL 的增强版本。

其次，我们引入了一种用于比较不同实现的统计验证方法：Python（参考实现）和 HIP/GPU（生产实现）使用不同的随机数生成器和矩阵操作，这些作为明确的实验变量。我们应用 Kolmogorov-Smirnov 检验来区分系统性的编解码器差异与实现引起的方差，并识别出 K 路径是 Jensen 不等式通过 softmax 非线性放大分数方差的直接特征。

第三，我们在所有场景和维度下对最终方案进行比较，并推导出场景特定的建议。重尾数据对任何基于特征基的方法都是灾难性的：异常值破坏了样本协方差的稳定性，校准后的基系统性地错位，并且任何预算增长都无法恢复损失。在结构化场景下，当独立的 K 和 V 特征基能够提供真正压缩时，SQ 胜出；在所有测试场景中，注水法退化为均匀分配。我们还刻画了有效语义维度 \(d_{\mathrm{eff}}\) 的自校准特性，它根据可用的校准预算自适应调整，而非恢复数据的真实秩——这一性质解释了令人惊讶的胜利和非单调的缩放行为。

## 1 引言

大规模 Transformer 推理受限于内存带宽：在长上下文生成中，KV-cache 访问占主导地位，直接减少缓存大小可转化为延迟和吞吐量的提升。对键和值进行量化是标准方法，越来越多的研究表明，将每个元素量化至 2-4 比特而不产生显著的精度损失是可行的[5, 4, 3]。

评估 KV-cache 压缩时的一个反复出现的挑战是，难以将观察到的质量差异归因于特定的算法选择。对真实大语言模型 (LLM) 流量的评估混淆了分布特性、硬件效应和算法假设，使得理解一种方法 *何时* 以及 *为何* 成功或失败变得困难。本文引入了一种用于受控评估 KV-cache 量化方案的方法论：一组六类合成统计场景，每类旨在隔离压缩流水线的一个结构假设；同时，还引入了一个统计框架来区分系统性的算法差异与实现噪声。完整的评估以面向 AMD GPU 的开源 HIP/C++ 基准测试形式发布；任何新的压缩方案只需实现一个方案接口，而无需修改数据生成、度量指标或统计验证基础设施，即可在同一场景下进行评估。

我们将该方法论实例化到两个代表性系列上。TurboQuant (TQ)[5] 是数据无关的：随机化 Walsh-Hadamard 变换将每个 token 的能量均匀分布到所有 \(d\) 个维度上，使得分析性的 Beta 分布码本可统一应用；可选的 1 比特 QJL 残差草图校正量化误差的主导方向。SpectralQuant (SQ)[2] 是数据自适应的：它从代表性的 token 校准每个注意力头的特征基，通过注水法将比特集中到高方差的语义维度上，并使用仅在顶部 \(d_{\mathrm{eff}}\) 维度上应用的选择性 QJL 草图来校正残差。这两个系列代表了数据依赖性谱系的两端，共同检验了我们场景设计的全部假设范围。

合成场景并非留作测试集。其目的是诊断性的：揭示压缩方案在何种条件下达成其全部潜力或灾难性失败，并解释每种结果背后的机制。每个场景单独针对一个假设——分布形状、K/V 子空间对齐或特征值衰减——以便失败可以被归因和理解，而不仅仅是观测到。

主要发现是：(1) 对八个 QJL 变体的完整消融确定了三个非支配方案，其余被淘汰；(2) 统计验证框架揭示出 K 路径 QJL 方差被 softmax (Jensen 不等式) 指数级放大，而 V 路径方差则不然——这是一个在仅关注精度的评估中不可见的区别；(3) 在特征基校准失败的重尾数据上，TQ 占优；(4) 在足够预算的结构化场景下，只要 K 和 V 在独立的代表性集上分别校准，SQ 胜出；(5) 在所有测试场景中，注水法退化为均匀分配，并且有效语义维度 \(d_{\mathrm{eff}}\) 根据可用的校准预算自校准，而非恢复数据的真实秩。

## 2 背景

标准多头注意力计算，对于每个头：

\[
\mathbf{T} = \mathrm{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d}}\right)\mathbf{V},
\qquad(1)
\]

其中 \(\mathbf{K},\mathbf{V} \in \mathbb{R}^{S \times d}\) 是过去 \(S\) 个 token 的键和值缓存，\(\mathbf{Q} \in \mathbb{R}^{N_{q} \times d}\) 是查询批次。

**TurboQuant.** 给定每个头随机符号 \(s \in \{-1,+1\}^d\)，TQ 将每个键旋转为 \(y = \mathrm{WHT}(k \odot s)/\sqrt{d}\)，使得对于单位球面上的任意 \(k\)，\(y\) 近似服从 Beta\(((d-1)/2,(d-1)/2)\) 分布。一个单独的分析性 Lloyd-Max 码本被拟合到该分布，并共享于所有维度。可选地，WHT 域残差 \(e = y - \hat{y}\) 使用随机化 Hadamard 变换 (RHT) 进行草图化：\(b = \mathrm{sign}(\mathrm{WHT}(e \odot s_2))\)，使用固定符号向量 \(s_2\)，并在解码时添加修正 \(\Delta\hat{y} = \frac{\sqrt{\pi/2}}{d}\|e\| \, s_2 \odot \mathrm{WHT}(b)\)。K 上的 QJL 改善了注意力分数估计 \(\langle q,k \rangle\)；V 上的 QJL 改善了值重建 \(\hat{V}\)。内积 \(\langle q,k \rangle\) 是嵌入维度 \(d\) 上的求和，因此 \(k \in \mathbb{R}^d\) 的 QJL 草图有理论依据 (Zandieh et al. 2024[6])。注意力输出 \(T_j = \sum_i a_i V_{ij}\) 是对 token 数 \(S\) 求和，因此 \(V_i \in \mathbb{R}^d\) 的 QJL 草图并不针对相关的内积，而是一种启发式修正。

**SpectralQuant.** SpectralQuant 的核心思想是将比特从许多低方差维度重新分配到少数高方差维度。考虑一个 \(d=128\) 维的头，预算为 \(b=2\) 比特：均匀方案为每个 token 分配 \(2 \times 128 = 256\) 比特，将预算薄薄地分散到所有维度，无论其重要性如何。如果键向量位于一个秩为 4 的子空间附近，则只有 4 个方向携带有意义的信号。SQ 从校准数据中识别这些方向，并比如说为 4 个语义维度各分配 8 比特（总共 32 比特），然后以每维 2 比特（248 比特）量化剩余的 124 个尾部维度，总共仅 280 比特——与均匀方案的存储相当，但关键维度的量化精细了 4 倍。等价地，SQ 可以在一小部分存储成本下达到高预算均匀方案的精度。

给定校准键 \(K_{\mathrm{cal}} \in \mathbb{R}^{n_{\mathrm{cal}} \times d}\)，SQ 计算经验协方差 \(C_K = K_{\mathrm{cal}}^{\top} K_{\mathrm{cal}} / n_{\mathrm{cal}}\)，通过分块幂迭代提取顶部 \(d_{\mathrm{eff}}\) 个特征向量 \(\mathbf{U}_K \in \mathbb{R}^{d \times d_{\mathrm{eff}}}\)，并将每个键投影为 \(z_k = \mathbf{U}_K^{\top} k\)。单独的特征基 \(\mathbf{U}_V\) 从值 token 校准。有效维度是参与比 \(d_{\mathrm{eff}} = \mathrm{round}(\mathrm{tr}(C)^2 / \|C\|_F^2)\)。注水法通过最小化 \(\sum_k \lambda_k \cdot 4^{-b_k}\)，满足 \(\sum_k b_k = B\)，\(b_k \ge b_{\mathrm{tail}}\) 来分配比特。尾部维度以 \(b_{\mathrm{tail}}\) 比特均匀量化。K 上的 QJL 在 \(z\) 空间中作为非对称分数估计器应用：

\[
\mathrm{score}(q,k) = \langle q_Z, z_k \rangle + \frac{\sqrt{\pi/2}}{m} \|e_z\| \, q_Z^{\top} S \, \mathrm{sign}(z_k - \hat{z}_k)^{\top} S,
\]

其中 \(S \in \{-1,+1\}^{m \times d_{\mathrm{eff}}}\)，\(q_Z = \mathbf{U}_K^{\top} q\)。

**误差度量.** 令 \(T\) 和 \(\hat{T}\) 表示参考和重建的注意力输出。相对均方误差为 \(\mathrm{relMSE}_T = \|T - \hat{T}\|_F^2 / \|T\|_F^2\)。我们使用与预算无关的 2D 误差度量：

\[
d_2 = \sqrt{(1 - \mathrm{cosine}(T, \hat{T}))^2 + \mathrm{NF}^2}, \quad
\mathrm{NF} = \frac{\mathrm{relMSE}_T}{1 + \mathrm{relMSE}_T},
\qquad(2)
\]

它将无界的 relMSE 压缩到 \([0,1)\) 并与方向误差 \(1 - \mathrm{cosine}(T, \hat{T})\) 结合。\(d_2 \in [0, \sqrt{2}]\)；越小越好。我们还报告隔离度量：\(\mathrm{relMSE}_K\) 是仅量化 K 时 (V 精确) T 的 relMSE，\(\mathrm{relMSE}_V\) 是仅量化 V 时 (K 精确) 的 relMSE，用以隔离每个缓存对输出误差的贡献。每 token 缓存度量 \(\mathrm{ptcosine\_K} = \mathrm{cosine}(K, \hat{K})\) 和 \(\mathrm{ptcosine\_V} = \mathrm{cosine}(V, \hat{V})\) 独立于 Q 测量每 token 的重建质量，用于编解码器验证。

总的来说，六个标量度量——\(\mathrm{snr\_err}_K\)、\(\mathrm{dir\_err}_K\)、\(\mathrm{snr\_err}_V\)、\(\mathrm{dir\_err}_V\)、\(\mathrm{snr\_err}_T\)、\(\mathrm{dir\_err}_T\)——定义了一个六维误差空间中的点。我们将每个方案-场景-预算配置表示为这样一个点，并使用 2D 投影 (散点图) 结合 Kolmogorov-Smirnov 检验和能量距离来丰富分析，提供几何直觉和统计推断。这个几何框架在本文的早期版本[1]中引入；当前版本修正了量化器二分搜索中的一个错误 (在 `quantise_scalar` 中的差一错误)，该错误影响了那个版本中预算 \(b \ge 3\) 的所有结果，并扩展了分析以包含 TQ 与 SQ 的完整比较以及 SQ 校准敏感性研究。

## 3 实验设置

所有实验使用 \(n_{\mathrm{cal}} = 512\) 个校准 token，每个单元 200 次试验，预算 \(b \in \{2,\ldots,7\}\)，序列长度 \(S \in \{64,\ldots,4096\}\)，查询批次大小 \(N_q \in \{1,\ldots,512\}\)，嵌入维度 \(d \in \{64,128,256\}\)。

### 3.1 合成场景

我们设计了六个场景，每个场景隔离了 SQ 的一个假设。

表 1：六个合成场景及其针对的 SQ 假设。A1: 次高斯校准。A2: 共享 K/V 子空间 (现已放宽：独立 \(\mathbf{U}_K\), \(\mathbf{U}_V\))。A3: 特征值衰减足够用于注水法。
### 3.2 方案

我们在完整消融 (第 4 节) 中评估了 16 个方案，并随后缩减至 6 个最终候选方案：

所有 QJL-K 方案使用 \(b-1\) 比特进行均方误差 (MSE) 量化，1 比特用于草图 (预算中性)。对于 SQU-QKV，\(m=64\) 个投影作用于 \(d_{\mathrm{eff}} \approx 8\) 个语义维度。

#### 存储复杂度

对于具有头维度 \(d\)、包含 \(S\) 个 token 的 KV 缓存，每个缓存的存储为：

TQ-KV: \(b \times d \times S\) 比特
SQU-KV: \((b_{\mathrm{sem}} \times d_{\mathrm{eff}} + b_{\mathrm{tail}} \times (d - d_{\mathrm{eff}})) \times S\) 比特

当 \(d=128\)、\(d_{\mathrm{eff}}=8\)、\(b_{\mathrm{tail}}=2\) 且名义预算 \(b_{\mathrm{sem}}=b\) 时：TQ 每 token 使用 \(128b\) 比特；SQU 每 token 使用 \(8b + 240\) 比特。在 \(b=2\) 时，两者均使用 256 比特/ token (等存储)。在 \(b=3\) 时，TQ 使用 384 比特/ token，而 SQU 仅使用 264——节省 31%。这个存储差距随 \(b\) 增长，这就是为什么在 \(b \ge 3\) 时等存储比较总是有利于 SQ，即使 TQ 在质量上胜出。

### 3.3 实现与验证

我们使用 HIP/C++ (AMD MI100, `hipcc -O3`) 和 Python (参考实现) 实现所有方案。Python 和 HIP 使用不同的随机种子 (PCG64 对比 mt19937)，这使得种子成为第 5 节统计验证的一个明确实验变量。

## 4 结果 I: TQ 消融

我们评估了所有 8 个 TQ 变体 (Plain-KV, TQ-KV, TQ-QKV, TQ-FKV, TQ-KQV, TQ-QKQV, TQ-KFV, TQ-FKFV) 在 6 个场景和 \(d \in \{64,128,256\}\) 上的表现。图 1 显示了三个代表性预算下 \(d_2\) 作为 \(d\) 的函数。

> 参见图注
图 1: 所有 8 个 TQ 变体的 \(d_2\) 对比嵌入维度 (\(d \in \{64,128,256\}\)，200 次试验，预算 2, 4, 6)。TQ-KV 和 TQ-QKV 在所有场景中占优。FULL 变体 (橙色/黄色) 始终比其 RHT 对应物差。V 路径 QJL 变体 (KQV, FKFV) 始终在每个场景上造成损害。Plain-KV 在 lowrank_aligned 场景中随 \(d\) 退化 (Beta 码本失配随 \(D\) 增长)。

三个发现淘汰了八个方案中的五个。

#### FULL 投影变体被支配

TQ-FKV 和 TQ-FKFV 应用 \(d \times d\) 的随机 Rademacher 投影。在 \(D\) 维向量上使用 \(m=D\) 个投影时，修正方差为 \(\mathrm{Var} \propto \|q\|^2 \|k\|^2 / D\)——理论上与 RHT 相同，但缺乏 Walsh-Hadamard 变换的结构化低方差特性。实践中，FULL-K 在每个场景和预算下均一致地差于 RHT-K (图 1)。我们保留 TQ-QKV