Google Open-Sources TorchTPU to Bring PyTorch Natively to TPUs
Google has open-sourced TorchTPU, its system for running PyTorch natively on Google TPUs, with the code now on GitHub. The release was detailed on Oct. 20, 2026. Google Cloud Blog
TorchTPU is a PyTorch backend for TPUs, Google's AI accelerator chips. It runs PyTorch models and operations on TPU hardware, making TPU a direct option alongside GPUs and CPUs for training, serving, and deploying models. GitHub repository
For PyTorch developers, the interface stays familiar. Data is managed as standard PyTorch tensors on a tpu device, with no wrapper types to learn. Google Developers Blog Code written for eager execution, where commands run line by line for quick testing, runs on TPUs as it is. TorchTPU supports native PyTorch Eager Mode. Google Cloud
Underneath, TorchTPU implements PyTorch ATen kernels for TPU. ATen is the library of basic operations inside PyTorch, so existing operations can shift to TPU hardware without a separate programming model. The current release is version 0.2.0, installable via pip as torch_tpu==0.2.0, and it supports Python 3.11 and newer.
TorchTPU also works with torch.compile(), which converts code into optimized graphs for faster repeated runs. For that path it includes a compilation cache for XLA computations, the optimized instructions TPUs run. That gives teams two modes in one stack. Use eager mode for iteration and debugging. Use compiled graphs with caching for throughput.
Large-scale training is covered with Data Parallelism and Tensor Parallelism, two standard ways of splitting work across chips. The repository provides examples for Llama, Qwen, ResNet, DLRM, minGPT and UNet models, covering dense language models, image-focused convolutional networks, recommendation models, and diffusion-style UNet work.
For inference, TorchTPU runs PyTorch models natively inside the widely used vLLM and SGLang serving engines without model rewrites. Google Cloud Blog That avoids maintaining a separate TPU-only branch for serving.
Google describes TorchTPU as a new engineering stack built to run PyTorch natively on TPU infrastructure with peak efficiency. Google Developers Blog
The broader context here is that the shift is about friction, not raw capability. PyTorch on TPU existed through other paths, but those paths had rough edges. Wrapper tensor types, graph-only limits, and forked serving stacks added extra work.
In my view, this changes hardware choices by making them cheaper to test. If TPU becomes a device setting rather than a separate stack, a team standardized on PyTorch, ATen operations, Data and Tensor Parallelism, and vLLM or SGLang can try TPU capacity with less rebuilding. That does not remove performance tuning, XLA compilation behavior, or capacity planning. It lowers the cost of trying, and over time engineers tend to adopt what is easy to run.
