With TorchTPU, TPU is now a native backend for PyTorch, removing the need to port code to torch_xla or to switch to JAX. This makes it possible for Hugging Face libraries — including Transformers, Diffusers and Kernels — to run directly on TPUs: pick a model from the Hub, send it to device "tpu", and the model runs on the accelerator.

What is a TPU?

TPU stands for Tensor Processing Unit. TPUs are Google’s custom AI accelerators designed for the matrix math that neural networks perform. Google has been designing TPUs for about ten years and uses them to train and serve its own models, including Gemini. TPUs are available through Google Cloud, from single chips up to pods with thousands of chips. They are most useful when substantial compute is needed for training or large-scale serving.

The long-standing software challenge was that the TPU ecosystem grew around TensorFlow and JAX, so PyTorch users had extra hurdles to use TPUs effectively.

The old route: torch_xla

Until now, running PyTorch on TPUs typically required PyTorch/XLA (torch_xla). It worked, but had drawbacks: tensors were lazy, which meant code execution timing could be non-obvious; minor changes could trigger expensive graph recompilations; and debugging was harder to reason about.

Hugging Face previously provided Optimum TPU to enable inference and training of Hugging Face models on TPU; with TorchTPU providing native support, Optimum TPU is no longer required to run Transformers and Diffusers on TPU.

What TorchTPU provides

TorchTPU registers TPU as a new PyTorch backend. Developers can write normal PyTorch code while TorchTPU provides Pallas kernels tuned for TPU hardware, supports distributed training, and integrates with torch.compile for additional speedups.

Installing on a TPU VM is a single command:

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

The extra index supplies the CPU build of torch; without it pip would download the much larger CUDA build, which is unnecessary for TPU usage.

Your favourite libraries, now on TPU

Transformers

Transformers v5 was designed so model definitions are the single source of truth, enabling models to run anywhere — and "anywhere" now includes TPU. To load a model onto a TPU, pass device_map="tpu" when calling from_pretrained and the model will be loaded directly on the TPU. By default models run in eager mode; you can also opt into torch.compile for better performance or run in a distributed setup.

Diffusers

TPU is a first-class device in Diffusers. For prototyping or debugging, eager mode only requires pipe.to("tpu"). For deployment, compiled mode adds a compile call. Your pipelines and code remain the same, so there’s no separate TPU fork to maintain.

Example (StableDiffusion3Pipeline with dtype torch.bfloat16):

import torch from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16) pipe.to("tpu")

eager mode: that’s the only change

optional: compiled mode

pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

The first call compiles a graph; subsequent calls with the same image size reuse that graph.

Large models that do not fit on a single chip can be sharded during loading by passing a TensorParallelConfig to from_pretrained, so the transformer is distributed across chips and the full model never needs to fit in one place. Many pipelines support this out of the box; others can be adapted with a tp_plan. The Diffusers TPU guide provides full instructions.

Example performance results on a Cloud TPU v6e-4 (TP=4, bf16, batch 1) are reported for tasks like MiniMax-H3 (video + audio) and Qwen-Image-2.1 at various resolutions and step counts, with some cases showing substantial speedups in compiled mode versus eager mode (e.g., 1024×1024, 40 steps: 3.5 s compiled vs. 8.0 s eager in one example).

Kernels

The Kernels Hub hosts ready-to-use TPU kernels that can be loaded via the Kernels library to accelerate models. TPU kernels are written in Pallas, which provides low-level control over the hardware. Kernels from the Tokamax library can also be used.

The article shows a ragged-dot kernel from Tokamax integrated into a Mixture-of-Experts model by registering the kernel and using it as the experts forward implementation for a Qwen model. The performance benefit grows with batch size; at batch size 64, decoding was 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B when using a Hub kernel versus the default implementation. The piece also compares decode throughput on a TPU v6e-8 between a Hub kernel and the default implementation.

Testing and ongoing work

Hugging Face is continuously improving TPU support by running the transformers test suite on TPU. They are working to add a nightly CI on TPU and to surface results in the Transformers CI dashboard. Currently, 96.5% of tests that run on TPU pass. Each failure is turned into a small repro script and an issue filed on the TorchTPU repository: so far 34 issues have been reported and 19 fixed. Running on TPU also revealed bugs in Transformers itself (e.g., in StaticCache and Qwen3 tensor-parallel plans) which were fixed upstream.

The goal is to detect breakages on TPU proactively so users encounter fewer surprises.

Getting started

TPUs are no longer a side quest for PyTorch users. Install TorchTPU, pick a model from the Hub, and send it to "tpu" with device_map or .to("tpu"). Further resources include the TorchTPU GitHub, the TPU kernels on the Hub, and Cloud TPU documentation. If you hit problems, open an issue and report your findings.

Hugging Face encourages the community to try the new TPU support and share feedback.