基于 PyTorch 和 vLLM 的大规模解耦推理

主要收获

PyTorch 与 vLLM 已有机结合，旨在加速尖端的生成式 AI 应用，包括推理、后训练以及智能体系统。
“预填充/解码解耦”（Prefill/Decode Disaggregation）是提升大规模生成式 AI 推理效率（涵盖延迟与吞吐量）的关键技术。
预填充/解码解耦已在 Meta 内部推理堆栈中启用，服务于 Meta 的大规模流量。通过 Meta 与 vLLM 团队的通力合作，Meta 的 vLLM 解耦实现已在性能上超越了 Meta 原有的内部 LLM 推理堆栈。
Meta 的优化成果与可靠性增强功能正逐步向上游 vLLM 社区贡献。

在我们之前的文章 PyTorch + vLLM 中，我们分享了 vLLM 加入 PyTorch 基金会这一激动人心的消息，并重点介绍了 PyTorch 与 vLLM 之间的多项整合成果及未来规划。核心举措之一便是大规模的“预填充-解码 (P/D) 解耦”推理，旨在满足 Meta LLM 产品在延迟预算内的吞吐量提升需求。在过去两个月中，Meta 的工程师致力于实现 vLLM 的内部 P/D 解耦整合，并在 TTFT（首字延迟）和 TTIT（迭代字延迟）指标上均优于 Meta 现有的内部 LLM 推理堆栈。vLLM 与 llm-d 及 dynamo 有着原生集成。在 Meta 内部，我们开发了加速 KV 缓存传输至服务集群拓扑及架构的抽象层。本文将重点介绍 Meta 对 vLLM 的定制化改进及其与 vLLM 上游版本的集成。

解耦预填充/解码

在 LLM 推理中，第一个 Token 依赖于用户提供的输入提示词，而随后的所有 Token 则以自回归方式逐一生成。我们将第一个 Token 的生成称为“预填充（Prefill）”，后续 Token 的生成称为“解码（Decode）”。

尽管执行的操作本质上相同，但预填充和解码展现出截然不同的特征。值得注意的特征包括：

预填充
- 计算密集型（Compute bound）
- 延迟受 Token 长度和批处理大小限制
- 每个请求仅发生一次
解码
- 内存密集型（Memory bound）
- 效率受批处理大小限制
- 在整体延迟中占主导地位

预填充/解码解耦旨在将预填充与解码解耦至不同的主机，解码主机将请求重定向至预填充主机进行首个 Token 的生成，后续工作则由自身处理。我们旨在独立扩展预填充和解码推理，从而实现更高效的资源利用，并同时改善延迟与吞吐量。

vLLM 集成概览

目前，TP（张量并行）+ 解耦在预填充和解码两侧均得到支持。通过 TCP 网络实现最优 P/D 解耦服务，主要依赖三个核心组件：

代理库（Proxy library）
Python KV 连接器（Python kv connector）
C++ 解码 KV 连接器与预填充 KV 连接器，通过 TCP 连接

我们通过路由层处理转发，负责负载均衡，并以点对点（P2P）模式连接预填充节点与解码节点，从而降低网络开销。预填充节点和解码节点可根据各自的工作负载独立扩缩容。因此，在生产环境中，无需手动维护预填充与解码的比例。

组件

服务代理（Service Proxy）

服务代理附加在解码服务器上。它将请求转发至远程预填充主机，并协调解码与预填充 KV 连接器之间的 KV 缓存传输。我们利用 Meta 内部的服务路由解决方案，根据服务器负载和缓存命中率在所有预填充主机之间进行负载均衡。

服务代理首先将传入请求转发至选定的预填充主机，同时建立多个流式通道，通过底层的 Meta C++ 连接器从该主机获取 KV 缓存。
获取的远程 KV 缓存将首先被复制到临时 GPU 缓冲区，等待 vLLM KV 连接器将其注入到适当的 KV 块中。

vLLM Python KV 连接器

我们基于 vLLM v1 KV 连接器接口实现了异步 KV 连接器方案。KV 连接器在主模型执行流之外并行执行 KV 缓存传输操作，并确保双方的 GPU 操作互不争抢。通过这种方式，我们实现了更快的 TTIT/TTFT；具体优化细节见下文。

在预填充侧
- 在每一层完成注意力计算后，Python KV 连接器会将特定请求的 KV 缓存保存到临时 CPU 缓冲区，这些保存操作通过底层的 Meta C++ 连接器执行。通过这种方式，我们确保了主流模型执行完全不会被阻塞。
- 当 KV 缓存保存完成后，它会立即流式传输至远程解码主机。
在解码侧
- 远程 KV 缓存获取并复制到临时 GPU 缓冲区后，Python KV 连接器会开始将远程 KV 缓存注入到由 vLLM 分配的本地 KV 缓存块中。这同样通过底层的 Meta C++ 连接器在其独立的 C++ 线程和 CUDA 流中完成。
- 当 KV 注入完成后，Python KV 连接器会将请求释放回 vLLM 调度器，该请求随后被调度在下一轮迭代中运行。
错误处理
- 我们还实现了一个通用的垃圾回收器，用于清理从远程获取的闲置 KV 缓存缓冲区，以避免 CUDA 显存溢出（OOM）问题。这涵盖了以下极端情况：
  - 被抢占的请求、被取消/终止的请求，远程获取可能已完成但本地注入已被中止的情况。

Meta C++ 连接器

由于 KV 传输操作涉及大量的 IO，我们选择使用 C++ 实现，以便更好地并行化数据传输并精细调整线程模型。所有实际的 KV 传输操作，如网络流传输、本地 H2D/D2H（主机到设备/设备到主机）、KV 注入/提取，均在各自的 C++ 线程中配合独立的 CUDA 流完成。

预填充 C++ 连接器

在每一层模型注意力计算完成后，KV 缓存被卸载到 DRAM 上的 C++ 连接器。随后，它会针对特定请求和层将 KV 缓存流式传回解码主机。

解码 C++ 连接器

从代理层接收请求及其路由的预填充主机地址后，它会建立多个流式通道以获取远程 KV 缓存。它将获取的 KV 缓存缓冲在 DRAM 上，并异步注入到预分配的 GPU KV 缓存块中。

优化

加速网络传输

多网卡（Multi-NIC）支持： 多个前端网卡链接至最近的 GPU，优化了解码与预填充 KV 连接器之间的连接。
多流 KV 缓存传输： 单个 TCP 流无法饱和网络带宽。为了最大化网络吞吐量，KV 缓存被切片并使用多个流并行传输。

优化服务性能

粘性路由（Sticky Routing）： 在预填充转发中，来自同一会话的请求始终被引导至同一个预填充主机。这显著提高了多轮对话场景下的前缀缓存命中率。
负载均衡： 我们利用 Meta 内部服务路由根据各主机的利用率有效地分配工作负载。结合粘性路由，这在维持 90% HBM（高带宽内存）利用率的同时，实现了 40%-50% 的前缀缓存命中率。

微调 vLLM

更大的块大小： 虽然 vLLM 建议每个 KV 缓存块包含 16 个 Token，但我们发现，在 KV 缓存注入和提取期间，这些较小的块在 CPU 和 GPU 之间传输会导致大量小内核启动，从而产生显著开销。因此，我们采用了更大的块大小（例如 128、256）以提高解耦性能，并配合必要的内核端调整。
禁用解码前缀缓存： 解码主机从 KV 连接器加载 KV 缓存，使得前缀哈希计算成为调度器中不必要的开销。在解码侧禁用它有助于稳定 TTIT（迭代字延迟）。

改善 TTFT（首字延迟）

提前返回第一个 Token： 代理层从预填充层接收响应并立即将第一个 Token 返回给客户端。与此同时，引擎解码第二个 Token。我们还复用了来自预填充的标记化提示词（Tokenized prompt），消除了解码侧额外的分词步骤。这确保了 P/D 解耦方案的 TTFT 尽可能接近预填充主机的 TTFT。

增强 TTIT（迭代字延迟）

独占原始类型使用： 我们观察到，在 vLLM 调度器与工作节点之间传输数据时，Python 原生 pickle dump 序列化张量的时间比序列化整数列表慢三倍。这经常导致随机的调度器进程挂起，从而对 TTIT 产生负面影响。最佳实践是避免在 KVConnectorMetadata 和 SchedulerOutput 中创建张量或复杂对象。
异步 KV 加载： 我们将 KV 加载操作与 vLLM 模型解码步骤并行化。这防止了等待远程 KV 的请求阻塞已在生成新输出 Token 的请求。
最大化 GPU 操作重叠： 由于 KV 传输操作主要是复制/IO 操作，而主流模型前向执行是计算密集型的，我们成功地在其独立的 CUDA 流中完全重叠了 KV 传输操作与主模型前向执行。这使得 KV 传输不会导致额外的延迟开销。
避免 CPU 调度争用： 我们不是同时调度所有层级的 KV 注入（本质上是索引复制操作），这可能会导致模型前向传递期间的内核调度争用，而是按照与模型前向传递同步的顺序，逐层调度 KV 注入。
非阻塞复制操作： 所有复制（主机到设备/设备到主机）操作均以非阻塞方式运行。我们还解决了一个问题，即在默认 CUDA 流中运行的主模型前向传递会无意中阻塞来自非阻塞 CUDA 流的其他复制操作。

性能结果

我们使用 Llama4 Maverick 在通过 TCP 网络连接的 H100 主机（每主机 8 张 H100 卡）上进行了基准测试。评估采用输入长度 2000，输出长度 150。

在相同的批处理大小下，我们发现解耦 (1P1D) 可以提供更高的吞吐量。

在相同的 QPS 工作负载下，由于 TTIT 曲线更加平滑，解耦 (1P1D) 在控制整体延迟方面表现更好。

然而，我们也注意到，当负载非常大时，TTFT 的曲线会更剧烈地回归，这可能有多种原因（正如我们在下一节探讨内容中所述）。

TCP 连接使网络成为瓶颈。
1P1D 设置给预填充侧带来了更大的负载压力，因为我们的评估是在预填充负载较重的工作（2000 个输入 vs 150 个输出）下完成的。理想情况下，需要更高的预填充与解码比率。

使用 PyTorch 和 vLLM 实现大规模解耦推理