Skip to content
Neural Daily

ToolsNews

TorchTPU brings native PyTorch and Hugging Face to TPUs

TorchTPU is a new PyTorch backend that makes the TPU just another device, like a GPU. Hugging Face announced that Transformers, Diffusers and Kernels now run on it natively, so a Hub model only needs to be sent to "tpu".

By: Mikola Máté Patrik · · 3 min read

Hugging Face TorchTPU banner: the Hugging Face emoji and the PyTorch logo as cartoon characters cheering behind a glowing Google TPU chip
Image: Hugging Face · Press kit · Official Hugging Face blog image · Source

What changes

Until now, running PyTorch code on TPUs meant PyTorch/XLA (torch_xla) or learning JAX. torch_xla worked, but its tensors were lazy: code did not always run when you expected, small changes triggered graph recompilations and debugging was hard. TorchTPU (GitHub) is a native PyTorch backend instead, so the TPU becomes just another device (Hugging Face).

It ships with Pallas kernels tuned for TPUs, distributed training support and a torch.compile integration. On a TPU VM it installs with one command: pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu. The extra index pulls the CPU build of torch, so you skip the much larger CUDA build you do not need.

What already runs on it

  • Transformers v5: pass device_map="tpu" and the model loads straight onto the TPU. It runs in eager mode by default, can be sped up with torch.compile and works in distributed setups.
  • Diffusers: pipe.to("tpu") is enough for prototyping, and deployment adds one compile call. There is no TPU fork to maintain. Models too big for one chip are sharded across chips at load time with TensorParallelConfig.
  • Kernels: the Kernels Hub has ready-to-use TPU kernels written in Pallas, and kernels from the Tokamax library can be plugged in too.
  • Optimum TPU: no longer required for Transformers and Diffusers now that support is native.

The numbers

In Hugging Face's measurements, a Tokamax ragged-dot kernel pays off on MoE models at large batch sizes: at batch 64, decoding is 21.5x faster on Qwen3-30B-A3B and 12.8x faster on Gemma 4 26B-A4B than the default implementation, on a TPU v6e-8.

Hugging Face chart: decode throughput on TPU v6e-8, the Tokamax ragged_dot Hub kernel is 21.5x faster on Qwen3-30B-A3B and 12.8x faster on Gemma 4 26B-A4B at batch 64
Hugging Face · Press kit · Official Hugging Face blog image · Source

Source: Hugging Face. Decode throughput on a TPU v6e-8 (TP=8), Hub kernel vs. the default grouped_mm implementation.

In Diffusers on a Cloud TPU v6e-4, with each model split across four chips, Qwen-Image-2.1 renders a 1024×1024, 40-step image in 3.5 seconds compiled and 8.0 seconds eager. MiniMax-H3 produces a 5.2-second video with audio in 188 and 262 seconds respectively.

How stable is it

Not everything is green yet. Of the Transformers tests that run on TPU, 96.5% pass. Hugging Face has reported 34 issues to TorchTPU so far, 19 of them already fixed, and it also found bugs in Transformers itself (StaticCache, Qwen3 tensor-parallel plans) that were fixed upstream. A nightly TPU CI is in the works, with results going to the Transformers CI dashboard.

Why it matters for developers

TPUs used to be a detour for PyTorch users; now they are roughly a device swap. If you train on Google Cloud or serve models at scale, it is worth comparing TPU pricing against GPU options, since you no longer have to port your code. The full Diffusers recipe is in the TPU guide. The 96.5% pass rate also means you should run your own model's tests before going to production.

    Produced with the help of AI, reviewed by an editor. How we make our articles

    Reviewed and edited by the Neural Daily desk. About

    Related articles