想用 TPU 跑 PyTorch，以前得先「翻译」一遍模型
Google 的 TPU 一向以性价比著称，但对 PyTorch 用户有个老毛病：模型想上 TPU，基本得重写成 XLA 图，或者套一层 torch_xla 适配器，调试报错经常让人摸不着头脑。2026 年 10 月 20 日，Google 在 PyTorch Conference 2026 上宣布开源 TorchTPU（Apache 2.0），直接把 TPU 注册成 PyTorch 原生设备——以后 torch.device('tpu') 和 cuda、cpu 平级，模型不用改一行就能跑。项目地址在 官方仓库，技术细节见 Google 官方博客。
核心机制：eager 调试，compile 上生产
TorchTPU 是 Google 与 Meta PyTorch 团队共建的 PyTorch 原生 TPU 栈。它的设计思路很直白：开发阶段你用 eager 模式照常写代码、照常打断点；要上生产时，一行 torch.compile(backend='tpu') 把模型交给编译器，剩下的事框架全包。
编译链路拆开看
model.to('tpu')：把 TPU 当普通设备用，算子调度走 PyTorch 原生路径torch.compile(backend='tpu')：经 TorchDynamo 捕获计算图，转 StableHLO，再由 OpenXLA 编译下发- Pallas 内核：热点算子用 Pallas 手写内核，榨干 TPU 的矩阵乘单元
- 安装只有一条命令：
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu，要求 Python 3.11+
这套链路的意义在于：XLA 那套东西还在，只是被藏进了 backend='tpu' 后面。你不再需要学 XLA 的编程模型，PyTorch 怎么写，TPU 就怎么跑。
推理生态：主流引擎发布即适配
开源同一时间，推理侧的几个主流项目已经把 TorchTPU 接了进去：
推理侧三家
- vLLM + Inferact：两家共建 TPU 后端，Pallas 内核打底，continuous batching、量化、多芯片并行一起上
- SGLang + RadixArk：共建版本用 RadixAttention 做 KV 缓存复用；值得一提的是，SGLang-TorchTPU 已经在 Meta 的生产环境里跑着了
- llm-d 1.0（CNCF Sandbox）：由 Red Hat、IBM 牵头，500 多位贡献者、25 个组织参与的分布式推理项目，把 TorchTPU 纳入了 TPU 支持路线
还有个彩蛋：vLLM 同步推出了一个免费 TPU 算力池，专门给同时为 vLLM 和 TPU 生态做贡献的开发者用——想白嫖 TPU 跑实验的可以留意。
训练与周边：从大模型训练到端侧部署
一条线全覆盖
- TorchTitan：在 Qwen 3 8B 上跑出 52% MFU（谷歌自报），验证规模是 8000 张 Ironwood TPU，大规模训练已经能用
- VeRL：RL/GRPO 这类强化学习训练流程直接支持
- Ray：分布式调度照常用
- Hugging Face：Transformers v5 里传
device_map='tpu'就能用，官方每晚都在 TPU 上跑完整测试套件 - LiteRT：
litert-torch.convert()一键把模型转到端侧 - TPU agent suite（预览版）：帮你做迁移和调优的智能体工具包
这一套下来，从训练、微调、强化学习到推理服务、端侧部署，TPU 第一次有了和 CUDA 差不多的全家桶体验。
性能数字怎么看：全部来自谷歌自家披露
官方博客给了一组很漂亮的数字，先说清楚：以下全部是谷歌自家披露，第三方还没复核，看看就好，别当采购依据：
- Qwen 3.5 跑在 Ironwood（TPU 第七代）上，推理成本 2.01 美分/百万 tokens（SemiAnalysis InferenceX 基准，谷歌自报）
- 2.8T 参数的 Kimi K3，用 Pallas megakernel 做单用户推理，700+ tok/s（谷歌自报）
- Zenteiq 用 GKE Inference Gateway，每 1000 张 TPU v6e 每天生成的 tokens 从 14.7 亿涨到 59.7 亿，吞吐提升 4 倍（谷歌自报）
数字漂亮归漂亮，真实业务里还受批大小、网络、模型结构影响。TPU 的价格优势一直存在，TorchTPU 解决的是「用起来麻烦」的问题，而不是「算力更便宜」的问题——这两件事别搞混。
同一周，亚马逊也出手了
有意思的时间线：10 月 19 日，也就是 TorchTPU 开源前一天，亚马逊发布了 TorchNeuron——Trainium 芯片的 PyTorch 原生支持。两家在同一周动手，目标很一致：打破英伟达 CUDA 的生态绑定，把自家芯片变成 PyTorch 的「一等公民」。
对用户的实际意义很直接：已经在 Google Cloud 上用 TPU 的团队，TorchTPU 几乎是零迁移成本；在 AWS 上跑 Trainium 的，可以去评估 TorchNeuron。两家芯片的 PyTorch 体验，2026 年底算是同时补上了最关键的一块拼图。
