Google Open-Sources TorchTPU, Putting Native PyTorch on Its TPUs

Google announced on October 20, 2026, at PyTorch Conference 2026 that TorchTPU, its native PyTorch stack for Tensor Processing Units (TPUs), is now available as open source on GitHub, the company said on the Google Cloud blog. The software was co-developed with the PyTorch team at Meta and is designed to run existing PyTorch workloads on TPUs with minimal code changes.

With the release, TPUs join GPUs and CPUs as a native option for PyTorch, Google said, spanning training, serving, and deployment through tools including vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face, and LiteRT. Developers can debug models in eager mode and then move to production-level performance through torch.compile to optimize TPU utilization and reduce memory overhead.

Jana van Greunen, Director of PyTorch Engineering at Meta, said in the announcement that TorchTPU allows Meta to co-develop hardware flexibly with its core machine learning systems team while supporting major LLM, Ads, and Recommendation systems workloads. “It enables us to deploy our PyTorch workloads to Google TPUs, provide direct feedback to the engineering teams, and build the ecosystem flexibility everyone needs,” she said.

Hugging and Serving Integrations

Google partnered with Hugging Face so that open-weight models on the Hugging Face Hub run natively on TorchTPU. In Transformers v5, deploying a model to TPUs is done by passing device_map="tpu", and the Diffusers library runs natively on TPUs with support for eager mode, torch.compile, and tensor parallelism across multiple accelerators. Google said it also built dedicated CI/CD pipelines for continuous nightly testing of the Hugging Face Transformers suite on TPU infrastructure. Jeff Boudier, VP of Product at Hugging Face, said the collaboration puts the platform’s hardware-choice commitment into practice.

For serving, the native PyTorch vLLM backend on TorchTPU was developed in partnership with Inferact. Google said it supports continuous batching, quantization, and multi-chip parallelism on TPU hardware using custom Pallas kernels, while preserving upstream worker models, schedulers, and API servers so existing inference stacks require no changes. Simon Mo, CEO and co-founder of Inferact, said the integration brings first-class vLLM support to TPUs.

A native PyTorch backend for SGLang, developed with RadixArk, targets structured and agentic workloads. According to the announcement, SGLang uses RadixAttention to reuse KV cache across requests so repeated prompt prefixes skip recomputation, and the TorchTPU version employs a disaggregation transfer backend that also serves as the I/O backend for HiCache. RadixArk CTO and co-founder Banghua Zhu said TorchTPU has accelerated the company’s work as it rolls out SGLang support across different silicon.

Google is also releasing llm-d 1.0 with Red Hat, IBM, and a community of more than 500 contributors and 25 organizations as a Cloud Native Computing Foundation Sandbox project. The announcement describes llm-d as connecting the PyTorch serving ecosystem with Kubernetes-based production deployment: its router uses prefix-cache-aware routing to send each request to the replica already holding its context, and it splits compute-heavy prefill from memory-heavy decode. Version 1.0 adds multi-tenant flow control, router high availability, stable v1 APIs, and session-aware routing for agents, and on Google Kubernetes Engine a common llm-d endpoint can run on either GPUs or TPUs.

Prof. Sashikumaar Ganesan, founder of Zenteiq, said in the announcement that GKE Inference Gateway powered by llm-d delivered over 4× higher aggregate inference throughput for the company’s synthetic data generation, increasing daily generated token capacity from 1.47 billion to 5.97 billion tokens per 1,000 TPU v6e accelerators.

Training and Orchestration

TorchTitan, the PyTorch-native reference implementation for distributed training, now includes what Google describes as a production-grade TorchTPU integration built on torch.compile, TorchTitan’s SimpleFSDP, specialized Pallas kernels, and graph-level optimizations. TorchTitan provides TorchTPU recipes for pre-training dense models such as Llama 3 70B and 405B, Qwen 3 8B and 32B, and MuseGlimmer 30B, along with the Qwen 3 235B mixture-of-experts model and the FLUX.1 12B diffusion transformer. Google said the integration has been validated on up to 8,000 Ironwood TPUs (TPU v7x) and that TorchTPU training through TorchTitan reached up to 52% Model Flops Utilization on Qwen 3 8B, with near-linear scaling across up to 8,192 Ironwood TPUs.

Google also upstreamed a TPU backend and an end-to-end reference recipe to VeRL, a reinforcement learning framework. The GRPO recipe, built on a disaggregated trainer-sampler architecture and demonstrated on mathematical reasoning datasets, is intended to let developers run VeRL workloads on TPUs while preserving their existing hardware-agnostic model and algorithm code.

Ray adds native TorchTPU support for orchestrating PyTorch data processing, training, reinforcement learning, and serving across CPUs, GPUs, and TPUs within a single workflow. Edward Oakes, chief architect at Anyscale, said the companies worked to make TPUs a first-class accelerator in Ray, adding: “TorchTPU embeds Google TPUs directly in that industry-standard stack.”

Prior Architecture and Availability

TorchTPU was publicly introduced on April 7, 2026, in a Google Developers Blog post describing an “Eager First” design built on PyTorch’s PrivateUse1 interface, with Debug, Strict, and Fused eager modes backed by a shared compilation cache. Google said at the time that Fused Eager delivers a 50% to 100+% performance increase over Strict Eager, and that torch.compile on TorchTPU captures FX graphs with Torch Dynamo and lowers them through XLA and StableHLO, with custom kernels written in Pallas and JAX. The post said the stack supports Distributed Data Parallel, FSDPv2, and DTensor out of the box and is architected to handle divergent execution across ranks. It also listed the launch of a public GitHub repository among TorchTPU’s 2026 roadmap items.

The torch_tpu repository is now public under the google-pytorch GitHub organization and licensed under Apache 2.0. Its README lists custom ATen kernels, a compilation cache for XLA computations, distributed training support for data and tensor parallelism, torch.compile integration, debugging and benchmarking utilities, and examples for Llama, Qwen, ResNet, DLRM, minGPT, and UNet. The project supports Python 3.11 and later, and the torch-tpu 0.2.0 package was released on PyPI on October 19, 2026, with Linux x86-64 wheels for CPython 3.11 through 3.14. A user guide is published at torchtpu.readthedocs.io, and the README notes that the project’s interactive notebooks must run from a TPU VM, with a v6e GCE instance the easiest option.

In the same announcement, Google said LiteRT, formerly TensorFlow Lite, is joining the PyTorch ecosystem: the LiteRT-Torch library converts a PyTorch model with a single litert-torch.convert() call, and converted models run on Android, iOS, macOS, Windows, Linux, web, and IoT devices. A TPU agent suite intended to automate onboarding, diagnostics, and inference performance tuning for PyTorch models is in early preview on GitHub.

The announcement directs developers to Google Cloud’s TPU documentation to begin running PyTorch on TPUs.