TPU 是 Google 的 AI 加速器，专为神经网络常用的矩阵运算设计，适合需要大量算力的训练和大规模推理。过去在 PyTorch 里用 TPU 很别扭：要么走 torch_xla，要么去学 JAX。torch_xla 的惰性张量让代码不按预期执行，小改动可能触发图重编译，调试让人头疼。TorchTPU 是一个 PyTorch 原生后端，把 TPU 变成普通设备。Hugging Face 生态里的 Transformers、Diffusers 和 Kernels 现在都能原生跑在 TPU 上，旧的 Optimum TPU 适配层不再需要。

安装只需一条命令：pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu。多加的 index 是为了拿 CPU 版 torch，避免拉取用不到的 CUDA 包。在 Transformers v5 中，加载模型到 TPU 只需要 device_map=”tpu”，示例代码用 Qwen/Qwen3.5-9B。Diffusers 也把 TPU 当作一等设备：原型阶段 pipe.to(“tpu”) 即可，部署阶段加一次 pipe.transformer.compile(backend=”tpu”)。模型太大放不进单芯片时，传入 TensorParallelConfig，加载时就会切分到多芯片。

TorchTPU 自带为 TPU 调优的 Pallas kernels，支持分布式训练和 torch.compile。Kernels Hub 提供现成的 TPU kernel，可以像插件一样插进模型，例如把 Tokamax 的 ragged-dot kernel 注册到 MoE 模型的专家实现。性能提升来自官方测试：在 TPU v6e-8 上，batch size 64 时解码吞吐比默认实现快 21.5×（Qwen3-30B-A3B）和 12.8×（Gemma 4 26B-A4B）。在 Cloud TPU v6e-4 上，Qwen-Image-2.1 生成 1024×1024 图像，编译模式 3.5 秒，eager 模式 8.0 秒。

Hugging Face 团队用 Transformers 测试套件持续检查 TPU 兼容性，正在着手增加夜间 CI。目前 TPU 上能跑的测试 96.5% 通过。他们已向 TorchTPU 仓库报告 34 个问题，其中 19 个已修复。这个过程中还发现并修复了 Transformers 自身的几处 bug，例如 StaticCache 和 Qwen3 的 tensor-parallel 计划。

用 TorchTPU，TPU 不再是 PyTorch 用户的支线任务。装上它，从 Hub 选一个模型，发送到 “tpu”，完事。