工程 / 工程
双向 KV 缓存传输如何将多轮智能体工作负载的首字延迟压缩至原来的三分之一
在多轮智能体应用中,每一轮对话都需要将完整的历史上下文重新发送给模型处理,随着对话增长,预填充节点反复重算 KV 投影的开销急剧攀升,首字延迟(TTFT)随之恶化。传统的预填充-解码(P-D)分离架构仅支持从预填充节点向解码节点单向传输 KV 缓存,解码节点生成回复后积累的 KV 块无法被下一轮的预填充节点复用,造成双重浪费:一是解码节点生成的 token 对应的 KV 需重新计算,二是预填充节点缓存命中率低,极端情况下需对整段对话历史全量重算。AWS 在 vLLM 中引入双向 KV 传输机制,通过扩展 NIXL KV 连接器并配合具备 KV 缓存感知能力的路由器,实现解码节点将已缓存的 KV 块元数据返回给路由器,供下一轮预填充节点直接读取复用。在搭载 AWS EFA 网络的 Amazon EC2 P5en 实例上,Qwen3-32B 的 TTFT 最高降低 1.6 倍,GLM-5-FP8 最高降低 3.3 倍。
多轮智能体工作负载的核心挑战在于上下文的持续累积。每当智能体完成一轮推理、调用工具或整合历史结果后,客户端需将包含系统指令、历史提示、工具返回值及模型回复的完整对话历史重新提交。在预填充-解码分离架构下,预填充节点负责处理提示并构建 KV 缓存,解码节点则读取这些 KV 块并生成回复。然而传统架构中 KV 缓存仅能从预填充节点流向解码节点,形成单向瓶颈。
单向传输带来两类具体低效。其一,解码节点在生成回复后持有包含提示与新生成 token 的完整 KV 缓存,但下一轮的预填充节点无法获取这些数据,只能重新计算上一轮回复 token 的 KV 投影。其二,解码节点长期维护会话状态,而预填充节点不保留,导致活跃会话在预填充侧的缓存极易被驱逐。一旦完全缓存未命中,预填充节点须对整段对话历史全量重算,随着对话轮次增加,这一冗余开销将持续消耗加速器算力并拉高 TTFT。
为解决上述问题,AWS 团队扩展了 NVIDIA 推理传输库(NIXL)的 KV 连接器,使预填充节点能够像解码节点消费预填充侧 KV 块一样,主动读取解码节点上的远程 KV 块。具体流程为:解码节点在生成结束后保留符合条件的 KV 块,并在流式响应中返回其块元数据;路由器缓存该元数据并与对话 ID 绑定;下一轮请求到来时,预填充节点先匹配本地缓存,再检查可复用的远程 KV 块,当可复用 token 数量达到可配置阈值后,通过异步 NIXL READ 从解码节点拉取 KV 块,仅对剩余后缀执行计算。
整套机制的协调依赖一个轻量级的 KV 缓存感知路由器。该路由器作为客户端、预填充节点与解码节点之间的有状态控制面组件,在内存中维护对话 ID 到解码节点 KV 传输参数的映射关系。路由器仅存储块位置元数据而非 KV 张量本身,元数据直接从响应流中提取,无需额外查询解码节点。首轮对话走标准 P→D 流程,路由器保存元数据;后续轮次路由器凭对话 ID 检索元数据,供预填充节点复用,生成结束后再刷新元数据以备下一轮使用。
在块保留策略上,解码节点默认在生成结束后立即释放请求对应的 KV 块,而在双向传输模式下,这些块将被保留一段可配置的生存时间(默认 480 秒),以供后续轮次的预填充节点读取。此外,系统还提供 KV 重计算阈值参数,只有当可复用的 token 数量超过该阈值时,才会触发跨节点传输,避免因传输少量 KV 块而引入不必要的网络开销。
性能测试在搭载 AWS EFA 低延迟 RDMA 网络的 Amazon EC2 P5en GPU 实例上以 1P1D 部署方式进行。EFA 支持 GPUDirect RDMA,可在预填充与解码节点的 GPU 显存之间直接传输 KV 块,无需经过 CPU 中转。结果显示,双向 KV 传输使 Qwen3-32B 的 TTFT 最高降低 1.6 倍,GLM-5-FP8 最高降低 3.3 倍。该特性最初在 AWS Trainium 实例集群上构建验证,随后作为加速器无关能力贡献至 vLLM 开源社区,GPU 侧的测试结果证明其收益同样适用于 GPU 工作负载。
要点
- 双向 KV 传输通过让解码节点将已缓存的 KV 块元数据返回给路由器,使下一轮预填充节点可直接复用,避免对历史上下文的重复计算。
- KV 缓存感知路由器是实现跨轮次 KV 复用的关键控制面组件,仅存储块位置元数据,开销极低,可按需适配不同部署环境。
- 在 EC2 P5en 实例上,Qwen3-32B 的 TTFT 最高降低 1.6 倍,GLM-5-FP8 最高降低 3.3 倍,对话轮次越多、历史越长,收益越显著。
- 该特性已作为加速器无关能力贡献至 vLLM 开源社区,同时适用于 GPU 和 AWS Trainium 加速器,并通过 EFA RDMA 实现高效的跨节点 KV 传输。
- 可配置的块保留 TTL 与 KV 重计算阈值参数,使用户能够在显存占用与传输收益之间灵活权衡,适应不同业务场景。
原始标题:Improve TTFT for multi-turn agentic workloads with bi-directional KV transfer in vLLM
本文由 DataHub 基于公开来源整理,用于信息发现与摘要阅读;具体事实、数据和后续更新以原始来源为准。