谷歌开源TorchTPU，PyTorch原生支持TPU，CUDA护城河松动

在圣何塞举行的 PyTorch Conference 2026 上，谷歌宣布开源 TorchTPU，一套让 PyTorch 原生运行在 TPU 上的完整软件栈。此后在 Hugging Face Transformers v5 中部署模型到 TPU，只需传入 device_map=“tpu” 一个参数。这被视为谷歌软件路线的转向，也是英伟达 CUDA 生态近十年来遭遇的最正面挑战。

PyTorch 自 2016 年诞生起与 CUDA 深度绑定，英伟达长期投入确保其在该平台性能最佳。谷歌 TPU 十年来主要服务 TensorFlow 和 JAX，PyTorch 用户需经 PyTorch/XLA 转接，付出额外工程成本。TorchTPU 使 TPU 与 GPU、CPU 并列成为 PyTorch 原生选项。项目由谷歌与 Meta 联合开发，vLLM、SGLang、Hugging Face、Ray、TorchTitan 等同步接入，23 页技术报告已发布在 arXiv。

此前 PyTorch/XLA 需侵入式修改源码、编译延迟严重、不支持即时执行；TorchAX 则逐算子翻译为 JAX。TorchTPU 于 2025 年年中立项，原则是不依赖 JAX、即时执行扎实、迁移只需将 .to(“cuda”) 改为 .to(“tpu”)。它基于 PrivateUse1 接口注册为 C++ 后端，开发者拿到普通 PyTorch 张量。即时执行分三档：Strict 逐算子异步分发；Debug 同步执行排查错误；DeferAndFuse 先攒算子再融合计算，速度比 Strict 快 50% 以上并降低峰值显存。三档共享分层编译缓存。使用 torch.compile 时，TorchDynamo 捕获计算图后映射到 StableHLO，交由 XLA 生成 TPU 程序，而非走 Inductor，原因是复用即时模式基础设施、XLA 稳定、Triton 难以移植到 TPU。分布式方面，Process Group 层实现 TPU 集合通信，DDP、FSDP2、DTensor 开箱即用，支持非对称写法。TorchTPU 公开后将取代 PyTorch/XLA。

性能方面，谷歌在 TPU 上测试 Hugging Face 下载量最高的约 5000 个模型，截至报告撰写共 4917 个、覆盖 500 多个架构族，98% 跑通不崩溃，657 个代表性架构中 97% 输出与 GPU 数值一致。30B 参数 MuseGlimmer 模型单颗 TPU v7 上 5120 序列长度预填充达每秒 2.57 万 token，MFU 59%；4 颗 TPU v7 训练 MFU 38%；TorchTitan 同款实现训练 MFU 36%。Qwen3 235B MoE 在 64 颗 TPU v7 上吞吐达 JAX 方案 MaxText 的约 94%；从 1024 扩展到 8192 颗芯片，每 TensorCore 吞吐仅降约 6%，MFU 稳定在 27% 至 29%。Qwen 3 8B 训练 MFU 最高 52%。推理侧，TPU v7 搭配 vLLM-TorchTPU 在 8K 输入、1K 输出场景下每百万 token 成本 2.01 美分，较 9 月首次亮相的 3.6 美分降低约四成。FLUX 训练单步时间从 11.49 秒降至 2.78 秒；DLRM v2 在 32 颗 TPU v7x 上与 JAX 基线持平。TorchTPU 已在部分谷歌云大客户生产环境承担推理和训练任务。

生态方面，vLLM 与 SGLang 后端保留上游调度器和 API 服务器；谷歌联合 Red Hat、IBM 发布 llm-d 1.0，在 GKE 上让 GPU 与 TPU 共用推理入口。训练侧 TorchTitan 提供预训练配方，VeRL 合入 TPU 后端，Ray 支持跨硬件调度。LiteRT 加入 PyTorch 生态。TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU 和 TorchTitan 四个核心仓库同步开源。Meta PyTorch 工程总监 Jana van Greunen 表示 Meta 不能被单一硬件架构束缚，TorchTPU 已支撑其 LLM、广告和推荐系统的关键生产场景。

项目节奏紧凑：2025 年年中谷歌内部启动；同年 12 月路透社首次披露代号；2026 年 4 月谷歌开发者博客公开架构，项目负责人在巴黎 PyTorch Conference Europe 登台；9 月进入私有 Beta 并在 SemiAnalysis 推理基准中首次亮相。