AWS Machine Learning

通过 vLLM 中的双向 KV 传输改善多轮智能体工作负载的 TTFT

AWS将双向KV缓存传输功能贡献给vLLM,让预填充节点复用解码节点保留的KV,减少多轮智能体对话的重复计算,在EC2 P5en实例上Qwen3-32B的TTFT最多降1.6倍,GLM-5-FP8最多3.3倍,但结果随工作负载和配置而变。

Flow diagram contrasting the standard prefill-to-decode KV path (blue) with the bi-directional additions numbered one through six (red)
配图来源 · AWS Machine Learning

多轮智能体工作负载在智能体进行推理、调用工具并整合先前结果时建立在不断累积的上下文之上。每当出现新一轮对话,客户端都会重新发送对话历史,包括系统指令、先前的提示、工具结果和模型响应,导致不断增长的上下文被反复处理。在 Prefill-decode (P-D) 分离架构中,每个请求被拆分到专用节点上:prefill 节点处理提示并构建其键值 (KV) 缓存,然后 decode 节点检索这些 KV 块并生成响应。

然而,传统的分离式服务只在单个方向上传输 KV 缓存:从 prefill 节点到 decode 节点。这种单向流动给多轮工作负载带来了两个低效问题。首先,在生成响应后,decode 节点会同时持有提示和新生成 token 的 KV 缓存条目。在处理下一轮时,prefill 节点无法检索这些条目,因此必须为之前的响应 token 重新计算 KV 投影。其次,decode 节点维护会话状态而 prefill 节点不维护,这增加了 prefill 节点上活跃会话的缓存被逐出的概率。在完全缓存未命中的情况下,prefill 节点必须为整个对话历史重新计算 KV 投影,包括系统指令、先前的用户请求、工具结果以及模型生成的响应。随着对话的增长,这种冗余工作会消耗更多的加速器容量并增加首 token 时间 (TTFT)。

在构建对话助手和智能体应用时,更长的交互历史可能通过重复的 prefill 计算增加延迟并消耗加速器容量。在这篇文章中,我们展示 双向 KV 缓存传输 如何让 prefill 和 decode 节点跨轮次复用先前已计算的上下文,从而提高加速器利用率,并帮助应用在对话增长时保持响应迅速的 TTFT。

在使用 Amazon Elastic Compute Cloud (Amazon EC2) P5en 实例的 1P1D 部署上进行的多轮基准测试中,双向 KV 传输将 Qwen3-32B 的 TTFT 最多降低了 1.6 倍,将 zai-org/GLM-5-FP8 的 TTFT 最多降低了 3.3 倍(参见性能数据部分)。结果因工作负载和配置而异。这些实例通过 AWS Elastic Fabric Adapter (AWS EFA) 连接,这是一种在 AWS 各服务器之间使用的低延迟远程直接内存访问 (RDMA) 网络。在 支持的实例类型上,AWS EFA 支持 GPUDirect RDMA 和 Neuron Direct RDMA,可以在 prefill 和 decode 节点上的 GPU 或 AWS Trainium 设备内存之间直接移动 KV 块。

我们最初在通过 EFA 连接的 AWS Trainium 实例集群上构建并验证了该特性,并将其作为一项与加速器无关的能力贡献给了 vLLM。本文展示的结果来自配备 EFA 的 Amazon EC2 P5en GPU 集群,表明这些优势同样适用于 GPU。本文介绍了双向 KV 传输架构、路由器元数据流、KV 块保留与重计算控制、性能结果以及如何复现该基准测试。

双向 KV 传输

双向 KV 传输是一种机制,其中 decode 节点在对话的后续轮次中将先前计算的 KV 返回给 prefill 节点。这避免了在 prefill 节点上重复计算共享上下文,并受一个可配置的 KV 重计算阈值控制。它缓解了前文描述的两种计算冗余场景。为使其生效,KV 必须在产生它的那一轮之后仍然存活,并在由具备 KV 缓存感知能力的路由器协调下向反方向传输。

面向双向传输的 NIXL KV 连接器增强

我们扩展了 NVIDIA Inference Xfer Library (NIXL) 的 KV 连接器 ,让 prefill 节点能够消费远程 KV 块,类似于 decode 节点目前已经可以消费 prefill 节点提供的 KV 块的方式。这并非一个简单的改动,因为两种节点类型在架构上不同:prefill 节点为新提示 token 计算 KV 表示,并且能够提供完整的提示 KV 块。decode 节点缺乏(或效率不足以)生成完整 KV 表示的能力,只能提供先前缓存的 KV 块。因此,prefill 和 decode 节点仍各自保留其用于计算远程和本地 KV 块元数据的角色与逻辑。不过,实际的 KV 传输机制(NIXL READ)对于 P->D 和 D->P 两个方向是通用的。

下图首先以蓝色显示现有流程。路由器将请求发送到 prefill 节点,prefill 节点计算 prompt 的 KV 并返回其块元数据,路由器再将请求和元数据转发给 decode 节点以执行标准的 P→D NIXL READ。带编号的红色路径展示了双向 KV 传输的新增内容:(1) decode 节点在生成完成后保留符合条件的 KV 块,(2) 在流式响应中返回其元数据,(3) 路由器缓存该元数据以供下一个 prefill 请求使用。随后 prefill 节点 (4) 在应用本地缓存匹配后检查可远程复用的 KV,(5) 当可复用 token 数量达到阈值时排队一个异步接收,(6) 在计算剩余后缀之前执行 D→P NIXL READ。

Flow diagram contrasting the standard prefill-to-decode KV path (blue) with the bi-directional additions numbered one through six (red)

图 1:标准的从预填充到解码的流程(蓝色)以及双向 KV 传输的新增部分(红色)

感知 KV 缓存的路由器

vLLM 包含连接器侧的双向 KV 传输支持,而在多轮对话间协调 KV 复用则需要具备 KV 缓存感知能力的路由器。vLLM GitHub 提供了一个 示例代理 它实现了本文所使用的流程,你可以根据自身环境进行调整。

KV-cache感知路由器是一个轻量级的有状态控制平面组件,位于客户端、prefill节点和decode节点之间。它维护一个从客户端提供的到内存映射的映射。 conversation_id 解码的 kv_transfer_params。该路由器仅存储块位置元数据,而不存储 KV 张量,并且从响应流中获取该元数据,而不是单独查询解码节点。

下图展示了路由器如何在对话轮次之间管理 KV 元数据。在第一轮中,缓存查找未命中,请求遵循标准的 P→D 流程。路由器随后从响应流中保存解码(decode)的元数据。在后续轮次中,路由器使用该 conversation_id 因此,prefill 可以复用来自 decode 的 KV。然后在生成之后刷新元数据,用于下一轮对话。

Two stacked sequence diagrams focused on the KV-cache-aware router. On turn one, the router records a cache miss, sends the request to Prefill, receives Prefill metadata, sends Decode the request and Prefill metadata, saves Decode metadata from the stream, and returns the response. On a later turn, the router pops saved Decode metadata, sends it to Prefill, receives updated Prefill metadata, sends Decode the request and Prefill metadata, refreshes the Decode metadata entry, and returns the respo

图2:跨对话轮次的KV缓存感知路由器元数据流

区块保留 TTL

通常,decode 实例会在生成完成后立即释放请求的块。在双向 KV 传输模式下,它会在一段可配置的生命周期内保留这些块(decoder_kv_blocks_ttl,默认 480 秒),并发布其位置和过期时间。对话的 KV(包括提示前缀和已生成的响应)随后保持可用,供下一轮使用。对于在生存时间(TTL)到期之前到达的轮次,预填充节点会复用 KV。较晚到达的轮次则回退到正常的预填充。由于此 TTL 是固定的且不会续期,因此需要根据预期的轮次间隔时间对其进行调整,以在 KV 复用与加速器内存占用之间取得平衡。一个自然的后续改进是将基于心跳续期的租约扩展到解码侧。这样,解码实例就可以在对话空闲时立即回收块,而不是在固定计时器结束时回收。

KV 重计算阈值

KV 复用并不总是比本地重计算更快。对于较短的可复用前缀,互连和协调的开销可能超过重新计算这些 token 所需的时间。 kv_recompute_threshold 设置(默认 64 token)控制这一决策。当需要计算的 token 数低于阈值时,预填充节点会进行重计算;当达到或超过阈值时,则从解码节点拉取 KV。部署方可根据模型计算时间、KV 大小和互连性能来调整此值。

性能数据

为了进行评估,我们在使用 EFA 连接的 Amazon EC2 P5en 实例上的 1P1D 部署中使用了 Qwen3-32B 和 GLM-5-FP8 模型。为了模拟预填充节点已逐出其本地缓存,或路由器选择了不具备该对话 KV 块的其他预填充节点的场景,我们通过以下方式禁用了预填充节点上的前缀缓存: --no-enable-prefix-caching.

在两个模型上,开启双向 KV 传输都能减少多轮对话的预填充计算,而且收益通常随提示长度增长而增加。在该功能关闭(OFF)时,预填充实例每轮都会重新计算整个增长后的上下文,因此 TTFT 急剧上升。Qwen3-32B 从 2k 时的 158 ms 升至 24k 时的 778 ms,GLM-5-FP8 从 2k 时的 500 ms 升至 20k 时的 2,334 ms。开启(ON)时,解码实例保留的 KV 被复用而非重计算,使 TTFT 曲线平坦得多。Qwen3-32B 在相同范围内仅从 122 ms 升至 460 ms(在最长提示下约降低 1.6 倍),而 GLM-5-FP8 从 358 ms 升至 702 ms,在长提示下比重计算路径最多低 3.3 倍。预填充时长分解说明了原因。

双向传输会增加少量 KV 传输时间(Qwen3-32B 约 30 ms 到 130 ms;GLM-5-FP8 约 29 ms 到 110 ms),但它省去的预填充计算量远超这一 KV 传输时间。因此,开启传输时的总预填充时长(Qwen3-32B:128 ms 到 414 ms;GLM-5-FP8:316 ms 到 541 ms)始终远低于仅重计算路径(Qwen3-32B:153 ms 到 744 ms;GLM-5-FP8:500 ms 到 2,331 ms)。更大的 GLM-5-FP8 模型显示出更大的绝对和相对差距,因为其更高的单 token 重计算时间使得随着模型和上下文规模的扩大,KV 复用的收益更加显著。

Qwen3-32B

下图绘制了 Qwen3-32B 的 TTFT 与提示长度的关系,对比了双向 KV 传输开启与关闭的情况。

Line chart of Qwen3-32B time to first token versus prompt length with bi-directional KV transfer on and off

图 3:双向 KV 传输开启和关闭时 Qwen3-32B 的 TTFT 与提示长度对比

下一张图将相同运行的预填充时长分解为重计算和 KV 传输两个组成部分。

Bar chart breaking Qwen3-32B prefill duration into recompute and KV-transfer components

图 4:Qwen3-32B 预填充时长分解为重计算和 KV 传输组成部分

两幅图表合在一起显示,在功能开启时 TTFT 几乎保持平稳,而关闭时随着上下文增长 TTFT 陡峭上升,因为较小的 KV 传输时间被其省去的重计算远远抵消。

zai-org/GLM-5-FP8

GLM-5-FP8 也呈现相同的两种视图。第一幅图绘制了功能开启与关闭情况下 TTFT 随提示长度的变化。

Line plot for zai-org/GLM-5-FP8 on Amazon EC2 P5en instances with AWS EFA, comparing time to first token with bi-directional KV transfer enabled and disabled across input prompt lengths.

图 5:zai-org/GLM-5-FP8 在 KV 复用开启和关闭时的 TTFT

第二幅图将预填充时长分解为其重计算和 KV 传输组成部分。

Grouped bar plot for zai-org/GLM-5-FP8 on Amazon EC2 P5en instances with AWS EFA, comparing prefill duration with bi-directional KV transfer enabled and disabled and showing the KV-transfer component.

图 6:zai-org/GLM-5-FP8 预填充时长分解

GLM-5-FP8 显示出比 Qwen3-32B 大得多的差距,在长提示下 TTFT 最多降低约 3.3 倍,因为其更高的每 token 重计算时间使得随着上下文长度增长,KV 复用的收益更加显著。

如何运行基准测试

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

结论

在这篇文章中,我们展示了双向 KV 缓存传输如何在智能体的对话历史增长时降低 TTFT。在 P-D 分离部署中,该功能让预填充节点跨轮次复用解码节点保留的 KV 块,从而减少重复计算。可配置的块保留和重计算阈值决定了何时进行复用。在使用通过 EFA 连接的 Amazon EC2 P5en 实例的 1P1D 部署中,在测试的最长提示长度下,我们观察到 Qwen3-32B 的 TTFT 降低了 1.6 倍,zai-org/GLM-5-FP8 降低了 3.3 倍。结果取决于工作负载和部署配置。要在您的环境中实现此功能,请使用 vLLM v0.21.0 或更高版本并搭配 NixlConnector,在预填充和解码服务器上启用 bidirectional_kv_xfer ,使用支持 KV 缓存感知的路由器,并遵循 GitHub 重现指南.

相关资源

致谢

我们感谢 vLLM 社区在维护此功能方面给予的支持。


关于作者

原始出处

AWS Machine Learning

内容说明

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

机器翻译 · 请以原文为准