# 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. - **Source:** [PyTorch Blog](https://pytorch.org/blog) - **Published:** 2026-11-10 - **Topic:** [technology](https://brifing.si/c/technology) - **Original article:** - **On Brifing:** https://brifing.si/s/7c076358-244d-4e8e-9d8a-d9a1490b5af9 ## Optimizing 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. ## Megakernel 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. ## Pipelining 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. ## Collective 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. ## Removing 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. ## Hardware 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. ## Speculative 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. ## Performance 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. ## Cost 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. ## Model 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. ## Expert 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. ## Future 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. ## Sources covering this story - https://pytorch.org/blog/optimizing-muse-spark-open-on-tpu-v7-with-sglang-torchtpu/ — --- _Summary written by Brifing from the original reporting. Copyright in the original article belongs to its publisher; read the full piece at the link above._