Google and Meta Release TorchTPU, Letting PyTorch Run on 8,192 TPU Chips

Google and Meta have released TorchTPU, an open-source PyTorch backend for Google TPUs. It exposes TPUs through familiar interfaces such as .to("tpu"), eager execution, torch.compile, and PyTorch distributed APIs. The project began inside Google in mid-2025, and its source is available on GitHub.

A backend PyTorch can dispatch directly

PyTorch/XLA uses lazy tensors to collect operations before sending them to XLA for compilation. That architecture can require code changes around compilation boundaries, adds compilation latency, and diverges from the imperative execution model developers use on CPUs and GPUs. Support for newer PyTorch distributed primitives has also tended to arrive later.

TorchTPU registers a C++ backend through PyTorch’s PrivateUse1 extension slot and makes eager execution the default. Its standard runtime path removes the JAX dependency, while torch.compile captures programs through TorchDynamo rather than PyTorch/XLA’s lazy tensors. XLA and StableHLO remain part of the compilation stack beneath the PyTorch interface.

A basic migration requires two primary changes:

• Replace the cuda device with tpu.
• Replace the nccl distributed backend with tpu_dist.

The team reports testing compatibility against the 5,000 most-downloaded Hugging Face models, including Qwen, DeepSeek, Gemma, Llama, Mistral, Whisper, CLIP, ViT, and BERT-family models. Compatibility still depends on operator coverage, model shape, custom kernels, and distributed configuration, so teams should test their exact workloads.

Strict, debug, or fused

TorchTPU offers three eager execution modes through the same dispatcher:

• Strict Eager, the default, dispatches one Aten operation at a time while allowing the CPU and TPU to work asynchronously.
• Debug, enabled with TORCH_TPU_LAUNCH_BLOCKING=1, makes dispatch synchronous so stack traces point more directly to failing operations.
• DeferAndFuse, enabled with TORCH_TPU_DEFER_AND_FUSE=1, accumulates operations and compiles them in fused groups. Google reports gains of 50% or more over Strict Eager on some workloads, along with lower peak high-bandwidth memory use.

DeferAndFuse builds a graph until the program reaches a materialization point such as .to("cpu") or .item(). TorchTPU converts that graph to StableHLO, fingerprints it, and checks a three-level compilation cache: process-local storage, host-shared memory, and optional remote storage such as Google Cloud Storage. Repeated shapes can then reuse compiled programs. A repeated-operations heuristic detects patterns such as stacked transformer blocks and can trigger compilation before an explicit materialization point.

Precision follows the MXU

TorchTPU defaults matrix multiplication to medium precision, using BF16 multiplication with FP32 accumulation. BF16 preserves FP32’s exponent range while using fewer mantissa bits, which aligns with the native matrix units in Google TPUs. CUDA’s PyTorch default is highest, so code that relies on implicit precision settings can behave differently after migration.

TorchTPU emits a performance warning when code selects high or highest. Developers comparing GPU and TPU results should set precision explicitly, check numerical tolerances, and include the selected mode in performance reports.

Scaling across 8,192 chips

Google evaluated weak scaling by training the 235-billion-parameter Qwen3 mixture-of-experts model with TorchTitan on TPU v7 Ironwood. Weak scaling increases the global batch with the chip count, testing whether each chip maintains similar throughput as the cluster grows from one slice to eight.

Model FLOPs utilization, or MFU, measures how much of the hardware’s theoretical floating-point capacity performs model computation. It remained between 26.9% and 28.8% across the test. Fully sharded data parallelism, or FSDP, shards model state across devices. Hybrid sharded data parallelism, or HSDP, combines sharding within a slice with replication across slices. XLA’s scheduler overlaps HSDP all-reduce communication during the backward pass with computation for the next layer, limiting the throughput loss from cross-slice networking.

A separate result for Llama 3 405B reports more than 500 tokens per second per chip and over 50% MFU on 1,024 TPU v7x chips, using BF16 and an 8,000-token context window.

Serving gains hinge on kernels

Google has integrated TorchTPU with vLLM and SGLang for inference. A SemiAnalysis benchmark cited by the project reports a cost of 2.01 cents per million tokens on InferenceX 8K/1K, down from 3.6 cents in an earlier September result. On AgentX, the reported cost was 0.83 cents per million tokens, along with the benchmark’s leading interactivity result.

The serving improvements rely on Pallas kernels that fuse a mixture-of-experts layer into one single-program, multiple-data program. A sequence-on-lane layout also changes how the key-value cache stores attention state under eight-way tensor parallelism, cutting the reported memory waste by half. Actual cost and latency will vary with model architecture, sequence lengths, batching, quantization, utilization, and cloud pricing.

Three routes to custom kernels

Developers can write TPU-specific operations in Pallas, use prebuilt kernels from Tokamax, or author portable kernels in Helion. The Pallas route binds a JAX-based kernel to PyTorch through torch.tpu.pallas.jax_op:

@torch.tpu.pallas.jax_op("pallas::add_vectors")
def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
    out_shape = jax.ShapeDtypeStruct(x.shape, x.dtype)
    return pl.pallas_call(add_vectors_kernel, out_shape)(x, y)

x = torch.tensor([1.0, 2.0, 3.0], device="tpu")
y = torch.tensor([4.0, 5.0, 6.0], device="tpu")
z = add_vectors(x, y)

Pallas remains the direct path for TPU-tuned kernels even though ordinary TorchTPU execution no longer requires JAX. Tokamax provides kernels for operations including grouped matrix multiplication for mixture-of-experts models, paged attention, Kimi Delta Attention, and multi-head latent attention. Helion provides a PyTorch accelerator-agnostic domain-specific language for kernels intended to run across hardware backends.

Limits to budget for

• Narrow or irregular workloads: TPU matrix units can be underused by small or awkwardly shaped operations that distribute efficiently across many GPU streaming multiprocessors.
• Collective performance: Distributed collectives default to the safer Skinny mode. The project reports an end-to-end penalty of 30% to 40% in some workloads until developers enable @spmd_safe or use the compilation handshake.
• Process model: A TPU device accepts one user-space process, so reinforcement-learning configurations that colocate several processes on one device require a different deployment design.
• Backend maturity: Operator coverage, debugging behavior, precision, and distributed scaling should be validated against the current release before committing production workloads.

A practical migration sequence

1. Check the current operator and model support in the project documentation.
2. Change device placement from cuda to tpu and distributed communication from nccl to tpu_dist.
3. Begin with Strict Eager, then use blocking debug mode to isolate unsupported or failing operations.
4. Benchmark DeferAndFuse and torch.compile with representative shapes, batch sizes, and sequence lengths.
5. Set matrix-multiplication precision explicitly and compare numerical results with the existing implementation.
6. Profile collectives, memory use, compilation-cache hit rates, and custom-kernel requirements at the intended cluster size.

TorchTPU gives PyTorch teams a direct way to evaluate Google TPUs while retaining familiar model code, FSDP, TorchDynamo, and serving frameworks. Google has published the TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan repositories, making the backend, distributed stack, inference integrations, and training examples available for inspection. Its value will depend on workload-specific compatibility and performance, but accelerator selection no longer requires a wholesale move away from PyTorch.