TorchTPU: Google Bridges PyTorch and TPUs for Scalable AI Workloads

The frontier of machine learning has expanded beyond single accelerators to distributed systems spanning thousands of chips. As models scale to run on clusters of O(100,000) chips, the software powering these models must meet new demands for performance, hardware portability, and reliability. At the center of this challenge lies the integration between popular frameworks like PyTorch and specialized hardware like Google's Tensor Processing Units (TPUs).

Google has addressed this challenge head-on with TorchTPU, a new integration that enables PyTorch to run natively on TPU hardware at Google scale. The engineering team, led by Technical Lead Claudio Basile and including Kat Ko, Ben Wilson, Lee Howes, and others from Google's Core ML division, has built a stack that prioritizes usability, portability, and performance.

Understanding the TPU Architecture

To appreciate TorchTPU's significance, one must first understand the hardware it targets. A TPU system is not merely a chip but an integrated network. Each host connects to multiple chips, which in turn communicate via Google's Inter-Chip Interconnect (ICI), forming a highly efficient 2D or 3D Torus topology. This design enables massive scale-up without traditional networking bottlenecks.

Within each chip, execution is divided between specialized units:

• TensorCores: Single-threaded units optimized for dense matrix math
• SparseCores: Handle irregular memory access patterns like embeddings and gather/scatter operations

These specialized features make TPUs uniquely powerful for certain machine learning workloads, particularly those requiring massive matrix operations and efficient handling of sparse data patterns.

The TorchTPU Philosophy: PyTorch, But on TPUs

The core principle driving TorchTPU is simple: it should feel like PyTorch. Developers should be able to take existing PyTorch scripts, change their device initialization to "tpu", and run their training loops without modifying core logic. This approach dramatically lowers the barrier to entry for leveraging Google's TPU infrastructure.

"Our mandate was to build a stack that leads with usability, portability, and excellent performance," explains the team. "We wanted to enable developers to migrate existing PyTorch workloads with minimal code changes while giving them the APIs and tools to extract every ounce of compute from our hardware."

Engineering the TorchTPU Stack

The engineering team implemented an "Eager First" philosophy for TorchTPU, rather than requiring immediate static graph compilation. This approach integrates deeply with PyTorch using the framework's "PrivateUse1" interface, allowing developers to work with familiar PyTorch Tensors on TPUs without wrappers or subclasses.

Three Eager Execution Modes

TorchTPU offers three distinct eager modes to support different stages of the development lifecycle:

1. Debug Eager: Dispatches one operation at a time and synchronizes with the CPU after each execution. While inherently slow, this mode is invaluable for debugging shape mismatches, NaN values, and out-of-memory crashes.

2. Strict Eager: Maintains single-op dispatch but executes asynchronously, mirroring the default PyTorch experience. This allows simultaneous CPU and TPU execution until synchronization points in the user's script.

3. Fused Eager: The breakthrough mode that uses automated reflection on operation streams to fuse steps on the fly into larger, computationally dense chunks before handing them to the TPU. By maximizing TensorCore utilization and minimizing memory bandwidth overhead, Fused Eager delivers a 50% to 100+% performance increase over Strict Eager with no setup required from users.

All three modes leverage a shared Compilation Cache that can operate on a single host or be configured persistently across multi-host setups, reducing compilation time as TorchTPU learns from workloads.

Static Compilation with XLA

For users seeking peak performance, TorchTPU integrates natively with PyTorch's torch.compile interface for full-graph compilation. The stack captures the FX graph using Torch Dynamo but routes to XLA (Accelerated Linear Algebra) as the backend compiler rather than Torch Inductor.

This was a deliberate architectural decision. XLA is rigorously battle-tested for TPU topologies and natively understands how to optimize the critical overlap between dense computation and collective communications across the ICI. TorchTPU's translation layer maps PyTorch operators directly into StableHLO, XLA's primary Intermediate Representation for tensor math.

The team also ensures extensibility doesn't compromise performance. TorchTPU natively supports custom kernels written in Pallas and JAX through a @torch_tpu.pallas.custom_jax_kernel decorator, allowing engineers to write low-level hardware instructions that interface directly with the lowering path.

Distributed Training and MPMD Support

Preserving flexibility at scale was a key focus for TorchTPU. The framework supports PyTorch's distributed APIs including Distributed Data Parallel (DDP), Fully Sharded Data Parallel v2 (FSDPv2), and DTensor out of the box. Many third-party libraries that build on PyTorch's distributed APIs work unchanged on TorchTPU.

A significant improvement over PyTorch/XLA (TorchTPU's predecessor) is support for MPMD (Multiple Program Multiple Data) execution. While XLA traditionally optimized for pure SPMD (Single Program Multiple Data) code, real-world PyTorch applications often have slight divergence between ranks—such as rank 0 handling additional logging or analytics.

"TorchTPU is architected to carefully support divergent executions and will isolate communication primitives where necessary to preserve correctness, at minimal cost," explains the team. "This approach ensures that using PyTorch on TPU feels natural to existing developers while preserving XLA's ability to overlap communication and computation."

Hardware Awareness and Optimization

While portability is a priority, TorchTPU acknowledges that optimal model design may differ across hardware. For example, models often hardcode attention head dimensions to 64, while current-generation TPUs achieve peak matrix multiplication efficiency at dimensions of 128 or 256.

TorchTPU facilitates a tiered workflow: establish correct execution first, then use guidelines to identify and refactor suboptimal architectures or inject custom kernels for optimal hardware utilization. This balance between portability and hardware-specific optimization represents a pragmatic approach to the challenge of cross-platform AI development.

Roadmap for 2026

The TorchTPU team has laid a solid foundation but continues to address several challenges:

• Reducing recompilations triggered by dynamic sequence lengths and batch sizes through advanced bounded dynamism within XLA
• Building a comprehensive library of precompiled TPU kernels to reduce first-execution latency
• Planned initiatives include:
  • Launch of a public GitHub repository with extensive documentation and tutorials
  • Integration with PyTorch's Helion DSL to expand custom kernel capabilities
  • First-class support for dynamic shapes through torch.compile
  • Native multi-queue support for asynchronous codebases
  • Deep integrations with ecosystem pillars like vLLM and TorchTitan

Significance for the AI Ecosystem

TorchTPU represents more than just another hardware integration—it embodies a philosophical approach to AI infrastructure development. By prioritizing the developer experience while maintaining hardware-specific optimizations, Google has created a template for how specialized AI hardware can integrate with existing ecosystems.

As models continue to scale, the ability to efficiently leverage thousands of accelerators will become increasingly critical. TorchTPU's approach of "PyTorch, but on TPUs" lowers the barrier to entry while providing paths to optimal performance, potentially accelerating the development of next-generation AI systems.

The integration also signals Google's commitment to making its TPU infrastructure accessible to the broader AI community, complementing existing offerings like Google Cloud TPUs while providing a more native experience for PyTorch developers.

For developers interested in exploring TorchTPU, Google has indicated that a public GitHub repository with documentation and tutorials will be available later in 2026. The TPU Developer Hub will serve as the central resource for staying updated on TorchTPU developments.

The evolution of AI infrastructure will continue to present challenges in balancing developer productivity with hardware-specific optimizations. TorchTPU offers one approach to this balance, and its development will be closely watched by both the research and industrial AI communities.