在规模化环境下使用 vLLM x Mooncake 提供 Agentic 工作负载

Agentic 工作负载产生大量可复用前缀

Agentic 工作负载会生成长上下文且高度复用,使得前缀缓存至关重要。 Codex/SWE‑bench Pro 追踪显示,在第 30 回合时上下文长度约为 80 K token,且可超过 180 K token,而每回合仅新增几百到几千个 token。 在 610 条追踪(每条追踪的中位回合数为 33)中,数据集表现为:

  • 94.2 % 的缓存命中率
  • 131:1 的输入‑输出 token 比率
  • 平均每回合上下文增长约 2 242 token
  • 中位上下文增长从 12 K 增至 80 K token 每条追踪
  • 回合间延迟从中位数 5.2 s 到 P99 81.4 s 不等 这些数字证实大部分输入都是可复用的前缀,缓存它们即可消除冗余的预填充。

使用 Mooncake Store 的分布式 KV 缓存池解决前缀共享

集成 Mooncake Store 可提供共享 KV 缓存,提升命中率并实现跨实例复用。 Mooncake 是一个用于 KV 缓存传输和分布式存储的开源库。 vLLM 已经通过 MooncakeConnector 使用 Mooncake 实现了预填充‑解码的解耦。 新的集成构建了一个分布式 KV 缓存池,多个 vLLM 实例嵌入 Mooncake 客户端并连接到集群范围的 Mooncake Store 主节点。 主节点管理 KV‑块元数据、服务发现以及客户端健康状态;工作节点在 GPU HBM 与分布式 DRAM/SSD 池之间通过 RDMA 传输 KV 块。 在调度器端,vLLM 对提示 token 块进行哈希,查询 Mooncake 主节点以获取匹配的 KV 缓存块,并利用结果指导调度。 在工作节点端,每个 GPU 工作节点运行 Mooncake 客户端,将 GPU KV 缓存内存注册为 RDMA 缓冲区,并通过后台 I/O 线程移动数据。

设计亮点:无 SM 的 GPUDirect RDMA、异步传输、MultiConnector

实现采用 GPUDirect RDMA 实现零拷贝,在后台 I/O 线程上执行传输,并通过 MultiConnector 与 PD 解耦配合工作。 KV 块移动利用 GPUDirect RDMA,避免了暂存缓冲区和 SM 消耗,并且可以聚合多个 RNIC 以获得更高带宽。 所有 RDMA 操作均在专用的后台 I/O 线程上运行,使传输路径完全异步,防止主 CPU 路径出现阻塞。 MultiConnector 接口将子连接器串联;预填充实例将 KV 块存入分布式池,并在缓存命中时检索;解码实例则将块写入池,以供预填充实例可见。

在 Codex 追踪上的性能结果与扩展性

在真实的 Agentic 追踪上,缓存命中率从 1.7 % 提升至 92.2 %,带来 3.8 倍吞吐量、46 倍更低的 TTFT、8.6 倍更低的延迟;扩展到 60 台 GB200 GPU 时仍保持 >95 % 的命中率并实现近线性吞吐增长。 在 Codex Agentic 追踪实验(1P1D,12 台 GB200 GPU)中,分布式 KV 缓存池使吞吐量提升 3.8 倍,P50 TTFT 降低 46 倍,端到端延迟降低 8.6 倍。 这些提升来源于缓存命中率从 1.7 %(仅系统提示被缓存)上升至 92.2 %(几乎整个前缀被缓存)。 为了评估扩展性,使用了从 Codex 工作负载衍生的合成数据集,并在节点之间进行轮询路由。 系统在所有规模下均保持超过 95 % 的缓存命中率,并且吞吐量从 12 台到 60 台 GB200 GPU 几乎呈线性增长,表明分布式池消除了跨实例的未命中,从而避免了性能下降。

未来工作

计划的扩展包括分布式磁盘离线、混合模型支持、缓存感知路由以及数据通路优化。 未来工作将把存储层次扩展到 NVMe SSD 和分布式文件系统,以提供更大的缓存容量。 将加入对混合模型(具有混合注意力机制)的支持,允许针对不同层采用不同的缓存策略。 缓存感知路由将与 KV 缓存池共同设计请求路由器,将回合直接调度到持有相关前缀的实例,以最大化本地命中,未命中时再回退到分布式池。 数据通路优化将探索 NVIDIA 多节点 NVLink 以及类似 DualPath 的同步 KV 加载(同时来自预填充和解码实例),以提升整体带宽。

致谢

vLLM 与 Mooncake Store 的集成受到 vLLM‑Ascend 先前工作的启发。 感谢来自蚂蚁集团的 Lei Chao 为最初实现所作的贡献,以及来自 Inferact 的 Liu Zijing 提供的 Agentic 追踪和分析。 特别感谢 Approaching.AI 的 Lu Jiahao、Zhang Zuoyuan、Tang Zihan、Yang Ke;华为的 Zhao Pengbo、Duan Fuqiao、Xu Tianyu;阿里云计算的 Ding Tianchen、Shang Xuchun、Yi Xingrui、Ma Teng;蚂蚁集团的 Ning Yunxiao、Zhu Dejiang、Zheng Shoujian;以及 9#AISoft 的 Ren Feng 提供的技术反馈。 同时感谢更广泛的 vLLM 与 Mooncake 社区的支持,以及 Inferact 团队的紧密合作。

Sources