Google 开源 TorchTPU：PyTorch/XLA 被弃用，将于 11 月 15 日归档
10 月 20 日，Google 在 GitHub 上放出了 TorchTPU，这是面向自家 TPU 的原生 PyTorch 后端。如果你已经在 TPU 上跑 PyTorch，最该看的是 TorchTPU FAQ 里的这句话：“PyTorch/XLA is deprecated and replaced by TorchTPU, which is now public.”（PyTorch/XLA 已废弃，由 TorchTPU 接替，现已公开。）
PyTorch/XLA 的 README 现在开头也是同样的说明：“This project is deprecated and replaced by TorchTPU.”（本项目已废弃，由 TorchTPU 接替。）里面还写明仓库将于 2026 年 11 月 15 日归档。也就是说，还在用 torch_xla 的团队只剩不到四周，之后仓库就会变成只读。
Google 这次交付了什么
Google 在 PyTorch Conference 2026 上宣布了 TorchTPU，并称这是与 Meta 的 PyTorch 团队共同打造的（Google Cloud 博客）。按 Google 的说法，TPU 现在和 GPU、CPU 一样成为 PyTorch 的原生目标设备：调试时用 eager 模式（算子逐条立即执行，方便定位问题），上生产要速度则用 torch.compile。仓库采用 Apache 2.0 协议，PyPI 上最新公开版本为 0.2.0，要求 Python 3.11 及以上（README）。
官方功能页列出已验证可用的芯片：TPU v5e、v5p、v6e 和 v7。TPU v8i 与 v8t 只在近期路线图里，目前还不支持。
推理部署也一并跟上。Google 的文章给出了两个链接：与 Inferact 合作构建的 vLLM 后端，与 RadixArk 合作构建的 SGLang 后端；并称 Hugging Face Transformers v5 只要设置 device_map="tpu" 就能部署到 TPU（Google Cloud 博客）。
训练方面，Google 称用 TorchTitan 在 Qwen 3 8B 上最高达到 52% 的模型算力利用率（MFU，即实际算力占峰值算力的比例），扩展到 8192 颗 Ironwood TPU 时性能接近线性增长。这些数字来自 Google 自家测试，不是独立第三方基准。
TechCrunch 把这解读为 Google 追随 Amazon，一起冲击 Nvidia 对 AI 开发者的掌控；而就在前一天，AWS 刚发布了自家的 Trainium PyTorch 后端。那个我在 TorchNeuron 那篇里写过。
迁移从改一个字符串开始
FAQ 里对现有业务的迁移说法是：.to("cuda") → .to("tpu")。
安装时要用 torch 的 CPU 版本，因为 README 提醒：默认安装会拉取 CUDA 版本（见 README）：
pip install torch_tpu==0.2.0 \
--extra-index-url https://download.pytorch.org/whl/cpu
下面是 VM 快速上手文档里的代码，已做精简：
import torch
device = torch.device("tpu") # 原来是 "cuda"
a = torch.ones((5, 5)).to(device)
b = torch.randn((5, 5), device="tpu")
result = (a + b).cpu() # 回传到主机再打印
print(result.sum())
这份指南在 TPU VM 上运行，示例用的是 v6e 实例。文档还提到，较新的 TPU 硬件需要用 GKE，而不是 VM（见指南）。
默认 eager 执行，以及会卡住的地方
执行模型页面解释了那一行改动背后发生了什么。在默认的严格 eager 模式下，每个算子都会变成一个单算子图，以 -O1 编译并缓存，因此缓存能很快预热。torch.compile 会把整个 FX 图交给 TorchTPU，并以 -O2 编译。中间还有两个环境变量：
# 错误会在引发它的算子处暴露
TORCH_TPU_LAUNCH_BLOCKING=1 python train.py
# 把延迟的算子融合成更大的图
TORCH_TPU_DEFER_AND_FUSE=1 python train.py
文档称，使用 TORCH_TPU_DEFER_AND_FUSE 相比默认 eager 最高可提速 50%，且无需改动模型。同一页还列出了哪些情况会让设备停顿。像 if my_tensor.all() 这种依赖数据的 Python 分支，会把值复制回 CPU 并清空队列；在编译或融合模式下，它们还会打断图。输出形状依赖数据的算子（比如 torch.masked_select）也会造成同样的问题。
形状增长不会每一步都触发重新编译。动态形状页面用一个 KV 缓存不断增长的解码循环做了演示。当 dynamic=False 时，每一步都会重新编译。在默认的 dynamic=None 下，增长的那一维会被当作符号处理，按文档示例，只有前三个步骤需要支付编译时间。
底层机制
TorchTPU 在底层会按 128 的倍数对齐编译桶（默认一个桶，最多八个），并在桶内做填充：
step = torch.compile(decode_step, backend="tpu", dynamic=None)
# 或者提前声明范围
torch._dynamo.
mark_dynamic(kv_cache, 2, min=16, max=128)
这些限制都写在同一页文档里：每次 trace 最多 8 个 bucket，不能使用 dynamic=True，这条路径下不支持 torch.nonzero 这类输出形状，也不能用自定义 Pallas kernel。这些“有界动态”的可执行文件不会存进 AOTAutograd 缓存，所以每次进程启动都得重新编译。
如果你现在还在用 PyTorch/XLA
FAQ 把两者的界线说得很清楚：PyTorch/XLA 依赖 JAX 的运行时组件，还有自己的一套 API；而 TorchTPU 是原生后端，不依赖 JAX 运行时。FAQ 同时强调，TorchTPU 不是要取代 JAX。
第一步，先把代码里所有碰过旧栈的地方找出来：
git grep -n -e "torch_xla" -e "xla_device" -e "xm\."
然后拿每一处命中对照功能清单。清单里标为“已验证”的有：eager 和编译模式混跑、前向和反向传播、Pallas kernel、SimpleFSDP。FSDPv2 和 SymInts 属于实验性，其中 FSDPv2 还标注了性能偏低。
11 月 15 日一过，PyTorch/XLA 仓库就变成只读：代码仍然能看、能 clone，但没人能再提新 issue 或 pull request。
TPU agent suite 接入 Claude Code 和 Antigravity
第二个发布是 TPU agent suite，Google 的博客把它称为预览版。它是一套给编程 agent 用的技能和子 agent 包，只有一个入口——一个叫 tpu-agent 的技能，由它把请求分发给对应的专职 agent，分别负责上手引导、性能剖析、硬件检查、带出处的 TPU 资料查询以及 kernel。
目前支持 Claude Code 和 Antigravity，Codex 和 Cursor 标注为“即将支持”。这个仓库暂时不接受外部贡献。
运行安装脚本之前，先看看它到底写了什么。按照安装指南，给 Claude Code 做项目级安装会做这几件事：在 .claude/skills/ 下添加技能，在 .claude/agents/ 下为每个 agent 放一个 prompt，生成一个 .claude/settings.json 预先授权 ~/.tpu-agent-suite 路径，再往你的 CLAUDE.md 里写一段受管理的配置块。对 Antigravity 则是写进 AGENTS.md。加上 --no-workspace-memory 可以跳过那段配置块。
安装请在分支上进行，提交前先看 git status 和 git diff。有个坑要注意：如果在 worktree 里跑 agent，新建的 git worktree 只会检出已被跟踪的文件。
未提交的项目级文件会留在你安装时所在的目录里。指南里用 --scope user 安装就是为了避开这一点。我自己在做 Hermes IDE，所以请结合这一点来权衡。目前这套工具支持的两种 agent 都是 Hermes 内置的，每个任务都能有自己的 hermes/<name> 分支和工作树（worktree），因此 Claude Code 和 Antigravity 在同一个 TPU 端口上跑时，不会互相碰到对方的文件（README）。我本人还没跑过这套 TPU agent 工具。Hermes 与 Google、Anthropic 或任何 agent 厂商都没有隶属关系。Hermes 2.3.4 的 agent 选择器：已安装 Claude Code 和 Codex，列出 Antigravity CLI 及其他支持的 agent，另有 Custom agent。
你会迁移吗？我还没跑过 TorchTPU，上面这些限制正是真正做移植时才会暴露问题的地方。这次弃用意味着，用 PyTorch/XLA 的团队没法再观望了。你现在就在 TPU 上跑 PyTorch 吗？11 月 15 日归档这个日期进你的路线图了吗？欢迎在评论区聊聊。
如果这篇文章对你有用——如果你确实用 vLLM 或 SGLang 把推理服务迁到了 TPU 上，下一个问题就是这次迁移到底值不值：延迟、每次请求的 token 数、每个回答的成本，以及能证明输出没有变化的评测。《LLM 可观测性口袋指南》讲了怎么挑追踪和评测工具来回答这些问题。