Systems Engineering Playbook: Optimizing Qwen 3.5-397B MoE on Ironwood (TPU7x)

Google published a systems engineering playbook on the Google Developers Blog describing how engineers served the 397B-parameter Qwen 3.5 Mixture-of-Experts (MoE) model on Ironwood TPUs (TPU7x). The core result was a modular JAX/Pallas optimization stack that delivered up to a 4.7x inference speedup on prefill-heavy workloads. That combination of model scale, MoE routing, and TPU serving is the practical problem the post is solving: get a very large expert model onto Ironwood without leaving performance on the table.

On the technical side, hardware sharding constraints made a clean, conventional partition strategy hard to use at this size. The team used a hybrid Data Parallelism and Expert Parallelism (DP+EP) topology so data and expert placement could be split across devices instead of forcing a single parallelism axis. They also built custom low-level communication fusions, including a hierarchical reduce-scatter, aimed at the expensive part of MoE serving: cross-device token routing. The stack is modular and written in JAX/Pallas, so the optimizations sit close to the kernel and collective path rather than only in higher-level graph rewrites.

For engineers and builders, the useful detail is not just the 4.7x number but where it shows up. Prefill-heavy workloads are the prompt-processing phase that often dominates latency and cost when contexts are long or batch composition is uneven. A hybrid DP+EP layout plus fused collectives is a concrete pattern for anyone trying to place MoE models on TPUs when pure expert or pure data parallel sharding hits device limits. JAX/Pallas as the implementation surface also signals that production wins at this scale still come from low-level control of memory layout and communication, not only from model quality.

In market terms, this is Google documenting how to run a large open MoE family model on its Ironwood TPU generation rather than only showcasing first-party models. Qwen 3.5 at 397B parameters is a high-profile MoE target; showing a 4.7x prefill path on TPU7x is as much a platform credibility claim as a research note. It positions Ironwood as a viable serving substrate for third-party MoE inference, not only training or Google-internal stacks, and puts pressure on competing accelerator stacks that advertise MoE throughput without publishing similar topology and collective detail.

The practical takeaway is to treat MoE serving as a systems problem of topology plus collectives, not only of FLOPs. If you are evaluating TPU inference for large MoE models, the playbook’s DP+EP topology and hierarchical reduce-scatter-style fusions are the pieces to map onto your own sharding and prefill path. Watch next for whether similar modular JAX/Pallas patterns land in reusable libraries or only in one-off serving stacks, and whether the same hybrid approach extends cleanly beyond the prefill-heavy regime the post highlights.