AREX Feed Article
Google 开源 TorchTPU，TPU 成为 PyTorch 原生设备
10 月 20 日，在圣何塞举行的 上，Google Cloud 在中宣布 TorchTPU 已在 GitHub 开源。这套与 Meta 的 PyTorch 团队联合开发的软件栈把 TPU 接进 PyTorch 的设备体系：模型代码把 .to("cuda") 换成 .to("tpu")，就能在 TPU 上以 eager（即时执行）模式运行，不必再经过 torch_xla 那套懒张量接口。
仓库 创建于 2026 年 3 月 9 日， 上的 torch-tpu 0.2.0 于 10 月 19 日上传；README 给出的安装命令是 pip install torch_tpu==0.2.0，要求 Python 3.11 以上，标注的许可证为 Apache 2.0。Google 工程师署名的 在 10 月 16 日发出第一版，10 月 20 日更新到第二版，讲这套栈的实现和第一批性能数字。
从 torch_xla 的懒张量，到 torch.compile 的整图编译
PyTorch 跑 TPU 的旧路径是 PyTorch/XLA，2020 年推出，用懒张量（lazy tensor）记录计算图，开发者需要在代码里手动标出可编译区域。论文回顾这段历史时列了四个问题：要侵入式修改源码、编译延迟严重、不支持 PyTorch 原生的 eager 执行、对 PyTorch 新增的分布式原语跟进慢。
TorchTPU 换了一个接入点：按 PyTorch 的扩展机制注册成 C++ 后端（PrivateUse1），张量是普通的 at::Tensor，设备名是 tpu。代码不在 pytorch/pytorch 主仓，而在 google-pytorch 组织下。编译器仍是 XLA：torch.compile(backend="tpu") 用 TorchDynamo 抓 FX 图，把整张图降成 StableHLO（XLA 的中间表示），再交给 OpenXLA 生成 TPU 可执行文件，不走 GPU 上的 Inductor/Triton 那条路。
论文把“运行时不依赖 JAX”列为设计原则之一。eager 执行分三档：Strict 逐算子异步下发、Debug 同步执行便于定位、DeferAndFuse 把连续算子推迟并融合，论文称最后一档比 Strict 快 50% 以上，同时降低 HBM（高带宽内存）峰值占用。
精度默认值上还有一处直接差异。TPU 的矩阵乘单元以 bfloat16 为原生输入、FP32 累加，TorchTPU 因此默认用 medium 精度模式；论文的对照表里，PyTorch 在 CUDA 上的默认是纯 FP32。要更接近 FP32 的结果，high 模式把每次乘拆成 3 遍、highest 拆成 6 遍，算力开销约为 3 倍和 6 倍，显式选择高精度还会触发一次性能警告。
执行模型上还有几处硬约束。一颗 TPU v7 芯片向 PyTorch 暴露两个设备（对应两颗 TensorCore），论文里“单芯片推理”的例子实际用了进程数为 2 的 FSDP（全分片数据并行）；一颗芯片只能被一个用户态进程访问，TorchTPU 因此采用一进程一设备，强化学习训练里的 colocation 无法开箱支持，论文称正在研究放宽这一限制的办法。PyTorch 的视图与原地写入语义，靠视图分解和逆变换翻译成不可变的 StableHLO，对广播维度或重叠维度的原地写仍不支持。动态形状同样没有完成：XLA 编译器与 PjRT 运行时假定张量形状在编译期已知，TorchTPU 的有界动态形状还在开发，当前在 TPU 上调用 torch.compile 会自动设 dynamic=False，遇到新形状仍可能重新编译，论文把这一项列为临时限制。
8,192 颗芯片上的 27.1%：扩展数据由 Google 自己给出
图来自论文 Figure 3，测试对象是 Qwen3 235B MoE（混合专家），用 TorchTitan 训练，全局批大小随芯片数同步放大，从 1,024 颗芯片时的 32,768 加到 8,192 颗时的 262,144。每颗 TensorCore 的吞吐（论文记为 TPS，即每秒 token 数）从 2,236 降到 2,104，MFU（模型浮点运算利用率）从 28.8% 降到 27.1%，单步时间从 29.3 秒变成 31.2 秒。论文把这称为弱扩展：总吞吐随芯片数近似线性增长，效率基本持平。
同一篇论文里还有几组数字。Llama 3 405B 在 1,024 颗 TPU v7x、8K 上下文、BF16 下达到每芯片 500 tokens/s，MFU 50% 以上。Google 的博客另称，TorchTitan 路径在 Qwen 3 8B 上录得最高 52% MFU；同一段里还说，最多 8,192 颗 Ironwood 芯片上保持近线性扩展。论文里 8,192 颗芯片那一组跑的是 Qwen3 235B，与 8B 那项不是同一组数据。
论文另给了一个只用标准 PyTorch 模块的例子。30B 的 Muse Glimmer 以 nn.Linear、nn.RMSNorm、F.scaled_dot_product_attention 这类模块和算子实现，不用模型专用算子，在单颗 TPU v7 芯片上（两颗 TensorCore，进程数为 2 的 FSDP）prefill（预填充）达到 25.7k TPS、MFU 59%，序列长度 10,240 时降到 22.4k TPS、51%。论文说，这些例子说明纯 PyTorch 代码的性能足够好；要达到峰值性能，仍要靠 Pallas kernel（用 Pallas 语言写的 TPU 算子）和生产级的 vLLM、SGLang、TorchTitan 栈。上面这些数字都由 Google 自己测得，论文是预印本，未经同行评审。
论文还称，TorchTPU 已在部分大型 Google Cloud 客户的生产环境中用于推理和训练。论文引用的外部评测是 SemiAnalysis 的 InferenceX：Google 称，TPU v7 配 vLLM-TorchTPU 在 8K/1K 场景把每百万 token 成本做到 2.01 美分，低于 9 月首次提交时的 3.6 美分。
vLLM 与 SGLang 各有一个保留上游接口的 TPU 后端
vLLM 的 TPU 后端放在 vLLM 官方组织下（），由 Google 与 Inferact 合作开发。它保留上游的 worker、scheduler 和 API server，替换底层执行，并带一批为 TPU 写的 Pallas kernel；支持 continuous batching（连续批处理）、量化和多芯片并行。
SGLang 侧对应 ，合作方是 RadixArk。SGLang 用 RadixAttention 在请求间复用 KV cache（键值缓存），重复的提示词（prompt）前缀可以跳过重算；TPU 后端保留上游结构，并让一个分离式（disaggregation）传输后端兼任 HiCache 的 I/O 后端。
推理引擎要落到集群还需要路由层。llm-d 1.0 在 10 月 20 日发布（），Google 称它由 Red Hat、IBM 和社区共同推进；它按前缀缓存（prefix cache）的状态把请求送到已经持有上下文的副本，并把计算密集的 prefill 和访存密集的 decode（解码）拆到不同 worker 池。在 Google Kubernetes Engine（GKE）上，GPU 和 TPU 可以挂在同一个 llm-d 端点后面，按负载挑硬件。
列的能力清单更长：FP8 KV cache，FP8 与 4 位（MXFP4、NVFP4）权重，MTP、EAGLE-3、DFlash 等投机解码方法，以及把 KV cache 换出到主机内存的两级前缀缓存。这些是厂商文档的说法，不是第三方测试结果。
编排和端侧也在同一批发布里。Ray 加了原生 TorchTPU 支持，可以在一个工作流里把数据处理、训练、强化学习和 serving（推理服务）按阶段派到 CPU、GPU 或 TPU；LiteRT（原 TensorFlow Lite）带 LiteRT-Torch 加入 PyTorch 生态，用一次 litert-torch.convert() 转换基于 torch.export() 的模型，产物可跑在 Android、iOS、Windows、Web 和 IoT 设备上。
Hugging Face 的 96.5% 通过率，和 34 个已提交的问题单
在 Transformers v5 里，加载模型时传 device_map="tpu" 就落到 TPU 上；Diffusers 里是 pipe.to("tpu")，需要时再调 compile，模型太大可以传 TensorParallelConfig 在加载时切到多颗芯片。Hugging Face 在中说，TorchTPU 带来原生支持之后，原先为 TPU 准备的 Optimum TPU 不再需要了。
Hugging Face 还公布了测试通过率。它把 Transformers 的测试套件放到 TPU 上跑，在能运行的用例里通过率 96.5%；截至博客发布，已向 TorchTPU 仓库提交 34 个问题单（issue），其中 19 个已修复。跑测试还找出了 Transformers 自身的几处缺陷，例如 StaticCache 和 Qwen3 的张量并行方案，并在上游修掉。
CI 的进度两边说法不同。Google 博客说，已经为 Transformers 在 TPU 上搭了专门的 CI/CD 流水线做持续夜测；Hugging Face 的说法是“正在添加”TPU 上的夜间 CI（nightly CI），测试面板本身还标着“进行中”。
Hugging Face 产品副总裁 Jeff Boudier 在博客里说：“Hugging Face 致力于做一个开放、中立的平台，让 AI 构建者选择最适合自己的模型、框架和硬件。硬件多样性是让开源 AI 对所有人可及的前提。与 Google 合作把原生 TorchTPU 支持带进 Transformers 等核心库，是把这一承诺落到实处。”
Hugging Face 还测了 kernel 的收益：在 TPU v6e-8 上，把 Tokamax 的 ragged-dot kernel 接进 MoE 模型，批大小为 64 时，Qwen3-30B-A3B 的解码速度是默认实现的 21.5 倍，Gemma 4 26B-A4B 是 12.8 倍。这是 Hugging Face 自己的测量，对应的 kernel 发布在它的 Kernels Hub 上。
Meta 的理由，与仍未归档的 PyTorch/XLA
Meta 的 PyTorch 工程总监 Jana van Greunen 在 Google 博客里给出的理由是算力规模和硬件多样性：“在 Meta，推进 AI 前沿需要极大的算力规模和深度的硬件多样性。我们不能被单一硬件架构限制。TorchTPU 是很大的一步——它让我们能同核心机器学习系统团队一起灵活地共同开发硬件，同时在关键生产用例上支撑主要的 LLM、广告和推荐系统负载。”按她的说法，Meta 因此能把 PyTorch 负载部署到 Google TPU 上，并把问题直接反馈给工程团队。
Google 与 Meta 的 PyTorch 团队还把 TPU 后端上游进了 TorchRec：高基数 embedding（嵌入向量）查表交给 SparseCore，稠密矩阵运算留在 TensorCore。论文称，在 32 颗 TPU v7x 上跑 DLRM v2，单步 18.3 毫秒，与 Google 自己的 JAX 基线持平，损失（loss）曲线也一致。
PyTorch/XLA 的替代关系已经写明：仓库 里 4 月 22 日的提示说，TorchTPU 公开后“将取代 PyTorch/XLA”。论文给出的 6～12 个月路线图里，几项直接关系到迁移能否落地：有界动态形状投入生产、预编译的 ATen 动态算子库、面向 TPU v8 的 FP8/FP4 量化，以及通过 Ray 把强化学习训练与 vLLM/SGLang 采样后端接通。
PyTorch/XLA 仓库尚未归档，2025 年 10 月提出的原生 TPU 后端 RFC（）仍处于开启状态。