Optimizing Muse Spark Open on TPU v7 with SGLang-TorchTPU
Achieving high-throughput inference for 816B parameter models using custom megakernel pipelining on TPU v7 hardware.
1Optimizing Large Model Inference
Serving Meta's 816B-parameter Muse Spark Open model requires extreme efficiency on TPU v7 hardware. By utilizing SGLang with the TorchTPU backend, developers implemented a megakernel architecture. This approach consolidates 62 transformer layers into a single Pallas program, ensuring that weight loads and data transfers are continuously pipelined behind compute operations.
2Megakernel Architecture Benefits
The megakernel significantly reduces overhead by eliminating boundaries between layers that previously stalled weight streams from high-bandwidth memory. For single-request workloads, this design halves the time per output token compared to native paths. The performance gain is most pronounced at low concurrency, effectively optimizing the model for interactive agentic coding tasks.
3Pipelining Data Movement
Data movement is the primary bottleneck in large-scale inference. The megakernel schedules DMA engines to fetch weights and KV cache data just before the compute units require them. By prefetching across layer boundaries and using a ring of tiles for dense weights, the system sustains high throughput despite the massive parameter count.
4Collective Folding Strategy
Collective folding optimizes tensor parallelism by allowing matvec operations to begin as soon as partial results arrive from other devices. Instead of waiting for a full all-gather operation, the kernel processes slices incrementally. This technique effectively hides the latency of inter-device communication, which is critical for maintaining high performance in MoE models.
5Removing Attention Bottlenecks
Traditional online softmax implementations create serial dependencies that prevent pipelining. By leveraging the QK-norm properties of Muse Spark, the kernel eliminates the need for a running maximum score. This allows attention chunks to be processed independently, enabling full overlap of DMA transfers and compute stages during long-context inference sequences.
6Hardware and Memory Layout
Each TPU v7 host contains four chips, with each chip featuring two TensorCores and 192 GiB of HBM. The system uses a sequence-on-lane KV layout to maximize memory efficiency. By utilizing L2 HiCache on the host's DRAM, the architecture provides a secondary cache tier that further improves session-aware data eviction and retrieval.
7Speculative Decoding Integration
The system employs DFlash2 speculative decoding to accelerate token generation. A smaller draft model predicts sequences, which the full 816B model verifies in a single forward pass. This verify step is highly optimized, running within 1.5–1.7 times the time required to simply stream model weights from HBM to the matrix units.
8Performance Benchmarking Results
Testing on the AgentX 1M-context benchmark demonstrates significant throughput capabilities. At low interactivity, the system achieves 433 output tokens per second per request. As concurrency increases, the throughput scales to 25,906 tokens per second per chip. All tested points finish requests within 47 seconds at the p90 latency threshold.
9Cost Efficiency Comparison
The inference setup is evaluated against GB200 GPU baselines using external owning costs. With TPU v7 chip-hours priced at $1.21 and GB200 GPU-hours at $1.86, the optimized TPU implementation provides a competitive cost-per-token profile. These figures reflect the total operational expenditure for running large-scale mixture-of-experts models in production environments.
10Model Architecture Details
Muse Spark Open features 62 decoder layers, all utilizing mixture-of-experts routing. Each token is routed to 8 of 256 experts. The model employs hybrid attention: 46 layers use sliding-window attention over 2,048 tokens, while 16 layers attend to the full 1M-token context, balancing local efficiency with long-range dependency requirements.
11Expert Weight Quantization
To manage the massive parameter footprint, the model uses NVFP4 expert weights. This microscaling format applies one scale per 16 values, significantly reducing memory bandwidth pressure. By combining this quantization with FP8 compute, the system maximizes the utilization of the TPU's matrix units, which peak at 4,614 TFLOPs in FP8 mode.
12Future Scaling Implications
The success of the megakernel approach on TPU v7 suggests a shift toward more integrated, whole-step kernel design for large language models. By minimizing kernel boundaries and maximizing hardware-level pipelining, developers can extract significantly more performance from existing silicon. This methodology provides a blueprint for scaling future MoE models on distributed hardware.
Original article: PyTorch Blog — Optimizing Muse Spark Open on TPU v7 with SGLang-TorchTPU