跳到正文
AWS Machine Learning Blog· Sunita Nadampalli·· 13 小时前精选AI 评分64

vLLM 双向 KV 传输如何降低多轮智能体工作负载的 TTFT

Improve TTFT for multi-turn agentic workloads with bi-directional KV transfer in vLLM

AI 导读

AWS介绍了 vLLM 的双向 KV cache 传输机制,让 Prefill 节点在多轮智能体工作负载中复用 Decode 节点保留的上下文,减少重复计算并降低 TTFT。

推荐理由

文章串联了双向 KV 传输的架构、路由元数据与阈值配置,并用两组模型的 TTFT 数据展示其在长上下文多轮场景中的收益。

正文 · AI 翻译

多轮智能体工作负载会随着智能体推理、调用工具并整合先前结果,逐步累积上下文。每增加一轮,客户端都会重新发送对话历史,包括系统指令、先前的提示、工具结果和模型响应,导致不断增长的上下文被反复处理。在预填充-解码(P-D)分离架构中,每个请求会分配给专用节点处理:预填充节点处理提示并构建键值(KV)缓存,然后解码节点获取这些 KV 块并生成响应。

然而,传统的分离式服务只会单向传输 KV 缓存:从预填充节点传输到解码节点。这种单向传输会给多轮工作负载带来两种低效情况。首先,生成响应后,解码节点会保留提示和新生成令牌对应的 KV 缓存条目。处理下一轮时,预填充节点无法获取这些条目,因此必须重新计算先前响应令牌的 KV 投影。其次,解码节点会维护会话状态,而预填充节点不会,因此预填充节点上活跃会话的缓存被驱逐的概率会增加。如果缓存完全未命中,预填充节点就必须重新计算整个对话历史的 KV 投影,包括系统指令、先前的用户请求、工具结果和模型生成的响应。随着对话不断增长,这种重复工作会消耗更多加速器容量,并增加首令牌时间(TTFT)。

在构建对话助手和智能体应用时,较长的交互历史可能会因重复的预填充计算而增加延迟并消耗加速器容量。本文将介绍双向 KV 缓存传输如何让预填充节点和解码节点跨轮次复用先前计算的上下文,从而提高加速器利用率,并帮助应用在对话不断增长时维持较短的 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,可在预填充节点和解码节点之间直接传输 GPU 或 AWS Trainium 设备内存中的 KV 块。

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

双向 KV 传输

双向 KV 传输是一种机制:在对话的后续轮次中,解码节点会将之前计算的 KV 返回给预填充节点。这样可以避免预填充节点重新计算共享上下文,并由可配置的 KV 重计算阈值控制。它可以缓解前文所述的两种计算冗余场景。要实现这一机制,KV 必须在生成它的轮次结束后继续保留,并在 KV 缓存感知型路由器的协调下反向传输。

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

我们扩展了 NVIDIA Inference Xfer Library (NIXL) KV 连接器,使预填充节点能够使用远程 KV 块,类似于解码节点目前使用预填充节点提供的 KV 块。这并非简单的改动,因为这两种节点在架构上有所不同:预填充节点会为新的提示词 token 计算 KV 表示,并能提供完整的提示词 KV 块。解码节点缺乏生成完整 KV 表示的能力(或效率),只能提供先前缓存的 KV 块。因此,预填充节点和解码节点仍各自承担计算远程和本地 KV 块元数据的职责并运行相应逻辑。不过,实际的 KV 传输机制(NIXL READ)在 P->D 和 D->P 两个方向上是通用的。

下图首先以蓝色展示现有流程。路由器将请求发送给预填充节点,预填充节点计算提示词 KV 并返回其块元数据,随后路由器将请求和元数据转发给解码节点,执行标准的 P→D NIXL READ。编号的红色路径展示了新增的双向 KV 传输流程:(1) 解码节点在生成后保留符合条件的 KV 块;(2) 在流式响应中返回这些 KV 块的元数据;(3) 路由器缓存这些元数据,以供下一次预填充请求使用。随后,预填充节点会 (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 缓存感知型路由器是一个轻量级、有状态的控制平面组件,位于客户端、预填充节点和解码节点之间。它会维护从客户端提供的 conversation_id 到解码节点的 kv_transfer_params 的内存映射。路由器只存储块位置元数据,不存储 KV 张量,并从响应流中获取这些元数据,而不是单独查询解码节点。

下图展示了路由器如何在多个对话轮次中管理 KV 元数据。第一轮中,缓存查找未命中,请求按标准 P→D 流程处理。随后,路由器保存从响应流中获取的解码节点元数据。在后续轮次中,路由器使用 conversation_id 检索这些元数据,使预填充节点能够复用解码节点的 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 response. Purple arrows show the Decode-to-Prefill and Prefill-to-Decode KV reads; KV tensors do not pass through the router.

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

块保留 TTL

通常,解码实例会在生成完成后立即释放请求的块。在双向 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 传输都会减少多轮对话中的预填充计算量,且提示越长,收益通常越大。关闭该功能时,预填充实例每轮都会重新计算整个增长后的上下文,因此 TTFT 会急剧上升。Qwen3-32B 的 TTFT 从 2k 时的 158 ms 升至 24k 时的 778 ms;GLM-5-FP8 则从 2k 时的 500 ms 升至 20k 时的 2,334 ms。启用该功能后,解码实例保留的 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:启用和禁用 KV 复用时 zai-org/GLM-5-FP8 的 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 复用带来的收益也更大。

如何运行基准测试

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

结论

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

相关资源

致谢

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


作者简介

Sunita Nadampalli

Sunita Nadampalli

Sunita 是 AWS 的首席工程师和 AI/ML 专家。她负责 AWS Trainium 和 Graviton 针对 AI/ML 工作负载的软件性能优化。她热衷于开源软件开发,致力于交付高性能且可持续的软件解决方案。

来源:AWS Machine Learning Blog · aws.amazon.com