Google 开源 TPU Sync：面向 TPU 的 KV cache 传输 C++ 库
Google Cloud 发布开源 C++ 库 TPU Sync（Apache-2.0），把 TPU HBM、主机 DRAM、共享存储和数据中心网络当作统一内存层级，用异步 DMA 传输 KV cache 和模型权重。
MarkTechPost· Michal Sutter·2026-11-13 18:32· 1 小时前AI 评分52
Google 发布 TPU Sync：面向 TPU 的开源 KV Cache 传输方案
AI 导读
Google Cloud 发布开源 C++ 库 TPU Sync（Apache-2.0），把 TPU HBM、主机 DRAM、共享存储和数据中心网络当作统一内存层级，用异步 DMA 传输 KV cache 和模型权重。
正文 · 原文
Google Cloud 发布了 TPU Sync，这是一个开源 C++ 库，用于在 Cloud TPU 上搬运 KV cache 和模型权重。Google TPU Sync KV cache 传输库将 TPU HBM、主机 DRAM、共享存储和数据中心网络视为一个统一的内存层级。它针对的是一个日益突出的瓶颈：搬运模型状态，而非计算模型状态。Google 报告称，在 235B 至 480B 参数的 MoE 模型上，KV 传输延迟降低超过 90%。
TL;DR
- 库：原生 C++ 核心，提供面向 PyTorch 和 JAX 的 Python wheel。已在 235B 至 480B MoE 模型上测试。
- 运行于：根据 README，支持 Cloud TPU v7x（Ironwood）和 v6e。基准测试使用 TPU v7x（2x2x1）、PCIe Gen5 和 200 Gbps 网卡。
- 性能：在 Google 内部实验室测试中，KV 传输延迟和首 token 时间（TTFT）大幅降低。
- 最佳：平均 TTFT 最高降低 90%（480B P/D 服务在低并发下，以及 64K 上下文前缀缓存）。
- 最差：未公布 RL 数据。README 称目前尚不推荐用于一般用途。
- 结论：最佳：Qwen3 235B 的 KV 传输从数秒降至约 230 ms。最差：所有数字均为自报，且处于早期阶段。
什么是 TPU Sync？
TPU Sync 是 Google 推出的开源 C++ 库，通过异步 DMA 在内存层级之间搬运 TPU 分块张量。它处理 2 项任务：推理期间的 KV cache 状态，以及强化学习（RL）中从训练器到采样器的权重广播。代码仓库采用 Apache-2.0 许可。
TPU Sync 如何工作？
该库覆盖 TPU HBM、锁页主机 DRAM、共享存储和多主机 DCN。存储使用实验性 POSIX 后端，包括本地 SSD 和 Google Cloud Managed Lustre。Google 列出了 4 项设计选择：
- 保持布局的直接 DMA：它就地读取 JAX PJRT 缓冲区或 PyTorch 张量。数据保留其原生 TPU 分块布局，无需 Python 序列化或 bounce buffer。
- 与计算重叠的传输：原始 PJRT 缓冲区别名让传输可在计算期间进行。内置 DMA 节流可防止大块拷贝延迟内核启动。
- 数据路径中无 XLA 编译：传输和重分片在 XLA 图之外以原生 C++ 运行。
- 高带宽 PCIe 和 DCN：每台 Ironwood 主机最高可达 223 GB/s，而 PCIe Gen5 上限为 256 GB/s。DCN 流式传输在 1 块网卡上达到约 192 Gbps，在 2 块网卡上约 350 Gbps。
在此之上是用于前缀查找的全局 KV 存储、传输控制器和 KV cache 管理器。
分离式 prefill/decode 服务如何使用它？
prefill 节点通过异步 DMA 将 KV 块卸载到锁页主机 RAM。控制器将它们以主机到主机方式流式传输到指定的 decode 节点。到达的块在解码继续进行的同时移入 decode 节点 HBM。
TPU Sync 基准测试显示了什么？
Google 于 2026 年 7 月在 TPU v7x（2x2x1）上进行了内部测试。技术栈为 tpu-inference v0.24.0、JAX 0.10.2 和 libtpu 0.0.42.1。Google 团队表示实际结果会有所不同。
对于 P/D 分离式服务：
- KV 传输开销下降超过 90%。Qwen3 235B 的 KV 传输延迟中位数从数秒降至约 230 ms。
- 在低并发下，235B 的平均 TTFT 下降 60%，480B 下降 90%。
- 在流量高峰期间，235B 的 token 输出吞吐量提升 30% 至 70%。
关于前缀缓存：
- 当 MoE 权重占据大部分 HBM 时，有效缓存容量提升了一个数量级。
- 在 64K 上下文工作负载上，平均 TTFT 最多下降 90%。
- 吞吐量提升超过 50%，P99 的 token 间延迟下降 30%。
TPU Sync 支持哪些框架？
TPU Sync 通过 TorchTPU 与 PyTorch 集成，并原生支持 JAX。在 vLLM-TorchTPU 中，它是 vLLM 的 KVConnector 接口背后的 TPU KV 连接器。在 SGLang-TorchTPU 中，它是 P/D 传输后端和分层缓存 I/O 后端。在 GKE 上，llm-d 1.0 将其用于前缀缓存亲和性路由，并实验性地溢出到 Managed Lustre。Inferact 和 RadixArk 共同参与了 vLLM 和 SGLang 集成的工程开发。Google 表示 Meta 和 Wayfair 正在评估该架构。
TPU Sync 与 GPU 侧替代方案相比如何？
| 特性 | Google TPU Sync | NVIDIA Dynamo KVBM + NIXL | LMCache | Mooncake |
|---|---|---|---|---|
| 目标硬件 | Cloud TPU v7x、v6e（README） | NVIDIA CUDA GPU；NIXL 中部分支持 ROCm（NIXL） | CUDA 和 ROCm（repo） | CUDA、HIP、Ascend、MUSA 等（repo） |
| 内存层级 | HBM、固定主机 DRAM、POSIX 存储（实验性）、DCN | GPU、CPU、本地 SSD、远程存储（KVBM 文档） | GPU、CPU DRAM、本地磁盘、远程后端 | DRAM 和 SSD/NVMe（Mooncake Store） |
| 网络传输 | DCN、多网卡条带化 | UCX、LIBFABRIC、GDS 插件 | NVLink、RDMA、TCP | TCP、RDMA、EFA、NVMe-oF、NVLink、CXL |
| 引擎集成 | vLLM-TorchTPU、SGLang-TorchTPU、llm-d、JAX | vLLM、TensorRT-LLM、SGLang（KVBM） | vLLM、NVIDIA Dynamo | vLLM、SGLang |
| 许可证 | Apache-2.0 | Apache 2.0（部分 wheel 模块为 NVIDIA 专有） | Apache 2.0 | Apache-2.0 |
| 核心语言 | C++ 及 Python 绑定 | C++ 及 Python 和 Rust 绑定 | Python 包，含 C/C++ 和 Rust 代码 | 未披露 |
| 标称带宽 | 每 v7x 主机最高 223 GB/s（PCIe D2H/H2D） | 未披露 | 未披露 | 在 8×400 Gbps RoCE（网络）上最高 190 GB/s |
| RL 权重传输 | 是（实验性） | 未披露 | 未披露 | 是，通过 SGLang（RDMA） |
这些带宽数字测量的是不同路径。TPU Sync 的数字是主机 PCIe DMA。Mooncake 的数字是 RoCE 上的网络传输。
关键要点
- TPU Sync 是 Google 的 Apache-2.0 C++ 库，用于在 TPU 上移动 KV 缓存和权重。
- 它通过直接、保持布局的 DMA，每台 Ironwood 主机最高可达 223 GB/s。
- 在 Google 的测试中，Qwen3 235B 的 KV 传输从数秒降至约 230 ms。
- 它可接入 vLLM-TorchTPU、SGLang-TorchTPU 以及 GKE 上的 llm-d。
- RL 和存储分层仍为实验性，所有数字均为自报。
查看技术博客、GitHub 仓库、vLLM-TorchTPU 和 SGLang-TorchTPU。另外，欢迎在 Twitter 上关注我们，别忘了加入我们的 150k+ML SubReddit 并订阅我们的 Newsletter。等等！你在 Telegram 上吗？现在你也可以在 Telegram 上加入我们了。
需要与我们合作推广你的 GitHub 仓库、Hugging Face 页面、产品发布或网络研讨会等？联系我们
文章 Google Releases TPU Sync: Open-Source KV Cache Transfer for TPUs 首次发布于 MarkTechPost。
来源：MarkTechPost · marktechpost.com
来源：RadarAI · aihot.news