PQCache: Product Quantization-Based KVCache for Long Context LLM Inference

发表时间: 2024-07 · arXiv:2407.12820 (SIGMOD 2025)

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

HAILIN ZHANG, XIAODONG JI, YILIN CHEN, FANGCHENG FU, XUPENG MIAO, XIAONAN NIE, WEIPENG CHEN, BIN CUI
Peking University, China; Purdue University, USA; Baichuan Inc., China

速读

一句话结论
本文提出了一种基于乘积量化（Product Quantization）的键值缓存（KVCache）管理系统 PQCache，通过将长上下文大语言模型推理中的选择性注意力转化为近似向量检索问题，在 InfiniteBench 评测集上相比现有方法提升了 4.60% 的得分，同时保持了极低延迟。

要解决什么问题
在大语言模型（LLM）长上下文推理中，随着输入序列变长，KVCache 的显存占用呈线性暴涨，极易突破单张 GPU 物理上限（如 7B 模型在 128K 上下文、128 批次下需 1TB 显存）。为绕开显存瓶颈，现有的“选择性注意力”方案主要分两类，但均存在致命卡点：一是 KVCache 丢弃法（如 H2O、SnapKV、PyramidKV），假设早期权重低的 Token 后续也不重要并将其丢弃，但这在真实场景中会导致关键信息永久丢失，严重损害生成质量；二是 KVCache 卸载法（如 SPARQ、InfLLM），将全量 KVCache 存放在 CPU 内存，生成新 Token 时通过代理分数检索相关键值对，但这要么引入无法与 GPU 计算掩盖的巨大通信延迟，要么依赖不切实际的块级连续性假设导致效果大跌。这个卡点直接决定了模型能否在有限 GPU 资源下，既不掉精度，又以人类可接受的阅读速度完成超长文本推理。

怎么做的
PQCache 的核心思路是将 KVCache 的存储与调用视为典型的近似最近邻检索问题，利用乘积量化（Product Quantization, PQ）技术对 Token 的键（Key）进行降维压缩和聚类，在极低通信开销下实现高精度 Token 召回。该方法由三个关键部件构成。首先是预填充阶段的异步构建。模型计算所有输入 Token 的 KVCache 时，将其异步卸载到 CPU 内存。在 CPU 上，PQCache 将每个注意力头的键向量划分为 $m$ 个子空间并进行 K-Means 聚类，生成聚类中心和极小体积的 PQ 编码。为防止 CPU 聚类拖慢 GPU，系统设计了自适应迭代策略，令最大聚类迭代次数 $T_{max}$ 满足聚类时间等于 GPU 计算时间：

$$T_{max} = \frac{\gamma_2 \cdot s^2 + \beta_2 \cdot s + \alpha_2 - \alpha_1}{\beta_1 \cdot s}$$


其中 $s$ 为序列长度，$\alpha, \beta, \gamma$ 为拟合系数。其次是解码阶段的近似检索。体积固定的聚类中心常驻 GPU，下一层的 PQ 编码与当前层计算并行预取。生成新 Token 时，GPU 计算查询向量与聚类中心的内积，结合 PQ 编码还原历史 Token 的近似注意力分数，精准定位得分最高的 $k$ 个目标 Token。最后是块级 GPU 缓存。为避免每次都从 CPU 拉取这 $k$ 个 Token 的完整键值对，PQCache 在 GPU 显存中维护了基于最近最少使用或最不经常使用策略的块级缓存。高频访问的 Token 块留在 GPU 中，未命中部分才通过 PCI-e 异步拉取，最大化掩盖通信延迟。

效果如何
实验基于 Llama-3.1-8B（128K 上下文）、Mistral-7B-Instruct-v0.2（32K 上下文）及 Llama-3.1-70B 模型展开，硬件为单张 RTX 4090 24GB GPU 搭配双路 Intel Xeon CPU。对比基线涵盖两大路线：代表 KVCache 丢弃路线的 H2O、SnapKV、PyramidKV，以及代表 KVCache 卸载路线的 SPARQ、InfLLM，并设置了精确检索基线 Oracle。量化结果表明，在 InfiniteBench 评测集中，当仅使用 1/10 的 Token 参与计算且额外通信量仅为键显存的 1/64 时，PQCache 平均得分比现有最佳基线高 4.60%，与无压缩全量计算表现几乎一致。在 LongBench 评测集中，当打乱问题位置（不放在末尾）时，依赖位置假设的 SnapKV 和 PyramidKV 性能大幅衰减，而 PQCache 依然稳健，得分反超基线 7.10%。在 131,000 长度的大海捞针测试中，PQCache 实现了近乎完美的检索准确率。系统效率上，单步解码延迟远低于 SPARQ，达到快于人类阅读速度的实用水平。该方法的代价与局限在于：需消耗额外 CPU 资源完成聚类；当输入序列极短时，GPU 计算时间不足以掩盖 CPU 聚类时间，自适应策略会削减迭代次数从而影响量化精度；此外，对于超长输出或多轮对话场景，仅基于输入构建的 PQ 结构可能无法捕捉新 Token 特征，需引入定期重建机制。

主要贡献

随着大型语言模型（LLMs）领域的不断发展，推理过程中的上下文长度正在稳步增长。由于GPU内存有限，LLM推理中token的中间表示——键值缓存（KVCache），现在已经成为主要的内存瓶颈。当前的方法通过选择性地确定合适的键和值来进行LLM中的自注意力计算以解决此问题。然而，这些方法要么无法保持模型质量，要么导致很高的服务延迟。

本文从数据管理社区中普遍存在的高级嵌入检索技术中汲取灵感，将KVCache的存储和检索视为一个典型的嵌入检索（Embedding Retrieval）问题。作者提出了PQCache，该方法采用乘积量化（Product Quantization, PQ）来管理KVCache，在保持模型质量的同时确保低服务延迟。在预填充（Prefilling）阶段，对每个LLM层和头的token键应用PQ。在自回归解码（Decoding）阶段，使用PQ编码（codes）和质心（centroids）来近似识别重要的前置token，然后获取相应的键值对进行自注意力计算。通过对重叠计算（overlapping）和缓存（caching）的精心设计，该方法将两个阶段的任何额外计算和通信开销降至最低。广泛的实验表明，PQCache实现了有效性和效率的双赢，在InfiniteBench上比现有方法提高了$4.60\%$的分数，并在预填充和解码阶段均保持了低系统延迟。

不同批处理大小、模型大小和序列长度下的KVCache内存大小和理论CPU-GPU传输延迟
Fig. 1 证实了随着序列长度的增加，KVCache的内存消耗急剧增加，远超单个GPU的内存容量。

信息检索与带有选择性注意力的LLM推理之间的比较
Fig. 2 展示了信息检索过程与LLM选择性注意力推理过程的高度一致性。

背景知识与关键观察

大型语言模型推理

LLM推理概述与自注意力机制：大型语言模型（LLM）由输入词汇嵌入、堆叠的Transformer层和用于输出的token分类器组成。自注意力模块是Transformer层的关键组件，促进不同token之间的交互和信息聚合。在每一层中，形状为$(n, s, d)$的输入被分别投影并转置为查询（query）、键（key）和值（value），形状均为$(n, h, s, d_h)$，通常满足$d = h * d_h$。注意力机制将查询和键相乘，应用下三角因果掩码以限制查询只能访问前面的键，并执行softmax以获得形状为$(n, h, s, s)$的注意力分数。然后使用这些分数对值进行加权求和，产生形状为$(n, h, s, d_h)$的输出。在长上下文推理中，查询、键和值之间操作的$O(s^2)$空间复杂度是不切实际的。因此，现代LLM通常采用FlashAttention【索引21，FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning + 2024 + ICLR】【索引22，FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness + 2022 + NeurIPS】，它利用分块矩阵乘法和softmax，仅使用$O(s)$的空间复杂度进行注意力计算，但时间复杂度仍为$O(s^2)$。为了减轻内存和计算负担，通常应用分组查询注意力（GQA），它为键和值使用较少数量的头$h_{kv}$。

预填充与解码阶段的KVCache生成：在LLM推理期间，模型的每次执行都会以自回归的方式生成一个新token。LLM的第一次遍历称为“预填充”，随后的遍历称为“解码”。在预填充阶段，自注意力模块计算所有输入token的查询、键和值，并将键值对作为KVCache存储以供后用。由于因果掩码的存在，后面的token不会影响前面的token，因此在自回归解码阶段，只需要计算最后生成的token的查询、键和值。此过程利用KVCache中先前token的键和值，并计算形状为$(n, h, 1, s)$的注意力分数。同时，新生成的键和值被添加到KVCache中。因此，KVCache的内存消耗随序列长度线性扩展，导致在长上下文输入和输出场景中出现内存瓶颈。
LLM推理概述

乘积量化 (Product Quantization)

乘积量化原理与应用：PQ【索引44，Product Quantization for Nearest Neighbor Search + 2011 + IEEE TPAMI】被提出用于促进高效的近似最近邻搜索（ANNS），即在给定查询嵌入向量的情况下，从大量候选者中检索相关的嵌入向量。PQ包含两个阶段：索引构建和搜索。在构建期间，PQ将每个候选嵌入划分为$m$个分区，本质上是将原始嵌入空间分解为$m$个独立的子空间。每个子空间分别进行K-Means聚类以对子嵌入进行分组，产生$2^b$个质心。然后为每个嵌入分配$m$个编码，每个编码有$b$位，指示其所属的质心。这些紧凑的PQ编码允许使用质心重建近似嵌入，从而显着降低内存需求。在搜索期间，查询嵌入被划分为$m$个分区。计算每个分区内与质心的相似度，并通过PQ编码进行聚合，从而绕过与每个嵌入进行完整相似度计算的需要。
PQ构建和搜索的概述

GPU-CPU 内存层次结构

现代深度学习任务中的GPU-CPU内存利用：现代深度学习任务严重依赖GPU来执行计算密集型操作。GPU-CPU结构形成了一个典型的内存层次结构：更昂贵的GPU内存为计算提供更快的内存I/O速度，而通过PCI-e或NVLink连接的CPU内存提供较低的带宽。这种硬件设置允许在GPU内存不足时利用可用的CPU内存以提高效率。随着模型参数的增加和对中间结果存储（如KVCache）需求的增长，CPU经常被用来分担内存负载。此背景下的主要挑战是如何有效地调度内存I/O（即GPU-CPU通信）并结合GPU计算，以有效地隐藏相关的开销。

方法细节

概览

选择性注意力的理论基础：我们设计PQCache将所有KVCache保留在CPU中，并选择性地获取相关的键值对进行自注意力计算。在长上下文推理场景中，整个KVCache对于内存层次结构中的注意力计算和I/O通信来说都太大了。根据先前的研究【索引59，Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time + 2023 + NeurIPS】【索引112，H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models + 2023 + NeurIPS】，注意力分数是衡量先前token重要性或相关性的合适指标。注意力分数通常遵循幂律分布，表明一小部分token比大多数其他token更重要。因此，我们可以仅包含那些具有高分数的token进行自注意力计算。遵循先前的工作，我们还将初始token和最近的token（称为局部token）包含在注意力计算中。
注意力分数的分布

PQCache的整体工作流程：识别具有最高分数的Top-$k$键的过程从根本上构成了一个近似最近邻搜索（ANNS）操作。基于上述观察，我们设计了PQCache，将所有KVCache卸载到CPU，并在解码阶段仅获取相关token的键值对。由于计算所有先前token的精确注意力分数涉及昂贵的I/O通信，我们利用轻量级的乘积量化（PQ）方法。在PQCache中，我们在预填充阶段构建PQ，并在解码阶段利用PQ。具体操作步骤如下：
* Step 1：在预填充阶段，计算每一层所有输入token的键和值。键和值的形状均为$(n, h_{kv}, s, d_h)$。它们随后被异步卸载到CPU，该过程可以与后续计算重叠。
* Step 2：在CPU上执行PQ构建。将每个头的键向量维度$d_h$划分为$m$个维度为$d_m$的子空间，其中$d_m = d_h / m$。对于划分后的向量，分别进行K-Means聚类，产生形状为$(m, 2^b, d_m)$的质心和形状为$(s, m)$的PQ编码。
* Step 3：在解码阶段，在计算前一个Transformer层时，预取PQ质心和编码。质心体积小，可在整个推理过程中保留在GPU上。
* Step 4：在GPU上执行PQ搜索：首先在查询和PQ质心之间进行矩阵乘法，然后与PQ编码聚合以获得所有token的近似分数。随后使用近似分数识别Top-$k$相关的token。
* Step 5：使用PQ分数，从CPU或GPU缓存中获取近似Top-$k$个token的键值对。
* Step 6：使用检索到的token继续进行选择性自注意力计算。对于新生成的token，当它们从局部token的滑动窗口中被驱逐时，将根据它们最近的质心被分配PQ编码。
PQCache的概览

复杂度分析

预填充与解码阶段的复杂度：在预填充期间，注意力计算保持不变。利用FlashAttention，空间复杂度为$O(sd)$，时间复杂度为$O(s^2 d/h + s d^2)$。使用PQCache后，额外的K-Means聚类过程的平均时间复杂度为$O(s h_{kv} m d_m 2^b T)$，其中$T$是K-Means迭代次数，其复杂度与输入序列长度呈线性关系。在解码阶段，原始LLM的时间复杂度为$O(sd + d^2)$。使用PQCache时，首先将查询与PQ质心相乘，然后使用PQ编码进行收集和归约，接着找到Top-$k$最大分数，最后计算选择性注意力及后续FFN。总体时间复杂度变为$O(2^b d^2 / (h m) + h_{kv} m s + k d + d^2)$。PQ质心和编码的内存复杂度为$O(h_{kv} m s + h_{kv} 2^b \cdot d_h)$。在长上下文场景中，当前的乘数$h_{kv} m$远小于原始的$d$（例如在7B模型中，$h_{kv}=8$, $m=2$, $d=4096$）。由于$h_{kv} m$尺寸较小，允许PQ结构的通信与GPU计算重叠，从而实现更高效的解码。

预填充阶段

GPU计算与CPU聚类的并行与自适应迭代：在预填充阶段，GPU计算、GPU到CPU的卸载通信以及PQ构建可以并发执行。每一层中，输入被投影为查询、键和值，键和值随后异步传输到CPU，同时GPU继续执行注意力和前馈网络操作。获取键和值后，CPU上的独立进程启动多个聚类进程进行PQ构建。为了实现按需同步，我们在下一个token解码阶段的相同Transformer层等待PQ构建完成。由于注意力计算时间随序列长度呈二次方增长，而通信和聚类时间呈线性增长，当序列长度足够长时，计算可以完全覆盖其他两个操作。然而，当序列长度不足时，聚类持续时间可能超过GPU计算。为了解决这个问题，我们提出了一种自适应K-Means方法，限制聚类的最大迭代次数。K-Means聚类时间拟合为线性曲线：
$Time_{clus} = \alpha_1 + \beta_1 \cdot s T$
GPU计算时间拟合为二次关系：
$Time_{comp} = \alpha_2 + \beta_2 \cdot s + \gamma_2 \cdot s^2$
确保聚类不干扰计算的最大迭代次数$T_{max}$应满足$Time_{clus} = Time_{comp}$，由此得出：
$T_{max} = \frac{\gamma_2 \cdot s^2 + \beta_2 \cdot s + \alpha_2 - \alpha_1}{\beta_1 \cdot s}$
在实践中，我们也会对$T_{max}$进行裁剪，以确保迭代次数既不会太小也不会太大。
PQCache与顺序调度的比较
预填充阶段的执行时间

# Algorithm 1: PQCache prefilling phase
Input: Tokenized query X
Output: Output token; KVCache; PQ structures

1: ProcComm = ProcPQ = KVCache = PQ = {}
2: L_x = X.length
   # Prefilling computation on GPU
3: for i <- 0 to L - 1 do
4:     Q, K, V = Projection_i(X)
5:     ProcComm.append(AsyncGPU2CPU(K, V))
6:     X = AttnFFN_i(Q, K, V)
7: end
8: Token = Classifier(X)
   # Launch PQ construction on CPU
9: for i <- 0 to L - 1 do
10:    K, V = ProcComm[i].sync()
11:    KVCache.append((K, V))
12:    ProcPQ.append(AsyncPQConstruct(K, L_x))
13: end
   # Wait for PQ construction in the next decoding phase
14: for i <- 0 to L - 1 do
15:    Centroids, Codes = ProcPQ[i].sync()
16:    PQ.append((Centroids, Codes))
17: end
   # Some decoding logic here
18: return Token, KVCache, PQ

解码阶段

PQ结构预取与块级缓存机制：在解码阶段，构建的PQ结构被每一层的注意力模块利用。PQ质心直接存储在GPU上，因为它们消耗的内存固定且微不足道。随着前一层计算的进行，下一层的PQ编码可以并行预取。我们将整个KVCache分为三个部分：初始token、中间token和局部token。初始和局部token的键值对直接参与注意力计算并存储在GPU上，而中间token的键值对存储在CPU上。在每个Transformer层中，局部token首先驱逐最早的token，生成其PQ编码，并异步将其键和值卸载到CPU。新生成的键和值随后被添加到局部token中。从CPU接收PQ编码后，进行PQ搜索以识别Top-$k$个token。唯一不能重叠的通信是检索Top-$k$相关token，因为它依赖于前面的PQ搜索。为了容纳那些在推理过程中始终重要的关键token，我们设计了一个具有最近最少使用（LRU）或最不常用（LFU）驱逐策略的块级GPU缓存。在识别出Top-$k$个token时，我们首先检查缓存中是否可用。为了最小化缓存查找开销，我们利用块级粒度而不是token级粒度。我们将所有token划分为块，并将常用的块存储在GPU缓存中。每次检索期间，我们使用包含最多Top-$k$个token的Top-$k_{cache}$个块来更新缓存。

# Algorithm 2: PQCache decoding phase
Input: Last generated token X
State: InitKV; MidKV; LocalKV; PQ codes; PQ centroids
Output: Output token

   # First pre-fetch first layer's PQ codes
1: Proc = AsyncCPU2GPU(PQCodes[0])
2: for i <- 0 to L - 1 do
       # Evict previous token from local tokens
3:     EvictK, EvictV = LocalKV.pop(0)
4:     NewCode = PQArgNearest(EvictK)
       # On CPU, MidKV and PQ will be updated
5:     AsyncGPU2CPU((NewCode, EvictK, EvictV)) 
       # Get Q, K, V for last token
6:     Q, K, V = Proj_i(X)
       # Get KV for current attention
7:     LocalKV.append((K, V))
8:     Codes = Proc.sync()
9:     Centroids = PQCentroids[i]
       
10:    if i != L - 1 then 
           # Pre-fetch next layer's PQ codes
11:        Proc = AsyncCPU2GPU(PQCodes[i + 1])
12:    end
       
13:    Codes.append(NewCode)
14:    TopkToken = PQSearch(X, Centroids, Codes)
15:    TopkKV = SyncFetchKV(TopkToken)
16:    AllKV = InitKV + TopkKV + LocalKV
       # Continue attention and FFN computation
17:    X = AttnFFN_i(Q, AllKV)
18: end
19: Token = Classifier(X)

实验环境

实验结果

补充细节

超参数调整指导：在PQCache中，检索集大小（Top-$k$）越大通常性能越好，但受限于GPU内存。对于PQ设置，建议从$m$和$b$的乘积开始调整，因为它代表总向量空间的大小$(2^b)^m = 2^{mb}$。最佳规模既不能太小以至于损害向量表示，也不能太大而使聚类过程过载。

多请求与多GPU场景：PQCache可以直接扩展以服务多个请求，通过增加批处理大小维度进行并行计算。如果CPU成为瓶颈，可以通过增加廉价的CPU资源来解决。这同样适用于多GPU场景。

长输出序列与多轮对话：对于长输出场景，基于输入的PQ结构可能无法捕获新生成的token信息。解决方案是定期重建PQ。在多轮对话中，可以采取每次用户输入时使用所有先前token重新进行预填充，或者对每个输入分别执行K-Means并将输出的PQ编码添加到相应输入的策略。

与LLM推理/服务系统的结合：PQCache旨在减少注意力计算的token数量，这与现有的LLM推理系统（如Orca、vLLM的PagedAttention、DistServe的预填充-解码解耦）是正交且兼容的，未来有望在现代LLM推理系统中得到实际应用。

结论

本文提出了 PQCache，这是一种用于有效且高效的长上下文 LLM 推理的系统-算法协同设计方法。该方法结合了嵌入检索技术 PQ，以减少内存和计算负担，并利用 PQ 编码和质心促进注意力模块中重要 token 的高效近似最近邻（ANN）搜索。通过细致的重叠和缓存设计，PQCache 将开销降至可忽略的水平。广泛的实验表明，与现有方法相比，PQCache 在 InfiniteBench 上将模型质量提高了 $4.60\%$，同时保持了较低的系统延迟，为下一代 LLM 推理建立了一个重要里程碑。

参考文献引用汇总