Hugging Face ❤️ TorchTPU: TPUs, Natively

What it claims

TorchTPU is a native PyTorch backend that adds tpu as an ordinary device, so Transformers, Diffusers and the Kernels library run on Google Cloud TPUs without torch_xla. Hugging Face says Optimum TPU is no longer required for running Transformers and Diffusers models. The backend ships Pallas kernels tuned for TPUs, distributed training support and a torch.compile integration. The post contrasts it with PyTorch/XLA, where lazy tensors deferred execution, small code changes could trigger graph recompilation, and debugging was hard.

Evidence

The speedups come from swapping in a Hub kernel for Mixture-of-Experts layers. The kernel is a Tokamax ragged-dot kernel, measured on a TPU v6e-8 against the default experts implementation. At batch size 64, decode throughput is 21.5x higher on Qwen3-30B-A3B and 12.8x higher on Gemma 4 26B-A4B. The post says the gain grows with batch size. It does not give figures for smaller batches.

On test coverage, 96.5% of the Transformers tests that run on TPU pass. The Hugging Face team has filed 34 issues on the TorchTPU repo, and 19 are fixed. TPU runs also exposed bugs in Transformers itself, in StaticCache and the Qwen3 tensor-parallel plans, which were fixed upstream. Diffusers sample images were produced on a Cloud TPU v6e-4 with TP=4, bf16 and batch 1. No timings are given for them.

How to use it

Install on a TPU VM. The extra index pulls the CPU build of torch instead of the much larger CUDA build:

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

Transformers (v5) loads directly onto the chip with device_map="tpu" and runs eager by default:

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="tpu", # will load model directly on the TPU
)

Diffusers uses pipe.to("tpu") for eager mode. Compiled mode adds one call:

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

The first call compiles. Later calls at the same image size reuse the graph. For models too large for one chip, passing a TensorParallelConfig to from_pretrained shards the transformer across chips as it loads. Pipelines without built-in support need a tp_plan.

MoE kernels register through ALL_EXPERTS_FUNCTIONS.register(...) and are selected with experts_implementation=.... The example loads tengomucho/ragged-dot-tpu via get_kernel(repo_id, version=1, trust_remote_code=[repo_id]).

Limits

The 3.5% of failing tests and 15 open issues mean some Transformers paths still break on TPU. The pass rate covers only the tests that run on TPU, and the post does not say how many tests are excluded. Nightly TPU CI is planned but not yet running, and the dashboard is marked work in progress. The headline speedups are for one kernel at batch 64 against the default implementation. They are not a comparison with GPUs or with torch_xla. Kernel loading requires trust_remote_code for a third-party repo. Compiled Diffusers graphs assume static shapes (dynamic=False), so changing the image size triggers a recompile.