TorchTPU: Running PyTorch Natively on TPUs at Google Scale

Google unveiled TorchTPU, introducing an integrated stack that runs PyTorch workloads natively on TPU.

The core goals are usability, portability, and performance. It aims for minimal changes—down to just switching initialization to tpu—from existing PyTorch scripts, while also targeting performance suited to TPU's large-scale distributed infrastructure.

It adopts an Eager First approach, prioritizing developer experience.

• Debug Eager: Synchronizes with the CPU on every operation, suitable for debugging
• Strict Eager: Runs asynchronous single operations, behaving close to the default PyTorch experience
• Fused Eager: Automatically bundles operation flows into larger chunks for execution, claiming a 50%~100%+ performance improvement over Strict Eager

The compilation path also supports full-graph optimization based on torch.compile. After PyTorch Dynamo captures the FX graph, it uses XLA as the backend instead of Torch Inductor, and operations are lowered to StableHLO. Custom operators support Pallas and JAX kernels, with Helion support in progress.

On the distributed training side, it states native support for DDP, FSDPv2, and DTensor. Whereas existing PyTorch/XLA excelled at pure SPMD, TorchTPU is designed to also handle MPMD execution—where code differs slightly per rank—accommodating asymmetric tasks such as logging.

On the hardware side, it noted that model architecture also needs adjustment to leverage TPU efficiency. For example, it explained that an attention head dimension of 128 or 256 may be more advantageous for TPU tensor core efficiency than 64.

A 2026 roadmap was also presented.

• Release of a public GitHub repository, documentation, and architecture tutorials
• Helion DSL integration
• Enhanced dynamic shape support for torch.compile
• Multi-queue support
• Ecosystem integration with vLLM, TorchTitan, and others
• Validation of linear scaling up to Pod-scale infrastructure