谷歌开源TPU Sync：把KV缓存与权重传输时延降低逾90%，服务推理与强化学习
谷歌云于11月9日开源TPU Sync，这是面向TPU集群的模型状态层库，用异步DMA在HBM、主机内存、共享存储与数据中心网络之间直接搬移张量。官方称，在前缀缓存与分离式推理负载中，相对未使用TPU Sync，KV传输时延可降低逾90%，重型开源MoE（2350亿至4800亿参数）的首字时延平均可降约90%。已对接vLLM、SGLang与llm-d。
导读
2026年11月9日，谷歌云杰出工程师Fenghui Zhang与高级产品经理Weiren Yu发文，介绍开源库 TPU Sync。它被设计为TPU集群上的模型状态层：既搬推理时存在海量KV缓存里的上下文状态，也搬强化学习中从训练器更新并广播到采样器的参数状态。
官方的核心观点是：加速器的HBM、主机DRAM、共享存储和数据中心网络（DCN）应当作为一套协同设计的存储层次来用。TPU Sync是原生C++库，以异步DMA直接按TPU的分块布局搬运张量，并与正在进行的计算重叠。对推理，它把TPU计算与网络反压解耦，把主机DRAM、本地SSD和Google Cloud Managed Lustre变成自动分层的KV缓存；对强化学习，它向上层框架暴露底层DMA与DCN原语，用于把更新后的权重从训练器流到采样器。
宣称的效果与架构要点
官方称，通过就地访问设备缓冲区、而不是拷贝或锁定它们，TPU Sync在前缀缓存和分离式服务负载中，能把KV传输时延稳定降低逾 90%；对参数量在 2350亿至4800亿 的重型开源MoE模型，首字时延（TTFT）平均可降低约 90%（对比未使用TPU Sync的情况）。这些数字来自谷歌2026年7月在TPU v7x上的内部测试，需按脚注中的配置理解。
TPU Sync直接接入现有的PyTorch（经TorchTPU）与JAX栈：作为vLLM、SGLang和llm-d分离式服务的TPU KV缓存连接器，同时为Pathways和多控制器JAX（McJAX）的强化学习流水线提供底层传输原语。
架构上的四点被着重写出：
- 保持布局的直接DMA与就地缓冲区别名：就地访问JAX PJRT缓冲区或PyTorch张量，在TPU HBM与页锁定（`DmaMap`注册）的主机内存之间单次硬件DMA搬移切片，并在HBM、主机DRAM、存储和网络上保持原生TPU物理分块布局，避免Python序列化、布局转换和用户态中转缓冲。
- 与计算重叠的传输：不为静态KV缓存和权重缓冲区加独占锁或做片上拷贝，而是创建不持有PJRT使用锁的原始别名，使块级传输可与就地计算并发，并带有DMA节奏控制，以免大块传输拖延内核启动。
- 数据路径不经过XLA编译：传输与感知分块的跨拓扑重分片在原生C++中、在XLA图外执行。
- 高带宽PCIe DMA与DCN流式传输：在Ironwood（TPU v7x）主机上，聚合的设备到主机与主机到设备吞吐最高可达 223 GB/s（PCIe Gen5硬件上限为256 GB/s）；单张200 Gbps网卡上DCN流式传输最高约 192 Gbps，双200 Gbps网卡做NUMA感知条带化时约 350 Gbps。
上层还有全局KV存储（按提示前缀哈希映射物理缓存块、LRU淘汰与分层降级）、TPU Sync控制器（生成传输计划、跨节点协调、租约与恢复）以及本地KV缓存管理器（跟踪块状态，管理HBM、主机DRAM，包括`/dev/shm`热启动持久化，以及实验性的POSIX存储后端）。
分离式推理与前缀缓存
在预填充/解码（P/D）分离部署中，TPU Sync的三步流水线是：预填充节点处理完初始提示后，立刻经异步D2H DMA把KV块搬进页锁定主机内存；控制器再做主机到主机的高带宽传输，把上下文流到指定解码节点；块到达解码节点主机内存后，再异步H2D写入TPU HBM，并与正在进行的解码计算重叠。
前缀缓存方面，TPU Sync把前缀缓存从HBM（第一层）扩展到分布式主机内存池（第二层），并可用实验性POSIX后端把更旧的KV块异步溢写到本地SSD或共享的Managed Lustre（第三层以上），再配合GKE上llm-d的前缀亲和路由，让副本跨层找回缓存前缀，而不必重新计算提示。
官方给出的部分内部测试结果（脚注：2026年7月，TPU v7x 2x2x1拓扑，每芯片PCIe Gen5 x16，每主机一或两张200 Gbps网卡，tpu-inference v0.24.0，JAX 0.10.2，libtpu 0.0.42.1）：
- 对Qwen3 235B，KV传输时延中位数从数秒降到约 230毫秒，降幅逾90%
- 在受控、低并发测试中，235B模型的TTFT平均降低约 60%，480B模型平均降低约 90%
- 在持续流量高峰中，235B的token输出吞吐平均提高约 30%到70%，随流量分布而变化
- 在64K上下文的前缀缓存负载上，平均TTFT最多可降约90%，吞吐提高逾50%，P99词间时延降低约30%
生态与获取方式
TPU Sync已作为开源提供，并带有PyTorch（`tpu-sync-torch`）与JAX（`tpu-sync-jax`）的Python包，集成到vLLM-TorchTPU、SGLang-TorchTPU以及GKE上的llm-d。合作方包括vLLM-TorchTPU维护者Inferact、SGLang-TorchTPU维护者RadixArk。官方还提到，Meta等前沿机构以及Wayfair等企业，已在评估或采用相关的PyTorch-on-TPU服务栈与多层架构。代码与文档在GitHub仓库 `google/tpu-sync`。
来源
Google Cloud Blog 2026-11-09，Fenghui Zhang与Weiren Yu，Introducing TPU Sync: Fast KV cache and weight transfer for inference and RL on TPU。文中性能数字来自谷歌内部实验室测试，实际效果随系统配置、网络带宽、并发与负载变化。