谷歌开源TorchTPU：PyTorch原生跑上TPU，一行参数不再绑定CUDA

一行参数带来的变化

这条新闻最精简的概括，就是下面这一行代码device_map="tpu"。

在圣何塞举办的 PyTorch Conference 2026 上，谷歌宣布 TorchTPU 正式开源。它是一整套软件栈，能让 PyTorch 直接在 TPU 上原生运行。对开发者而言，在 Hugging Face Transformers v5 里部署模型时，不用再折腾繁琐配置，也不用加额外的转换层，只要把这一个参数传进去，代码就能在 TPU 上跑。

改动看起来不大，实质是谷歌软件战略的一次明显转向，也让英伟达 CUDA 生态近十年来第一次面对这么直接的冲击。

过去十年，人们谈 AI 芯片竞争，常常只比硬件参数。可真正难被复制的是软件生态。PyTorch 是全球最主流的 AI 框架，2016 年问世以来一直和 CUDA 绑得很紧。英伟达为此投入了大量工程师，保证 PyTorch 在自家 N 卡上的体验最好。华尔街分析师甚至把 CUDA 称作英伟达挡住对手最硬的一面盾。

谷歌 TPU 的硬件性能不弱，但多年来的主要服务对象是自家的 TensorFlow 和 JAX。PyTorch 用户想用 TPU，得先过 PyTorch/XLA 这层“转接”。工程量上去了，性能还有折损，不少开发者因此放弃。

TorchTPU 把这道门槛拆掉了。按谷歌官方的说法，TPU 从今以后和 GPU、CPU 平起平坐，成了 PyTorch 的原生选项之一。

这个项目是谷歌和 PyTorch 创始方 Meta 一起做的。vLLM、SGLang、Hugging Face、Ray、TorchTitan 等主流开源项目也在同一时间完成了接入。arXiv 上同步放出了一份技术报告，里面写明了技术细节。

不再靠“转接头”：原生支持是怎么实现的

PyTorch 跑在 TPU 上其实早有先例。2020 年谷歌就做过基于惰性张量的 PyTorch/XLA。不过技术报告里，谷歌少见地承认了这套方案的问题：得侵入式改源码来划出可编译区域，编译延迟很重，PyTorch 原生的即时执行模式也不支持。分布式新原语的支持同样偏慢。

之后的 TorchAX 想把 PyTorch 算子一个个翻译成 JAX 再执行，说到底还是让 PyTorch 去适应 JAX 的运行方式。

TorchTPU 在 2025 年年中立项，当时定下了几条核心原则：运行时不依赖 JAX，即时执行必须得到扎实支持，代码迁移应当像把.to("cuda") 改为 .to("tpu") 一样简单。

落地上，TorchTPU走的是PyTorch给第三方硬件留的PrivateUse1接口，在分发器里注册成一个新的C++后端。开发者手里拿到的还是标准PyTorch张量，既没有子类，也没有多套一层包装。

灵活性和性能要兼顾，TorchTPU因此给出三种即时执行模式：

• Strict模式：默认走这一档，算子逐个分发，和CPU异步执行，用起来跟GPU上的常规体验差不多。
• Debug模式：改成同步执行，主要用来查形状错误和NaN值。
• DeferAndFuse模式：这是TorchTPU新引入的混合玩法。先把一串算子缓存住，等到需要取值时再自动融合成大块计算任务，交给TPU去跑。跟Strict模式比，速度快了50%以上，峰值显存占用也压了下来。

三种模式共用一套分层编译缓存，同一串算子只编译一次，之后反复使用。

要榨性能，开发者可以切到torch.compile。这一步谷歌做了个刻意的取舍：用TorchDynamo捕获计算图之后，绕开PyTorch默认的Inductor编译器，把算子映射到StableHLO，再交给XLA生成TPU程序。

报告里讲了这么选的三条理由：一是即时模式搭好的基础设施能直接复用；二是XLA经过多年验证，稳定性够；三是Inductor的成功很大程度上靠Triton，而Triton很难移植到TPU架构上。这也从侧面说明，PyTorch和GPU的绑定已经扎到了编译器底层。

分布式这块，TorchTPU在Process Group层做了TPU集合通信，DDP、FSDP2和DTensor都能开箱即用。它还支持让rank 0进程顺手处理日志等额外工作的“不对称”写法，只支持SPMD的PyTorch/XLA做不到这一点。按官方文档的说法，TorchTPU公开后将正式取代PyTorch/XLA。

技术成绩单：覆盖率与性能

技术报告亮出了三张关键答卷。

头一张是覆盖面。谷歌在TPU上持续测Hugging Face下载量最高的约5000个模型，到报告撰写时一共测了4917个，横跨500多个架构族。其中98%的模型能正常跑起来不崩溃。657个代表性架构里，97%的输出结果和GPU数值对得上。

第二张是未经专门优化的性能。谷歌用PyTorch标准模块写了个30B参数的MuseGlimmer模型，没调用任何定制内核。单颗TPU v7上，5120序列长度的预填充速度达到每秒2.57万token，模型利用率（MFU）为59%；4颗TPU v7训练时，MFU是38%。就连原本为GPU写的TorchTitan实现，稍加修改迁过来，训练MFU也能到36%。

第三张是扩展规模。Qwen3 235B MoE模型上，用64颗TPU v7时，PyTorch方案的吞吐量达到谷歌自家JAX方案MaxText的约94%。芯片从1024颗扩到8192颗（也就是一个完整的TPU v7 Pod），每个TensorCore的吞吐量只降了约6%，MFU稳在27%到29%之间，各规模下的损失曲线几乎重合。官方博客还提到，Qwen 3 8B的训练MFU最高能到52%。

说白了，PyTorch在TPU上的表现跟“亲儿子”JAX之间的差距，已经缩到个位数百分比。

推理侧的数字同样抓眼。SemiAnalysis InferenceX基准测试显示，TPU v7搭配vLLM-TorchTPU，在8K输入、1K输出的场景下，每百万token的成本降到2.01美分，比9月首次亮相时的3.6美分低了约40%。

报告还列了几个案例：FLUX训练从开箱即时模式切到完整编译优化，单步时间从11.49秒降到2.78秒；Meta的DLRM v2推荐模型在32颗TPU v7x上的表现，跟谷歌高度优化的JAX基线完全持平。眼下，TorchTPU已经在部分谷歌云大客户的生产环境里承担推理和训练任务。

生态“全家桶”：从训练到端侧

比技术本身更有分量的是生态的完整度。

推理侧，和Inferact合作的vLLM后端、和RadixArk合作的SGLang后端都保留了上游原有的调度器和API服务器，现有推理栈不用改。谷歌还拉上Red Hat、IBM发布了llm-d 1.0，在GKE上GPU和TPU能共用同一个推理入口。

训练侧，TorchTitan给出了稠密模型、MoE和扩散模型的预训练配方。强化学习框架VeRL合入了TPU后端，Ray则支持把同一工作流的不同阶段分别调度到CPU、GPU或TPU上。

端侧，前身是TensorFlow Lite的LiteRT正式加入PyTorch生态。TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU和TorchTitan四个核心仓库同步开源。

站台方的态度也有意思。Meta PyTorch工程总监Jana van Greunen表示，Meta不能被单一硬件架构绑住，TorchTPU已经撑起其LLM、广告和推荐系统的关键生产场景。考虑到Meta是英伟达的超级大客户，这番话的象征意味很足。

CUDA 的护城河会被填平吗？

框架层这道墙一拆，问题就来了：英伟达手里还剩下什么牌？

得说清楚一点：TorchTPU 填平的是框架层面的落差，硬件本身的物理限制它管不了。报告里写得很明白，只改那一行并不等于就能拿到峰值性能。.to("tpu")原因在于，模型的形状、分片方式和算子选择，原先多半是照着 GPU 调的。TPU 的 TensorCore 说到底是个单线程处理器，形状得足够宽才能把矩阵单元喂饱；GPU 那套思路正相反，靠并行处理一大堆小内核吃饭。

迁移的人还得留意几处细节差别：TorchTPU 默认用 BF16 做乘法、FP32 做累加，PyTorch 在 CUDA 上默认走的是完整 FP32；默认的集合通信编译方式会吃掉 30% 到 40% 的端到端性能，得额外加标注去优化；每颗 TPU 只允许一个进程访问，强化学习里训练和推理放在一起跑的做法，眼下还没法开箱即用。“能跑”和“跑得好”之间，差的还是调优那部分活。

另外，动手的不只谷歌。同一场会议上，AWS 也讲了 Trainium 转原生 PyTorch 的过程：其 Neuron SDK 已经不再支持通过 PyTorch/XLA 训练，改走原生的 TorchNeuron。

PyTorch 正变成事实上的硬件抽象层，PrivateUse1 接口则成了各家的芯片接进来的标准插槽。等到所有硬件都能原生跑同一份 PyTorch 代码，比拼的重点就会从“生态能不能用”挪到“每美元换来多少吞吐”。谷歌在发布稿里反复强调的，正是这个指标。

变量当然还在。报告里大部分性能数据是谷歌自己测的，还得靠社区在更多模型和场景里独立验证。TPU 能不能以足够的产能、合适的价格让普通开发者用上，也决定着这次开源最终能掀起多大浪。谷歌说，未来 6 到 12 个月的重点是生产环境加固、动态形状支持，以及面向下一代 TPU 的 FP8、FP4 量化支持。

不过方向已经摆在那了。过去十年，选 PyTorch 基本等于选 CUDA。从今天起，这个等式不成立了。