vLLM 如何通过 push-based KV transfer 改善 TTFT 和 TPOT
Push-based KV transfer to improve TTFT and TPOT for disaggregated inference in vLLM
AWS 介绍了 vLLM 的 push-based NIXL connector,通过将 KV 传输协调和准备移出关键路径,改善分离式推理的 TTFT 和 TPOT。
文章拆解了 push 模式如何将 KV 传输准备移出关键路径,并以基准结果和路由取舍说明其适用场景。
vLLM 是一款用于大型语言模型(LLM)的开源推理和服务引擎。LLM 推理分为两个不同阶段:预填充阶段处理输入提示并构建键值(KV)缓存;解码阶段利用该缓存逐个生成输出 token。这两个阶段在计算、内存和扩展特性方面各不相同。预填充-解码(P-D)分离会在不同的引擎实例上运行这两个阶段,使每个阶段都能独立扩展和优化。这可以提升负载下的资源利用率、吞吐量和延迟表现。不过,在开始生成 token 之前,还必须将 KV 缓存从预填充实例传输到解码实例,因此,高效的 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 连接器的默认模式是拉取式:预填充实例完成提示的 KV 缓存计算后,路由器会将源元数据(包括预填充块 ID)转发给解码实例。随后,解码实例分配目标块、准备 NIXL 传输,并发起 NIXL READ,从预填充实例拉取 KV 块。由于这些协调和设置步骤在预填充完成后、解码开始前依次进行,因此会延长 TTFT 关键路径,如下方时间线所示。
图 1:拉取模式下的 KV 传输时间线;路由器转发、解码端分配和 NIXL READ 在生成首个 token 的路径上串行运行
本文介绍一种推送式 NIXL 连接器,可将预填充后的传输协调和设置工作移出关键路径。如果你使用 vLLM 运行分离式预填充/解码推理,推送模式可以通过缩短 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 和解码吞吐量:
- 它通过让生产者和消费者的工作重叠,消除了拉取模式的串行化过程。在生产者仍在计算 KV 时,解码实例就会预先注册目标块。生产者的块一旦就绪,传输便立即开始,而无需等待消费者侧完成一轮 ID 发现往返,从而将这轮往返从首个 token 的关键路径中移除。
- 它将传输设置所需的时间从解码实例移开。设置工作在预填充侧的后台线程中运行。该线程准备 NIXL 描述符并提交传输,因此解码实例在延迟关键的 token 生成阶段无需耗费这些周期。
最终效果是降低 TTFT,并提高解码实例的输出 token 吞吐量。
推送连接器协调三个参与方:解码实例预先注册其目标块,预填充实例匹配每个请求并发起传输,专用后台线程则将 KV 写入操作移出引擎的前向传递循环。下图展示了这些部分如何协同工作。
图 2:基于推送的 KV 传输架构,展示解码实例、预填充实例和后台传输线程如何协同工作
实现细节
该协议包含三个逻辑步骤。
1. 注册(解码侧)。当解码实例为需要远程预填充的请求分配本地块时,会准备一份注册信息,说明预填充实例应将 KV 写入何处:包括解码实例自身的标识,以及它刚刚预留的块位置。在同一个引擎步骤中,工作线程会通过 NIXL 通知将这份注册信息发送给预填充实例。由于注册在分配时完成,连接器在预填充实例完成预填充之前就知道目标块的位置。
2. 汇合(预填充侧)。来自解码器的注册信息和生产者侧的预填充完成事件可能以任意顺序到达,连接器会处理这两种情况。如果注册信息先到,生产者会将其保留,直到自身的预填充完成。如果预填充先完成,生产者会保留该请求的块,直到注册信息到达。后发生的事件会触发传输。连接器按 ID 配对请求,并能处理重新发出的请求,确保正确关联被抢占或重试的请求。
3. 传输(预填充侧)。请求匹配后,生产者会与解码实例建立连接(如果尚未建立),并发起 NIXL WRITE,将 KV 块直接复制到解码器注册的内存中,同时定位到每个目标张量并行 rank。生产者和消费者可以使用不同的内部块布局。传输时,两侧各自将块位置映射到自己的物理布局,因此两侧可以使用不同的块大小。每次 WRITE 都会携带完成通知;数据到达后,NIXL 会将通知发送给解码实例。解码器会统计这些通知,每个向其写入数据的生产者 rank 对应一个通知。只有在所有预期通知都到达后,解码器才会将 KV 视为已接收;此时请求便会继续进入解码阶段。
线程模型
Push connector 在每个 prefill worker 上使用专用后台传输线程,将模型执行与 KV 传输准备分离。请求开始时,decode 主循环分配目标 KV 块,并将其注册到 prefill 实例;与此同时,prefill 主循环计算 prompt KV 缓存。目标注册和 prompt KV 都准备就绪后,prefill 主循环便将 KV 块交给后台线程,并迅速返回模型执行。
后台线程建立或复用 NIXL 连接,将请求的 KV 块 ID 映射为内存描述符,把连续块合并为更大的传输区域,然后向 decode 实例提交 NIXL WRITE。NIXL 完成传输后,会向 decode 实例发送完成通知,随后 decode 实例开始生成 token。将描述符准备和传输提交移出对延迟敏感的主循环,可避免这些逐请求处理工作拖慢模型执行。下图展示了单个请求中主循环与后台线程之间的这些交接过程。
图 3:基于 push 的线程模型中,单个请求的主循环与后台线程交接过程
Pull 与 push 的耗时对比
下图总结了前文所述的 pull 和 push 传输流程,并突出显示它们对 TTFT 的影响。尽管两种模式使用相同的 NIXL 和 EFA/RDMA 数据路径,但 pull 会在 prefill 完成后进行元数据中继和 decode 端设置,而 push 则让 decode 端设置与 prefill 并行进行。因此,push 将这些操作移出 prefill 后的关键路径,使解码能够更早开始。
图 4:Pull 与 push 传输流程对比,突出显示 push 如何让 decode 端设置与 prefill 并行进行
两种模式还涉及路由方面的权衡。Pull 模式允许路由器在 prefill 完成后选择 decode 实例,因此在负载感知路由和服务质量管理方面更灵活。Push 模式则更早选择 decode 实例,以便让目标端设置与 prefill 并行进行。作为回报,在负载相当的情况下,prefill 实例可以更早释放其 KV 块,因为只要 prompt KV 和目标注册都准备就绪,传输便会立即开始。
性能数据
在评估中,我们使用 Qwen3-32B 模型,在通过 EFA 连接的 Amazon EC2 P5en 实例上进行 1P1D 部署。为了单独评估 P→D 传输性能,我们将基准测试配置为通过 --no-enable-prefix-caching 禁用前缀缓存。此配置会使 decode 节点在每个请求中都从 prefill 节点获取完整的 KV 缓存。否则,重叠的 prompt 可能会复用 decode 节点的本地缓存并绕过传输,从而导致对比结果产生误导。
下图对比了 push 和 pull KV 传输的 TTFT。我们测试了 1k–16k 的输入长度,以及每秒 1–8 个查询(QPS)的请求速率。在低负载和短 prompt 下,两种模式的表现几乎相同。随着 QPS 和输入长度增加,pull 模式的性能急剧下降,在 8 QPS、16k 输入长度时达到约 1,744 ms。Push 模式则保持在低得多的水平,约为 540 ms,降幅约为 3 倍。差距扩大的原因在于,pull 模式会在 decode worker 上以内联方式、在其引擎循环中同步提交 NIXL 描述符,从而使这项工作阻塞 token 生成。Push 模式将其移至后台线程,避开了关键路径。
图 5:不同输入长度和请求速率下,push 与 pull KV 传输的 TTFT 对比(1P1D 拓扑,128 个输出 token,解码节点禁用前缀缓存)
此图展示了相同条件下的 TPOT。其趋势与 TTFT 相似。在低负载下,push 和 pull 的表现接近,约为 5.6–5.8 ms。随着负载和输入长度增加,pull 在 4–8 QPS、16k 输入长度时上升至约 10.4 ms,而 push 保持在 7 ms 左右。原因在于,push 将描述符准备和传输提交工作从解码工作线程中卸载。token 生成不会受到资源争用,因此即使在高并发和长上下文下,每个 token 的延迟也保持平稳。
图 6:不同输入长度和请求速率下,push 与 pull KV 传输的 TPOT 对比(1P1D 拓扑,128 个输出 token,解码节点禁用前缀缓存)
如何运行基准测试
如需复现我们的 1P1D 基准测试配置(一个预填充节点和一个解码节点,使用通过 EFA 连接的 Amazon EC2 P5en 实例),请参阅复现我们的结果指南。该指南涵盖环境设置、启动 vLLM 预填充和解码服务器、运行基准测试以及收集结果。此功能也可在其他受支持的配置上运行,但性能结果会因硬件、模型、工作负载和网络配置而异。
结论
本文展示了基于 push 的 KV 连接器如何将预填充计算与解码端的块分配和注册重叠执行,从而将预填充后的路由器往返通信以及解码端传输设置移出关键路径。这降低了 TTFT 和 TPOT,尤其是在请求速率较高时。在使用通过 EFA 连接的 Amazon EC2 P5en 实例进行的 Qwen3-32B 基准测试中,采用 1P1D 部署时,与 pull 模式相比,push 模式的 TTFT 最高提升 3 倍,TPOT 提升 30%。结果会因工作负载和配置而异。要使用 push 模式,请将预填充服务器和解码服务器都配置为使用 NixlPushConnector,而不是 NixlConnector。
当可以接受提前选择解码节点,且优先考虑降低 TTFT 时,应选择基于 push 的连接器。对于需要双向 KV 传输,或为了实现负载感知路由而将解码实例的选择推迟到预填充完成后的工作负载,基于 pull 的连接器仍然更合适。vLLM 同时支持基于 push 和 pull 的 KV 连接器,因此你可以选择最适合部署需求的方案。
作为一项后续改进,我们正在探索将后台 NIXL 线程模型引入 pull 模式连接器,把 NIXL 设置移出关键路径,使 pull 模式工作负载也能受益于相同的优化。我们已提交一个拉取请求以对此进行探索。
相关资源
致谢
感谢 vLLM 社区对维护此功能提供的支持。
作者简介
Sunita Nadampalli
Sunita 是 AWS 的首席工程师和 AI/ML 专家。她负责 AWS Trainium 和 AWS Graviton 针对 AI/ML 工作负载的软件性能优化。她热衷于开源软件开发,致力于提供高性能、可持续的软件解决方案。
Nicolò Lucchesi
Nicolò 是 Mistral 的研究工程师,负责领导推理基础团队。他是 vLLM 维护者,拥有 LLM 和计算机视觉快速推理方面的经验。
来源:AWS Machine Learning Blog · aws.amazon.com





