CommVQ: Commutative Vector Quantization for KV Cache Compression

发表时间: 2025-07 · arXiv:2506.18879 (ICML 2025)

原文: https://arxiv.org/abs/2506.18879

Junyan Li 1, Yang Zhang 2, Muhammad Yusuf Hassan 1, Talha Chafekar 1, Tianle Cai 3, Zhile Ren 4, Pengsheng Guo 4, Binazir Karimzadeh 4, Colorado J Reed 4, Chong Wang 4, Chuang Gan 1

速读

一句话结论 本文提出了一种名为 CommVQ 的键值缓存压缩方法，通过引入与旋转位置编码可交换的向量量化技术，在几乎不损失模型精度的前提下实现了极低比特的缓存压缩，使得 LLaMA-3.1 8B 模型能够在一张 RTX 4090 显卡上支持高达 128K 的上下文长度。

要解决什么问题 大语言模型在处理长上下文时，自注意力机制所依赖的键值缓存会随着上下文长度的增加呈线性增长，最终成为限制 GPU 显存的核心瓶颈。以 LLaMA-3.1 8B 为例，FP16 权重仅占 16GB，但在批处理大小为 2、上下文 128K 时，键值缓存会消耗 88GB，导致单张 80GB 顶级显卡也无法独立推理。为了缓解显存压力，现有的量化方法通常对缓存中的每个标量进行独立压缩，虽然能减少显存占用，但在极端压缩率（如 2-bit 或 1-bit）下会丢失大量信息，导致模型生成质量严重下降。此外，如果在注意力计算前强行对高度压缩的缓存进行解码，又会引入巨大的额外计算开销，拖慢推理速度。因此，如何在维持极低显存占用的同时，既保住模型精度，又不引入难以承受的计算延迟，是当前长上下文模型推理面临的关键卡点。

怎么做的 本文的核心思路是放弃传统的标量独立量化，转而采用向量量化将整个键值向量作为一个整体进行压缩，并巧妙利用旋转位置编码的数学性质重构注意力计算流程，从而彻底绕开解码带来的计算瓶颈。整个方法由加性量化编码器和可交换码本两个关键部件构成。首先，轻量级编码器会将多维的键或值向量映射为低比特离散表示，需要时通过与学习到的码本进行矩阵乘法来重建原始向量：

$$ \hat{t}_i = s_i C $$

这种加性量化极大地压缩了显存，但如果直接将解码后的向量代入自注意力机制，会引入与码本大小成正比的巨大计算开销。为了解决这个卡点，作者针对键缓存设计了与旋转位置编码矩阵可交换的特殊码本。由于位置编码矩阵是由多个二维旋转子矩阵构成的分块对角矩阵，作者在对应的二维子空间内约束码本矩阵，使其满足乘法交换律：

$$ R_m^i C = C R_m^i $$
这一性质是整个方法的内核。在标准的自注意力计算中，查询向量与键向量的内积原本需要逐个处理带有不同位置编码的解码向量。利用交换律，作者将解码与位置编码的顺序对调，使得注意力得分的计算可以重写为：
$$ \alpha_i = \sum_{j,l} (q^j R_t^j) C_K^{jlT} R_i^{jT} [s_i^j = l]^T $$
在这个重写后的公式中，查询向量与码本的乘积项完全脱离了目标词位置的依赖。这意味着模型可以在处理当前查询向量时，一次性预计算出该项，并在所有历史上下文中高效复用。对于值缓存，作者则通过简单调换注意力权重与码本的矩阵乘法顺序来降低复杂度。为了进一步逼近 1-bit 的极限压缩率，方法还引入了残差量化和分组共享机制，并通过 EM 算法在校准集上迭代优化码本，最终在极低比特下保住了向量的表征精度。

效果如何 实验主要在支持 128K 上下文的 LLaMA-3.1-8B-Instruct 模型上展开，并在 LLaMA-2-7B 和 Mistral-8B 上进行了泛化验证。码本训练数据来自 FineWeb-Edu，硬件测试基于 H100-80GB 和消费级 RTX 4090 显卡。对比基线点名了三种路线：KIVI代表非对称量化、KVQuant代表非均匀量化、VQLLM代表残差向量量化。在 LongBench 长文本基准测试中，2-bit 设置下的 CommVQ 实现了几乎无损的压缩，平均分达到 47.98（全精度基线为 48.05），在显存占用更低的情况下击败了 KIVI。在极限的 1-bit 量化下，CommVQ 的优势彻底显现，LongBench 平均分保持在 44.94，而同等压缩率的 VQLLM 暴跌至 27.42，KIVI 跌至 16.70。在考验复杂推理的 GSM8K 数学任务中，1-bit 的 CommVQ 依然维持了 66.57% 的准确率，而所有基线方法在此设置下均崩溃至 3% 以下。工程实现上，基于 Triton 编写的算子将 128K 上下文的推理显存从 60GB 锐减至 20GB，使得单张 RTX 4090 即可流畅运行，且得益于交换律设计，计算延迟比暴力解码方案降低了 9.6 倍。不过，该方法也存在一定代价和局限：它需要在部署前使用校准数据离线运行 EM 算法训练编码器和码本；此外，尽管 1-bit 压缩远超基线，但在 InfiniteBench 中对精度要求极高的数字检索任务上，准确率仍从全精度的 99.49% 下降到了 62.37%，说明极限压缩对细粒度信息的保留仍有折损。

主要贡献

随着大型语言模型（LLMs）上下文长度的不断增加，用于加速自回归生成的键值（KV）缓存的内存占用呈线性增长，最终成为GPU内存的主要瓶颈。例如，在批处理大小为2的情况下，运行128K上下文的LLaMA 3.1 8B模型仅KV缓存就需要88 GB内存，这使得单张GPU（如H100-80GB）无法在不卸载KV缓存的情况下进行推理。现有的KV缓存量化方法（如INT4或更低比特）通常独立处理每个标量，在极端量化（如1-bit或2-bit）时会导致严重的信息丢失和模型性能下降。

为解决上述问题，本文提出了交换向量量化（CommVQ），这是一种专为长上下文LLMs量身定制的高效且准确的KV缓存量化方法。其核心创新点包括：
1. 向量级量化：不同于以往独立量化KV缓存中每个标量的工作，CommVQ利用加性量化（Additive Quantization）技术，使用学习到的密码本将KV缓存中的每个向量作为一个整体压缩为低比特宽度的表示。
2. RoPE交换密码本设计：为了将加性量化高效地集成到自注意力机制中，本文创新性地设计了与旋转位置嵌入（RoPE）矩阵满足交换律的密码本。这允许在计算冗长的键-查询乘积时，预先计算并重用中间结果，从而大幅降低解码的计算开销。
3. 极低比特推理能力：通过EM算法训练密码本并结合Triton内核实现，CommVQ在2-bit量化下实现了近乎无损的KV缓存压缩（大小减少 $87.5\%$），在1-bit量化下仍保持极高的准确率，使得在单张RTX 4090 GPU上运行128K上下文的LLaMA-3.1 8B模型成为可能。

背景知识与设计原则

自注意力与KV缓存：
在LLM推理的预填充（Prefilling）阶段，给定输入提示的隐藏状态 $X \in \mathbb{R}^{N \times d}$，计算查询（Q）、键（K）和值（V）矩阵，并输出自注意力结果：

$$\begin{aligned} \begin{array}{rl} Q = X W_Q, \quad K = X W_K, \quad V = X W_V \\ \operatorname{Self-Attn} = \operatorname{Softmax}\left( \frac{Q K^T}{\sqrt{d}} \right) V \end{array} \end{aligned}$$


计算出的K和V矩阵被缓存。在解码（Decoding）阶段，给定当前输入隐藏状态 $x \in \mathbb{R}^{1 \times d}$，KV缓存被更新为 $K \gets \mathsf{Concat}(K, x W_K)$ 和 $V \gets \mathsf{Concat}(V, x W_V)$，并复用这些缓存计算当前的自注意力输出。

旋转位置嵌入（RoPE）：
RoPE被广泛应用于LLMs（如LLaMA、Mistral等）以编码位置信息：

$$ q_m \gets q_m R_m, \quad k_m \gets k_m R_m $$


其中 $q_m, k_m \in \mathbb{R}^{1 \times d}$ 是第 $m$ 个Token的查询和键向量，$R_m \in \mathbb{R}^{d \times d}$ 是RoPE矩阵。由于RoPE是一个仅在 $2 \times 2$ 对角块上有非零值的稀疏矩阵，可以将 $k_m$ 划分为多个二维子向量 $k_m^i = (k_{ix}, k_{iy})$。对应的 $2 \times 2$ 对角子矩阵 $R_m^i$ 如下：

$$\begin{aligned} R_m^i = \left( \begin{array}{cc} \cos m \theta_i & -\sin m \theta_i \\ \sin m \theta_i & \cos m \theta_i \end{array} \right) \end{aligned}$$
关键Observation（交换律）：该 $2 \times 2$ 旋转矩阵 $R_m^i$ 满足以下交换律属性。设 $C \in \mathbb{R}^{2 \times 2}$ 定义为 $C = \left( \begin{array}{ll} x & y \\ -y & x \end{array} \right)$，则有：
$$ R_m^i C = C R_m^i $$
这一矩阵乘法下的交换律属性是本文优化计算复杂度的核心基础。

方法细节

KV缓存的加性量化学习
本文采用受加性量化【1，Additive quantization for extreme vector compression+2014+CVPR】启发的每Token量化方案，分别对键和值向量进行独立量化。
* 编码过程：给定第 $i$ 个Token的 $d$ 维键或值向量 $t_i \in \mathbb{R}^d$，使用编码器 $E$ 将其编码为长度为 $N_c$ 的二进制序列 $s_i \in \{0, 1\}^{N_c}$，即 $s_i = E(t_i)$。编码器 $E$ 由线性层、激活函数和输出线性层组成，并使用Gumbel-softmax使其端到端可导。量化后的向量 $s_i$ 被拼接并存储为量化KV缓存 $S$。
* 解码过程：加载KV缓存时，对于每个Token的 $s_i$，通过与密码本 $C \in \mathbb{R}^{N_c \times d}$ 进行简单的矩阵乘法来重建解码张量 $\hat{t}_i$：

$$ \hat{t}_i = s_i C $$


解码后的键和值随后参与自注意力计算。编码器 $E$ 和密码本 $C$ 通过梯度下降优化，以最小化原始张量 $t_i$ 和解码张量 $\hat{t}_i$ 之间的MSE损失。

Figure 1. 用于KV缓存压缩的向量量化图示。KV缓存中对应于一个Token的每个向量首先被编码为低比特宽度的表示，从而显著减小存储大小。此过程独立应用于KV缓存中的所有向量。需要时，加载压缩的表示并使用密码本将其解码回原始KV缓存。
Figure 1. 用于KV缓存压缩的向量量化图示。KV缓存中对应于一个Token的每个向量首先被编码为低比特宽度的表示，从而显著减小存储大小。此过程独立应用于KV缓存中的所有向量。需要时，加载压缩的表示并使用密码本将其解码回原始KV缓存。
  • KV缓存缩减分析：量化后的层KV缓存大小变为 $B \times N \times N_c \times 2 \times 1$ bits。缩减率 $RR = 1 - \frac{N_c}{16d}$。KV缓存中表示每个标量的平均比特宽（Avg. bit）为：

    $$ \mathbf{Avg. bit} = 16 - 16 \times RR = \frac{N_c}{d} $$

  • 计算复杂度问题：传统的解码后自注意力计算在推理步骤 $t$ 的复杂度为 $O((2d+1)N + 2dN_c N)$。第二项 $2dN_c N$ 代表键和值的KV缓存解码开销。由于 $N_c$ 通常很大（例如1-bit量化时为1024），额外的向量量化解码开销比普通自注意力高出 $N_c$ 倍，这构成了显著的计算瓶颈。

基于交换密码本的效率优化
为了解决上述巨大的计算开销，本文利用交换密码本重新设计了自注意力计算。
* 交换律带来的计算重用：在计算注意力得分 $\alpha_i = q R_t (s_i C_K R_i)^T$ 时，由于 $R_i$ 随 $i$ 变化，无法提取公因式。但如果设计密码本 $C_K$ 使其与 $R_i$ 满足交换律，公式可重写为 $(q R_t) C_K^T R_i^T s_i^T$。此时，大块计算 $(q R_t) C_K^T$ 独立于 $i$，可以在不同 $i$ 之间重用，从而大幅节省计算量。
* 设计交换密码本：由于RoPE矩阵是块对角的，问题被分解到二维子空间。令 $\mathcal{C}_K^j = \{C_K^{j0}, C_K^{j1}, \dots, C_K^{j(N_{c'}-1)}\}$ 为键向量子空间 $j$ 的密码本集合，其中 $N_{c'}$ 是量化级别的数量，每个 $C_K^{jl}$ 是一个满足公式7形式的 $2 \times 2$ 矩阵，因此满足交换律 $R_i^j C_K^{jl} = C_K^{jl} R_i^j$。子向量 $k_i^j$ 的量化表示为二维向量 $s_i^j$。解码键表示为：

$$ \hat{k}_i^j = \sum_{l=0}^{N_{c'}-1} [s_i^j = l] C_K^{jl} $$
其中 $[s_i^j = l]$ 是二维布尔指示向量。基于此，注意力得分计算被重写为：
$$\begin{aligned} \begin{array}{l} {\alpha_i = \sum_j (q^j R_t^j) (\hat{k}_i^j R_i^j)^T} \\ {= \sum_j (q^j R_t^j) \big( \sum_l [s_i^j = l] C_K^{jl} R_i^j \big)^T} \\ {= \sum_{j,l} (q^j R_t^j) R_i^{jT} C_K^{jlT} [s_i^j = l]^T} \\ {= \sum_{j,l} (q^j R_t^j) C_K^{jlT} R_i^{jT} [s_i^j = l]^T} \end{array} \end{aligned}$$
最后一步应用了交换律，使得 $(q^j R_t^j) C_K^{jlT}$ 可跨 $i$ 重用。
* 密码本的学习与残差优化：通过最小化重建误差 $\sum_i \|\hat{k}_i^j - k_i^j\|^2$ 来学习密码本。这通过类似EM的算法解决（E步更新分配，M步更新密码本）。为了提高压缩率，将连续的 $g$ 个子向量分入一组并共享量化值 $s$。为了提高量化精度，迭代地在新密码本的量化误差张量上应用聚类算法 $R$ 次（类似于残差向量量化【Barnes等，1996】）。最终 $d$ 维向量的平均量化比特为：
$$ \mathbf{Avg. bit} = \frac{R \log_2(N_{c'})}{g} $$
* 降低计算复杂度：对于值（Value）量化，保留原始方法但重新排序矩阵乘法：$\operatorname{Self-Attn} = \operatorname{Softmax}(A) S_V C_V$。这使得复杂度从 $O(dN_c N + dN)$ 降至 $O(N_c N + dN_c)$。整合上述所有优化后，最终的计算复杂度为：
$$ O((Rd + N_c + 1)N + d(N_c + R N_{c'})) $$
相比于未优化的计算，复杂度从普通自注意力的 $N_c$ 倍降低到了约 $\frac{R+1}{2}$ 倍，由于 $R$ 较小（如1-bit时 $R=11$），额外开销微乎其微。

实验环境

  • 数据集：使用FineWeb-Edu数据集的子集用于训练编码器和密码本。评估任务使用LongBench（128K最大长度）、InfiniteBench（128K最大长度）、Needle-in-a-Haystack（NIAH）和GSM8K。
  • 模型架构：主要评估模型为支持128K上下文的LLaMA-3.1-8B-Instruct。消融实验使用了LLaMA-2-7B（32K）和Mistral-7B-v0.3（32K）。
  • 硬件与软件配置：内存节省和延迟测量在H100-80GB GPU上进行；128K上下文的1-bit推理在单张RTX 4090 GPU上验证。代码提供了Triton内核实现以反映真实的内存节省。
  • 基线方法：对比了KIVI【14，Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024】、KVQuant【9，Kvquant: Towards 10 million context length llm inference with kv cache quantization+2024】和VQLLM【12，Residual vector quantization for kv cache compression in large language model+2024】。

实验结果

  1. 长上下文基准测试（LongBench）：

    • 实验内容：在8个任务上评估多文档QA、摘要和代码补全能力。
    • 结果与分析：CommVQ-2在大多数任务上实现了无损准确率（平均分47.98，接近FP16的48.05），比KIVI节省33%内存且得分更高。CommVQ-1大幅超越其他基线，平均分比VQLLM-1高出17.52%，在极低比特下最小化了精度下降（详见原论文Table 1）。

  2. 超长上下文基准测试（InfiniteBench）：

    • 实验内容：评估近乎无限输入长度的检索、QA等任务。
    • 结果与分析：在需要准确信息的检索任务（如R.PK, R.Num）中，基线方法在1-bit时完全失效（得分为0），而CommVQ-1仍保留了强大的能力（R.PK得分为99.15），证明其密码本设计有效保留了KV缓存信息（详见原论文Table 2）。

  3. 大海捞针测试（Needle-in-a-Haystack）：

    • 实验内容：测试模型在大量无关文本中检索特定小段信息的能力。
    • 结果与分析：Fig 2证实CommVQ-2成功保留了完整的检索能力，与FP16基线性能匹配。CommVQ-1在极高压缩率下的检索保真度显著优于KIVI-1。
      Figure 2. 使用LLaMA-3.1 8B模型进行的大海捞针测试。我们展示了FP16基线（左上）、KIVI-1（左下）、CommVQ-2（右上）和CommVQ-1（右下）的测试结果。我们方法的2-bit版本匹配了FP16的性能，而1-bit CommVQ变体优于KIVI，证明了在极端压缩下强大的检索保真度。

  4. 复杂推理测试（GSM8K）：

    • 实验内容：评估多步数学推理能力。
    • 结果与分析：CommVQ-2准确率（76.04%）仅比FP16下降0.23%，远超VQ-LLM-2（52.69%）。在1-bit下，基线方法崩溃至~2%，而CommVQ-1仍保持66.57%的准确率。

  5. 模型消融与域偏移鲁棒性：

    • 实验内容：在LLaMA-2-7B和Mistral-7B上测试；在数学、代码和合成UUID数据集上测试困惑度（PPL）。
    • 结果与分析：方法在不同模型上泛化良好。跨域测试显示PPL没有显著增加，证明密码本和编码器在面对未见过的域（如合成UUID）时具有强鲁棒性。

  6. 效率测试：

    • 实验内容：对比朴素实现与优化实现的延迟；测试Triton内核的实际显存占用。
    • 结果与分析：交换密码本优化使延迟加速了6.0倍至9.6倍。Fig 3证实了实际显存节省：在120K上下文下，FP16需要60GB，而CommVQ-1仅需20GB；在32K上下文下，FP16在批处理大小为8时OOM，而CommVQ可扩展至128。
      Figure 3. CommVQ (1-bit) 与 FP16 模型的每Token解码内存使用量比较。实验在 LLaMA3.1-8B-Instruct 模型上进行。

结论

本文提出了CommVQ，一种用于长上下文LLMs的新型KV缓存量化方法。通过结合向量量化和与RoPE满足交换律的密码本，CommVQ在保持高计算效率的同时显著减小了KV缓存大小。在长上下文基准测试上的评估表明，CommVQ优于现有的KV缓存量化方法，能够在有限的GPU内存约束下实现更高效、可扩展的长上下文LLM推理。

附录细节

A.1 交换密码本的编码与解码解释
将量化过程表述为聚类问题，聚类中心定义为：

$$\begin{aligned} c_{a,b} = \left[ \begin{array}{l} 1 \\ 0 \end{array} \right] \mathcal{C}_K^j[a] + \left[ \begin{array}{l} 0 \\ 1 \end{array} \right] \mathcal{C}_K^j[b] \end{aligned}$$


其中 $\mathcal{C}_K^j[a]$ 和 $\mathcal{C}_K^j[b]$ 是 $\mathcal{C}_K^j$ 中的第 $a$ 个和第 $b$ 个 $2 \times 2$ 子密码本。密码本构成了 $N_{c'}^2$ 个聚类中心。将 $k_i^j$ 量化为最近的聚类中心，使用 $s = \{a, b\}$ 作为量化表示。解码时，$\hat{k}_i^j = c_{a,b}$，这与正文中的公式18完全等价。

A.2 EM算法实现细节
使用FineWeb-Edu的子集作为校准集 $K$。E步将 $K$ 中的每个向量分配给最近的聚类中心。M步给定当前分配 $S$，推导了更新 $\mathcal{C}_K^j$ 的闭式解。定义参数向量 $\phi$、均值向量 $m$、对角计数矩阵 $S$ 和辅助常量矩阵 $T$。公式18的矩阵形式目标函数为 $\operatorname*{min}_\phi (T \phi - m)^T S (T \phi - m)$，其闭式解为：

$$ \phi^* = (T^T S T)^{-1} T^T S m $$


为了稳定拥有大量聚类中心（如4096个）的优化过程，采用了两种技术：

  1. 软聚类中心分配：计算权重矩阵 $W \in \mathbb{R}^{N \times N_{cc}}$，$W_{ij} = \frac{e^{-D_{ij}}}{\sum_k e^{-D_{ik}}}$，其中 $D$ 是L2距离矩阵。
  2. 温度退火：在早期迭代中需要平滑分布以防止死聚类中心，后期需要锐化分布以收敛。引入温度参数 $T$（随迭代指数衰减）：$W_{ij} = \frac{e^{-\frac{D_{ij}}{T}}}{\sum_k e^{-\frac{D_{ik}}{T}}}$。

A.3 密码本大小分析
存储密码本需要额外的GPU内存（FP16格式）。键密码本大小计算为 $2 \times 2 \times N_{c'} \times R \times \frac{d}{2} \times 2$ MB。对于LLaMA-3.1-8B，1-bit配置下值密码本为2.00 MB，键密码本为2.75 MB。相比于128K上下文所需的256 MB KV缓存，密码本的内存开销是恒定且微不足道的。

A.4 交换密码本配置的消融实验
通过调整量化级别 $N_{c'}$、残差量化次数 $R$ 和共享量化值的子向量组大小 $g$ 来控制压缩率。
* 在保持平均量化比特不变时，增大 $g$ 会降低量化误差（MSE），但会导致 $N_{c'}$ 增加从而提高计算复杂度。
* 在保持 $g$ 和 $N_{c'}$ 不变时，增大 $R$ 会持续降低量化误差，但代价是平均量化比特升高。
* Fig 4展示了在 $g=64$ 和 $N_{c'}=64$ 时，改变 $R$ 对MSE的影响。为了在精度、计算成本和压缩率之间取得平衡，主要实验设置 $g=64, N_{c'}=64$，1-bit量化时 $R=11$，2-bit量化时 $R=21$。
Figure 4. 在保持 $g=64$ 和 $N_{c'}=64$ 的情况下对 $R$ 进行的消融实验。

A.5 量化误差对比
通过计算原始KV缓存与解码KV缓存之间的MSE来对比量化误差。结果显示，CommVQ的加性量化方法显著优于KIVI中使用的非对称量化，特别是在1-bit极低量化比特下（KIVI为0.00380，CommVQ仅为0.00027）。