Google has open-sourced TorchTPU, the software stack that lets PyTorch run natively on its TPUs, the AI accelerators it rents out through Google Cloud. It announced the move on October 20, 2026, at the PyTorch Conference 2026. The project, developed together with Meta’s PyTorch team, is on GitHub under the Apache 2.0 license, and Google says it turns TPUs into a native PyTorch hardware option, just like GPUs and CPUs already are.
TorchTPU in 20 seconds
- Google open-sources TorchTPU under the Apache 2.0 license, developed with Meta’s PyTorch team.
- It runs PyTorch on TPUs with minimal code changes, in eager mode or with torch.compile.
- vLLM, SGLang, Hugging Face, Ray, and llm-d already integrate it.
- Google reports near-linear scaling up to 8,192 Ironwood TPUs.
The announcement, signed by Mark Lohmeyer, vice president and general manager of AI infrastructure and compute at Google Cloud, and Bill Jia, vice president of engineering, targets one of the classic barriers to using TPUs: most AI teams work in PyTorch, while Google’s software for its chips has traditionally revolved around JAX and TensorFlow. For anyone buying cloud capacity, the question is whether a workload can move to a different accelerator without being rewritten.
What changes compared to PyTorch/XLA
Until now, the way to use PyTorch on TPUs was PyTorch/XLA, which Google’s developer blog describes as TorchTPU’s predecessor and which only supported pure SPMD code (the same program on every node). TorchTPU plugs into PyTorch through the PrivateUse1 interface, the one the framework reserves for external hardware backends, so tensors on TPU are regular PyTorch tensors, with no wrappers or subclasses.
The workflow Google proposes is to debug in eager mode, operation by operation, and then move to torch.compile for production. In that second step, TorchTPU captures the graph with TorchDynamo and compiles it with XLA instead of TorchInductor, PyTorch’s default compiler. According to Google, there is also an eager mode with automatic operation fusion that delivers between 50% and more than 100% improvement over standard eager mode. For distributed training it supports DDP, FSDPv2, and DTensor, and allows custom kernels written in Pallas or JAX.
The version published on PyPI is 0.2.0, which requires Python 3.11 or later and installs with pip. Google points to its PyTorch on TPU documentation to get started.
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpuScaling on Ironwood and the inference stack on GKE
The performance figures Google provides come from its own tests. Using TorchTitan, PyTorch’s reference implementation for distributed training, the company says it reached 52% model FLOPs utilization (MFU) training Qwen 3 8B, with near-linear, stable scaling up to 8,192 Ironwood TPUs (TPU v7x), the generation Google introduced for large-scale inference. TorchTitan’s recipes include models such as Llama 3 at 70 billion and 405 billion parameters, the MoE-based Qwen 3 235B, and FLUX.1 12B.
For inference, Google has built native backends for vLLM, together with Inferact, and for SGLang, together with RadixArk. Above that sits llm-d 1.0, a project in the CNCF’s Sandbox stage involving Red Hat, IBM, and more than 500 contributors from 25 organizations, which adds prefix-cache-aware routing and separates the prefill and decode phases. On GKE, Google Cloud’s Kubernetes service, it lets teams serve models on GPUs or TPUs behind the same endpoint.
As an example, Google cites Zenteiq.ai, which its founder says has more than quadrupled aggregate inference throughput using GKE Inference Gateway and llm-d, going from 1.47 billion to 5.97 billion tokens per day for every 1,000 TPU v6e accelerators. Ray, Anyscale’s orchestrator, also now supports TorchTPU to distribute workloads across CPUs, GPUs, and TPUs.
An alternative to CUDA, with caveats
The move has a clear reading in the accelerator market. Much of NVIDIA’s advantage lies in CUDA and in PyTorch running better on its GPUs than on anything else. TPUs already compete on price and performance, as explored when looking at how Google’s TPUs challenge NVIDIA, but software remained a barrier. Meta, which uses Google’s TPUs, has an interest in not depending on a single vendor. Meta’s head of PyTorch engineering, Jana van Greunen, says in the announcement that the company cannot be limited by a single hardware architecture and that TorchTPU lets it deploy its PyTorch workloads on TPUs for language models, ads, and recommendation systems.
There are things still to prove outside Google’s own tests. Its roadmap includes reducing recompilations when sequence length or batch size change, adding full dynamic-shapes support in torch.compile, and expanding the library of precompiled kernels. These are known issues with XLA-based compilers and weigh on highly variable inference workloads. TorchTPU also only works with Google’s TPUs, which are accessed through Google Cloud, so the portability it offers is between GPUs and TPUs within PyTorch, not between clouds.
Frequently Asked Questions
What is TorchTPU?
It’s Google’s software stack for running PyTorch natively on its TPUs. It has been open-sourced under the Apache 2.0 license and developed together with Meta’s PyTorch team.
Does TorchTPU replace PyTorch/XLA?
Google describes PyTorch/XLA as its predecessor. TorchTPU integrates into PyTorch as a native backend, though Google has not formally announced the retirement of PyTorch/XLA.
What performance does TorchTPU deliver?
According to Google, 52% MFU training Qwen 3 8B with TorchTitan, and near-linear scaling up to 8,192 Ironwood TPUs. These are the company’s own figures.
Can TorchTPU be used outside Google Cloud?
The software is open source, but it only works on Google’s TPUs, which are available through Google Cloud.
Sources:
