TorchTPU makes Google TPUs "just another device" in PyTorch: 96.5% of tests pass, MoE decode up to 21.5x faster than the default path
Hugging Face’s native TPU backend turns a port into a device string. What the benchmarks do and do not show, and how to try it on a v6e today.
With TorchTPU, a native PyTorch backend, TPU is now just another device.
Comprehensive take
Pilot one model on a TPU v6e and measure your own cost per million tokens; the post shows portability, not economics.
Use cases
Run Transformers models on TPU without XLA
Set device_map="tpu" and run eagerly or with torch.compile; the example uses Qwen3.5-9B.
Serve diffusion models with tensor parallelism
Compile the pipeline transformer with backend="tpu" and shard with TensorParallelConfig; Qwen-Image-2.1 at 1024x1024 took 3.5 s compiled on v6e-4.
MoE decode on TPU v6e-8
Tuned kernels give 21.5x on Qwen3-30B-A3B and 12.8x on Gemma 4 26B-A4B at batch 64 versus the default grouped_mm.
Specs
| Library tests passing on TPU | 96.5% source |
|---|---|
| Issues filed / fixed | 34 / 19 source |
| MoE decode speedup (v6e-8, batch 64) | 21.5x Qwen3-30B-A3B; 12.8x Gemma 4 26B-A4B source |
| Qwen-Image-2.1 1024x1024, 40 steps (v6e-4) | 3.5 s compiled vs 8.0 s eager source |
| Install | pip install torch-tpu (repo lists 0.2.0, Python 3.11+) source |
| Licence | Apache 2.0 source |
Insights
- The benchmarks compare against HF-chosen baselines (default grouped_mm, eager); there is no GPU or PyTorch/XLA comparison.
- The repo shows version 0.2.0 and a "Dependencies: TBD" line, which signals early-stage software.
- Together with TorchNeuron, this is a pattern: Hugging Face is standardising non-NVIDIA backends behind the same Transformers and Diffusers APIs.
Axion Scorecard
Overall: 7.5
96.5% of TPU-runnable library tests pass; large MoE decode speedups against the default path. source
Apache 2.0 repo under google-pytorch; Pallas kernels distributed via the Kernels Hub. source
One device string in Transformers and Diffusers; pip install. source
Not yet scored — no cited data.
Not yet scored — no cited data.
Native backend replaces lazy-tensor XLA path. source
What changed
Hugging Face announced on 20 October 2026 that PyTorch models can now target Google TPUs through TorchTPU, a native PyTorch backend that adds "tpu" as a device. The post's framing: "With TorchTPU, a native PyTorch backend, TPU is now just another device." Previously you needed PyTorch/XLA, with its lazy tensors and recompilation, or Optimum TPU. The backend lives in the google-pytorch organisation on GitHub under Apache 2.0, with custom ATen kernels, a compilation cache, torch.compile integration and data and tensor parallelism per its README.
How to try it
The blog gives a one-line install: pip install torch-tpu with the PyTorch CPU wheel index to avoid pulling CUDA builds. The repo lists release 0.2.0 on PyPI, Python 3.11 or newer, and recommends a TPU VM (v6e on Google Compute Engine is the easiest). In Transformers you set device_map="tpu" (the example uses Qwen3.5-9B) and run eagerly or with torch.compile. In Diffusers you call pipe.to("tpu"), and for speed compile the transformer with backend="tpu", fullgraph=True and dynamic=False, with TensorParallelConfig for sharding.
The numbers
All figures are Hugging Face's own and unreproduced. The headline: with tuned kernels on a TPU v6e-8 at batch 64, decode ran 21.5x faster on Qwen3-30B-A3B and 12.8x faster on Gemma 4 26B-A4B than the default grouped_mm path. That baseline is HF's choice, so read it as "versus the naive path", not "versus the best GPU kernel". For diffusion on v6e-4 with TP=4 in bf16: Qwen-Image-2.1 at 1024x1024 and 40 steps took 3.5 s compiled versus 8.0 s eager, and MiniMax-H3 video at 1344x768 took 188 s compiled versus 262 s eager.
On robustness, 96.5% of library tests that can run on TPU pass. HF filed 34 issues on the TorchTPU repo and 19 are fixed; nightly TPU CI is still work in progress. Their stance: "If something breaks on TPU, we want to find out before you do."
What is not in the post
There is no GPU comparison and no comparison to PyTorch/XLA, so the post does not tell you whether TPU is cheaper or faster than your current setup. The first compiled call triggers compilation (reused for the same image size), some pipelines need a tp_plan for tensor parallelism, and kernel gains depend on batch size. The repo README still shows a "Dependencies: TBD" line, a hint of early-stage software, and we found no Google announcement or GA label to cite. Optimum TPU is described as no longer required, but formal deprecation is not stated.
Who should act
If you already have TPU credits or capacity, try it now: moving a Transformers or Diffusers script becomes a device string change rather than a port. If you are GPU-only, wait for third-party cost-per-token comparisons before planning a migration.
Axion take
This is the same pattern Hugging Face showed with TorchNeuron on AWS Trainium a day earlier: first-class PyTorch backends for non-NVIDIA silicon, shipped through the standard libraries. The portability story is real; the economics are not yet shown. Pilot one model, measure your own cost per million tokens, and decide then.