Google Open-Sources TorchTPU, a Native PyTorch Backend for Cloud TPUs

Google has open-sourced TorchTPU, a native PyTorch backend for Cloud TPUs that the company says replaces PyTorch/XLA, announced at PyTorch Conference 2026 on October 20, 2026. Google co-developed the stack with the PyTorch team at Meta. Existing PyTorch code now targets a TPU through a standard torch.device("tpu"), placing TPUs beside GPUs and CPUs as native PyTorch devices. Google recommends TorchTPU for all new workloads.

The stakes are straightforward. PyTorch dominates research and much of production AI, but TPU adoption outside Google's internal JAX ecosystem has historically required an impedance-matching layer, PyTorch/XLA, which imposed lazy-tensor execution and manual compilation hints. TorchTPU removes that friction at a moment when Google is competing with NVIDIA's CUDA ecosystem and AWS's Trainium stack for AI training and inference workloads.

What shipped, exactly?

TorchTPU is a software stack, not a model. The release spans four open repositories:

• TorchTPU — the core backend, v0.2.0 on PyPI, requiring Python 3.11+, Apache 2.0 licensed
• vLLM-TorchTPU — inference serving, built with Inferact, keeping upstream schedulers and API servers with custom Pallas kernels
• SGLang-TorchTPU — serving built with RadixArk, preserving RadixAttention prefix reuse
• TorchTitan — training recipes for Llama 3 70B/405B, Qwen 3 8B/32B, Qwen 3 235B MoE and FLUX.1 12B

Google tested the stack on TPU generations from v5e onward and scaled it to 8,192 Ironwood (TPU7x) chips. Serving supports FP8, MXFP4 and NVFP4 checkpoints.

How does TorchTPU execute PyTorch code?

TorchTPU plugs in through PyTorch's PrivateUse1 interface, the extension point for third-party accelerators. Tensors remain ordinary PyTorch tensors, with no subclasses or wrappers. The stack runs one process per device and preserves tensor views and in-place mutations.

Unlike PyTorch/XLA, which recorded lazy-tensor graphs and required developers to mark compilable regions manually, TorchTPU defaults to eager execution. It also handles divergent MPMD (multiple program, multiple data) execution, where PyTorch/XLA supported only pure SPMD (single program, multiple data) code.

The stack ships three eager modes:

• Strict (default) — dispatches each operator asynchronously with standard PyTorch semantics
• Debug — runs synchronously with tensor bounds checking
• DeferAndFuse — fuses runs of eager operators to cut HBM (high-bandwidth memory) traffic; the TorchTPU paper reports it can be 50% or more faster than Strict

For peak throughput, torch.compile(backend="tpu") captures the graph with TorchDynamo and lowers it through StableHLO to OpenXLA, which emits TPU machine code. Compiled executables are cached across processes, hosts and job restarts.

Custom Pallas kernels are callable as standard PyTorch functions, alongside prebuilt Tokamax kernels and Helion tile kernels. Distribution runs through the tpu_dist process group, supporting DDP, FSDP2 and DTensor, plus a topology-aware mesh for TPU interconnects.

How fast is it?

All figures below are self-reported by Google and its partners.

• Training scale: TorchTitan reached up to 52% MFU (model FLOPs utilization) training Qwen 3 8B on Ironwood, with near-linear scaling to 8,192 TPUs
• Tuning ladder: in a TorchTitan fine-tuning case, step time fell from 11.49 seconds out-of-box eager to 2.78 seconds fully tuned; DeferAndFuse alone delivered 1.15x, full torch.compile delivered 3.80x
• Model coverage: Google continuously tests roughly 5,000 Hugging Face models in both eager and compiled modes
• Fleet serving: Zenteiq grew daily output from 1.47B to 5.97B tokens per 1,000 TPU v6e chips — a gain Google attributes to llm-d routing on GKE, not TorchTPU alone

What integrations ship with it?

• Hugging Face: Transformers v5 deploys to TPU with device_map="tpu"; Diffusers supports eager, torch.compile and tensor parallelism
• Serving features: FP8 KV caches, FP8/MXFP4/NVFP4 checkpoints, speculative decoding with MTP and EAGLE-3
• Ecosystem: llm-d 1.0, Ray, TorchRec, VeRL (preview) and LiteRT-Torch for on-device export

How does it compare with PyTorch/XLA, TorchNeuron and JAX?

The competitive picture sharpens in comparison. TorchNeuron, AWS's counterpart for Trainium 2 and Trainium 3, remains in public beta at Neuron 2.33 and also uses the PrivateUse1 mechanism. JAX on TPU, Google's first-party framework path, stays maintained but sits outside PyTorch entirely. PyTorch/XLA is now formally being replaced by TorchTPU.

The compile paths differ. TorchTPU lowers through TorchDynamo to StableHLO to XLA. TorchNeuron uses the Neuron Compiler. PyTorch/XLA went directly through XLA with lazy tensors. On distributed APIs, TorchTPU supports DDP, FSDP2, DTensor, SPMD and MPMD, against PyTorch/XLA's pure-SPMD limitation. TorchTPU is the only option in this set with a disclosed license: Apache 2.0.

Installation

On a TPU VM, install the package and select the device:

# pip install torch_tpu
import torch
device = torch.device("tpu")
x = torch.randn((4, 4), device=device)

The repository README pins torch_tpu==0.2.0 with the CPU PyTorch wheel index; without it, pip pulls the CUDA build of torch. Full documentation lives at torchtpu.readthedocs.io.

The release signals where Google's accelerator software strategy now points: converge PyTorch and TPU through standard abstractions rather than a parallel framework, and let the ecosystem — vLLM, SGLang, Hugging Face, roughly 5,000 tested models — carry adoption. Whether self-reported 52% MFU and 3.80x compile gains hold up under independent benchmarks will determine how quickly PyTorch/XLA users migrate.