Introducing TorchTPU

TL;DR TorchTPU is an open source collaboration between Google and Meta that introduces a native and principled PyTorch backend for Google TPUs. It enables seamless out-of-the-box model execution and TPU migration while delivering industry-competitive inference and training performance through eager execution optimizations, torch.compile integration, and native hardware acceleration. Introducing TorchTPU Google TPUs have been used very successfully, both within and outside of Google, to train and deploy large ML models at massive scale. They are specifically designed for ML and, as such, they provide significant benefits in terms of performance, energy, and cost. While TPUs have been historically associated with the TensorFlow and the JAX ecosystems, we’re excited to introduce TorchTPU , a new effort started at Google in mid 2025, in collaboration with the PyTorch Team at Meta, to run PyTorch natively on TPUs. To program Google’s TPU at a level below the Python framework, developers utilize the XLA compiler stack. In 2020, Google introduced the first compiled version of PyTorch for TPU, PyTorch/XLA, based on Lazy Tensors. While foundational, this initial integration struggled to keep pace with the rapid, community-driven evolution of PyTorch. Crucially, PyTorch/XLA required intrusive source-code modifications to explicitly define compilable regions, suffered from severe compile-time latency, lacked support for PyTorch’s native imperative (“eager”) execution mode, and lagged in supporting emerging distributed computing primitives within the PyTorch ecosystem. In contrast with previous efforts such as PyTorch/XLA and TorchAX, which shoehorn PyTorch into JAX’s modus operandi, TorchTPU was built as a principled TPU backend for PyTorch: eager execution is a top priority, and existing PyTorch workloads run out of the box with trivial code changes (e.g., .to("cuda") → .to("tpu") ). More specifically, as we started TorchTPU we deliberately formulated a set of principles to follow: PyTorch Citizenship : No JAX dependencies in the runtime; Solid eager mode; Torch.compile() based on Torch Dynamo (not Lazy Tensor) and with similar performance to JAX compile; Wide support for PyTorch distributed idioms (DDP, FSDP, DTensor, ...) by providing TPU collectives at the lowest level in the PyTorch stack (e.g. Process Group API); Seamless integration into the PyTorch ecosystem; Community engagement. Focus on the User: Porting a workload to TPU should be as simple as changing model.to("cuda") to model.to("tpu") . (One cannot achieve peak TPU performance with that simple change, since a model’s shapes, sharding and ops may be skewed toward a GPU architecture; however, the model should run out of the box.) Provide user-friendly and actionable messages for runtime/compiler errors; Integrate with the tools and frameworks that the PyTorch community already uses (e.g., torch.profiler, DTensor, Ray). Improve the State of the Art : Build a solution that can compete with PyTorch on GPUs for large scale workloads; Introduce a JIT compilation approach for eager execution, with runtime optimizations that can bridge the performance gap between op-by-op execution and compiled execution; Improve the TPU programming stack by introducing bounded dynamism, by avoiding XLA recompilation due to shape changes, and by exploring dynamic TPU kernels that can be precompiled ahead of time. TorchTPU is in active development, it is used in production for both inference and training workloads for some of the largest Google Cloud customers, and can be found as an Open Source project in this GitHub repo . At the time of this writing, we have: Added support for both eager and compiled modes, including a new DeferAndFuse eager mode; Added support for ~450 Aten ops (skipping ops with structural or composite delegates), including forward and backward pass for all Core ATen ops required for the major deep learning architectures; Added ability to author custom kernels via JAX/Pallas and via Helion DSL; Validated built-in compatibility with Hugging Face’s Transformers v5, Diffusers, and Kernels libraries, with out-of-box execution on the 5,000 most downloaded Hugging Face models, including: LLMs: Qwen, DeepSeek, Gemma, LLama, Mistral, GPT, OPT, BLOOM, Falcon, StableML Audio/Speech: Whisper, Wav2Vec2, HuBERT, MMS Vision: DinoV2, SigLIP, CLIP, Swin Transformer, ConvNeXt, ViT, ResNet, DETR Text Encoders & Search: BERT, ALBERT & Electra, RoBERTa, DeBERTa, T5, MT5, Flan-T5 Validated DTensor-based sharding; Validated multi-host inference (vLLM and SGLang) and fine-tuning / training workloads (TorchTitan), with Ray-based orchestration; Validated integration with vLLM, SGLang, TorchTitan, Ray, and TorchRec. In the remainder of this document we discuss key aspects of the TorchTPU architecture and some initial results. Architecture Deep Dive A TPU Primer Google TPUs are specialized ASICs designed specifically to perform the massive tensor operations that are the foundation of deep learning. From a HW architecture perspective , a TPU chip contains (depending on the generation): One or more TensorCores , each comprising one or more matrix-multiply units (MXUs), a vector unit (VPU), and a scalar unit (see this link ). Based on the TPU generation, an MXU can include 256×256 multiply-accumulators, which are arranged as a systolic array . One or more SparseCores , which are specialized units for sparse operations (e.g., gather, scatter, embeddings) and for collectives. They can run in parallel with a TensorCore, enabling compute-compute and compute-communication overlaps. A large High Bandwidth Memory (HBM) – 288 GB at 8.6 TB/s in TPU 8i , which can be used to store ML data and share that with the CPU. Bidirectional Inter-Chip Interconnect (ICI) – 19.2 TBps in TPU 8t and 8i , allowing fast communication between TPU chips within user-configured TPU Slices , which are logical subgroups of a TPU Pod , comprising of a preconfigured number of TPU chips connected in a 3D torus (for TPU 7, TPUv 8t) or Boardfly (TPU 8i) configuration. Bidirectional Data-Center-Network (DCN) – 400 GBps in TPU 8t , allowing communication between TPU chips across Slices. TPUs are programmed with the XLA compiler stack through two different entry points: StableHLO (higher-level IR for ML ops) and Mosaic (lower-level IR for custom kernels). The actual TPU ISA is proprietary, subject to change, and not publicly available. Effectively, one can think of StableHLO and (even more so) of Mosaic as equivalent to PTX. See picture below for a high-level sketch. --> A C++ program can access a local TPU by instantiating a PjRT Client , whose API provides methods to compile MLIR (StableHLO or Mosaic) into a TPU executable, to launch that executable, and to copy data to-from the TPU. Note that TPU execution is asynchronous and the CPU can wait for a TPU result to be ready by waiting on the associated PjRTBuffer. Notable features of the TPU programming model: A TensorCore device can be thought of as a single-threaded processor with one or more massive matrix-multiply units (MXU). Consequently, to reach peak performance it is key that a model’s shapes be sufficiently wide to exploit the full potential of the MXU. This differs from GPUs where launching multiple smaller kernels in parallel can help maximize HW utilization. The XLA compiler can instruct a TensorCore to offload sparse and collective operations to SparseCore, for more efficient and/or parallel execution. This happens either automatically or through user-supplied MLIR annotation in a compiled StableHLO module. A TPU device can be accessed by only one user-space process. Following our goal of seamless integration in the PyTorch ecosystem, for distributed TorchTPU workloads we adopt a one-process-per-device model in contrast with JAX’s one-process-per-host model. However, because a TPU can be accessed by only one user-space process, RL colocation can’t be supported out of the box and we are investigating ways