TorchTPU 让 TPU 成为 PyTorch 的原生设备，Transformers、Diffusers 和 Kernels 改一个设备名就能跑
不再需要 torch_xla 或学 JAX，Hugging Face 生态里的模型 to("tpu") 即可。
// 要点
- TorchTPU 是 PyTorch 的原生 TPU 后端，自带为 TPU 调优的 Pallas 内核、分布式训练支持和 torch.compile 集成。
- Transformers 中传入 device_map="tpu" 即可把模型直接加载到 TPU；默认 eager 模式，也可用 torch.compile 提速。
- Diffusers 把 TPU 作为一等设备，超出单芯片的模型可通过 TensorParallelConfig 在加载时分片到多块芯片。
- 有了原生支持，运行 Transformers 和 Diffusers 不再需要 Optimum TPU。
开发者视角和昨天的 TorchNeuron 连起来看，这周 PyTorch 接连拿下了 Trainium 和 TPU 两块非 GPU 加速器。对做产品的人，推理不再只能比 GPU 报价；我的建议还是先拿一个已经在跑的推理服务原样迁过去，比单位成本，而不是为了新硬件改代码。