Replacing the PyTorch/XLA layer
Running PyTorch on TPUs used to mean either switching to JAX or going through PyTorch/XLA, introduced in 2020. The latter was built on lazy tensors: you had to mark compilable regions in your source, compile latency was high, PyTorch's native eager mode wasn't supported, and small changes could trigger whole-graph recompilation. TorchTPU, described on the official PyTorch blog, is a project Google started in mid-2025 with Meta's PyTorch team, aiming to make the TPU just another PyTorch device. Eager execution comes first, and existing code runs once you swap .to("cuda") for .to("tpu"); for speed you add torch.compile, built on TorchDynamo; distributed idioms such as DDP, FSDP and DTensor reach TPU collectives through the lowest-level process group API.
As of the post, the project reports support for about 450 ATen ops, covering forward and backward passes for the core ops major architectures need; out-of-the-box execution of the 5,000 most-downloaded Hugging Face models; validated multi-host inference with vLLM and SGLang and training with TorchTitan; and production use for training and inference by some of the largest Google Cloud customers.
The Hugging Face side
Hugging Face followed on October 20: Transformers v5, Diffusers and Kernels now support TPUs natively. Pass device_map="tpu" when loading a model, and the Optimum TPU library built for this purpose is no longer needed. Pallas kernels from the Kernels Hub can be plugged straight into models. For example, Tokamax's ragged-dot kernel applied to mixture-of-experts models on a TPU v6e-8 at batch size 64 made decoding 21.5x faster on Qwen3-30B-A3B and 12.8x faster on Gemma 4 26B-A4B; those numbers cover only those two models and that setup.
Hugging Face also listed what isn't finished. Of the Transformers tests that run on TPU, 96.5% pass. Every failure becomes a reproduction script; so far 34 issues have been reported to TorchTPU and 19 fixed, and the work also caught several bugs in Transformers itself, in StaticCache and the Qwen3 tensor-parallel plans.
Limits to know before migrating
TPU matrix units natively use BF16, and TorchTPU defaults to BF16 multiplication with FP32 accumulation; higher precision is available through 3-pass or 6-pass algorithms costing roughly 3x and 6x the compute. Each TPU can be accessed by only one user-space process, so distributed jobs are organized as one process per device, and the training-and-inference colocation common in reinforcement learning isn't possible yet. The project also notes that a one-line change only guarantees the model runs; models whose shapes and sharding lean toward GPU designs still need tuning to reach peak TPU performance.
It's the same direction as TorchNeuron, which AWS open-sourced for Trainium the same day: cloud providers are packaging their own chips as native PyTorch devices to lower the cost of leaving CUDA.
via: PyTorch Blog: Introducing TorchTPU, Hugging Face: Hugging Face ❤️ TorchTPU: TPUs, Natively