Google and Meta launch TorchTPU, a native PyTorch backend for TPUs

A native backend instead of a JAX workaround

TPUs have long been tied to the TensorFlow and JAX ecosystems. TorchTPU, a project started at Google in mid 2025 together with the PyTorch Team at Meta, aims to change that by letting developers run PyTorch directly on the hardware.

The authors describe it as a principled TPU backend for PyTorch. They contrast it with earlier efforts such as PyTorch/XLA and TorchAX, which in their words "shoehorn PyTorch into JAX's modus operandi". PyTorch/XLA, introduced by Google in 2020 and based on Lazy Tensors, required intrusive source-code modifications to define compilable regions, suffered from severe compile-time latency, lacked support for eager execution and lagged in supporting new distributed primitives, according to the post.

With TorchTPU, eager execution is a top priority. Porting a model should be as simple as changing .to("cuda") to .to("tpu"). The authors add that this change alone will not deliver peak TPU performance, because a model's shapes, sharding and ops may be skewed toward a GPU architecture, but the model should run.

Design principles and current state

The team set out three groups of principles: PyTorch citizenship, a focus on the user, and improving the state of the art. In practice this means no JAX dependencies in the runtime, a solid eager mode, and torch.compile based on Torch Dynamo rather than Lazy Tensor, with performance similar to JAX compile. Distributed idioms such as DDP, FSDP and DTensor are supported by providing TPU collectives at the lowest level of the PyTorch stack.

According to the post, TorchTPU is in active development and is used in production for inference and training by some of the largest Google Cloud customers. It is available as an open source project on GitHub. At the time of writing, the team lists the following:

• Support for eager and compiled modes, including a new DeferAndFuse eager mode.
• Support for ~450 Aten ops, including forward and backward pass for all Core ATen ops required for the major deep learning architectures.
• Custom kernel authoring via JAX/Pallas and Helion DSL.
• Validated out-of-box execution on the 5,000 most downloaded Hugging Face models, spanning LLMs such as Qwen, DeepSeek, Gemma, LLama and Mistral, speech models such as Whisper, vision models such as CLIP and ViT, and text encoders such as BERT and T5.
• Validated DTensor-based sharding, multi-host inference with vLLM and SGLang, and training with TorchTitan, using Ray-based orchestration.

How TPU hardware shapes the design

The post also gives a primer on the hardware. TPUs are specialized ASICs with TensorCores built around matrix-multiply units, and some generations add SparseCores for sparse operations and collectives. The post cites 288 GB of High Bandwidth Memory at 8.6 TB/s in TPU 8i, and 19.2 TBps of inter-chip interconnect in TPU 8t and 8i.

These traits explain several design choices. A TensorCore behaves like a single-threaded processor with massive matrix units, so models need sufficiently wide shapes to reach peak performance. A TPU can be accessed by only one user-space process, so TorchTPU uses a one-process-per-device model, unlike JAX's one-process-per-host model. As a consequence, RL colocation cannot be supported out of the box, and the team says it is investigating workarounds.

The XLA compiler and PjRT runtime assume shapes are known at compile time, which is a limitation for PyTorch workloads. TorchTPU is adding bounded dynamism, where only the maximum bounds of a dimension are known at compile time. Collective ops are lowered to XLA as StableHLO ops rather than using a library such as NCCL. Because StableHLO does not allow in-place mutations, TorchTPU provides familiar PyTorch view semantics, including mutations, through view decomposition and inversion.

What it means for PyTorch users

For teams that already work in PyTorch, TorchTPU offers a path to try Google TPUs without rewriting code for JAX. The authors say the goal is a solution that can compete with PyTorch on GPUs for large scale workloads, and describe the work as ongoing. The source text provided here breaks off while discussing the autoload mechanism and does not include the benchmark results the authors mention, so performance claims remain the project's own until those figures are reviewed.