SGLang-TorchTPU: PyTorch-Native Production Serving on TPUs
SGLang-TorchTPU brings PyTorch-native production serving to TPUs, built with Google and Meta.
Extract
TL;DR SGLang-TorchTPU is open source today, built by the SGLang team at RadixArk in collaboration with Google and Meta. Launch models include Muse Spark Open, GLM-5.3 Flash, GLM-5.3, Kimi K3, DeepSeek V4, Ling 3.1 Flash, Qwen3.8, Gemma 4 and MiniMax H3. The plugin brings the full SGLang feature set to TPU, ready for agentic workloads. It is already proven in large-scale production at Meta. RL training on TPU works with SGLang-TorchTPU and TorchTitan in Miles. Contributors can apply for free TPU compute from RadixArk’s community TPU capacity pool. Introduction Today we are launching SGLang-TorchTPU, a native PyTorch backend that brings SGLang’s full serving stack to Google Cloud TPU. SGLang’s serving engine and model implementations are built on PyTorch, and with TorchTPU the same codebase serves on both GPUs and TPUs, so you can pick the hardware that best fits each workload and its deployment economics. TPU-specific components remain isolated in the plugin. New models and features added to SGLang could therefore become available on TPU through a single implementation. How it works Figure 1: How the sglang-tpu plugin fits into SGLang. SGLang-TorchTPU is an out-of-tree plugin. Installing the sglang-tpu package registers a TPU platform with SGLang through the sglang.srt.platforms entry point, and SGLANG_PLATFORM=tpu_oot selects it. The scheduler, RadixAttention and API server are SGLang’s own. The plugin fills in the parts SGLang already asks each hardware platform for: the attention backend, the KV cache pools, the MoE kernels and the graph runner. Each TPU core runs its own PyTorch process through TorchTPU, as each GPU does on CUDA. The model runs under torch.compile and lowers to XLA. There is no graph capture: instead, the runner pads each forward to a fixed set of token-count buckets for prefill and batch-size buckets for decode. Each bucket compiles once, and later steps dispatch to the cached program by shape. Attention and MoE run on Pallas kernels, registered as torch.library custom ops that Dynamo treats as opaque calls inside the compiled graph: ragged paged attention over a paged KV cache, and a grouped matmul over the sorted expert groups for MoE. When KV has to leave the device, for KV cache offloading to host memory or for prefill-decode disaggregation across hosts, it can move over TPU Sync, Google’s open-source TPU DMA library. Offload copies whole KV pages between HBM and a pinned host pool by direct PJRT DMA. For disaggregation the …
Explained
What happened
The SGLang team at RadixArk, working with Google and Meta, open-sourced SGLang-TorchTPU, a native PyTorch backend that runs SGLang's full serving stack on Google Cloud TPUs. It is already used in large-scale production at Meta.
- Out-of-tree plugin. Installing the
sglang-tpupackage registers a TPU platform through SGLang's platform entry point, selected withSGLANG_PLATFORM=tpu_oot. The scheduler, RadixAttention prefix cache and API server stay SGLang's own. The plugin supplies only what every hardware platform must: attention backend, KV cache pools, MoE kernels and the graph runner. - Compile once per shape bucket. Each TPU core runs its own PyTorch process. The model runs under
torch.compileand lowers to XLA. Instead of graph capture, each forward pass is padded to a fixed set of token-count buckets for prefill and batch-size buckets for decode. Each bucket compiles once and later steps reuse the cached program. - Custom kernels as opaque ops. Attention (ragged paged attention over a paged KV cache) and MoE (grouped matmul over sorted expert groups) run as Pallas kernels registered as custom ops that the compiler treats as black boxes.
- KV movement with TPU Sync. KV pages can be offloaded to pinned host memory by direct DMA. For prefill/decode disaggregation, the decode side pulls each chunk as soon as prefill commits it, while the next chunk computes, so only the last chunk is left to move at the end. Pages the decode side already has in its prefix cache are never sent, and KV heads are resharded in flight when the two sides use different tensor-parallel sizes.
- Coverage and testing. Launch models include GLM-5.3, Kimi K3, DeepSeek V4, Qwen3.8 and Gemma 4, with speculative decoding, FP8 KV cache, context parallelism and multimodal input. CI runs 193 TPU tests across 15 models on single-host and multi-host TPU v7x, and the support matrix is generated from those tests.
Why it works
Accelerators differ mostly in a few hot spots: attention, expert routing, memory layout and how graphs are compiled. Keeping everything else (scheduling, caching, APIs) shared and isolating the hardware-specific parts behind a narrow interface means new models and features reach TPU through one implementation.
Shape buckets solve XLA's main pain point. XLA compiles a program per input shape, and serving traffic has endless shapes. Padding to a small fixed set trades a little wasted compute for no recompiles at runtime.
The bigger pattern
This is the hardware abstraction layer pattern, the same idea as database storage engines or video codec backends behind one API. For LLM serving it means hardware choice becomes a cost decision rather than a rewrite, which adds pricing pressure on GPUs.
The chunked, pull-based KV transfer is the same overlap trick used in pipelined data transfers: start moving data before the producer finishes.
What it means in practice
- If you serve open models, TPU is now a realistic option with the same serving features; compare cost per token on your own traffic.
- When adding hardware backends to any system, keep the plugin surface small and make CI on real hardware the source of truth for what is "supported".
- On compiled runtimes (XLA, TensorRT), plan input shape buckets explicitly to avoid recompiles under load.
- For disaggregated serving, overlap KV transfer with compute and skip pages the receiver already caches.