SGLang: Accelerating agentic and structured workloads
Developed in partnership with RadixArk, our native PyTorch backend for SGLang accelerates highly structured and agentic workloads on TPUs. SGLang uses RadixAttention to reuse KV cache across requests, allowing repeated prompt prefixes to skip recomputation entirely. SGLang-TorchTPU preserves upstream worker models, schedulers, and API servers, using a disaggregation transfer backend that also acts as the I/O backend for HiCache to optimize latency and throughput for complex multi-step generation.
“At RadixArk, TorchTPU has supercharged our development velocity as we roll out SGLang support across diverse silicon. It’s given us a painless path to native TPU acceleration, passing massive efficiency gains and flexible compute choices directly down to our customers.” – Banghua Zhu, CTO and co-founder, RadixArk
llm-d for distributed LLM inference
To run inference engines across large fleets, we are releasing llm-d 1.0 with Red Hat, IBM, and a community of more than 500 contributors and 25 organizations. As a Cloud Native Computing Foundation (CNCF) Sandbox project, llm-d connects two open-source communities: PyTorch, where vLLM and SGLang serve models, and Kubernetes, where distributed systems run and scale in production. The llm-d router sends each request to the replica that already holds its context, using prefix-cache-aware routing. It also splits compute-heavy prefill from memory-heavy decode. llm-d 1.0 makes that ready for production with multi-tenant flow control, router high availability, stable v1 APIs, and session-aware routing for agents. On Google Kubernetes Engine (GKE), you can have the choice of running on either GPUs or Google TPUs, with a common llm-d endpoint, and choose the hardware that fits each workload.
“At Zenteiq, we use GKE Inference Gateway powered by llm-d to accelerate synthetic data generation for our sovereign Scientific Foundation Models. By leveraging load-aware and prefix-cache-aware routing, we have achieved over 4× higher aggregate inference throughput, increasing daily generated token capacity from 1.47 billion to 5.97 billion tokens per 1,000 TPU v6e accelerators. The open-source foundation of llm-d, integrated with GKE, enables us to efficiently serve multiple model families at scale without developing and maintaining a custom inference scheduler.” – Prof. Sashikumaar Ganesan, Founder, Zenteiq.ai
Training at scale
Scaling PyTorch training across thousands of accelerators requires fault-tolerant, compile-aware distributed frameworks. We have integrated TorchTPU into the PyTorch community’s core training, reinforcement learning (RL), and orchestration frameworks.
TorchTitan: Distributed high performance training at scale
TorchTitan is the PyTorch-native reference implementation for scalable distributed training and features a production-grade integration with TorchTPU. By connecting to PyTorch’s torch.compile and TorchTitan’s compile friendly SimpleFSDP, specialized Pallas kernels, and graph-level optimizations, developers can build and train models on TPUs with high efficiency at scale.
TorchTitan provides recipes for pre-training large dense models (such as Llama 3 70B/405B, Qwen 3 8B/32B, and MuseGlimmer 30B), Mixture-of-Experts (MoE) such as Qwen 3 235B, and diffusion transformers (such as FLUX.1 12B), and has been validated running with up to 8000 Ironwood TPUs (TPU v7x). Beyond initial pre-training, TorchTitan serves as the scalable trainer inside post-training pipelines like VeRL for Group Relative Policy Optimization (GRPO) and RL alignment.To test large-scale efficiency on demanding open architectures, TorchTPU’s training paths via TorchTitan demonstrated up to 52% Model Flops Utilization (MFU) on Qwen 3 8B, showing near-linear scaling and stability across up to 8,192 Ironwood TPUs.
VeRL: Scalable post-training and reinforcement learning in preview
RL is a foundational technique for frontier labs that are advancing model capabilities, as well as for enterprises customizing models for domain-specific tasks. Yet, orchestrating distributed RL workloads remains one of the most complex challenges in modern AI infrastructure. Bridging that gap, VeRL has quickly become one of the most popular RL frameworks, supporting everything from verifiable math and coding rewards, to multi-turn agentic tool use at massive scale.
To help expand hardware choice for this thriving community, we upstreamed a TPU backend and an end-to-end reference recipe to VeRL. Built on a disaggregated trainer-sampler architecture designed to maximize training efficiency, our GRPO reference recipe on mathematical reasoning datasets demonstrates how developers can deploy VeRL workloads on TPUs while preserving their existing, hardware-agnostic model and algorithm code.
Ray: Orchestrating heterogenous compute
We’ve also added native TorchTPU support to Ray, enabling you to orchestrate PyTorch data processing, training, RL, and serving across CPUs, GPUs, and TPUs within a single workflow. With a few lines of resource configuration, you can schedule each stage of a heterogeneous AI pipeline onto the accelerator best suited for the job.
“At Ray Summit we talked about the PARK stack, PyTorch, AI models, Ray, and Kubernetes, because that’s what teams are building production AI on today. TorchTPU embeds Google TPUs directly in that industry-standard stack. We’ve also worked closely with the Google team to support TPUs as a first-class accelerator in Ray. Coupled with native PyTorch support, teams can now seamlessly develop, scale, and monitor their workloads on TPUs.” – Edward Oakes, Chief Architect, Anyscale
Deployment flexibility and operational efficiency
True portability requires standardizing deployment tools across different hardware profiles, including edge devices and cloud accelerators.
LiteRT joins the PyTorch Ecosystem
Many PyTorch models need to run on phones, laptops, browsers, and embedded devices. Today, LiteRT (formerly TensorFlow Lite), Google’s high-performance on-device AI inference, joins the PyTorch ecosystem.
With LiteRT-Torch, you convert a PyTorch model with a single litert-torch.convert() call. It builds on torch.export() and Core ATen operators, so your model code stays in PyTorch. For LLMs, LiteRT Torch exports Hugging Face Transformer models to run with LiteRT-LM. The converted models run on Android, iOS, macOS, Windows, Linux, web, and IoT devices, with hardware acceleration across CPUs (via XNNPack), GPUs (via open source ML Drift), and NPUs from major chipsets provided through a single API.
TPU agent suite in early preview
To expand Pytorch adoption on TPUs, we are announcing that the TPU agent suite is in preview on GitHub. Integrating through plug-ins with several coding harness-model combinations, the TPU agent suite helps automate onboarding, diagnostics, user questions, and inference performance tuning for many PyTorch models.
