How SGLang-TorchTPU Accelerates an 816B MoE Model on TPU v7

Introduction

Large-model inference is not limited by matrix multiplication. During low-batch decoding, a TPU may finish the available arithmetic quickly while waiting for weights to arrive from HBM, for KV data to be read, or for a collective communication step to complete. Separate kernels make the problem worse: every boundary can interrupt the flow of weights and prevent memory transfers from overlapping with computation.

A PyTorch engineering report describes how SGLang-TorchTPU addresses this problem while serving Meta’s Muse Spark Open. The system uses SGLang for scheduling, KV-cache management, and model serving, with TorchTPU as its PyTorch backend and Pallas kernels on the hot paths. The deployment uses one TPU v7 host with four chips, exposed as eight devices and sharded eight ways with tensor parallelism.

Muse Spark Open is an 816-billion-parameter mixture-of-experts model. Around 42 billion parameters are active for each token. All 62 decoder layers are MoE layers, and each layer routes a token to eight of 256 experts. The architecture mixes sliding-window attention over the most recent 2,048 tokens with full-context attention that can reach roughly one million tokens. That combination puts pressure on both weight movement and KV-cache capacity.

The main engineering choices

• A whole-step megakernel. Instead of launching several kernels and collectives for each layer, the new path runs all 62 layers and the final normalization inside one Pallas program per device. SGLang still performs surrounding operations such as embedding, drafting, the language-model head, and sampling, but the main decode or verify step becomes one continuous execution region.
• Pipelined movement. The kernel schedules dense and expert weights, KV reads, partial sums, and cross-device transfers ahead of the computation that consumes them. DMA engines can move data while the matrix units and vector units continue working. The objective is not merely to make each transfer faster, but to keep transfers from appearing on the critical path.
• Decode-specific optimizations. Collective folding reduces communication boundaries. A max-free attention loop pipelines KV chunks, while the expert loop is designed to keep register spills away from its common path. DFlash2 speculative decoding adds another layer: a draft model proposes several tokens, and the full model verifies them in one pass.
• More effective KV capacity. The megakernel uses a sequence-on-lane KV layout. TPU L2 HiCache adds host DRAM as another cache tier, with session-aware eviction intended to preserve useful data for long-running or prefix-sharing sessions.
• A faster prefill path. A fused MoE kernel and sequence parallelism target prompt processing, which has different bottlenecks from token-by-token decoding.

What the measurements show

On an earlier Muse Spark checkpoint, the DFlash2 verify path took 1.95 milliseconds per output token for one request with the megakernel, compared with 3.78 milliseconds for the native path. That is a 1.94x improvement. With two concurrent requests, the reported gain was 1.85x. The advantage narrows at higher concurrency because batching itself improves hardware utilization and makes kernel-boundary overhead less dominant.

The verify step also came within roughly 1.5 to 1.7 times the time required simply to stream the distinct weights for that step from HBM. This is an important signal: the implementation is approaching a bandwidth-limited regime rather than being dominated by per-layer launch and collective overhead.

The team also evaluated a one-million-token AgentX trace set based on agentic coding sessions. On one four-chip TPU v7 host, the run used FP8 compute and NVFP4 expert weights. Its performance curve ranged from highly interactive, lower-throughput operation to much higher aggregate throughput with a lower output-token rate per request. The published comparison uses external owning-cost assumptions and a specific GB200 baseline, so those numbers should be read as measurements of this configuration, not universal hardware rankings.

Why it matters

The broader lesson is that serving a large MoE model requires optimizing the schedule of the entire step. A faster matmul alone cannot solve a pipeline in which weights, KV blocks, and collective results repeatedly wait at kernel boundaries. By placing the 62-layer path inside one programmable kernel, SGLang-TorchTPU can decide when data should move, when it should be consumed, and how much work can be overlapped.

This approach is promising for sparse models and very long contexts, but it also raises the cost of engineering. Layouts must match TPU vector lanes, VMEM capacity must be managed carefully, and the kernel must remain compatible with paged KV caching and dynamic serving workloads. Whether the same strategy generalizes broadly will depend on the cooperation among the compiler, runtime, model implementation, and serving scheduler. The report nevertheless illustrates a useful direction: inference optimization is increasingly about eliminating idle time across the whole system, not just improving isolated operators.

Source: PyTorch Blog