AWS Machine Learning

基于推送的 KV 传输,用于提升 vLLM 中分离式推理的 TTFT 和 TPOT

AWS为vLLM引入推式NIXL连接器,将KV缓存传输协调移出关键路径,基准中TTFT最高改善3倍、TPOT改善30%。该结果基于P5en实例上的Qwen3-32B测试,随负载与配置而异。

Pull-mode timeline: router relay, decode-side allocation, and NIXL READ run serially on the path to first token.
配图来源 · AWS Machine Learning

vLLM 是一个面向大语言模型(LLM)的开源推理和服务引擎。LLM 推理有两个不同的阶段:prefill(预填充),处理输入提示并构建键值(KV)缓存;以及 decode(解码),使用该缓存逐个 token 地生成输出。这两个阶段在计算、内存和扩展特性上各不相同。 Prefill-decode (P-D) 分离(disaggregation) 将这两个阶段运行在单独的引擎实例上,使每个阶段可以独立扩展和优化。这可以在负载下改善资源利用率、吞吐量和延迟。然而,它也要求在 token 生成开始之前,将 KV 缓存从 prefill 实例传输到 decode 实例,这使得高效的 KV 传输对首个 token 时间(TTFT)至关重要。

vLLM 执行这种传输的方式之一是通过 NVIDIA Inference Xfer Library (NIXL) 连接器,它在支持远程直接内存访问(RDMA)的互连上移动 KV 块,例如 AWS Elastic Fabric Adapter(AWS EFA,AWS 中跨服务器使用的低延迟 RDMA 网络)。在 支持的实例类型上,AWS EFA 支持 GPUDirect RDMA 和 Neuron Direct RDMA。因此,NIXL 可以直接在 GPU 或 AWS Trainium 设备内存之间传输 KV 块,而无需经过主机内存中转。NIXL 连接器的默认模式是 基于拉取(pull-based)的:在 prefill 实例完成提示 KV 缓存的计算后,路由器将源元数据(包括 prefill 块 ID)转发给 decode 实例。decode 实例随后分配目标块,准备 NIXL 传输,并发起 NIXL READ 从 prefill 实例拉取 KV 块。由于这些协调和准备步骤在 prefill 之后、解码开始之前按顺序执行,它们延长了 TTFT 关键路径,如下面的时间线所示。

Pull-mode timeline: router relay, decode-side allocation, and NIXL READ run serially on the path to first token.

图 1:拉取模式 KV 传输时间线,其中路由器转发、decode 侧分配和 NIXL READ 在通往首个 token 的路径上串行运行

本文介绍一种 基于推送(push-based) 的 NIXL 连接器,它将 prefill 之后的传输协调和准备工作移出关键路径。如果您使用 vLLM 运行分离式 prefill/decode 推理,推送模式可以通过降低 TTFT 和 TPOT 来改善应用响应速度,尤其在高请求速率下。

在使用 EFA 连接的 Amazon Elastic Compute Cloud (Amazon EC2) P5en 实例的 1P1D 部署的基准测试中,对于 Qwen3-32B 模型,与拉取模式相比,推送模式实现了最高 3 倍的 TTFT 提升和最高 30% 的每输出 token 时间(TPOT)提升(参见“性能数据”一节)。结果取决于工作负载和配置。

我们在通过 EFA 连接的 AWS Trainium 实例集群上构建并验证了该功能,随后将其作为一项与加速器无关的功能贡献给 vLLM。本文的结果来自带有 EFA 的 Amazon EC2 P5en(GPU)集群,表明这些收益可以延续到 GPU 上。本文解释基于推送的架构、传输流程、线程模型、性能结果,以及如何复现该基准测试。

基于推送的 KV 传输

本节解释基于推送的设计:其架构、传输协议,以及使 KV 传输远离关键路径的线程模型。

架构概述

基于推送的传输在两个方面改善了 TTFT 和 decode 吞吐量:

  1. 它通过重叠生产者和消费者工作,消除了拉取模型的串行化。 decode 实例在生产者仍在计算 KV 时就预先注册其目标块。传输随后在生产者的块一就绪时立即开始,而不必等待消费侧的 ID 发现往返,从而将该往返从通往首个 token 的路径中移除。
  2. 它将传输准备时间从 decode 实例上移走。 该准备工作在 prefill 侧和后台线程上运行。它准备 NIXL 描述符并提交传输,因此 decode 实例不会在其延迟敏感的 token 生成阶段花费这些周期。

最终效果是 decode 实例上更低的 TTFT 和更高的输出 token 吞吐量。

推送连接器协调三个角色:decode 实例预先注册其目标块,prefill 实例匹配每个请求并发起传输,一个专用的后台线程将 KV 写入从引擎的前向传播循环中剥离出来。下图展示了这些部分如何协同工作。

Push-based KV transfer architecture overview

图2:基于推送的 KV 传输架构,展示 decode 实例、prefill 实例和后台传输线程如何协调

实现细节

该协议包含三个逻辑步骤。

1. 注册(解码侧)。 当解码实例为需要远程预填充的请求分配本地块时,它会准备一份注册信息,描述预填充实例应将 KV 写入的位置:包括自身的身份标识以及刚刚预留的块位置。在同一个引擎步骤内,工作进程通过 NIXL 通知将该注册信息传输给预填充实例。由于注册发生在分配时,连接器在预填充实例完成预填充之前就已知道目标块。

2. 会合(预填充侧)。 来自解码器的注册信息和生产者上预填充的完成可能以任意顺序到达,连接器可以处理这两种情况。如果注册信息先到达,生产者会保存它,直到自身的预填充完成。如果预填充先完成,生产者会保存请求的块,直到注册信息到达。无论哪个事件后发生,都会触发传输。连接器通过 ID 配对请求,并且匹配机制会考虑重新发出的请求,正确关联被抢占或重试的请求。

3. 传输(预填充侧)。 一旦请求被匹配,生产者会与解码实例建立连接(如果尚未存在),并发出一个 NIXL WRITE,将 KV 块直接复制到解码器已注册的内存中,并针对每个目标张量并行 rank。生产者和消费者可能使用不同的内部块布局。双方在传输时将块位置映射到各自的物理布局,从而允许两侧使用不同的块大小。每个 WRITE 都携带一个完成通知,NIXL 会在数据到达后将其传递给解码实例。解码器会统计这些通知,每个向其写入的生产者 rank 对应一个通知。只有在所有预期的通知都已到达后,解码器才将 KV 视为已接收,此时请求进入解码阶段。

线程模型

推送连接器通过在每个预填充工作进程上使用专用的后台传输线程,将模型执行与 KV 传输准备工作分离。在请求开始时,解码主循环分配其目标 KV 块并向预填充实例注册,同时预填充主循环计算提示词 KV 缓存。在目标注册和提示词 KV 都就绪后,预填充主循环将 KV 块交给后台线程,并迅速返回模型执行。

后台线程建立或复用 NIXL 连接,将请求的 KV 块 ID 映射到内存描述符,将连续的块合并为更大的传输区域,并向解码实例提交 NIXL WRITE。当 NIXL 完成传输后,它会向解码实例发送完成通知,解码实例随即开始生成 token。将描述符准备和传输提交移出对延迟敏感的主循环,可以防止这些针对每个请求的工作延迟模型执行。下图展示了单个请求在主循环和后台线程之间的这些交接过程。

Threading model showing the decode and prefill main loops handing KV blocks to a background thread that performs the NIXL transfer off the forward-pass loop

图 3:基于推送的线程模型中单个请求的主循环与后台线程交接过程

拉取与推送的时序比较

下图总结了前面几节所述的拉取和推送传输流程,并突出显示了它们对 TTFT 的影响。虽然两种模式使用相同的 NIXL 和 EFA/RDMA 数据路径,但拉取在预填充之后执行元数据中继和解码侧设置,而推送则将解码侧设置与预填充重叠。因此,推送将这些操作从预填充后的关键路径中移除,使解码能够更早开始。

Side-by-side timeline comparing pull and push KV transfer, showing push overlapping decode-side setup with prefill to shorten the path to the first token

图 4:拉取与推送传输流程对比,突出显示推送如何将解码侧设置与预填充重叠

两种模式还涉及路由上的权衡。拉取模式允许路由器在预填充完成后选择解码实例,这为负载感知路由和服务质量管理提供了更大的灵活性。推送模式更早选择解码实例,以便目标设置可以与预填充重叠。作为回报,在可比负载下,预填充实例可以更早释放其 KV 块,因为只要提示词 KV 和目标注册都就绪,传输就会立即开始。

性能数据

为了进行评估,我们在使用 EFA 连接的 Amazon EC2 P5en 实例上的 1P1D 部署中使用了 Qwen3-32B 模型。为了隔离 P→D 传输性能,我们将基准测试配置为禁用前缀缓存,使用 --no-enable-prefix-caching。此配置使解码节点为每个请求从预填充节点检索完整的 KV 缓存。否则,重叠的提示词可能会复用解码节点的本地缓存并绕过传输,从而产生误导性的比较结果。

下图比较了推送和拉取 KV 传输的 TTFT。我们扫描了 1k–16k 的输入长度和每秒 1–8 个查询(QPS)的请求速率。在低负载和短提示词下,两种模式的表现几乎相同。随着 QPS 和输入长度的增长,拉取模式急剧恶化,在 8 QPS 和 16k 时达到约 1,744 ms。推送模式保持在约 540 ms 的低得多的水平,约降低了 3 倍。差距出现的原因是拉取模式在解码工作进程上内联发布 NIXL 描述符,在其引擎循环内同步执行,从而阻碍了 token 生成。推送模式将这项工作移至后台线程,脱离关键路径。

Four-panel TTFT comparison of push and pull KV transfer for Qwen3-32B on Amazon EC2 P5en instances with EFA. Test configuration: 1P1D topology, 128 output tokens, and prefix caching disabled on the Decode node.

图 5:在不同输入长度和请求速率下,推送与拉取 KV 传输的 TTFT 比较(1P1D 拓扑,128 个输出 token,解码节点禁用前缀缓存)

此图显示了相同条件下的 TPOT。其模式与 TTFT 相似。在轻负载下,推送和拉取接近,约为 5.6–5.8 ms。随着负载和输入长度上升,拉取在 4–8 QPS 和 16k 时攀升至约 10.4 ms,而推送保持在 7 ms 附近。原因在于推送将描述符准备和传输提交从解码工作进程卸载。token 生成不会因资源匮乏而受影响,因此即使在高低并发和长上下文下,每个 token 的延迟也保持平稳。

Four-panel TPOT comparison of push and pull KV transfer for Qwen3-32B on Amazon EC2 P5en instances with EFA. Test configuration: 1P1D topology, 128 output tokens, and prefix caching disabled on the Decode node.

图 6:在不同输入长度和请求速率下,推送式与拉取式 KV 传输的 TPOT 对比(1P1D 拓扑,128 个输出 token,解码节点禁用前缀缓存)

如何运行基准测试

要使用通过 EFA 连接的 Amazon EC2 P5en 实例重现我们的 1P1D 基准测试配置(一个预填充节点和一个解码节点),请参阅以下指南中的详细说明: 重现我们的结果 指南。该指南涵盖环境设置、启动 vLLM 预填充和解码服务器、运行基准测试以及收集结果。该功能可以在其他受支持的配置上运行,但性能结果会因硬件、模型、工作负载和网络配置而有所不同。

结论

在这篇文章中,我们展示了基于推送的 KV 连接器如何将预填充计算与解码侧的块分配和注册重叠,从而将预填充后的路由器往返以及解码侧的传输设置从关键路径中移除。这降低了 TTFT 和 TPOT,在高请求速率下尤为明显。在我们使用通过 EFA 连接的 Amazon EC2 P5en 实例的 1P1D 部署中进行的 Qwen3-32B 基准测试中,与拉取模式相比,推送模式使 TTFT 提升高达 3 倍,TPOT 提升 30%。结果因工作负载和配置而异。要使用推送模式,请将预填充和解码服务器都配置为使用 NixlPushConnector 而不是 NixlConnector.

。当可以提前选择解码节点且以更低的 TTFT 为优先时,请选择基于推送的连接器。对于需要 双向 KV 传输 或者将解码实例的选择推迟到预填充之后以实现负载感知路由的工作负载,基于拉取的连接器仍然是更合适的选择。vLLM 同时支持基于推送和拉取的 KV 连接器,因此您可以选择最适合您部署的方式。

作为未来的改进方向,我们正在探索将后台 NIXL 线程模型引入拉取模式连接器,将 NIXL 设置移出关键路径,以便拉取模式的工作负载也能从相同的优化中受益。我们已经开启了一个 拉取请求 来探索这一方向。

相关资源

致谢

我们感谢 vLLM 社区对维护此功能的支持。

 


关于作者

原始出处

AWS Machine Learning

内容说明

原始发布及相关权利归来源方。

机器翻译 · 请以原文为准