The Real Problem With Serving Sparse MoE Models
Every engineering team that has tried to serve a frontier-class Mixture-of-Experts (MoE) model on specialized accelerators has hit the same wall: you can't brute-force your way to good performance. Qwen 3.5-397B-A17B is a perfect case study. It packs 397B total parameters into a weight footprint of roughly 400 GB, but only activates 17B parameters per token — a 4.3% routing activation ratio. That sparsity is the whole point: you get the intelligence of a 400B-class model at the serving cost of a ~20B one.
But that promise only holds if your systems engineering is airtight. Naive tensor-parallel sharding breaks immediately when a model has exactly 2 KV heads in its GQA layers and 512 experts in its MoE stack. You can't shard 2 heads across 8 devices. Replicating them wastes HBM. And on a chip with 192 GB of HBM (versus the 288 GB on Blackwell GB300 — a ~50% capacity gap), wasting HBM is fatal.
The solution isn't a single clever kernel. It's a modular, model-agnostic optimization methodology — a playbook of reusable building blocks (Batched RPA, Grouped GEMMs, SparseCore unpermutation) that can be ported to new architectures with near-zero friction. This article breaks down how that playbook was applied to Qwen 3.5 on Ironwood (TPU v7x), yielding ~3.1x decode and ~4.7x prefill speedups at the 512-concurrency tier.
For a full technical breakdown, see the original engineering report.
The Architectural Constraints That Break Naive Sharding
Qwen 3.5's hybrid layout is not a uniform Transformer stack. It uses 60 layers arranged in 15 repeating blocks, each with a 3:1 ratio of Gated DeltaNet (GDN) linear attention to standard GQA attention. Three mathematical structures coexist:
- GDN linear attention (75% of layers): 64 V-heads, 16 QK-heads, head dim 128. Maintains a constant-size recurrent state matrix per head, updated via the delta rule. Scales O(S) instead of O(S²).
- GQA (25% of layers): 32 query heads, exactly 2 KV heads, head dim 256, RoPE dim 64. Compresses KV cache but imposes brutal sharding constraints.
- MoE FFN: 512 small experts, intermediate dim 1024,
top_k=10routing plus 1 always-on shared expert.
Why TP=8 Fails
With only 2 KV heads, a tensor-parallel size of 8 forces fractional head sharding (2/8 = 0.25 heads per device) — physically impossible. Replicating the heads across 8 cores duplicates the KV cache on every device, capping real concurrency at ~200 instead of the planned 512.
The Fix: Hybrid Attention-DP + Expert-Parallel
The team co-designed a 8-way Attention Data Parallelism (DP=8) + 8-way Expert Parallelism (EP=8) scheme:
# Conceptual sharding topology for Qwen 3.5 on 8-device TPU mesh
# 주석: 8개 디바이스 메시에서 어텐션은 DP, MoE는 EP로 샤딩
sharding = {
"gdn_layers": "replicated", # Full 2 KV heads per device
"gqa_layers": "replicated", # Preserve local KV cache consistency
"moe_experts": "expert_parallel", # 512 experts / 8 devices = 64 per chip
}
# Cross-device token routing: All-Gather metadata, then Reduce-Scatter outputs
# 주석: 라우팅 메타데이터는 All-Gather, 전문가 출력은 Reduce-Scatter로 복귀
Attention weights are replicated across all 8 devices (each core sees the full 2 KV heads), eliminating intra-attention sharding communication. The 512 routed experts are then evenly distributed (64 per device), avoiding the 400 GB weight duplication that would otherwise be required.
Fusing Three Collectives Into Two
Under naive EP, preparing for local MoE compute required three separate All-Gathers: expert indices, topk weights, and token hidden states. Since expert indices (int) and topk weights (float) share identical shapes [1024, 10], they were stacked, bitcast, and packed into a single dense 32-bit integer blob — one All-Gather instead of two. That halved the routing metadata collective latency.
Pallas Kernels, SparseCore Co-Design, and the Roofline
Once the sharding topology was correct, the remaining wins came from hand-scheduled kernels that bypass the standard XLA lowering path.
Three High-Leverage Optimizations
| Optimization | Technique | Measured Impact |
|---|---|---|
| Coarse-grained RPA indexing | KV page size 16 → 256 (--block-size=256) | Decode step latency at C=512: 428µs → 283µs (33.8% faster) |
| SparseCore unpermutation | Offload token unpermutation + local reduction to SparseCore | HBM reads 20→10, writes 15→5 |
| Fused GDN + causal 1D conv | Register-level sliding window in VPU registers | Eliminated 6 redundant HBM round-trips |
Register-Level Fusion in the GDN Path
The GDN recurrent update was previously compiled as an independent op, forcing intermediate convolution outputs to be written to and read from HBM. The team fused the causal 1D convolution (K=4) and the GDN recurrent state update into a single execution block, caching historical token states directly in VPU registers. They also transitioned the SSM state variables from Float32 to BFloat16, doubling VPU vector arithmetic throughput without compromising numerical convergence.
Roofline Validation
At Concurrency 64 with an 8K/1K layout, empirical throughput landed remarkably close to the first-principles roofline limits — meaning the low-level kernels are pushing the hardware near its physical execution bounds. Under high concurrency, gating and routing matrices are highly sensitive to low-precision accumulation errors, so a Numerical Verification Layer continuously audits FP8 scaling blocks, confirming zero deviation from the Float32 reference path.
Limitations, Warnings, and Where to Go Next
What This Playbook Does Not Solve
- HBM capacity is the hard ceiling. TPU v7's 192 GB per chip is ~50% smaller than Blackwell GB300's 288 GB. No amount of kernel fusion recovers that. The hybrid memory layout helps, but the fundamental capacity gap remains.
- Model-specific novelty still costs engineering time. The modular playbook reduces friction, but Qwen 3.5's GDN layers and
top_k=10routing still required bespoke Pallas kernels. It's not free. - Numerical correctness must be re-verified per model. FP8 accumulation is model-sensitive. What works for Qwen 3.5 may not transfer cleanly.
The Remaining Roadmap
Two tracks remain open: (1) fusing the top_k selection kernel directly on the VPU to eliminate the TensorCore→VPU serialization bottleneck, and (2) further reducing cross-device collective overhead through chunk-level pipelining.
If You're Building on This
Start with the sharding topology before touching kernels. Getting DP+EP right is worth more than any single micro-optimization. Then profile aggressively — bottlenecks hide in collectives, not just matmuls.