Google open-sources TorchTPU so PyTorch runs natively on TPUs

What Google released

Google has open-sourced TorchTPU, a native PyTorch backend that lets models written in PyTorch run on Google's Tensor Processing Units, or TPUs, the custom AI chips it rents out through Google Cloud. The company announced it at PyTorch Conference 2026 in a Google Cloud blog post on 20 October, saying TPUs now join GPUs and CPUs as a native option for PyTorch.

The code is on GitHub under the Apache 2.0 licence, with version 0.2.0 on the Python package index. Google says it built TorchTPU with the PyTorch team at Meta, which created PyTorch, and the two published a technical paper on arXiv describing the design. Alongside the core backend, Google released TPU versions of the vLLM and SGLang serving engines and TorchTPU support in TorchTitan, PyTorch's reference project for large-scale training.

The practical promise is simple: a PyTorch model that today runs on an Nvidia GPU should be able to target a TPU by changing the device name from cuda to tpu, then get faster with standard PyTorch tools rather than a rewrite.

Why PyTorch on TPUs was awkward

TPUs have mostly been the domain of Google's own frameworks, TensorFlow and later JAX. PyTorch, which most researchers and many companies use, could reach TPUs through an older bridge called PyTorch/XLA, introduced in 2020. According to the TorchTPU paper, that bridge recorded operations lazily into graphs, forced developers to edit their code to mark which regions could be compiled, suffered long compile times, lacked a true step-by-step eager mode and lagged behind PyTorch's distributed training features.

That gap mattered commercially. Developers who wanted to try TPUs often faced porting work, which kept them on GPUs. Google is now retiring the old path: the PyTorch/XLA repository is marked deprecated, tells users to start new TPU work with TorchTPU and migrate existing projects, and says it will be archived on 15 November 2026.

How TorchTPU runs PyTorch code

TorchTPU plugs into PyTorch through its standard extension point for new hardware, called PrivateUse1, so tensors on a TPU are ordinary PyTorch tensors rather than special wrappers. It runs one process per TPU device, as PyTorch does with GPUs, and the paper says it supports about 450 of PyTorch's core ATen operations, including the forward and backward passes needed for mainstream deep learning models.

It offers three eager modes for running code step by step. Strict, the default, sends each operation to the chip asynchronously with normal PyTorch behaviour. Debug runs synchronously with extra bounds checks. A new mode called DeferAndFuse batches runs of operations together to cut memory traffic, and the paper says it can be 50% or more faster than Strict.

For production speed, developers call torch.compile with the tpu backend. PyTorch's TorchDynamo captures the model as a graph, which TorchTPU lowers to StableHLO and hands to Google's XLA compiler to produce TPU machine code, with compiled programs cached across processes and restarts. Custom kernels written in Google's Pallas language or the Helion kernel language can be called like ordinary PyTorch functions, and distributed training uses PyTorch's own DDP, FSDP2 and DTensor tools.

The performance numbers

The headline figure in the TorchTPU paper is up to 52% model FLOPs utilisation, a measure of how much of the chip's theoretical maths capacity is doing useful work, when training Qwen 3 8B with TorchTitan, with near-linear scaling up to 8,192 of Google's Ironwood TPUs. The paper also reports Llama 3 405B on 1,024 Ironwood chips at 500 tokens per second per chip and over 50% utilisation, and a Qwen 3 235B mixture-of-experts model reaching 28% on 8,192 chips.

A detailed example shows how much tuning still matters. On a FLUX image-model training job, an untouched eager run took 11.49 seconds per step. DeferAndFuse cut that to 9.95 seconds, sharding with SimpleFSDP to 6.52, custom attention kernels to 5.55, full torch.compile to 3.02, and final model rewrites to 2.78, about four times faster than where it started.

The paper also ported a Muse Glimmer 30B training script originally written for GPUs, changing only the attention function because TorchTPU does not yet support PyTorch's FlexAttention, and reports 29% to 36% utilisation on four TPU v7 chips. Google says it continuously tests about 5,000 popular Hugging Face models in both eager and compiled modes.

Serving, training and the wider stack

Google paired the launch with integrations across the PyTorch ecosystem. Hugging Face Transformers version 5 can now place a model on a TPU by passing device_map tpu, and the Diffusers image library runs natively too. The new vLLM-TorchTPU engine was built with Inferact and SGLang-TorchTPU with RadixArk, and both support compact FP8, MXFP4 and NVFP4 model formats.

For fleets, Google points to llm-d 1.0, which also launched this week, to route inference requests across TPUs and GPUs on Kubernetes. Ray now treats TPUs as a first-class accelerator, the VeRL reinforcement learning framework has a TPU backend in preview, and Google's on-device runtime LiteRT is joining the PyTorch ecosystem with a one-call converter.

Meta's quote in the announcement is notable. Jana van Greunen, its director of PyTorch engineering, says TorchTPU lets Meta deploy PyTorch workloads to Google TPUs and supports major language model, ads and recommendation workloads on key production use cases.

What it means for the chip contest

TorchTPU arrives one day after Amazon opened TorchNeuron, its own native PyTorch backend for Trainium chips. Both moves attack the same advantage: Nvidia's dominance rests partly on the fact that PyTorch code simply works on its GPUs. If PyTorch runs natively and quickly elsewhere, switching chips becomes a pricing and availability decision rather than an engineering project.

Open inference stacks are moving the same way, from DeepSeek's newly open DeepInfer engine to llm-d, and TorchTPU slots Google's chips into that shared software layer instead of asking developers to adopt JAX.

Caveats

The performance numbers are Google's and Meta's own, measured on Google's hardware, and the FLUX example shows that out-of-the-box speed can be several times slower than a tuned setup. The paper lists known gaps: a TPU can be used by only one process at a time, which blocks some reinforcement learning setups that share a chip between training and sampling; FlexAttention is not supported yet; and collectives are compiled separately by default, which the paper says can cost 30% to 40% end to end unless developers opt into faster modes.

TorchTPU also only targets Google Cloud TPUs, tested from the v5e generation onward, so it does not reduce dependence on Google itself. The authors describe this as the start of a 6 to 12 month roadmap covering production hardening, FP8 and FP4 support on next-generation chips, and wider model coverage.