TorchTPU adds a native PyTorch backend for Google TPUs, eliminating the need for torch_xla or JAX. Installing it on a TPU VM requires a single pip command with a CPU-only torch index. Users can load any Hugging Face model onto a TPU with device_map=“tpu” and run it in eager mode without code changes. Diffusers pipelines also support TPU natively, requiring only pipe.to(“tpu”) for eager execution or an optional compile call for deployment. The Kernels Hub provides ready-to-use TPU kernels written in Pallas, which can accelerate MoE models by over 20× at large batch sizes, as demonstrated with Qwen3-30B-A3B. The Hugging Face team is continuously testing TPU support, with 96.5% of tests passing and upstream fixes for discovered issues, including bugs in StaticCache and Qwen3 tensor-parallel plans.
Context
PyTorch users previously needed torch_xla or JAX to run on TPUs; TorchTPU adds a native backend.
Approach / What changed
TorchTPU adds TPU as a native PyTorch backend, installable via pip with a CPU-only torch index. Users set device_map=‘tpu’ in Transformers or call pipe.to(‘tpu’) in Diffusers. It supports eager and compiled modes, distributed training, and Pallas kernels.
Takeaways
1
TorchTPU allows running PyTorch models on TPUs with device_map=“tpu” without code changes, supporting eager and compiled modes.
2
Diffusers pipelines work on TPU natively with pipe.to(“tpu”) and optional torch.compile for performance.
3
TPU kernels from the Hub can accelerate MoE models significantly, e.g., 21.5× faster decoding on Qwen3-30B-A3B at batch size 64.
FAQ
How do I install TorchTPU?
Run pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu on a TPU VM. The extra index provides the CPU-only torch build.
What models work with TorchTPU?
Any model from the Hugging Face Hub works. Load it with device_map='tpu' in Transformers or pipe.to('tpu') in Diffusers.
How do I use TPU kernels?
Load a kernel from the Kernels Hub using get_kernel, then register it with ALL_EXPERTS_FUNCTIONS for MoE models. The bigger the batch, the bigger the speedup.
A bytecode interpreter is a compiler's output plus a loop that runs it. This explainer maps its five parts — bytecode, dispatch loop, runtime state, and the host boundary — and why the shape repeats.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.