TorchTPU adds TPU as a native PyTorch backend, turning TPUs into just another device for Hugging Face Transformers and Diffusers. Before this, running PyTorch on TPUs meant porting code to PyTorch/XLA (torch_xla) or learning JAX. torch_xla worked, but it had lazy tensors, graph recompilations triggered by small changes, and debugging that was hard to reason about. Hugging Face had built Optimum TPU to give its models inference and training support on TPUs; with TorchTPU, that layer is no longer required for Transformers and Diffusers.

TPUs are Google’s AI accelerators: chips built for the matrix math that neural networks spend most of their time on. Google has been designing them for about ten years and uses them to train and serve its own models, including Gemini. They are available on Google Cloud from a single chip to pods of thousands, and they are most useful when a lot of compute is needed for training or serving at scale. The catch has always been software: the TPU stack grew up around TensorFlow and JAX. TorchTPU changes that by adding TPU as a backend to regular PyTorch code, with Pallas kernels tuned for TPUs, distributed training support, and torch.compile integration. Installing on a TPU VM is one command: pip install torch-tpu –extra-index-url https://download.pytorch.org/whl/cpu. The extra index gives the CPU build of torch, avoiding the much larger CUDA build.

In Transformers, a model can be loaded on a TPU by passing device_map=”tpu”. The article’s example uses Qwen/Qwen3.5-9B, moves the inputs to the “tpu” device, and runs generate. Eager mode is the default, with torch.compile and distributed execution available. This is part of the Transformers v5 goal of having model definitions be the source of truth so they can run anywhere. In Diffusers, TPU is now a first-class device: eager mode only needs pipe.to(“tpu”), while compiled mode adds a compile call such as pipe.transformer.compile(backend=”tpu”, fullgraph=True, dynamic=False). Pipelines and code stay the same, with no TPU fork to maintain. For models too large for one chip, a TensorParallelConfig passed to from_pretrained shards the transformer as it loads; most popular pipelines support it out of the box, and the rest are one tp_plan away. Sample runs on a Cloud TPU v6e-4 with TP=4, bf16, and batch 1 include MiniMax-H3 for video and audio and Qwen-Image-2.1 image generations. A 1344×768, 49-step Qwen-Image-2.1 generation took 188s compiled versus 262s eager; a 1024×1024, 40-step generation took 3.5s compiled versus 8.0s eager.

The Kernels Hub provides ready-to-use TPU kernels, written in Pallas for low-level hardware control, and kernels can also come from Tokamax. The post demonstrates a ragged-dot kernel from Tokamax plugged into a Mixture-of-Experts model using get_kernel and ALL_EXPERTS_FUNCTIONS.register. Batch size matters: at batch size 64, decoding was 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B, comparing a Hub kernel with the default implementation on a TPU v6e-8.

Hugging Face is testing TPU support by running the Transformers library suite, with nightly CI on TPU planned and results destined for the Transformers CI dashboard. Of the tests that run on TPU, 96.5% pass; the article is explicit that things are not all green yet. Each failure is turned into a repro script and an issue on the TorchTPU repo: 34 issues so far, 19 already fixed. TPU testing has also caught bugs in Transformers itself, including StaticCache and Qwen3 tensor-parallel plans, fixed upstream. The closing guidance is to install TorchTPU, pick a model from the Hub, and send it to “tpu”, with links to TorchTPU on GitHub, TPU kernels on the Hub, and Cloud TPU docs.