在这篇文章中，我们将介绍 vLLM 0.11.0 中引入的全新 KV 缓存卸载功能。我们将重点探讨卸载到 CPU 内存（DRAM）及其对提高整体推理吞吐量的益处。在博客的后半部分，我们将深入探讨在优化主机到设备（Host-to-Device）和设备到主机（Device-to-Host）传输吞吐量以实现 KV 卸载方面的研究。

动机

服务 LLM 模型是一项计算复杂度极高的操作，其核心涉及计算称为 KV 数据的数据块。响应用户提示的初始步骤是计算与该提示相对应的 KV 值。这一阶段被称为请求处理生命周期中的预填充（prefill）阶段。预填充阶段（即每个提示计算 KV 值的过程）计算开销大，需要专门的加速硬件（如 GPU）才能快速完成。

为同一个提示计算的 KV 值可以被共享相同前缀的其他提示复用，从而避免重新计算。对于许多用例，缓存并重用 KV 值可以实现两大主要优势：

• 改善请求延迟（假设从缓存读取数据比重新计算 KV 数据更快）
• 提高单节点吞吐量（因为 GPU 核心的负载降低，从而允许处理更多并发请求）。

此外，即使对于请求之间没有公共前缀的工作负载，KV 缓存卸载也非常有用。具体而言，当处理大量并发请求时，GPU 可能会因空间不足而无法存储处理当前请求集所需的 KV 值。在这种情况下，推理引擎可能会抢占正在运行的请求，将其 KV 值从 GPU 内存中丢弃。随后，当该请求被重新调度进行处理时，其 KV 值将需要重新计算。通过在请求被抢占前将 KV 缓存卸载到更大的存储层（如 CPU DRAM），可以避免重新计算 KV 值的成本。

CPU 卸载

在本文中，我们重点讨论 KV 到 CPU 内存（DRAM）的卸载。由于以下综合原因，这种做法特别值得关注：

• CPU RAM 在各种部署环境中都广泛可用。
• 其容量通常超过 GPU 内存，允许存储更大的 KV 缓存。
• CPU RAM 和 GPU 内存之间的传输具有低延迟和高吞吐量的优势。
  结合前一点，这使得 CPU 卸载成为高效处理请求抢占的理想方案。
• CPU RAM 还是方便的过渡区域，用于进一步卸载到外部存储。
  在存储延迟较高的情况下，这一点尤为有利。

全新的卸载连接器

vLLM 连接器 API

vLLM 长期以来一直支持用于读写 KV 数据的 API，并与请求生命周期集成。该 API 被称为连接器 API（Connector API）。概括地说，vLLM 在处理任何请求之前都会查询此 API，从而允许从外部源导入 KV 数据。在 KV 数据计算完成后，vLLM 还会调用此 API 将新生成的 KV 值存储在外部目标中。

最初，连接器 API 是同步的。这意味着当 vLLM 在外部加载/存储 KV 数据时，vLLM 引擎会被阻塞，无法并行处理新的请求批次。vLLM 0.9.0 扩展了连接器 API，以支持异步加载和存储 KV 数据。卸载连接器正是利用了这一新的异步 API 来实现 KV 缓存卸载。

我们引入了卸载连接器，它允许异步卸载和加载 KV 数据。它公开了一个可插拔的后端 API，允许使用任何媒介进行卸载。此 API 简化了添加新卸载后端的流程。用户只需定义一个实现媒介间 KV 数据复制的传输函数即可。

卸载连接器内置了 CPU 后端，实现了 vLLM 中 KV 数据的原生 CPU 卸载。在本文的其余部分，我们将仅关注 CPU 卸载。

使用卸载连接器

要使用卸载连接器进行 CPU 卸载，只需在 vllm serve 命令中添加以下 CLI 标志：

--kv_offloading_backend native --kv_offloading_size <size_in_GB>

此 CLI 基于此 PR #24498，预计将包含在 0.14.0 版本中。

对于旧版本，可以使用以下 CLI 启用 CPU 卸载：

--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_cpu_blocks": <num_cpu_blocks>}}'

其中 num_cpu_blocks 是为 CPU KV 缓存分配的 CPU 块数量。

通过卸载连接器进行 CPU 卸载的优势

我们提出了两个不同的微基准测试。第一个衡量单个请求的首字延迟（TTFT），强调加速单个请求的服务速度；而第二个衡量系统处理多个并发请求的吞吐量，展示了卸载如何帮助处理更繁重的工作负载。

在我们的第一个基准测试中，我们测量处理单个预填充请求的延迟，比较 CPU 缓存加载与 GPU 计算 KV 值的时间。

图 1：单个请求 TTFT（Llama-3.1-8B-Instruct，NVIDIA H100）。

结果表明，从 CPU 加载 KV 值可使 TTFT 缩短 2 倍至 22 倍，具体取决于提示大小。我们基准测试的确切设置和代码出现在本博客的末尾。

请注意，KV 卸载的延迟（将 KV 数据从 GPU 复制到 CPU）不是用户可见的，即它不应影响响应时间。这是因为卸载也是异步完成的，用户的请求无需等待此传输完成即可结束。这意味着使用卸载连接器对缓存未命中时的 TTFT 影响极小。

接下来，我们对使用 CPU 卸载处理多个并发请求时的整体吞吐量进行了基准测试。我们提交了一批 10,000 个唯一请求（每个 512 个 token），并测量了在 CPU 缓存中不同命中率下实现的吞吐量。

我们测量处理这些请求的时间（省略预热 CPU 缓存的时间），并将其用于推导 token/s 吞吐量。为了专注于 CPU 缓存的效果，本次测试未利用 GPU 缓存。

图 2：并发请求吞吐量（Llama-3.1-8B-Instruct，NVIDIA H100，10000 个 512 token 的预填充请求）。

结果显示，吞吐量随着 CPU KV 缓存命中率的提高而增加。我们观察到吞吐量最高增加了 9 倍，尽管该提示大小的 TTFT 仅减少了 2 倍。这表明 KV 缓存卸载的主要增益在于最大化吞吐量。

卸载连接器的 vLLM 版本

请注意，卸载连接器的性能在 0.12.0 版本中得到了显著提升。例如，在使用 Llama-3.1-8B-Instruct 和 NVIDIA H100 GPU 进行测试时，我们观察到 TTFT 降低了高达 4 倍，吞吐量提高了 5 倍。我们将在讨论 vLLM 物理块大小的部分详细展开这些改进。

希望在即将发布的 0.14.0 版本中引入进一步的改进。
特别是：

• 允许将被抢占的请求从 CPU 加载回来（PR #29870）
• 修复卸载与模型计算之间的竞态条件（PR #31341）

我们在本文中的评估包含了这些改进。

评估 GPU-CPU 传输技术

在本文的其余部分，我们将深入探讨设计 CPU 卸载时的一些技术考量。具体而言，我们将介绍旨在通过最大化 GPU-CPU 吞吐量并最小化 GPU 和 CPU 核心开销来优化推理吞吐量的研究。

如上所述，在定义卸载连接器的后端时，核心组件是传输函数。对于 CPU 后端，该传输函数将数据从 GPU 内存复制到 CPU 内存（反之亦然）。它目前支持 CUDA 兼容设备（NVIDIA 和 AMD）。

CPU 后端实现的传输函数使用 cudaMemcpyAsync 函数，该函数利用了 GPU 上称为 DMA（直接内存访问）的硬件组件。此组件专为设备（GPU）和主机内存之间的高吞吐量数据传输而设计。此外，利用 DMA 执行传输意味着 CPU 和 GPU 核心的开销最小。由于我们的传输相对于模型计算是异步运行的，这一属性尤为重要。

在处理大的物理连续拷贝时，DMA 可提供最佳吞吐量。这意味着我们预期的卸载性能将根据 KV 数据布局而变化。具有更大 KV 数据块的 LLM 模型性能表现更好。

但是 DMA 的速度有多快？它与使用自定义 CUDA 内核等替代方案相比如何？

为了回答这些问题，我们创建了微基准测试 gpu_cpu_benchmark。
在此基准测试中，我们测试了两种在 GPU 和 CPU 之间复制数据的替代方案：

• 使用 DMA 复制 —— 通过 cudaMemcpyAsync。
• 使用利用 GPU 核心并使用原始指针复制 16 字节字长的自定义 CUDA 内核。这种方法利用了 GPU 核心提供的海量并行性，非常有效。另一方面，它会与 GPU 核心的主要任务产生更大的干扰。

我们的第一个测试测量了 1000 个块的单次传输吞吐量，测试块大小范围从 4KB 到 16MB。

图 3：单个 GPU -> CPU 传输吞吐量（NVIDIA H100，1000 个块的单次传输）。

图 4：单个 CPU -> GPU 传输吞吐量（NVIDIA H100，1000 个块的单次传输）。

结果证实 DMA 表现良好，但仅限于较大的块大小。对于较小的块大小，自定义内核能实现明显更好的吞吐量。然而，我们注意到自定义内核的结果波动较大，方差更大。

我们现在转向测试双向传输吞吐量，通过发出两个并发传输（一个用于读取，一个用于写入）来进行测试。在此测试中，我们将块大小固定为 2MB，并调整两个方向传输大小的比例。对于这两种复制机制，当两个方向传输的量大致相等时，都能达到峰值吞吐量。然而，尽管对于单向传输两者都可以达到约 50GB/s，但在双向传输中结果有所不同：

• DMA 达到 83.4 GB/s
• 自定义内核达到 68.5 GB/s

因此，要在两种方法之间做出决定，现在的问题仍然是：

• vLLM 使用的有效块大小是多少？
  这取决于所服务的模型以及 vLLM 配置。在下一节中，我们将针对目前常用的一些模型回答这个问题。
• 两种方法如何影响 GPU 模型计算性能？
  回想一下，卸载连接器旨在与 GPU 执行的模型计算工作并行地卸载/加载 KV 数据。在我们的评估中，我们将观察每种方法如何影响整体吞吐量。

更改 vLLM 的内存布局

在本节中，我们将描述 vLLM 中 GPU 内存布局的更改，使其采用更支持 KV 传输的格式（同时不牺牲计算速度）。

首先描述 vLLM 为其 KV 缓存使用的默认内存布局，并了解卸载 KV 数据时需要在 GPU 和 CPU 之间复制的碎片大小。这决定了 vLLM 中用于传输 KV 数据的有效物理块大小。

vLLM 以 token 块的形式分配 GPU 内存，默认为每个块 16 个 token。实际物理布局取决于所使用的注意力机制后端（如 FlashAttention、FlashInfer 等）和所服务的模型。当今最常见的模型是均匀模型（uniform models），由多个层组成，每层都有自己的 KV 缓存，但形状相同。vLLM 也支持混合模型（hybrid models），目前尚未针对卸载连接器进行优化。对于均匀模型，vLLM 为每一层分配自己的 KV 缓存，因此单个逻辑块的 KV 缓存被碎片化为 num_layers 个块，每层一个。此外，根据注意力后端的不同，每层块可以进一步碎片化为 2 个子块，一个用于 K（键缓存），一个用于 V（值缓存）。

这种碎片化对于模型计算性能来说意义不大，但对于 KV 卸载却是毁灭性的，因为它在 KV 缓存布局中造成了不必要的碎片化，导致有效块大小变小。为了克服这一问题，我们最近上游合并了 vLLM KV 缓存布局的更改，该更改创建了一个包含所有层 KV 数据的连续物理块。此更改有效地将物理块大小增加了 2*num_layers 倍，反过来又将卸载连接器的吞吐量提高了一个数量级。

下表总结了当今一些常用模型，比较了旧版（0.11.0）和新版（0.12.0）的物理块大小（假设 vLLM 使用 16 个 token 的块）。

请注意，新的 vLLM KV 缓存布局产生的物理块大小约为 0.5-2 MB，而旧布局仅为几 KB。结合我们从 GPU-CPU 微基准测试中获得的数据，我们预计 DMA 方法具有可比的性能，或者（取决于模型）略逊于自定义内核方法。

复制方法的端到端评估

在下一节中，我们使用两个 vLLM 微基准测试来比较卸载连接器的两种变体：

• 包含 DMA 传输函数的上游版本
• 使用我们 GPU-CPU 微基准测试中自定义内核的补丁版本。

我们特意选择展示卸载连接器最坏情况下的结果，即使用物理块大小相对较小（0.5 MB）的模型。

图 5：单个请求 TTFT（Llama-3.2-1B-Instruct，NVIDIA H100）。

对于单个请求基准测试，我们看到自定义内核产生了略好的 TTFT，对于 1K 的提示差异小于 1ms，对于 90K 的大提示差异最高达 15ms。考虑到 0.5 MB 块大小下 GPU-CPU 微基准测试的结果，这些结果符合预期。对于具有更大块大小的模型，两种变体的结果大致相同。

图 6：并发请求吞吐量（Llama-3.2-1B-Instruct，NVIDIA H100，10000 个 512 token 的预填充请求）。

然而，对于并发请求测试，我们看到 DMA 实现了比自定义内核更好的吞吐量。增益在 0 命中率时约为 5.5%，并在 80% 命中率测量时增加到约 15%。

这些结果的解释是：自定义内核方法会干扰模型计算，因为两者都利用了 GPU 核心。在 0% 命中率下，自定义内核方法实际上比根本不使用 CPU 卸载时产生的吞吐量低 6%。在 100% 命中率下，由于没有与 CPU 加载并行的模型计算，因此两种方法之间的差距缩小了。

我们要强调的是，我们展示的是 DMA 方法在最坏情况模型下的结果。最常见的模型具有更大的物理块大小，因此更倾向于使用 DMA。以 Llama-3.1-8B-Instruct 为例，DMA 在保持 TTFT 不变的情况下，比自定义内核获得了高达 32% 的吞吐量提升。

综上所述，我们看到 GPU 内存布局的更改使我们能够利用 DMA 进行 KV 传输，从而实现了更好的整体吞吐量。

评估设置和基准测试代码

为了评估 vLLM 的 CPU 卸载，我们使用了以下设置：

• 单个 Ubuntu 24.04.1 LTS 容器
• 内核 5.14.0-427.81.1.el9_4.x86_64
• Intel Xeon SapphireRapids 2.1Ghz（限制为 8 个核心）
• NVIDIA H100 80GB HBM3
• 500GB DRAM
• CUDA 版本：12.9
• vLLM 提交哈希 2a1776b7ac4fae7c50c694edeafc1b14270e4350
• Flash Attention 后端
• 禁用 GPU 前缀缓存（以便评估 CPU 命中）
• GPU 块大小 16 个 token
• CPU 块大小 16 个 token
• 禁用编码/解码

我们的基准测试代码可以在这里找到。

下一步是什么？

我们正在继续增强 vLLM 的原生 KV 卸载功能。我们的下一个里程碑是使 CPU KV 缓存能够作为存储卸载的中间层。

一如既往，我们的首要任务仍然是正确性和性能。我们邀请您尝试一下，分享您的结果，如果您遇到任何问题，请告诉我们。

加入讨论：在 vLLM Slack 的 #feat-v1-cpu-offloading 频道中分享您的用例和反馈。