AI Industry
Google Open-Sources TorchTPU to Run PyTorch Natively on TPUs
Oct 21, 2026·5 min read·By Rahish Saifi
Google on October 20, 2026 open-sourced TorchTPU, a native PyTorch backend for its Tensor Processing Unit (TPU) chips, at PyTorch Conference 2026. The stack was co-developed with Meta's PyTorch team and arrives one day after Amazon unveiled a similar path for Trainium. Outlets including TechCrunch framed the back-to-back releases as cloud rivals offering PyTorch developers alternatives to defaulting on Nvidia GPUs—without claiming an overnight shift away from CUDA.
Key facts
| Item | Detail |
|---|---|
| Announcement | Google Cloud and PyTorch Foundation posts — Oct. 20, 2026, PyTorch Conference 2026 |
| What shipped | TorchTPU open source on GitHub; TPUs join GPUs and CPUs as a native PyTorch device option |
| Porting goal | Change model.to("cuda") to model.to("tpu"); distributed backend tpu_dist; torch.compile(model, backend="tpu") for production |
| Ecosystem | Integrations with vLLM, SGLang, llm-d, TorchTitan, Ray, TorchRec, Hugging Face, LiteRT |
| Repos | TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan (Google / PyTorch Foundation) |
| Docs | torchtpu.readthedocs.io |
What Google announced
In a Google Cloud post titled "TorchTPU is now open source: Run native PyTorch on TPUs," executives Mark Lohmeyer and Bill Jia said Google is releasing TorchTPU so teams can run existing PyTorch workloads with minimal code changes, debug in eager mode, then move to torch.compile for production. Google positions TPUs alongside GPUs and CPUs as a first-class PyTorch target—not a separate JAX-only stack.
The PyTorch Foundation's companion post, "Introducing TorchTPU," traces the effort to mid-2025 at Google with Meta's PyTorch engineers. Unlike PyTorch/XLA (2020, built around Lazy Tensors), TorchTPU has no JAX dependency in the runtime and treats eager execution as a priority—addressing a long-standing complaint that TPU workflows felt indirect for PyTorch-first teams.
Google also previewed a TPU agent suite on GitHub to help automate onboarding, diagnostics, and inference tuning via coding-agent plug-ins, and highlighted LiteRT (formerly TensorFlow Lite) joining the PyTorch Ecosystem with litert-torch.convert() for deploying converted models across mobile, desktop, web, and IoT.
How it works
The porting story Google and the PyTorch Foundation emphasize is deliberately simple: swap the device string and keep the rest of your model code familiar.
model = MyModel()
model.to("tpu")
# distributed: backend="tpu_dist"
# production: torch.compile(model, backend="tpu")
Under the hood, TorchTPU supports roughly 450 ATen ops, according to the PyTorch Foundation post. Developers can stay in strict eager mode, use a new DeferAndFuse eager path that the PyTorch Foundation post says can be 50% or more faster than strict eager on some workloads, or compile with backend="tpu". Custom kernels can be written in Pallas or with the Helion DSL.
For inference, Google's developers blog notes that vLLM and SGLang previously reached TPUs through a JAX backend; the new engines sit on TorchTPU with Pallas kernels so teams can debug in PyTorch. vLLM documents installation via pip install vllm-torchtpu.
Google vs Amazon: two PyTorch escape routes from CUDA in two days
On October 19, 2026, Amazon open-sourced TorchNeuron so PyTorch can target AWS Trainium2/3 with a device change rather than a rewrite—our coverage is in Amazon's TorchNeuron Brings PyTorch to Trainium, Eyes Nvidia. A day later, Google answered with TorchTPU for its own accelerators.
Both moves target the same tension: PyTorch is widely used in AI application development, yet much of the ecosystem still assumes CUDA and Nvidia GPUs. TechCrunch and SiliconANGLE describe the parallel launches as hyperscalers courting PyTorch developers while Nvidia remains the default for many production fleets. Google's pitch adds depth on partner integrations—Meta, Hugging Face, Inferact (vLLM), RadixArk (SGLang)—plus the TPU agent suite and LiteRT bridge.
Ecosystem: Hugging Face, vLLM, SGLang, llm-d, LiteRT
Google's announcement reads less like a lone runtime drop and more like a coordinated platform push:
- Hugging Face — In Transformers v5, deploy to TPUs with
device_map="tpu"; Diffusers runs natively on TPUs; nightly CI of the Transformers suite on TPUs. Jeff Boudier, VP Product at Hugging Face, said in Google's post: "Hardware diversity is essential to making open-source AI accessible to everyone." - vLLM — Backend built with Inferact on TorchTPU.
- SGLang — Backend built with RadixArk.
- llm-d 1.0 — Released with Red Hat, IBM, and 500+ contributors from 25 organizations (CNCF Sandbox), routing across GPUs and TPUs on GKE.
- TorchTitan, Ray, TorchRec — Training and serving paths validated on TorchTPU; native TorchTPU support in Ray; VeRL TPU backend for RL post-training in preview.
- LiteRT — Cross-platform deployment from PyTorch via
litert-torch.convert().
Jana van Greunen, Director of PyTorch Engineering at Meta, quoted in Google's post: "We cannot be constrained by a single hardware architecture. TorchTPU is a massive step forward."
Google's numbers (attributed to Google's own tests)
Google and the PyTorch Foundation published several benchmarks. Treat them as vendor-reported unless independently reproduced.
| Claim | Source (Google / PyTorch Foundation) |
|---|---|
| 98% of 4,917 tested Hugging Face models run without crashing; 97% numerical parity with GPU on architecture representatives | PyTorch Foundation blog |
| Qwen3 235B MoE on 8,192 TPU v7x chips at 28% MFU; per-TensorCore throughput nearly flat from one to eight slices | PyTorch Foundation blog |
| On 64 TPU v7 chips: 1,828 TPS PyTorch/TorchTitan vs 1,937 TPS JAX/MaxText | PyTorch Foundation blog |
| TorchTitan validated up to 8,000 Ironwood TPUs (TPU v7x); up to 52% MFU on Qwen 3 8B with near-linear scaling to 8,192 Ironwood TPUs | Google Cloud blog |
| TPU v7 + vLLM-TorchTPU lowest cost per million tokens at 2.01 cents on 8K/1K (citing SemiAnalysis InferenceX) | PyTorch Foundation blog (Google claim citing SemiAnalysis) |
These figures illustrate Google's argument that native PyTorch can compete on throughput and cost on Ironwood-class pods—but they are not neutral third-party audits of every workload.
What it can't do yet
The PyTorch Foundation post is explicit about gaps worth planning around:
- FlexAttention is not supported yet.
- Dynamic shapes (bounded dynamism) remain in development;
torch.compilesetsdynamic=Falseon TPU for now. - RL colocation is not supported out of the box.
- FP8/FP4 via torchao is on the roadmap.
Peak performance may still require tuning beyond the one-line device swap—the PyTorch Foundation post notes that models should run out of the box, but peak TPU performance needs more than that one-line change.
How to try it
Google points developers to the open-source repositories (TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, TorchTitan) and documentation at torchtpu.readthedocs.io. For Hugging Face users on Transformers v5, Google highlights device_map="tpu". For serving, vLLM's TorchTPU package is installable via pip as documented on the vLLM blog.
For context on the wider hyperscaler custom-silicon race, see our write-up on Microsoft's Maia 300 roadmap and its multi-gigawatt AI compute plans.
Sources
- Google Cloud: TorchTPU is now open source: Run native PyTorch on TPUs — Oct. 20, 2026
- PyTorch Foundation: Introducing TorchTPU — Oct. 20, 2026
- Google Developers Blog: Announcing New vLLM and SGLang Inference Engines on TPU — Oct. 20, 2026
- vLLM Blog: Introducing vLLM TorchTPU — Oct. 20, 2026
- TechCrunch: Google has joined Amazon in attacking Nvidia's stronghold with AI app developers — Oct. 20, 2026
- SiliconANGLE: Google unveils TorchTPU in a push to break Nvidia's stranglehold on AI compute — Oct. 20, 2026