跳到正文
AWS Machine Learning Blog· Sunita Nadampalli·· 1 小时前AI 评分51

AWS 在 vLLM 推出 push 式 KV 传输,TTFT 最高提升 3 倍

Push-based KV transfer to improve TTFT and TPOT for disaggregated inference in vLLM

AI 导读

AWS 机器学习博客介绍为 vLLM 分离式推理新增的 push 式 NIXL KV 连接器,将解码侧注册与传输准备工作移出关键路径。在 Amazon EC2 P5en 实例加 EFA 的 1P1D 部署中,Qwen3-32B 在 8 QPS、16k 输入下 TTFT 从约 1744 ms 降至约 540 ms(约 3 倍提升),TPOT 提升最高 30%。

正文 · AI 翻译

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

vLLM 执行这种传输的方式之一是通过 NVIDIA Inference Xfer Library(NIXL)连接器,它通过支持远程直接内存访问(RDMA)的互连(例如 AWS Elastic Fabric Adapter,即 AWS EFA,AWS 中跨服务器使用的低延迟 RDMA 网络)移动 KV 块。在受支持的实例类型上,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 模型,与拉取模式相比,推送模式的 TTFT 提升高达 3 倍,每输出 token 时间(TPOT)改善高达 30%(参见性能数据一节)。结果取决于工作负载和配置。

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

推送式 KV 传输

本节讲解推送式设计:其架构、传输协议,以及让 KV 传输脱离关键路径的线程模型。

架构概览

推送式传输能够改善 TTFT 和 decode 吞吐量,原因有二:

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

总体效果是降低 TTFT 并提高解码实例的输出 token 吞吐量。

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

Push-based KV transfer architecture overview

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

实现细节

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

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

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

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

线程模型

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

后台线程建立或复用 NIXL 连接,将请求的 KV 块 ID 映射到内存描述符,把连续的块合并为更大的传输区域,并向 decode 实例提交 NIXL WRITE。当 NIXL 完成传输后,它会向 decode 实例发送完成通知,随后 decode 实例开始生成 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:基于 push 的线程模型中单个请求的主循环与后台线程交接过程

pull 与 push 的时序对比

下图总结了前几节所述的 pull 和 push 传输流程,并突出它们对 TTFT 的影响。尽管两种模式使用相同的 NIXL 和 EFA/RDMA 数据路径,pull 在 prefill 之后执行元数据中继和 decode 侧设置,而 push 则将 decode 侧设置与 prefill 重叠进行。因此,push 将这些操作从 prefill 之后的关键路径中移除,使解码能够更早开始。

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:pull 与 push 传输流程对比,突出 push 如何将 decode 侧设置与 prefill 重叠

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

性能数据

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

下图对比了 push 和 pull KV 传输的 TTFT。我们扫描了 1k–16k 的输入长度和每秒 1–8 个查询(QPS)的请求速率。在低负载和短提示词下,两种模式的表现几乎相同。随着 QPS 和输入长度的增长,pull 模式急剧退化,在 8 QPS 和 16k 时达到约 1,744 ms。push 模式保持在约 540 ms 的远低水平,约降低了 3 倍。差距的出现是因为 pull 模式在 decode 工作节点上内联提交 NIXL 描述符,即在其引擎循环内同步执行,这会阻塞 token 生成。push 模式则将其移至后台线程,脱离关键路径。

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:在 1P1D 拓扑(128 个输出 token,解码节点禁用前缀缓存)下,不同输入长度和请求速率下推送式与拉取式 KV 传输的 TTFT 对比

此图表展示了相同条件下的 TPOT。其模式与 TTFT 一致。在轻负载下,推送与拉取非常接近,约为 5.6–5.8 毫秒。随着负载和输入长度的增加,在 4–8 QPS 和 16k 时,拉取攀升至约 10.4 毫秒,而推送保持在 7 毫秒左右。原因是推送将描述符准备和传输提交的工作从解码工作进程中卸载出去。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:在 1P1D 拓扑(128 个输出 token,解码节点禁用前缀缓存)下,不同输入长度和请求速率下推送式与拉取式 KV 传输的 TPOT 对比

如何运行基准测试

要使用通过 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 Blog · aws.amazon.com