AI & Tech

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

October 21, 2026
briefly

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

TL;DR

  • Size: A software stack, not a model. 4 open repos: TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan. TorchTPU is v0.2.0 on PyPI, Python 3.11+, Apache 2.0.
  • Runs on: Google Cloud TPUs. Tested on generations from v5e, scaled to 8,192 Ironwood (TPU7x) chips. Serving supports FP8, MXFP4 and NVFP4 checkpoints.
  • Performance: Eager mode for debugging, torch.compile for production-level throughput.
  • Best: 52% MFU training Qwen 3 8B with TorchTitan on Ironwood (self-reported).
  • Worst: Out-of-box eager took 11.49s per step in a fine-tuning case, versus 2.78s fully tuned.

What is TorchTPU?

TorchTPU is an open-source PyTorch backend from Google and Meta that runs PyTorch models natively on TPUs. It plugs in through PyTorch’s PrivateUse1 interface. Tensors are ordinary PyTorch tensors, with no subclasses or wrappers. It runs one process per device and preserves tensor views and in-place mutations.

The predecessor, PyTorch/XLA, recorded lazy-tensor graphs. Developers often added code to mark compilable regions. PyTorch/XLA also only supported pure SPMD code. TorchTPU also handles divergent MPMD execution.

How does TorchTPU execute PyTorch code?

TorchTPU ships 3 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 traffic. The TorchTPU paper says it can be 50% or more faster than Strict.

For peak speed, torch.compile(backend="tpu") captures the graph with TorchDynamo. It lowers 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. Prebuilt Tokamax kernels and Helion tile kernels also work.

Distribution uses the tpu_dist process group. It supports DDP, FSDP2 and DTensor, along with a topology-aware mesh for TPU interconnects.

How fast is TorchTPU?

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

  • Training scale: TorchTitan reached up to 52% MFU on Qwen 3 8B. Google reports near-linear scaling up to 8,192 Ironwood TPUs.
  • Tuning ladder: In a TorchTitan fine-tuning case, step time fell from 11.49s to 2.78s. DeferAndFuse alone gave 1.15x. Full torch.compile gave 3.80x.
  • Model coverage: Google continuously tests about 5,000 Hugging Face models in eager and compiled modes.
  • Fleet serving: Zenteiq grew daily output from 1.47B to 5.97B tokens per 1,000 TPU v6e chips. That gain came from llm-d routing on GKE, not TorchTPU alone.

What integrations ship with TorchTPU?

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

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

Feature TorchTPU (repo) PyTorch/XLA (repo) TorchNeuron (AWS docs) JAX on TPU (Google docs)
Vendor Google + Meta Google / PyTorch AWS Google
Hardware Cloud TPU (v5e and later tested) Cloud TPU Trainium 2, Trainium 3 Cloud TPU
PyTorch device tpu xla neuron Not PyTorch
Backend mechanism PrivateUse1 Lazy-tensor graph recording PrivateUse1 Not applicable
Default execution Eager (Strict, Debug, DeferAndFuse) Lazy tensors Eager plus torch.compile Not disclosed
Compile path TorchDynamo to StableHLO to XLA XLA Neuron Compiler XLA
Distributed APIs DDP, FSDP2, DTensor, SPMD and MPMD Pure SPMD only FSDP, DTensor, DDP, TP Not disclosed
Custom kernels Pallas, Tokamax, Helion Not disclosed NKI Pallas
vLLM / SGLang serving vLLM-TorchTPU, SGLang-TorchTPU Not disclosed vLLM Neuron vLLM-JAX, SGLang-JAX
Status Open source, v0.2.0 Being replaced by TorchTPU Public beta (Neuron 2.33) Maintained
License Apache 2.0 Not disclosed Not disclosed Not disclosed

How do you install TorchTPU?

Install the package on a TPU VM, then pick the tpu device:

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

The repo’s README pins torch_tpu==0.2.0 with the CPU PyTorch wheel index. Without it, pip pulls the CUDA build of torch. Full guides live at torchtpu.readthedocs.io.

Key Takeaways

  • TorchTPU is open source under Apache 2.0, co-built with Meta.
  • It replaces PyTorch/XLA with eager-first, PrivateUse1-based execution.
  • torch.compile(backend="tpu") lowers through StableHLO and XLA.
  • TorchTitan hit 52% MFU on Qwen 3 8B, scaling to 8,192 Ironwood chips.
  • vLLM, SGLang and Hugging Face now run on TPU through standard PyTorch.

Check out the GitHub repo, paper, docs and Google Cloud announcement. All credit goes to the researcher of this project. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

The post Google Open-Sources TorchTPU, a Native PyTorch Backend for Cloud TPUs appeared first on MarkTechPost.