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

TorchTPU is a software stack, not a model. It is published as v0.2.0 on PyPI under Apache 2.0, requires Python 3.11 or later and runs on Google Cloud TPUs. The project spans four open repositories: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU and TorchTitan. Google tested it on TPU generations from v5e up to 8,192 Ironwood (TPU7x) chips. Serving supports FP8, MXFP4 and NVFP4 checkpoints.

The backend plugs in through PyTorch's PrivateUse1 interface. Tensors are ordinary PyTorch tensors with no subclasses or wrappers, and it runs one process per device while preserving tensor views and in-place mutations. Its predecessor, PyTorch/XLA, recorded lazy-tensor graphs, often required developers to add code to mark compilable regions, and supported only pure SPMD code. TorchTPU also handles divergent MPMD execution.

TorchTPU ships with three eager modes. Strict is the default and dispatches each operator asynchronously with standard PyTorch semantics; Debug runs synchronously with tensor bounds checking; DeferAndFuse fuses runs of eager operators to cut HBM traffic and can be 50% or more faster than Strict, according to the TorchTPU paper. For peak speed, 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, and prebuilt Tokamax kernels and Helion tile kernels also work. Distribution uses the tpu_dist process group and supports DDP, FSDP2 and DTensor, along with a topology-aware mesh for TPU interconnects.

All performance figures are self-reported by Google and its partners. TorchTitan reached up to 52% MFU training Qwen 3 8B on Ironwood, and Google reports near-linear scaling up to 8,192 Ironwood TPUs. In a TorchTitan fine-tuning case, step time fell from 11.49 seconds in out-of-box eager mode to 2.78 seconds fully tuned; DeferAndFuse alone gave 1.15x, and full torch.compile gave 3.80x. Google says it continuously tests about 5,000 Hugging Face models in eager and compiled modes. Zenteiq grew daily output from 1.47 billion to 5.97 billion tokens per 1,000 TPU v6e chips, a gain the report attributed to llm-d routing on GKE, not TorchTPU alone.

Integrations include Hugging Face Transformers v5, which deploys to TPU with device_map="tpu", and Diffusers, which supports eager, torch.compile and tensor parallelism. vLLM-TorchTPU was built with Inferact and keeps upstream schedulers and API servers while adding custom Pallas kernels. SGLang-TorchTPU was built with RadixArk and keeps RadixAttention prefix reuse. Serving features include FP8 KV caches, FP8, MXFP4 and NVFP4 checkpoints, and speculative decoding with MTP and EAGLE-3. TorchTitan includes recipes for Llama 3 70B/405B, Qwen 3 8B/32B, Qwen 3 235B MoE and FLUX.1 12B. The stack also lists llm-d 1.0, Ray, TorchRec, VeRL in preview and LiteRT-Torch for on-device export.

Installation requires the package on a TPU VM and selecting the tpu device. The README pins torch_tpu==0.2.0 with the CPU PyTorch wheel index; without it, pip pulls the CUDA build of torch.