TorchTPU is now open source: Run native PyTorch on TPUs

TorchTPU is an open-source native PyTorch stack that brings Google TPUs into the standard PyTorch ecosystem, enabling developers to run, debug, optimize, and deploy PyTorch workloads on TPUs with minimal code changes. Co-developed with Meta’s PyTorch team, it supports eager-mode debugging and production optimization via torch.compile, exposes standard PyTorch APIs, advanced diagnostics, and native distributed primitives, and integrates directly with popular libraries and tools like Hugging Face Transformers (device_map="tpu"), Diffusers, vLLM, SGLang, TorchTitan, Ray, and TorchRec. Hugging Face and Meta collaborations include CI/CD and upstream support so many Hub models and pipelines run natively on TorchTPU with minimal friction.

The stack targets both inference and large-scale training: vLLM’s native TPU backend (with Pallas kernels) provides continuous batching, quantization, and multi-chip parallelism for high-throughput serving; SGLang-TorchTPU uses RadixAttention and HiCache to reuse KV cache and cut recomputation in agentic, multi-step workloads; and llm-d 1.0 (a CNCF Sandbox project) supplies prefix-cache-aware routing, prefill/decode splitting, multi-tenant flow control, and GKE integration for TPU/GPU fleets. For training, TorchTitan delivers compile-aware distributed training recipes (Llama 3, Qwen, MuseGlimmer, MoE) and demonstrated near-linear scaling and up to ~52% model FLOPS utilization on Qwen 3 8B across thousands of Ironwood TPUs. Additional pieces include a VeRL TPU backend for RL, Ray orchestration for heterogeneous pipelines, LiteRT-Torch for on-device export, and a TPU agent suite in preview.