Google and Meta open-source TorchTPU: native PyTorch on TPUs, where porting is as simple as .to("cuda") to .to("tpu"), with native TPU backends for vLLM and SGLang

Background and the problem

PyTorch is the default framework for most AI apps, but many of its features were built around Nvidia GPUs, so changing chips usually meant changing code. Google's 2020 PyTorch/XLA relied on Lazy Tensors, required intrusive code changes, compiled slowly and lacked true eager mode. TorchTPU is a new native backend Google started in mid-2025 with Meta's PyTorch team so TPUs work out of the box like GPUs. AWS released the comparable TorchNeuron a day earlier.

Architecture and how it works

TorchTPU plugs into the PyTorch dispatcher as a PrivateUse1 backend with no JAX runtime dependency; torch.compile goes through Dynamo and lowers to StableHLO for XLA. Eager mode has three variants: async op-by-op by default, a synchronous debug mode, and the new DeferAndFuse, which batches and compiles groups of ops and is described as 50%+ faster than default eager. Collectives sit at the Process Group level so DDP, FSDP and DTensor work directly, and custom kernels can use Pallas, the Tokamax library or the experimental Helion DSL.

Benchmarks

All vendor-reported: 98% of about 4,917 popular Hugging Face models run without crashing and 97% of architecture representatives match GPU numerics; Qwen3 MoE 235B reaches 1,828 TPS on 64 TPU v7 chips versus 1,937 TPS for JAX/MaxText and keeps about 27% MFU at 8,192 chips; Google Cloud cites up to 52% MFU training Qwen 3 8B. Google also cites SemiAnalysis InferenceX at 2.01 cents per million tokens on 8K/1K for TPU v7 with vLLM-TorchTPU. No independent like-for-like comparison against CUDA has been published yet.

Getting started

On a TPU machine with the TorchTPU wheel, set the device to "tpu" and the process group backend to tpu_dist, or use torch.accelerator.current_accelerator() for device-agnostic code. In Hugging Face Transformers v5, pass device_map="tpu"; for serving use vllm-torchtpu or sglang-torchtpu. The preview TPU agent suite (TPU v6e and TPU7x) plugs into coding assistants such as Claude Code and asks for approval before running commands. Google notes a simple device swap guarantees the model runs, not peak performance.