Skip to content
audiotale.net
Sign in
NewsNews

Google open-sources TorchTPU, a native PyTorch backend for TPUs built with Meta

Existing PyTorch code can now target Google's TPUs as a normal device, with Hugging Face, vLLM and SGLang support, a day after Amazon's TorchNeuron.

audiotale.net team

· 7 min read

Written by AI from the sources linked below. How we use AI

On this page
  1. Key points
  2. What happened
  3. Background
  4. How it works
  5. What it means
  6. What we don't know yet
  7. What to watch next
  8. Sources

Google has released TorchTPU, a native PyTorch backend for its Tensor Processing Units (TPUs), as an open-source project on GitHub. Announced on 20 October 2026 and built together with the PyTorch team at Meta, it lets existing PyTorch code target a TPU as a normal device ("tpu" instead of "cuda"), without the detour through PyTorch/XLA or JAX that TPU users have had to take until now. It arrives one day after Amazon released its own native PyTorch backend for Trainium, and it is the clearest sign yet that the big cloud vendors are attacking Nvidia's lead at the software layer, not only the chip.

Key points

  • TorchTPU is a native PyTorch backend for Google TPUs, open source on GitHub, co-developed by Google and Meta's PyTorch team; the PyTorch blog says work began in mid-2025.
  • Porting is meant to be a device change: model.to("cuda") becomes model.to("tpu"), and the distributed backend nccl becomes tpu_dist. The authors note this gets a model running, not running at peak speed.
  • The PyTorch post says it has added about 450 ATen ops and validated out-of-the-box execution on the 5,000 most downloaded Hugging Face models.
  • Google says its new vLLM and SGLang engines for TPU are built on TorchTPU, and that SGLang-TorchTPU already runs in production at Meta.
  • Hugging Face reports that, of the Transformers tests that run on TPU, 96.5% pass.

What happened

On 20 October, at PyTorch Conference 2026, Google published TorchTPU's code and a set of announcements around it. The PyTorch blog, written by "The TorchTPU Team", describes it as "an open source collaboration between Google and Meta" providing a "native and principled PyTorch backend" for TPUs. The project is on GitHub as google-pytorch/torch_tpu; the README points to version 0.2.0 on PyPI and a user guide at torchtpu.readthedocs.io. Hugging Face's guide gives the install as pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu, where the extra index is there so pip fetches the small CPU build of torch rather than the large CUDA build.

Three groups of partners shipped alongside it. Hugging Face says Transformers, Diffusers and its Kernels library now run natively on TPU: a model loads with device_map="tpu", and a Diffusers pipeline moves with pipe.to("tpu"). Google's developer blog says it launched native PyTorch backends for the two best-known open-source inference servers, vLLM (built with Inferact, the company founded by vLLM's creators) and SGLang (built with RadixArk). TechCrunch adds that Google updated LiteRT, its runtime for phones, laptops and embedded devices, and launched a "TPU agent suite" to help developers migrate apps and fix errors. We could not open Google's own launch post for those two items, so they rest on TechCrunch's report.

Meta is a user, not just a collaborator. Zongwei Zhou, Meta's PyTorch-TPU lead, is quoted on Google's blog saying SGLang-TorchTPU "is running in production at Meta today".

Background

PyTorch grew up on Nvidia GPUs. Meta created it in 2016 and handed it to the Linux Foundation in 2022, and TechCrunch cites an estimate that it underpins about 85% of AI apps. That makes the framework itself part of Nvidia's moat: if your code and your tools assume CUDA, a cheaper chip is not much use.

TPUs have had the opposite problem. Google's own chips grew up with TensorFlow and JAX. The PyTorch blog explains that in 2020 Google introduced PyTorch/XLA, built on "lazy tensors". It required intrusive changes to mark compilable regions, suffered from long compile times, lacked proper eager execution and lagged behind PyTorch's distributed features. Hugging Face describes the same experience from the user side: lazy tensors meant code did not always run when you expected, and small changes could trigger recompilations.

The project itself is not new. Reuters reported in December 2025 that Google was working on a "TorchTPU" effort with Meta's help, and Google's developer blog described its design in April 2026. What changed on 20 October is the open-source release and the surrounding ecosystem work. Our coverage of the previous day's Amazon TorchNeuron release shows the parallel: two of the three big clouds now ship their own native PyTorch device.

How it works

According to the PyTorch post, TorchTPU registers as a PrivateUse1 device in PyTorch's dispatcher, so tensors on a TPU are ordinary PyTorch tensors, with no wrapper types. The design follows three stated principles: behave like PyTorch (no JAX dependency in the runtime, solid eager mode, torch.compile based on Dynamo), keep the user's job simple, and improve the state of the art.

It offers three eager modes:

  • Strict (the default) sends one operation at a time and runs it asynchronously, as PyTorch does on a GPU.
  • Debug (TORCH_TPU_LAUNCH_BLOCKING=1) is synchronous, slow and meant for tracking down shape errors and crashes, much like CUDA_LAUNCH_BLOCKING.
  • DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) is new: it defers operations, fuses them into larger chunks and compiles them, which the PyTorch post says can be 50% or more faster than Strict while using less memory. Google's April post put the gain at 50% to over 100%.

For peak performance, torch.compile captures the graph with Dynamo and then compiles through XLA via StableHLO rather than through Torch Inductor. Custom kernels can be written in Pallas or JAX, and Helion support is listed. A three-tier compilation cache (in-process, shared memory on one machine, and an optional persistent store such as Google Cloud Storage) avoids recompiling.

Two details matter to anyone porting a model. First, precision: TorchTPU defaults to a "medium" mode in which matrix multiplies take bfloat16 inputs with 32-bit accumulation. Asking for high or highest precision costs roughly three or six times the matrix-unit compute and triggers a warning. Second, shapes: the compiler wants shapes known at compile time, and Google says it is adding "bounded dynamism" to limit recompilation when they change. Google's April post adds that models often hard-code attention head dimensions of 64 while current TPUs reach peak efficiency at 128 or 256.

type: bar
title: TorchTPU: reported scale of validation
unit: count
source: PyTorch blog, "Introducing TorchTPU"; Hugging Face, "Hugging Face ❤️ TorchTPU"
Hugging Face models validated out of the box | 5000
ATen ops supported | 450
Transformers issues reported on TorchTPU | 34
Of those, already fixed | 19

What it means

For PyTorch developers, the pitch is that a TPU stops being a separate programming model. Hugging Face's examples show a Qwen model loading straight onto a TPU and Stable Diffusion 3.5 running with a one-line device change. Its own benchmark of a ragged-dot Mixture-of-Experts kernel on a TPU v6e-8 reports decoding 21.5 times faster than the default implementation on Qwen3-30B-A3B at batch size 64, and 12.8 times faster on Gemma 4 26B-A4B. That compares a tuned kernel with the default path, not TPU with GPU.

For inference teams, Google's developer blog says both new engines support prefix caching, speculative decoding, disaggregated prefill and decode, and tiered KV-cache offload, and cites a lowest cost per million tokens of 2.01 cents on the public SemiAnalysis InferenceX benchmark, for Qwen 3.5 on its Ironwood TPU. That is Google's figure on a benchmark it chose to highlight; we have not seen an independent comparison against Nvidia hardware. Google also says Inferact used a Pallas "megakernel" to serve the 2.8-trillion-parameter Kimi K3 model at more than 700 tokens per second for a single user.

For Nvidia, nothing changes overnight. TechCrunch puts it plainly: developers will not automatically drop GPUs just because moving is easier. But the strategy is visible: Google on 20 October, Amazon on 19 October, each making its chip a first-class PyTorch device and each recruiting Hugging Face, vLLM and SGLang to do so. Amazon also said it is upstreaming code to make PyTorch less tied to Nvidia hardware by default, according to TechCrunch.

What we don't know yet

  • Independent performance. The speed claims come from Google, Meta and partners. The PyTorch post says TorchTPU aims to compete with PyTorch on GPUs for large workloads; independent head-to-head results on the same models were not available to us.
  • Training maturity. Google says TorchTPU is used in production for inference and training by some of its largest cloud customers, but names none. Recreating tuned training recipes without code changes is a harder test than loading a model.
  • Gaps. The authors say TPUs allow only one process per device, so reinforcement-learning setups that colocate workloads do not work out of the box. Hugging Face says its TPU CI is "not all green yet".
  • Availability beyond Google Cloud. TorchTPU targets TPUs, which are rented through Google Cloud; this release does not change that.

What to watch next

  • Whether the nightly TPU results on the Transformers CI dashboard that Hugging Face says it is working towards actually appear, and how the 96.5% pass rate moves.
  • Independent InferenceX and MLPerf-style comparisons of vLLM and SGLang on TPU against GPUs.
  • Whether Amazon's and Google's upstream changes make other PyTorch backends (for AMD, Intel and others) easier to maintain.
  • Which named customers confirm moving PyTorch training to TPUs.

Sources

SaveStart the discussion

Written by

audiotale.net team

The official account of the people who run audiotale.net: announcements, conversation starters and the newsroom's drafts. Posts written with AI say so. Every other account here is a real person.

Read next