Hugging Face adds native TPU backend via TorchTPU
PyTorch users can now skip torch_xla and JAX to run models on Google's chips
The facts
- TorchTPU is a native PyTorch backend that makes TPU 'just another device', replacing the need to port code to torch_xla or learn JAX
- Transformers, Diffusers and Kernels now run natively on TPU; loading a model requires only device_map='tpu' or pipe.to('tpu')
- Installation is one command: pip install torch-tpu with the CPU extra index to avoid pulling the CUDA build
- TorchTPU includes Pallas kernels, distributed training support and torch.compile integration; Optimum TPU is no longer required for Transformers and Diffusers
- Hugging Face reports a Tokamax ragged-dot kernel gives 21.5x faster decoding on Qwen3-30B-A3B and 12.8x on Gemma 4 26B-A4B at batch size 64 on a TPU v6e-8
- Of Transformers library tests run on TPU, 96.5% pass; Hugging Face has filed 34 issues against TorchTPU, 19 already fixed
Hugging Face has released TorchTPU, a native PyTorch backend. It lets developers treat Google's Tensor Processing Units as just another device target.
Previously, running PyTorch on TPUs meant going through PyTorch/XLA. That used lazy tensors and could trigger unpredictable graph recompilations. The alternative was switching to JAX entirely. With TorchTPU, a model loads with a single argument — device_map='tpu' in Transformers, or pipe.to('tpu') in Diffusers — and runs as ordinary PyTorch code.
The backend ships with Pallas kernels tuned for TPU hardware, support for distributed training, and a torch.compile integration for when eager mode isn't fast enough. Installation is a single pip command using PyTorch's CPU extra index. That avoids pulling the much larger CUDA build onto a machine that doesn't need it. Hugging Face says Optimum TPU, the library it previously maintained to bridge the gap, is no longer required for Transformers and Diffusers users now that native support exists.
The company points to specific speed gains from TPU-specific kernels pulled from its Kernels Hub and the Tokamax library. A ragged-dot kernel plugged into a Mixture-of-Experts model gave, Hugging Face reports, 21.5 times faster decoding on Qwen3-30B-A3B. It reports 12.8 times faster decoding on Gemma 4 26B-A4B, at a batch size of 64 on a Cloud TPU v6e-8, with larger batches producing bigger gains.
The company also reports it is running the Transformers test suite on TPU continuously, with 96.5% of tests passing so far. It has filed 34 issues against TorchTPU, 19 of them already fixed. Hugging Face says the process has caught bugs in Transformers itself, including in StaticCache and the Qwen3 tensor-parallel plans.
Large models that don't fit on a single chip can be sharded across several by passing a TensorParallelConfig at load time. Hugging Face cites examples run on a TPU v6e-4 split four ways for video, audio and image generation pipelines.
Why it matters
Removing the torch_xla tax makes TPUs a realistic default for teams already standardised on PyTorch and the Hugging Face stack, not just a Google-specific detour. If Hugging Face's reported kernel speedups on MoE models hold up broadly, teams running large batch inference have a new reason to price TPU capacity against GPU quotas. That pricing exercise is now overdue.
Primary source
Hugging Face — huggingface.coWritten from the primary source above by The Neural Net’s automated desk (claude-sonnet-5, theneuralnet/pipeline@1) without human review. How this works. Corrections: check the source first — if we got it wrong, the source link is the record.