PuzzleKV：基于页面的低秩分解用于KV缓存压缩

# PuzzleKV：基于页面的低秩分解用于KV缓存压缩
来源：https://arxiv.org/html/2608.23843
###### 摘要

大型语言模型（LLM）的长上下文推理日益受限于键值（KV）缓存所需的内存。KV缓存压缩通过减少先前token的存储成本来解决此问题。在现有方法中，低秩压缩尤为引人注目，因为它在降维空间中表示每个token。先前的低秩方法通常从模型权重中推导出固定的投影空间、从校准激活中构建固定空间，或在大范围缓存区域上构建共享基底。这些表示可能无法捕获细微但重要的信息。我们将每个注意力头的KV缓存划分为固定长度的逻辑页面，并观察到单个页面内存在显著的低秩结构。基于此观察，我们提出了PuzzleKV，一种无需训练和校准的方法，它将每个已完成的页面视为独立的压缩单元。PuzzleKV在每一层和每个KV头内对页面进行分解，直接在稠密和分解后的页面上计算注意力，并在自回归解码过程中增量压缩新符合条件的页面。跨模型、上下文长度和基准的实验证明了PuzzleKV在匹配存储预算下的有效性。在约占用原始KV缓存存储量的60%时，PuzzleKV在所有评估模型和基准设置下均达到了完整KV性能的96%以上，并在RULER上较全局SVD有显著提升，在LongBench上表现具有竞争力。为实现更激进的压缩比，PuzzleKV可与量化进一步结合，仅使用18.7%的原始存储空间即可保留超过93%的完整KV性能。

## 1引言

大型语言模型越来越多地依赖少样本示范、思维链以及交织的推理与行动来解决复杂任务\[16 (https://arxiv.org/html/2608.23843#bib.bib1),18 (https://arxiv.org/html/2608.23843#bib.bib2),23 (https://arxiv.org/html/2608.23843#bib.bib3)\]。随着这些输入的增长，更长的上下文窗口允许LLM访问更多任务相关信息，但也对推理施加了显著的内存压力。

自回归解码器将先前token的注意力键和值存储在键值（KV）缓存中，以避免在每个生成步骤重新计算它们\[17 (https://arxiv.org/html/2608.23843#bib.bib6),9 (https://arxiv.org/html/2608.23843#bib.bib7)\]。对于一个具有$L$层、$H_{kv}$个KV头、头维度$d$和上下文长度$T$的模型，单个请求在批处理前会持有$2LH_{kv}Td$个KV元素。例如，Qwen3-32B\[22 (https://arxiv.org/html/2608.23843#bib.bib8)\]有64层、8个KV头，头维度为128，因此在BF16精度下，一个128K token序列大约需要32 GiB的KV缓存空间。由于此占用空间随上下文长度和推理批处理大小线性增长，KV缓存压缩对于高效的长上下文推理至关重要。

现有方法通过token淘汰、量化或低秩分解来减少KV缓存存储。淘汰会丢弃选定的token\[27 (https://arxiv.org/html/2608.23843#bib.bib15)\]，量化则以更低的精度存储键和值\[12 (https://arxiv.org/html/2608.23843#bib.bib19)\]，而低秩方法则在低维空间中表示键和值。后者的主要区别在于压缩空间的来源：一些方法分解模型投影权重\[2 (https://arxiv.org/html/2608.23843#bib.bib9),26 (https://arxiv.org/html/2608.23843#bib.bib10)\]，通常需要校准或针对模型进行离线处理；而另一些则从推理过程中产生的KV缓存构建或更新基底\[28 (https://arxiv.org/html/2608.23843#bib.bib13),3 (https://arxiv.org/html/2608.23843#bib.bib14)\]。无论哪种情况，单个基底都在广泛的缓存区域共享，这有利于主导整体重建的方向，但可能遗漏稀疏但对任务至关重要的信息。

为了在更细粒度上捕获这种局部结构，我们从PagedAttention用于管理KV缓存内存的固定大小块\[9 (https://arxiv.org/html/2608.23843#bib.bib7)\]中获得灵感，并将逻辑页面重新解释为压缩单元。这一选择是有充分依据的：我们的分析表明，在不同模型、层和KV头中，单个键和值页面内存在显著的低秩冗余（图2 (https://arxiv.org/html/2608.23843#S3.F2)）。以页面作为压缩单元，使得一个抽象概念能够同时管理低秩表示和增量缓存更新，同时页面的统一形状和独立性允许它们的分解在GPU上以批处理方式执行。

基于此设计，我们提出了PuzzleKV，一种无需训练和校准的基于页面的低秩KV缓存压缩方法。PuzzleKV在每一层和每个KV头内独立分解已完成的页面，同时将一个小的sink区域和最近窗口保持为稠密形式。其混合注意力内核直接处理稠密和分解后的页面，无需重建历史KV缓存。在自回归解码过程中，新符合条件的页面被增量转换为因子存储。

我们的主要贡献如下：

- • 我们提出了PuzzleKV，一种无需训练和校准的方法，它独立地对每一层和每个KV头中已完成的KV页面进行因子化，通过页面内的低秩表示保留每个token。
- • 我们实现了PuzzleKV，包括批量页面分解、直接在稠密和分解后的页面上计算注意力（无需重建），以及在解码过程中增量页面转换。
- • 我们在Qwen3-8B和Llama-3.1-8B-Instruct模型上，针对RULER和LongBench基准评估了PuzzleKV。在约占用完整KV存储的60%时，PuzzleKV在所有评估设置下均达到了完整KV性能的96%以上，在RULER上相对于全局SVD有显著提升，并在LongBench上达到了有竞争力的性能。结合逐因子INT4量化，它仅使用18.7%的原始存储空间即可达到超过93%的完整KV性能。

## 2相关工作

KV缓存压缩对于内存高效的LLM推理至关重要，现有方法大致可分为三类：低秩分解、量化和淘汰\[10 (https://arxiv.org/html/2608.23843#bib.bib23)\]。

#### 低秩KV缓存压缩。

低秩方法通过在低维空间中表示键和值来减少KV缓存存储。Palu\[2 (https://arxiv.org/html/2608.23843#bib.bib9)\]和LoRC\[26 (https://arxiv.org/html/2608.23843#bib.bib10)\]分解投影矩阵，允许模型缓存低维中间表示。ECKVH\[24 (https://arxiv.org/html/2608.23843#bib.bib11)\]和EigenAttention\[13 (https://arxiv.org/html/2608.23843#bib.bib12)\]从校准数据集上收集的激活构建固定的压缩基底。OjaKV\[28 (https://arxiv.org/html/2608.23843#bib.bib13)\]在线更新序列级低秩基底，而xKV\[3 (https://arxiv.org/html/2608.23843#bib.bib14)\]在预填充阶段利用跨层冗余。PuzzleKV则独立地对每一层和每个KV头中每个已完成的页面进行因子化。

#### KV缓存量化。

量化通过以更低的数值精度存储键和值来减少KV缓存内存，使用的方案包括非对称低比特量化\[12 (https://arxiv.org/html/2608.23843#bib.bib19)\]、角度分量的极坐标变换\[7 (https://arxiv.org/html/2608.23843#bib.bib21)\]和失真感知向量量化\[25 (https://arxiv.org/html/2608.23843#bib.bib22)\]。PuzzleKV则通过低秩分解降低每个KV页面的维度，并且与量化互补，因为量化可以应用于其低秩因子。

#### KV缓存淘汰。

淘汰方法通过仅保留预计将继续有用的token来减少KV缓存存储，但被淘汰的token在未来推理中将永久丢失。H2O\[27 (https://arxiv.org/html/2608.23843#bib.bib15)\]保留heavy-hitter和最近token，StreamingLLM\[21 (https://arxiv.org/html/2608.23843#bib.bib16)\]保留注意力沉点和局部窗口，SnapKV\[11 (https://arxiv.org/html/2608.23843#bib.bib18)\]从提示时的注意力中选择重要位置。与淘汰方法不同，PuzzleKV以压缩形式保留每个token。

## 3方法论

图1：PuzzleKV概览。工作流程涵盖预填充阶段的页面划分和页面级分解，随后是解码阶段的混合分页注意力和增量缓存更新。### 3.1设计原则：页面作为局部子空间

低秩KV压缩的一个关键设计选择是token共享低秩基底的粒度。一个跨越整个提示的单一基底可最大化重用，但序列本身跨越了具有不同主导方向的区域。为每个KV页面分配其自己的紧凑基底，可以让每个区域都由与其最相关的方向表示，从而产生细粒度的低秩表示，这些表示共同覆盖了序列中更丰富的方向集合，同时保持每个页面强烈的低秩性。当单个页面本身具有低秩性时，这一设计会带来回报，我们通过检查Qwen3-8B和Llama-3.1-8B-Instruct中不同秩和页面大小的KV页面重建精度来确认这一点。如图2 (https://arxiv.org/html/2608.23843#S3.F2)所示，在两种模型和所有评估的页面大小中均观察到页面级的低秩结构。在$P=32$时，尽管不同层和KV头所需的秩有所不同，但在所有评估的层和头中，它们都保持在完整页面秩之下。这些结果共同证实了页面级低秩结构是所评估模型、页面大小、层和KV头中的一致属性。

### 3.2预备知识：低秩KV缓存

#### KV缓存与注意力。

在每个Transformer层中，键和值的投影产生$K,V \in \mathbb{R}^{H_{kv} \times T \times d}$，其中$H_{kv}$是KV头的数量，$T$是序列长度，$d$是头维度。在自回归解码过程中，每个新的键-值对沿着序列维度追加到缓存。为清晰起见，省略头索引，查询$Q$的注意力计算为

$$\operatorname{Attn}(Q,K,V) = \operatorname{softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right)V$$  (1)

#### 截断SVD。

对于一个矩阵$X \in \mathbb{R}^{n \times d}$，其奇异值分解为$X=U\Sigma W^{\top}$，其中$\Sigma$中的奇异值按降序排列，其秩为$r$的截断近似保留前$r$个分量，$\widehat{X}_{r}=U_{r}\Sigma_{r}W_{r}^{\top}$。其中$U_{r} \in \mathbb{R}^{n \times r}$，$\Sigma_{r} \in \mathbb{R}^{r \times r}$，$W_{r} \in \mathbb{R}^{d \times r}$。根据Eckart-Young-Mirsky定理，$\widehat{X}_{r}$在所有秩至多为$r$的矩阵$\widehat{X}$中最小化$\|X-\widehat{X}\|_{F}$。

### 3.3方法：PuzzleKV

图1 (https://arxiv.org/html/2608.23843#S3.F1)总结了PuzzleKV的工作流程，它跨越了预填充阶段（左）和解码阶段（右），两者操作在一个混合KV缓存上。

在预填充期间，PuzzleKV：1）计算标准稠密KV缓存；2）将每一层和每个KV头划分为固定大小的页面；3）对每个已完成的页面进行因子化，同时保持注意力沉点（attention sink）和局部窗口为稠密形式。结果是一个混合缓存，其中沉点和最近的页面保持稠密，而历史页面存储为低秩因子。在解码期间：4）每个查询在稠密和分解后的页面上联合进行注意力计算，并通过在线softmax合并其部分结果，无需重建历史缓存；5）然后，新的键-值对被追加到最近的稠密页面，一旦该页面填满并离开局部窗口，它就被因子化并移至低秩存储。步骤4和5构成了稳定的解码循环：混合注意力内核读取由增量更新持续维护的混合缓存，在整个生成过程中以压缩形式保留每个token。

图2：跨模型、层、KV头和页面大小的页面级低秩结构。左图：不同秩和页面大小下的重建精度。右图：$P=32$时达到80%重建精度所需的秩，各KV头以细线表示，其均值以粗线表示。#### 页面级低秩KV缓存构建（步骤1, 2, 3）。

PuzzleKV将每个已完成的KV页面视为独立的低秩子空间，同时以压缩形式保留每个token。在预填充期间，它计算标准的稠密KV缓存，并将每一层和每个KV头中的token划分为包含$P$个token的固定大小页面。然后，PuzzleKV将两个区域保持为稠密形式：一个初始的*sink页面*，保存最初的几个token，这些token吸引了不成比例的注意力份额，保持未压缩的成本很低\[21 (https://arxiv.org/html/2608.23843#bib.bib16)\]；以及一个移动的*局部窗口*，包含最近的页面，这些页面的token被访问最频繁，并且在解码期间其当前页面仍在填充中。每个其他已完成的页面都通过截断SVD独立地进行因子化，将每个符合条件的稠密页面替换为两个因子：

$$\widehat{X}_{p} = L_{p}^{(X)} R_{p}^{(X)}$$  (2)

$$L_{p}^{(X)} = U_{p,r_X}^{(X)}, \quad R_{p}^{(X)} = \Sigma_{p,r_X}^{(X)} \left( W_{p,r_X}^{(X)} \right)^{\top}$$

其中$X \in \{K, V\}$，对于键$r_X = r_K$，对于值$r_X = r_V$。由此产生的缓存包含稠密的沉点和最近页面，以及因子化的历史页面，不丢弃任何token。不包括以稠密形式存储的页面，一个键值页面对的因子化存储比为

$$\rho_{\mathrm{page}} = \frac{(r_K + r_V)(P + d)}{2Pd}$$  (3)

为了高效地获得这些因子，PuzzleKV避免对每个$P \times d$页面调用单独的SVD。因为KV页面是短而宽的（$P < d$）且具有相同的形状，它将它们批处理起来，并从较小的$P \times P$ Gram矩阵$G_{p} = X_{p}X_{p}^{\top}$中恢复主导的左奇异子空间，取$L_{p}^{(X)} = Q_{p,r_X}$作为其前$r_X$个特征向量，取$R_{p}^{(X)} = Q_{p,r_X}^{\top}X_{p}$。由于$X_{p}X_{p}^{\top}$的特征向量就是$X_{p}$的左奇异向量，这正好得到了上述截断SVD的因子，同时映射到批处理的GPU操作上。

#### 在混合KV页面上的注意力（步骤4）。

PuzzleKV通过一个自定义的混合注意力路径在稠密和分解后的页面上计算注意力，直接从存储的因子评估每个分解后的页面，从不重建历史稠密缓存。这避免了显式重建可能创建的瞬态稠密缓冲区，先前的工作指出这是内存节省损失的一个来源\[14 (https://arxiv.org/html/2608.23843#bib.bib24)\]。对于查询$q$和分解后的页面$\widehat{K}_{p} = L_{p}^{(K)} R_{p}^{(K)}$，$\widehat{V}_{p} = L_{p}^{(V)} R_{p}^{(V)}$，PuzzleKV计算页面局部softmax状态

$$s_{p} = \frac{\left( q \left( R_{p}^{(K)} \right)^{\top} \right) \left( L_{p}^{(K)} \right)^{\top}}{\sqrt{d}}$$