TorchTPU Is Open Source: Native PyTorch on Google TPUs Explained

Google open-sourced TorchTPU on October 20, 2026, giving PyTorch developers a native backend for Google's Tensor Processing Units. The release brings familiar model code, debugging and serving tools to TPUs, reducing the software work involved in evaluating an alternative to GPUs.

The important change is not that PyTorch can run on TPUs for the first time. It already could. The change is how closely that experience now follows ordinary PyTorch development.

What TorchTPU changes from PyTorch/XLA

Google's Cloud TPU documentation describes TorchTPU as the replacement for PyTorch/XLA and recommends it for new workloads. The older approach used lazy-tensor graph recording, which could require developers to identify compilable regions explicitly.

TorchTPU executes eagerly by default, so operations behave more like the step-by-step code developers inspect while debugging. For performance, torch.compile captures the computation and lowers it through OpenXLA and StableHLO into TPU machine code. Distributed work uses familiar PyTorch interfaces and one process per device.

That distinction is useful. Removing a JAX dependency from a developer's model workflow does not mean removing XLA compilation from the system. TorchTPU changes the interface developers work with while retaining a compiler underneath.

In its launch announcement, Google says it co-developed the backend with Meta's PyTorch team. The release connects the backend to training, inference and orchestration tools rather than asking users to adopt a completely separate framework.

How to start with TorchTPU

The official repository lists version 0.2.0 as its public release, Python 3.11 or later, and an Apache 2.0 license. Its installation command uses the CPU PyTorch package index to avoid pulling in an unnecessary CUDA build:

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

This belongs on a suitable TPU machine, not an ordinary laptop that happens to have PyTorch installed. Open-source software does not include free cloud accelerator capacity.

For Transformers, the relevant model-loading option is device_map="tpu". Diffusers pipelines can move to the TPU device too. Those familiar entry points make a first compatibility experiment easier, but they should not be treated as a guarantee that every model or custom operation works unchanged.

Hugging Face support is useful, but incomplete

Hugging Face's release covers Transformers, Diffusers and TPU kernels. It says 96.5% of the tests that run on TPU pass, with 34 issues reported and 19 already fixed. That percentage is a test-suite result, not the share of all Hub models that work or a production reliability score.

Hugging Face also describes its nightly TPU testing dashboard as work in progress. Teams should inspect the actual model and configuration they intend to run rather than treating ecosystem support as uniform.

For a developer maintaining a custom model, this is where the change becomes practical: the first investigation can stay closer to the existing PyTorch implementation. Unsupported operations, memory limits and compilation behavior still need attention, but the model does not have to begin life as a separate TPU-specific rewrite.

What changes for vLLM serving

The separate vLLM TorchTPU release uses a platform plugin while retaining upstream scheduling, model runners, KV-cache management and the OpenAI-compatible server. TPU-specific kernels sit underneath that familiar serving interface.

Its requirements are more specific than the base backend: the launch guide calls for Python 3.12 on a Cloud TPU VM. Ironwood, or TPU v7x, is recommended and tested in continuous integration; v6e, v5e and v5p support is described as experimental.

Startup also includes compilation for the configured shapes, which can take minutes. The project provides a cache path so compiled artifacts can survive server restarts. A migration test should therefore measure both initial startup and warmed-up serving, not quietly exclude one from an operational comparison.

Portability is the beginning of the cost comparison

The useful promise is a lower cost of trying different hardware. That is different from a promise that every workload becomes cheaper once it reaches a TPU.

A team evaluating the release should hold model quality, context length and latency requirements constant, then measure throughput, memory use, compilation overhead and actual infrastructure cost. The effort required to maintain custom kernels belongs in that comparison too.

TorchTPU makes hardware choice a more realistic engineering experiment. Whether it becomes a production choice depends on how much of the existing workload remains familiar after the first successful run.

I’m Paras. I work at the intersection of AI, engineering and distribution.

I help technical founders explain infrastructure launches and engineering decisions in clear stories that developers want to read and share.

Book a 30-minute call.