谷歌开源 TPU Sync：在 TPU 之间高速搬运 KV 缓存的新软件库

KV 缓存（键值缓存）是语言模型在读取提示词时建立的工作记忆：对于每一个已处理的词元（token），模型都会保存中间数值，这样在生成每个新词时就不必重新计算。对话或文档越长，这个缓存就越大，而它通常存放在加速器芯片上稀缺且昂贵的高带宽内存（HBM）中。

谷歌近日发布了 TPU Sync，这是一个开源的原生 C++ 库，用于在其 TPU（张量处理单元）之间搬运并管理这种缓存。据 TechBytes 报道，它概述了谷歌云博客上由工程师 Fenghui Zhang 和 Weiren Yu 撰写的文章，这个库还会在强化学习流程中同步模型参数。SemiAnalysis 指出，TPU Sync 以前名为 TPU-raiden，而 Open Source For You 则以 TPU Raiden 之名介绍该项目，并称其采用 Apache 2.0 许可证。

为什么 KV 缓存需要搬运？

Open Source For You 解释说，语言模型的推理分为两个阶段：预填充（prefill），模型读取整个提示词；解码（decode），模型逐个词元地生成回答。分离式服务把这两个阶段放在不同的硬件池上运行，因此预填充结束后，KV 缓存必须被送到负责解码的芯片上。SemiAnalysis 指出，没有传输库的话，TPU 的外部服务只能停留在聚合模式，尽管谷歌多年来一直在内部为 Gemini 使用分离式服务。

第二个需求是卸载。SemiAnalysis 解释说，当一块芯片的 HBM 装不下所有用户的缓存时，数据必须从 HBM 移到主机的 DRAM，再移回来。据 The Next Platform 2026 年 11 月 9 日报道，谷歌 AI 与计算基础设施副总裁兼总经理 Mark Lohmeyer 在电话会议上表示，在智能体类工作负载中，TPU 很大一部分时间处于闲置状态。他说，TPU Sync 会在智能体暂停时把工作记忆移出芯片，在智能体恢复时再还原，这样加速器就能腾出来处理其他任务，并降低每个智能体工作负载的成本。The Next Platform 还补充说，由于 TPU Sync 是软件，谷歌可以把它回溯移植到更早几代的 TPU 系统上。

谷歌宣称的性能如何？

以下数字来自日期为 2026 年 11 月 9 日的 TechBytes，该媒体将其归于谷歌云博客。我们未能读到谷歌的原文，因此这些是厂商公布、由二手媒体转述的数字。

• 据 TechBytes 报道，在参数规模为 2350 亿至 4800 亿的开放权重混合专家（Mixture of Experts）模型上，键值传输延迟降低超过 90%。
• 据 TechBytes 报道，首个词元时间（TTFT，即回答的第一个字出现前的等待时间）在前缀缓存和分离式服务的基准测试中，与标准配置相比平均下降 90%。
• 据 TechBytes 报道，raw_transfer 组件在 Ironwood TPU v7x 主机上，设备与主机之间双向的聚合吞吐量最高可达 223 GB/s，而 PCIe Gen5 的物理上限为 256 GB/s。
• 据 TechBytes 报道，在数据中心网络上，该软件在一块 200 Gbit/s 网卡上可持续约 192 Gbit/s，在两块网卡并采用 NUMA 感知条带化时可达 350 Gbit/s。

TPU Sync 内部是如何工作的？

据 TechBytes 介绍，TPU Sync 依靠保持数据布局的直接内存访问（DMA，一种无需处理器参与即可复制内存的硬件机制），并配合原地缓冲区别名技术。该库直接与 JAX 的 PJRT 缓冲区和 PyTorch 张量协作，从而避免了 Python 序列化和中间拷贝，而且传输发生在 XLA 编译图之外。SemiAnalysis 把同一思路描述为零拷贝传输，即通过提取 PJRT 缓冲区的原生硬件描述符来实现。

据 TechBytes 称，运行时由三个模块组成。全局键值存储相当于一个注册表，把提示词哈希映射到物理存储块，并负责最近最少使用的淘汰。集群控制器负责执行规划和租约管理。本地缓存管理器跟踪设备上各块的状态，管理锁定的主机内存，并在 /dev/shm 的共享内存中保留一份热启动副本。内存层级涵盖 HBM、锁定的主机 DRAM、本地 SSD 以及 Google Cloud Managed Lustre。

在集成方面，TechBytes 列出了通过 TorchTPU 面向 PyTorch 的官方绑定、面向 JAX 的官方绑定、面向 vLLM、SGLang 和 llm-d 服务栈的连接器，以及面向 Pathways 和多控制器 JAX 流程的底层原语。

有哪些限制，成熟度如何？

TechBytes 称，HBM 与主机 DRAM 之间的传输已可用于生产，而向 Managed Lustre 等共享存储和外部 POSIX 后端的分层存储仍处于实验性验证阶段。强化学习的权重同步模块同样是实验性的，要达到最高网络吞吐量，需要两块 200 Gbit/s 网卡并正确配置 NUMA 感知条带化。Open Source For You 的文章没有标注日期，它的说法更保守，称该项目仍在积极开发中，尚不适合一般的生产使用，因此不同来源对成熟度的描述并不一致。

SemiAnalysis 提供了竞争视角。其 InferenceX 预览版把聚合模式的 TPU v7 与聚合模式的英伟达 B200 和 B300 作比较，并称英伟达的 GB200 和 GB300 NVL72 系统目前在分离式服务曲线的一部分上领先。待基于 TPU-Sync 的分离式服务在外部技术栈中优化后，将推出与分离式 TPU v7 的对比，SemiAnalysis 预计 TPU 上的 Mooncake Store 支持将建立在 TPU-Sync 原语之上。

这对企业意味着什么？

对于在谷歌云上运行或租用长上下文助手和智能体的中国企业来说，搬运 KV 缓存变成了需要配置的事，而不是需要自己构建的事。实际步骤是：用自己的提示词测试与 vLLM 或 SGLang 的连接器，测量启用前后的首个词元时间，并在依赖之前确认哪些内存层级仍被标注为实验性。使用强化学习训练的团队应在自己的训练器和采样器拓扑上验证权重同步模块。已公布的任何数字都不能取代在您自身工作负载上的基准测试。