SGLang moves its serving stack onto TPU: one engine for GPU and TPU, already in production at Meta
Synopsis
The SGLang team, with Google and Meta, released SGLang-TorchTPU, an open-source backend that runs the same SGLang serving engine on both GPUs and TPUs. TPU had not been a native PyTorch device, so a PyTorch-centric serving stack could not be reused directly; the new backend keeps the scheduler, RadixAttention and API server unchanged and places TPU work underneath. It is already running in production at Meta and serves as the rollout engine for reinforcement learning with TorchTitan through Miles.
Interpretation
SGLang's full serving stack now runs on Google Cloud TPU with the same engine used on GPUs, so deployments can pick hardware by workload and cost. TPU was not a native PyTorch device, so a serving stack built around PyTorch could not simply carry over torch.compile and torch.distributed. The backend is running in production at Meta, whose engineers contributed compilation and disaggregated-serving work and brought production workloads that shaped the engine from the start.
Each TPU device runs its own PyTorch process, so SGLang's scheduler, RadixAttention and API server run unchanged while TPU work sits underneath them. TorchTPU makes TPU a native PyTorch device that runs eagerly or compiles through XLA, and SGLang-TorchTPU builds on it with TPU-tuned Pallas kernels. The KV cache lives in a paged pool laid out for Pallas ragged paged attention with FP8 KV by default, Pallas kernels for attention and MoE grouped matmul are called as PyTorch custom ops inside the torch.compile graph, and everything ships as the sglang-tpu plugin loaded with --device tpu.
The same PyTorch stack covers training too: SGLang-TorchTPU generates rollouts, TorchTitan trains, and Miles syncs weights between them. Reinforcement learning training faced the same GPU-or-TPU choice, and the trainer side could not necessarily reuse the same stack. Qwen3-4B was trained with GRPO on DAPO-Math-17k on a single 4-chip TPU v7x host against 6 H200 GPUs with the same recipe and hyperparameters and two seeds each; reward and response-length curves matched within seed-to-seed variation and AIME-24 stayed within eval noise.
Perspective
The demonstrated scope is serving and single-host RL training on TPU v7x: CI covers 15 models on single-host and multi-host TPU v7x with 193 TPU tests registered, and Miles TPU training is in preview. It is aimed at engineers and researchers deploying or training models within PyTorch workflows. The release plans to keep adding models and optimizing supported ones, starting with DeepSeek V4.1 Flash and Reflection Beam.
Miles TPU training is a small-scale experiment demonstrating an end-to-end RL workflow on a single host, with early learning behavior similar to the GPU baseline; planned next steps include larger dense and MoE models, multi-host training, faster weight transfer over ICI, closing the step-time gap with GPU, parity checks on longer runs, and bitwise-identical log-probs between training and inference. Some newest models and features remain experimental until their tests pass.