Google Demos Modular Optimization Approach for Giant MoE Models on TPU v7x Hardware

Google's systems engineering team has unveiled a scalable methodology for tuning mixture-of-experts (MoE) models on high-end accelerator hardware, using Alibaba's Qwen 3.5-397B as a case study. Rather than optimizing each new model family from scratch, the team built a modular framework of reusable kernel libraries and hardware-aware cost models that transfer across architectures with minimal friction.

A Pragmatic Alternative to One-Off Optimization#

Deploying a sparse model like Qwen 3.5-397B—which weighs 400 GB but activates only 17 billion parameters per token—demands careful orchestration of memory, routing, and compute on specialized hardware. Google's approach flipped the traditional workflow: instead of fitting a model to a platform through months of empirical tuning, the team pre-optimized a library of self-contained building blocks (batched RPA, grouped GEMMs, SparseCore operations) and ported them to new architectures with near-zero engineering overhead.

Across the April–June 2026 period, this strategy achieved approximately 3.1x throughput improvement for decode-heavy workloads and 4.7x for prefill-heavy workloads at 512 concurrent requests. The modular optimizations were then integrated into open-source serving frameworks vLLM and SGLang, creating a production-ready migration path for enterprise deployments.

The Architecture Challenge#

Qwen 3.5's hybrid design—60 layers arranged in a 3:1 alternating pattern of Gated DeltaNet (linear attention) and Grouped Query Attention (GQA) layers, paired with 512 MoE experts—created unusual hardware constraints. Standard tensor-parallel sharding approaches collapsed under these constraints: GQA's 2 key-value heads could not be cleanly sliced across 8 devices (2÷8 = 0.25 heads per device), and replicating them across all cores bloated the KV cache memory footprint.

Hybrid Data and Expert Parallelism#

To resolve the conflict, Google co-designed a hybrid parallelism scheme: 8-way attention batch sharding (replicating GQA and Gated DeltaNet weights locally for full-head computation) combined with 8-way expert parallelism in the MoE layers (distributing 512 experts as 64 per device). This avoided both head-splitting and parameter duplication, but required efficient cross-device token routing.

The team evaluated two routing topologies: all-to-all shuffling (lower redundancy, unpredictable latency) and token replication with collective aggregation (higher memory cost, deterministic latency). For production serving, they chose token replication and then optimized its collective pathways—fusing three separate all-gather operations into two and implementing a custom TPU-native hierarchical reduce-scatter that overlapped intra-chip and inter-chip communication.