The Pulse
RadixArk team serves Meta’s Muse Spark Open on one TPU host
PyTorch's SGLang-TorchTPU server runs Meta’s 816B-parameter Muse Spark Open on one TPU v7 host with four chips. The post reports faster decode results from a single-program kernel and describes changes to attention, expert processing and KV
AI.info Team · Written with AI · fact-checked ·
One host, eight devices
A RadixArk team serves Meta’s Muse Spark Open, an 816-billion-parameter mixture-of-experts model, on one TPU v7 host with four chips. The team's November 10 post on the PyTorch blog describes the system running SGLang on the TorchTPU PyTorch backend, with Pallas kernels handling performance-critical work.
The post says Muse Spark Open has about 42 billion active parameters per token. Its 62 decoder layers are all mixture-of-experts layers, and each layer routes a token to eight of 256 experts.
Each TPU v7 chip has two chiplets, which the framework exposes as separate devices; the host therefore has eight devices. The model is sharded across those devices with tensor parallelism, assigning one shard to each.
A single program for a serving step
The central change is a megakernel: one Pallas program on each device handles all 62 layers of a decode or verification step. SGLang continues to run the embedding, draft model, KV-cache write, language-model head and sampling around that program, the post says.
The megakernel takes 1.95 milliseconds per output token for one request, compared with 3.78 milliseconds on its previous “native” path, a 1.94-times speedup, the team reports. For two concurrent requests, it reports a 1.85-times speedup; the gain narrows as concurrency increases.
The post says the verification step runs within 1.5 to 1.7 times the time needed simply to stream its weights from high-bandwidth memory. To keep transfers from interrupting computation, the kernel pipelines weight loads, KV-cache reads and transfers among devices.
Changes to attention and expert processing
For attention, the team removed the running maximum used in the online softmax calculation, the team says. The post argues that Muse Spark’s query and key normalization, together with the model’s fixed query gain, bounds attention scores and makes the maximum unnecessary; without that dependency, KV-cache chunks can be processed in a pipeline.
In the serving system, the pipelined attention loop increased decode speed per request by 18% for two requests at a 500,000-token context, according to the team. A separate kernel test with random routing showed decode steps with many requests running 13% to 19% faster after the team split its expert-processing work into two loops.
DFlash2 speculative decoding was also used, in which a smaller draft model proposes tokens that the full model checks together. On GSM8K with real acceptance, the company reports that this reduced megakernel time per output token for one request from 3.93 milliseconds to 1.95 milliseconds, a 2.0-times speedup.
Agent workload results
For the 1-million-context version of SemiAnalysis’s AgentX coding benchmark, the post gives a range from 433 output tokens per second per request at 6,101 tokens per second per chip, to 25,906 tokens per second per chip at 95 tokens per second per request. Every TPU result shown finished requests within 47 seconds at the 90th percentile, the team says.
The post describes AgentX as based on 393 opt-in Claude Code sessions, replaying request lengths, shared prefixes, tool pauses and subagent structure with synthetic tokens replacing the original content. The comparison uses the benchmark’s 1-million-context trace set, where prompts reach about 990,000 tokens; throughput counts both input and output tokens, the team says.
The post identifies the TPU configuration as one host with four chips, using FP8 compute and NVFP4 expert weights. Its cost figures are InferenceX’s external owning costs: $1.21 per TPU v7 chip-hour and $1.86 per GB200 GPU-hour.
Each device has about 20 GiB left for KV cache after the model weights, the team says. The post describes a sequence-on-lane layout for FP8 cache data and L2 HiCache on TPU, with session-aware eviction, as parts of its approach to increasing cache capacity.