KVPR: Efficient LLM Inference With I/O-Aware KV Cache Partial Recomputation

发表时间: 2025-07 · ACL 2025 Findings

原文: https://aclanthology.org/2025.findings-acl.997

Chaoyi Jiang, Lei Gao, Hossein Entezari Zarch, Murali Annavaram / University of Southern California

速读

一句话结论 本文提出了 KVPR，一种通过在 GPU 上部分重计算 KV Cache 并与 CPU-GPU 异步传输相重叠的 I/O 感知大语言模型推理方法，在解码阶段将延迟降低了最高 35.8%，吞吐量提升了最高 46.2%。

要解决什么问题 大语言模型在自回归解码阶段需要缓存键值对（KV Cache）以避免重复计算，但随着批处理大小、序列长度和模型规模的增加，KV Cache 的显存占用会迅速超出 GPU 的物理容量。现有的高性价比方案是将 KV Cache 卸载到廉价且容量大的 CPU 内存中，但这会将系统瓶颈转移到 CPU 与 GPU 之间的 PCIe 总线带宽上。由于 PCIe 的传输延迟通常比 GPU 重新计算这些特征的延迟高出一个数量级（例如在 PCIe 4.0 下，传输延迟可能是计算延迟的数十倍），传统的异步流水线方法无法将漫长的 I/O 传输时间与短暂的 GPU 计算时间完全重叠，导致 GPU 存在大量空闲等待时间。另一方面，部分现有路线尝试让 CPU 直接分担注意力计算来掩盖传输开销，但这在单 CPU 挂载多 GPU 的分布式部署中会严重透支 CPU 算力，使其成为新的系统卡点。

怎么做的 核心思路是“部分重计算与异步传输重叠”。与其通过缓慢的 PCIe 总线把完整的 KV Cache 从 CPU 传到 GPU，不如只传一小部分体积更小的历史激活值（Activations）到 GPU，让 GPU 利用这些激活值重新计算出对应的局部 KV Cache；与此同时，CPU 通过 PCIe 并发传输剩余的那部分 KV Cache。因为激活值的体积远小于生成的 KV Cache，且 GPU 算力极强，这种做法能人为增加 GPU 的计算耗时并减少 PCIe 的传输数据量，从而让两者的耗时完美对齐，消除 GPU 的空闲等待。该方法由三个关键部件构成：分析器（Profiler）、调度器（Scheduler）和运行时（Runtime）。分析器负责收集当前系统的硬件特征，如 PCIe 传输带宽 $v_{com}$ 和 GPU 处理速度 $v_{gpu}$。调度器负责计算最优的切分点 $l$（即前 $l$ 个 Token 走重计算，后 $s'-l$ 个 Token 走直接传输）。假设当前序列长度为 $s'$，批大小为 $b$，隐藏层维度为 $h$，重计算前 $l$ 个 Token 的 KV Cache 所需的时间为：

$$t_{recomp}^i = \frac{4 \times b \times l \times h^2}{v_{gpu}}$$


处理单层的总耗时由激活值传输时间，加上“重计算时间”与“剩余 KV Cache 传输时间”两者的最大值构成：

$$t^i = \frac{M_{X^i[0:l]}}{v_{com}} + \max\left(t_{recomp}^i, \frac{M_{KV^i[l:s']}}{v_{com}}\right)$$
调度器通过求解 $\min t^i$ 的整数线性规划问题，自适应地得出每个生成步下的最优 $l$ 值。最后，运行时模块负责执行这一策略。为了防止 GPU 在等待所有注意力权重加载时无事可做，运行时设计了细粒度的流水线：优先将生成 KV Cache 必需的 $W_K$ 和 $W_V$ 权重传给 GPU 以立即启动重计算，同时在后台异步传输剩余的 $W_Q$ 和 $W_O$ 权重，确保重计算过程与权重加载完美掩盖。

效果如何 实验在配备单张 NVIDIA A100（40GB 显存）和 AMD EPYC 64 核 CPU 的服务器上搭建，两者通过带宽为 32 GB/s 的 PCIe 4.0 x16 连接。测试模型主要为 OPT-6.7B、OPT-13B 和 OPT-30B。实验分为两个主要场景并对比了不同路线的基线方法：在以降低延迟为目标的场景下（模型权重保留在显存，仅卸载 KV Cache），对比了 DeepSpeed Inference 和 Hugging Face Accelerate；在以最大化吞吐量为目标的场景下（权重和 KV Cache 均卸载到 CPU，采用列式调度），对比了 FlexGen；此外还对比了代表“CPU 辅助计算”路线的 FastDecode。量化结果显示，在延迟导向测试中（OPT-6.7B，输入 128 且生成 128 个 Token），KVPR 比 Hugging Face Accelerate 的延迟降低了 35.8%。在吞吐量导向测试中，KVPR 相比 FlexGen 在 OPT-13B 上实现了最高 46.2% 的吞吐量提升，并在解码阶段将 GPU 利用率从 85% 压榨到了 99%。在多 GPU 共享单 CPU 的并发测试中，KVPR 避免了 FastDecode 因透支 CPU 算力而导致的性能崩塌，展现出更好的扩展性。该方法的代价是增加了 GPU 的纯计算开销（在运行耗时占比中从 2.3% 增加到 13.3%）。作者承认的局限性包括：当前仅支持单卡或数据并行架构，尚未扩展到张量并行或模型并行等高级分布式系统以支持更大模型；此外，系统分析器目前仅在推理开始前进行静态采样，无法在多租户等硬件资源动态波动的环境中进行自适应调整。

A1 主要贡献

大型语言模型（LLM）的推理计算量巨大。为了降低自回归解码的成本，通常使用键值（KV）缓存来存储中间激活值，从而显著降低生成Token的计算开销。然而，KV缓存所需的内存增长迅速，经常超出GPU内存的容量。一种具有成本效益的替代方案是将KV缓存卸载到CPU内存中，这虽然缓解了GPU内存压力，但将系统瓶颈转移到了CPU和GPU之间有限的PCIe连接带宽上。现有方法（如计算与I/O重叠、CPU-GPU异构执行）试图解决这些问题，但受到过量数据移动和对CPU能力依赖的阻碍。随着KV缓存规模的增长或GPU计算能力的提升，完全隐藏PCIe通信延迟变得极具挑战性。

为了解决这一问题，本文提出了KVPR，一种高效的I/O感知LLM推理方法。其核心创新点在于：
1. KV缓存部分重计算与传输重叠：CPU首先传输一小部分激活值，GPU利用这些激活值开始重计算部分KV缓存。在GPU重计算部分KV缓存的同时，CPU并发地通过PCIe传输剩余的KV缓存。该方法将GPU重计算与KV缓存传输重叠，以最小化GPU空闲时间并最大化推理性能。
2. 全自动化的执行框架：KVPR集成了三个模块以实现全自动化：收集输入特征和系统硬件信息的分析器（Profiler）模块；利用线性规划优化计算和通信工作负载分配的调度器（Scheduler）模块；以及高效执行派生执行计划的运行时（Runtime）模块。

实验结果表明，与最先进的方法相比，KVPR在解码期间的延迟降低了最高 $35.8\%$，吞吐量提高了最高 $46.2\%$。

为了评估通信开销的影响，作者设置了一个使用NVIDIA A100 GPU的LLM推理服务系统（如图1所示）。表1展示了基于该系统的PCIe传输时间和GPU计算延迟对比。结果表明，PCIe延迟超过KV缓存重计算延迟一个数量级以上。因此，在KV缓存存储于CPU DRAM的系统中，漫长的传输时间会导致GPU空闲，严重影响推理效率。

图1：配备A100 GPU的LLM推理系统。
图1：配备A100 GPU的LLM推理系统。

表1：基于图1系统的不同KV缓存大小的PCIe延迟和计算延迟（Batch Size=32, Sequence Length=1024, FP16）

Model Hidden Dim KV Cache (MB) PCIe Latency (ms) Comp. Latency (ms)
OPT-6.7B 4,096 512 15.6 0.3509
OPT-13B 5,120 640 19.5 0.4388
OPT-30B 7,168 896 27.3 0.6143

A3 背景知识/关键Observation/设计原则

LLM推理过程
仅解码器（Decoder-only）LLM的推理过程采用自回归方法顺序生成Token。它包含两个阶段：预填充（Prefilling）阶段和解码（Decoding）阶段。
在预填充阶段，第 $i$ 个解码器层的输入表示为 $X^i \in \mathbb{R}^{b \times s \times h}$，其中 $i \in \{1, \dots, n\}$， $b$ 是批量大小（batch size）， $s$ 是提示词长度（prompt length）， $h$ 是输入嵌入维度。多头注意力（MHA）块通过对 $X^i$ 进行线性投影来计算一组查询（$Q$）、键（$K$）和值（$V$）：

$$Q^i = X^i \cdot W_Q^i, \quad K^i = X^i \cdot W_K^i, \quad V^i = X^i \cdot W_V^i$$


其中 $W_Q^i, W_K^i, W_V^i \in \mathbb{R}^{h \times h}$ 是投影矩阵。生成的 $K^i$ 和 $V^i$ 被存储在KV缓存中。
MHA中的自注意力分数计算如下：

$$Z^i = \mathrm{softmax}\left(\frac{Q^i (K^i)^T}{\sqrt{d_{\mathrm{head}}}}\right) \cdot V^i$$
其中 $d_{\mathrm{head}}$ 表示每个注意力头的维度。最后，对注意力分数应用线性投影以产生MHA块的输出：
$$O^i = Z^i \cdot W_O^i$$
其中 $W_O^i \in \mathbb{R}^{h \times h}$ 是投影矩阵。
MHA块之后是前馈神经网络（FFN），它包含两个全连接层并在其间应用非线性激活函数。它处理注意力输出 $O^i$ 以生成下一个解码器层的输入：
$$X^{i+1} = \sigma(O^i \cdot W_1^i) \cdot W_2^i$$
其中 $W_1^i \in \mathbb{R}^{h \times d_{\mathrm{FFN}}}$ 和 $W_2^i \in \mathbb{R}^{d_{\mathrm{FFN}} \times h}$ 是两个线性层的权重矩阵，$\sigma(\cdot)$ 表示激活函数。

在解码阶段，第 $i$ 个解码器层接收单个Token $x^i \in \mathbb{R}^{b \times 1 \times h}$。KV缓存通过将新计算的键值对与现有键值对拼接来进行更新：

$$K^i = \mathrm{concat}(K^i, x^i \cdot W_K^i)$$

$$V^i = \mathrm{concat}(V^i, x^i \cdot W_V^i)$$
解码阶段中剩余的注意力和前馈计算与预填充阶段相同。

LLM推理调度
针对大容量KV缓存存储在CPU DRAM并在需要时提取到GPU内存的系统，存在不同的调度策略：
* 逐行调度（Row-by-row schedule）：一次处理一个Batch，采用逐层执行。如果模型权重也被卸载到CPU，则单层的KV缓存和模型权重被传输到GPU，处理当前Batch，然后清除。此过程逐层重复，直到生成一个Token。该方法适用于以最小化延迟为首要目标的场景。
* 逐列调度（Column-by-column schedule）：通过增加有效批量大小（批次数乘以批量大小）来并行处理更多序列，以最大化吞吐量为目标，代价是延迟增加。模型权重被卸载到CPU以容纳大批量。单层的模型权重和KV缓存被传输到GPU并处理第一个Batch。随后，权重保留在GPU中，继续处理后续Batch。当一组Batch在第一层处理完毕后，再整体移动到第二层。

A2 方法细节

设计概述
为了缓解PCIe压力并提高GPU计算利用率，KVPR在将部分KV缓存传输到GPU的同时，在GPU上重计算剩余的部分KV缓存。KVPR包含三个主要模块：分析器（Profiler）、调度器（Scheduler）和运行时（Runtime）。用户配置包括性能目标（延迟或吞吐量）、数据参数（提示词长度、生成长度、批量大小）和模型信息。分析器模块收集系统统计信息（如PCIe带宽和GPU处理速度）。调度器模块利用这些信息，通过求解线性规划问题计算出最佳的KV缓存重计算分割点，旨在最大化计算与通信的重叠。运行时模块利用该执行策略处理用户输入，管理内存分配和数据传输流。
图2：KVPR设计概述。用户配置和分析信息提供给调度器，调度器计算出最佳的KV缓存重计算比例。然后运行时通过重叠数据传输和GPU计算来提高推理效率。

调度器模块（Scheduler Module）
* 带有KV缓存部分重计算的逐行调度：当性能目标是最小化延迟时，调度器启动逐行执行计划。在传统的卸载流水线中，KV缓存和模型权重均卸载到CPU，由于KV缓存体积大于MHA权重，它在异步传输中较晚到达GPU。在KVPR中，CPU不传输整个KV缓存，而是首先传输对应的输入激活值，GPU利用这些激活值重计算部分KV缓存，同时剩余的KV缓存通过PCIe异步传输到GPU。随后GPU将重计算的KV缓存与传输的KV缓存合并以执行MHA计算（如图3(b)所示）。
图3：两种卸载流水线的比较。(a) 带有异步数据传输的逐行调度传统卸载流水线。(b) 带有KV缓存部分重计算的逐行调度卸载流水线。
* 带有KV缓存部分重计算的逐列调度：当性能目标是最大化吞吐量时，采用逐列执行计划。该方法通过在多个Batch间复用模型权重来适应大批量推理。一旦Batch 0的KV缓存传输完毕，Batch 1的激活值就开始传输，同时GPU开始计算Batch 0的MHA。与逐行调度不同，逐列调度在同一层上处理多个Batch，因此对应于重计算KV缓存的激活值必须被存储，直到该Batch的生成完成（如图4所示）。
图4：旨在最大化吞吐量的带有KV缓存部分重计算的逐列调度卸载流水线。
* 确定最优KV缓存划分点：目标是确定最优划分点 $l$，即在GPU上重计算的KV缓存与从CPU传输的KV缓存之间的分割比例。给定当前序列长度 $s'$（大于提示词长度 $s$），第 $i$ 层传输到GPU的激活值表示为 $X^i[0:l]$，其中 $0 \leq l \leq s'$。后续Token的剩余KV缓存表示为 $K^i[l:s']$ 和 $V^i[l:s']$。
这些激活值的内存使用量为：

$$M_{X^i[0:l]} = b \times l \times h \times p$$

$$M_{KV^i[l:s']} = 2 \times b \times (s' - l) \times h \times p$$
重计算 $X^i[0:l]$ 的KV缓存需要：
$$K^i[0:l] = X^i[0:l] \cdot W_K^i$$
$$V^i[0:l] = X^i[0:l] \cdot W_V^i$$
GPU上的重计算需要浮点运算量为：
$$N_{KV^i[0:l]} = 4 \times b \times l \times h^2$$
因此，KV缓存的重计算时间 $t_{recomp}^i$ 为：
$$t_{recomp}^i = \frac{N_{KV^i[0:l]}}{v_{gpu}}$$
其中 $v_{gpu}$ 表示GPU处理速度。总处理时间 $t^i$ 为：
$$t^i = \frac{M_{X^i[0:l]}}{v_{com}} + \max\left(t_{recomp}^i, \frac{M_{KV^i[l:s']}}{v_{com}}\right)$$
其中 $v_{com}$ 表示激活值和KV缓存的数据传输速度。
目标是确定使总处理时间 $t^i$ 最小化的最优 $l$，这转化为一个线性规划问题：
$$\min \quad t^i$$
$$\mathrm{s.t.} \quad 0 \leq l \leq s \quad \forall i \in \{1, \dots, n\}$$
最优划分点 $l$ 依赖于生成过程中不断增加的当前序列长度 $s'$，因此必须自适应地确定。由于只有一个整数变量，求解该线性规划问题的计算开销可忽略不计。如果省略上述总时间公式中的第一项，该问题就简化为逐行调度的情况。

运行时模块（Runtime Module）
* 异步重叠：为了实现GPU计算与CPU-GPU通信的并发执行，运行时模块采用了包含六个进程的通信并行策略：权重加载、KV缓存加载、激活值加载、重计算激活值加载、KV缓存存储和激活值存储。通过结合双缓冲（double buffering）和预取（prefetching）技术，它可以同时加载下一层的权重，检索用于KV缓存重计算的激活值和下一个Batch的KV缓存，同时存储上一个Batch的缓存和激活值，并处理当前Batch。
* 锁页内存（Pinned memory）：为了优化数据传输，借鉴先前的工作【索引编号1，Sheng 等人 2023】和【索引编号2，Yu 等人 2024】，KVPR对传输到GPU的重计算激活值和权重使用锁页CPU内存。这避免了数据的换入换出，实现了更快且异步的传输。
* 隐藏KV缓存部分重计算：如果KV缓存和模型权重都被卸载，且传输的KV缓存大小小于模型权重大小，粗粒度的计算流水线可能会降低推理性能，因为重计算必须等待所有MHA权重（$W_Q, W_K, W_V, W_O$）完全加载后才能开始（如图5(a)）。由于KV缓存重计算仅需要 $W_K$ 和 $W_V$，KVPR实现了一种细粒度的MHA流水线，优先加载 $W_K$ 和 $W_V$。一旦这些权重可用，KV缓存重计算立即开始（如图5(b)）。随后使用 $W_Q$ 和 $W_O$ 进行MHA计算。这种方法有效地将KV缓存重计算与权重加载重叠，确保在最坏情况下（即受限于权重加载时），该方法的表现也不会差于基线。
图5：MHA层中不同粒度级别的卸载流水线比较。(a) 带有延迟KV缓存部分重计算的粗粒度卸载流水线。
图5：(b) 将KV缓存重计算与权重加载重叠的细粒度卸载流水线。

A4 实验环境

A4 实验结果

A5 结论

本文提出了KVPR，这是一种旨在加速KV缓存加载的高效CPU-GPU I/O感知LLM推理方法。KVPR通过利用KV缓存部分重计算，最小化了CPU和GPU之间的数据传输。通过将这种重计算与数据传输重叠，KVPR显著减少了GPU空闲时间并提高了整体推理性能。未来的工作可以将该方法扩展到容忍从远程网络存储加载KV缓存，或扩展到大型多GPU基础设施，以进一步增强其在多样化部署场景中的适用性和性能。
局限性：目前该方法仅限于单GPU和数据并行的多GPU推理，尚未扩展到模型或张量并行等高级分布式系统。此外，当前实现仅在推理开始时进行系统分析，假设硬件条件是静态的，未来加入动态分析和运行时自适应优化将增强其在异构或多租户环境中的鲁棒性。

A6 附录

Algorithm 1 KV Cache Partial Recomputation with Overlapping
for i = 1 to generation_length do
    for j = 1 to num_layers do
        for k = 1 to num_GPU_batches do
            // Load the weight of the next layer
            load_weight(i, j + 1, k) 
            // Load the activation for KV cache recomputation of the next batch
            load_activation_recompute(i, j, k + 1) 
            // Load the KV cache and activation of the next batch
            load_cache(i, j, k + 1) 
            load_activation(i, j, k + 1) 
            // Compute this batch
            compute(i, j, k) 
            // Store the KV cache and activation of the previous batch
            store_activation(i, j, k − 1) 
            store_cache(i, j, k − 1)

参考文献引用汇总