Google and Meta Announce Native PyTorch Support for TPUs with TorchTPU

Details

Google and Meta introduced TorchTPU, a native PyTorch backend for TPUs that removes JAX dependencies and allows execution via model.to("tpu"). The project plans to open-source TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan repositories at the October 2026 PyTorch Conference.
Validated on 4,917 of the top 5,000 Hugging Face models, achieving a 98% pass rate for popular models and 97% for numerical accuracy.
Performance benchmarks include 500 tokens/sec/chip with >50% MFU for Llama 3 405B on 1,024 TPU v7x chips.
Key optimizations include SimpleFSDP, Scan Higher-Order Op, and Pallas SplashAttention, which provide up to 3.80x speedup for FLUX.1-dev and 17.5% attention speedup.
Distributed training uses 2D+ DeviceMesh to assign Tensor/Expert Parallelism to high-speed ICI and Data Parallelism to slower DCN, with SparseCore offloading for All-Gather/Reduce-Scatter.
TorchRec now includes a native TPU backend for recommendation models, matching JAX baseline performance (18.3ms step latency) on DLRM v2 with 32-chip TPU v7x clusters.
Future roadmap includes FP8/FP4 quantization support for TPUv8, Ray Core orchestration for RL loops, and expanded support for Reasoning, Multimodal, and Diffusion models.

Key point