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

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

To serve the 397B-parameter Qwen 3.5 Mixture-of-Experts (MoE) model on Ironwood TPUs, engineers developed a modular JAX/Pallas optimization stack that achieved up to a 4.7x inference speedup for prefill-heavy workloads. The team bypassed severe hardware sharding constraints by deploying a hybrid Data Parallelism and Expert Parallelism (DP+EP) topology, paired with custom low-level communication fusions like a hierarchical reduce-scatter to optimize cross-device token routing. Finally, by executing hardware-aware custom kernels-such as Batched Ragged Page Attention and a fully-fused Gated DeltaNet (GDN) block-they successfully saturated HBM bandwidth and TensorCore MXUs to push system throughput near its theoretical roofline limits.

Deploying and serving a Mixture-of-Experts (MoE) model like Qwen3.5-397B on specialized hardware accelerators presents significant systems engineering challenges. Loading the 400 GB weight footprint into High Bandwidth Memory (HBM) and maximizing hardware utilization requires a disciplined, first-principles engineering methodology over empirical trial-and-error modifications. Crucially, as the landscape of open-weights models grows in complexity, engineering teams can no longer afford to spend months optimizing each new model family in isolation.

To solve this scalability challenge, our performance team has pioneered a modular, model-agnostic optimization strategy. Rather than tackling models as monolithic systems, we decompose them into self-contained, independent building blocks (such as Batched RPA, Grouped GEMMs, and SparseCore unpermutation) accompanied by hardware-aware cost models. When a new architecture arrives, these pre-optimized modules are ported with near-zero engineering friction. This allows our engineers to deliver state-of-the-art serving performance well ahead of initial projections, shifting our focus from localized model optimization to global, platform-level scalability.

This technical report details how we systematically applied this global optimization playbook to Qwen 3.5 MoE on the Ironwood (TPU v7x) platform. By leveraging our library of reusable JAX/Pallas kernels and targeting only Qwen 3.5's novel components-such as Gated DeltaNet (GDN) linear attention and Attention Data Parallelism-our team achieved significant performance uplift for both decode-heavy and prefill-heavy workloads. The optimizations discussed below allowed us to improve inference performance by approximately 3.1x for Decode-heavy and by approximately 4.7x for Prefill-heavy workloads (512 Concurrency tier) between April and June 2026.

Furthermore, by integrating these modular optimizations natively into open-source serving frameworks like vLLM and SGLang, we have neutralized legacy software barriers, providing a seamless, production-ready migration path for global enterprise workloads at scale.

Model Architecture Overview

The model consists of 397 billion total parameters, but leverages a highly sparse routing scheme that activates exactly 17 billion parameters per token per forward pass. This sparse configuration represents a 4.3% routing activation ratio, enabling the model to deliver the expressive capacity and intelligence of a 400B-class model while maintaining the inference footprint and execution speed of a much smaller 20B-class system.

The official model weights and configuration can be accessed directly via the Qwen3.5-397B-A17B Hugging Face Repository. For a comprehensive structural analysis of Qwen 3.5's hybrid linear attention and gating components, see the technical deep-dives on Qwen3.5: Nobody Agrees on Attention Anymore (Hugging Face Blog) and Gated DeltaNet for Linear Attention (Sebastian Raschka, PhD).

The network consists of 60 layers in total, hidden dimension D=4096, and a padded vocabulary size of 248,320 tokens. Rather than utilizing a uniform Transformer layer stack, Qwen 3.5 employs a highly customized hybrid layout composed of 15 repeating structural blocks. Each block is arranged in a 3:1 ratio:

This repeating sequence can be expressed as:

• 3 layers of Gated DeltaNet (GDN) linear attention
• 1 layer of Grouped Query Attention (GQA)

Key Mathematical Elements

The model's hybrid nature integrates three distinct mathematical formulations:

Gated DeltaNet (GDN) Linear Attention. Standard self-attention mechanisms scale quadratically with sequence length (O(S²)), creating a computational bottleneck for long-context generation. GDN solves this by computing linear attention, utilizing 64 linear attention heads for Values (V) and 16 heads for Queries and Keys (QK) with a head dimension of 128. Instead of constructing a pairwise softmax attention matrix, GDN maintains a constant-sized hidden state matrix per head (matching the dkdv key-value dimensions) that functions as a recurrent memory. At each token step t, the state matrix is updated using the delta rule:

S_t = S_{t-1} + σ(β_t) · (v_t - S_{t-1} · k_t) ⊗ k_t

Where q_t, k_t, and v_t are the query, key, and value vectors, and β_t is a learned gating parameter. This recurrent update is preceded by a causal 1D convolution (K=4) to capture local spatial dependencies. This recurrent formulation allows the context window to scale linearly (O(S)) in memory, keeping the recurrent state footprint constant.

Grouped Query Attention (GQA). To anchor the linear attention retrieval, the model uses standard GQA in 25% of its layers. GQA uses 32 query heads (Nq=32) and exactly 2 key-value (KV) heads (Nkv=2) globally, with a head dimension of 256 and a Rotary Position Embedding (RoPE) dimension of 64. This extreme GQA layout compresses the KV cache footprint during generation, but imposes strict hardware-level sharding constraints, as detailed in Section 3.

Mixture-of-Experts (MoE) Feed-Forward Network. The feed-forward network (FFN) layers are sharded into 512 small experts with an intermediate expert dimension of 1024. During execution, a router gate projects token representations and selects the top 10 routed experts via a softmax probability distribution. Crucially, the model also incorporates one shared expert path that is always executed, functioning as a common representation layer:

FFN(x) = SharedExpert(x) + Σ_{i ∈ top-k} (gate_i(x) · Expert_i(x))

This native multimodal MoE architecture natively processes text, image, and video inputs via an early-fusion training paradigm on trillions of multimodal tokens. The context window supports a native context length of 262,144 tokens, extensible to over 1,010,000 tokens using YaRN RoPE scaling.

Benchmarking Methodology

To systematically isolate, profile, and resolve compiler and kernel bottlenecks, the systems engineering team established a rigorous multi-dimensional evaluation matrix based on real-world, asymmetric workloads. Our benchmarking sweeps across asymmetric workloads designed to stress separate hardware execution subsystems:

The benchmark was executed on an enterprise-grade single-host cluster: vllm-project/tpu-inference. For the final optimized runs utilizing Attention DP, the server execution loop was configured with --max-num-batched-tokens=1024 and --max-num-seqs=64 per core (compared to --max-num-batched-tokens=8192 and --max-num-seqs=512 utilized in early tensor-parallel baselines).

Hybrid Sharding: DP+EP Topology

The specific architectural constraints of Qwen 3.5-namely, having exactly 2 KV heads in the GQA layers and 512 experts in the MoE layers-invalidate traditional uniform sharding approaches. In standard Attention Tensor-Parallel (TP) + Expert MoE configurations, attention weights are sliced and sharded across the device dimension. However, attempting to shard the GQA layers with a tensor parallelism size of 8 (TP=8) forces fractional head sharding (2/8 = 0.25 heads per device), which is physically impossible on hardware. Replicating the heads locally across 8 cores duplicates the physical KV cache memory footprint on every device, neutralizing the memory-saving benefits of GQA.

This memory redundancy severely restricts the HBM headroom available for active KV caches under high-load workloads. This capacity limitation forces the server engine to cap the actual achieved concurrency far below expected targets-limiting the system to roughly ~200 concurrent requests instead of the planned 512.

To eliminate this bottleneck, we co-designed a hybrid sharding scheme (PR #2577):

• 8-way Attention Batch Sharding (Data Parallelism, DP=8) combined with
• 8-way Expert Parallelism (EP=8) in the MoE layers

Replicating GQA and GDN weights across all 8 devices allows each core to process attention locally with the full 2 KV heads, preserving local KV cache consistency and eliminating intra-attention sharding communication. In the feed-forward MoE layers, we switch to Expert Parallelism (EP=8). The 512 routed experts are distributed evenly (64 experts per device), which avoids duplicating the 400 GB parameter footprint across all nodes while keeping collective payload sizes manageable.

Transitioning between Attention DP and MoE EP requires cross-device token routing.

Optimizing Expert Parallelism Communication

In designing our Mixture-of-Experts (MoE) routing layer, we evaluated two primary structural approaches to handle this cross-device transition:

• Option A: All-to-All dispatch - each device sends tokens directly to the device hosting the required expert
• Option B: All-Gather + local slice - each device gathers the full token batch, then selects its local expert slice

Because deterministic latency is critical for real-world serving, we opted for Option B and subsequently developed low-level communication fusions to optimize its collective pathways.

Under a naive Option B implementation, preparing for local MoE computation requires broadcasting three distinct pieces of data across the cluster to every device rank. Assuming a local tensor slice with a shape of [1024, 4096] for token hidden dimensions, we typically must perform three separate collective operations:

1. All-Gather for token hidden states ([1024, 4096])
2. All-Gather for expert indices ([1024, 10])
3. All-Gather for topk routing weights ([1024, 10])

Every collective communication call carries a fixed kernel launch and network synchronization latency penalty on the TPU. To optimize Expert Parallelism (EP) efficiency, we consolidated these three All-Gathers down to two in PR #2836. Because the expert indices (integers) and the topk weights (floats) share identical tensor shapes ([1024, 10]), we stack, bitcast, and pack them together along a new dimension into a single dense 32-bit integer array (blob). This allows us to run a single All-Gather across the data dimension (ShardingAxisName.MLP_DATA) for both routing metadata blocks, unpacking them locally and halving the routing metadata collective latency.

After expert execution, token outputs must return to their data-parallel ranks. A standard All-Reduce over the 8-device mesh is highly inefficient. We replaced this with a custom, TPU-native Hierarchical Reduce-Scatter written in Pallas/Mosaic (see PR #2679). The collective runs in two pipelined phases:

1. Phase 1: Within each pair of adjacent devices, perform a local reduce-scatter
2. Phase 2: Across device pairs, perform a global reduce-scatter

To prevent VMEM Out-of-Memory (OOM) errors, the data is sliced into 2 to 4 micro-batches. The kernel pipelines remote DMA transfers of micro-batch i while the TensorCore is performing vector additions for micro-batch i-1, hiding the communication latency behind the compute.

Roofline Analysis and Theoretical Bounds

To identify the theoretical bounds of our systems engineering and understand where execution stalls occurred, we conducted a first-principles roofline analysis for the Qwen 3.5 workload under a standard 8K/1K configuration at 64 concurrency. During the prefill phase, a batch of 64 prompts with 8,192 input tokens each yields 524,288 tokens processed in parallel. During the decode phase, the model processes 64 tokens per step (1 token per active request).

To translate these first-principles hardware constraints into plannable software engineering metrics, we modeled our standard evaluation workload (Concurrency 64, with an 8K/1K prefill-heavy and 1K/8K decode-heavy sequence length layout) using our end-to-end roofline model. This analysis establishes the absolute, application-level throughput bounds (in tokens-per-second per physical chip) for each serving phase:

Custom Pallas Kernel Optimizations

By authoring custom kernels using the JAX custom kernel language, Pallas, we bypassed the standard XLA lowering path to control VMEM layout, registers, and memory scheduling directly across the three primary execution tracks:

• Vector Processing Unit (VPU)
• TensorCore MXU
• SparseCore (SC)

Batched Ragged Page Attention (RPA)

Managing the KV cache for the 25% GQA layers requires dynamic memory allocation. We employ Ragged Page Attention (RPA) to index non-contiguous memory blocks in HBM (see PR #2632). Historically, a block size of 16 tokens was used to minimize memory fragmentation. However, on TPU, smaller block sizes result in massive indexing overhead, causing the Vector Processing Unit (VPU) to stall during the decode phase.

We resolved this by coarse-graining the indexing to a KV page size of 256 (enabled via the server command --block-size=256). This coarse-grained indexing reduced the decode step latency under Concurrency-512 from 428µs to 283µs, achieving a 33.8% kernel-level speedup.

To further saturate the memory bus, we designed batched RPA kernels. This design groups multiple decode streams together into a single compiled Pallas kernel (PR #2632), amortizing VPU instruction dispatch latency, breaking the data dependency stalls of sequential requests, and improving memory alignment.

SparseCore-Assisted Token Routing

The fine-grained routing factor of top_k=10 in Qwen 3.5 introduces non-power-of-two tensor dimensions. Permuting and unpermuting these arrays on the TensorCore previously resulted in heavily padded, unaligned HBM memory writes. We resolved this through a SparseCore-TensorCore co-design flow:

We authored a custom Pallas/Mosaic kernel that offloads token routing to the TPU's SparseCore (SC), a hardware unit optimized for indirect addressing (see PR #2137). The SC reads the routing indices, performs an indirect DMA gather of token embeddings directly from HBM, and writes them into a contiguous virtual buffer. This bypasses the materialization of heavily-padded, unaligned intermediate tensors in HBM, saving massive memory bandwidth.

Fused Gated DeltaNet (GDN) Block

In the GMM V2 kernel, we fused the SwiGLU activation functions directly into the main matrix multiplication loops (gating and up-projection are packed and processed in a single tile via dual DMA reads), avoiding register spills to HBM. Additionally, we implemented dynamic bounded slices to process the variable token payloads of each expert with minimal padding. We transitioned to 512 subchannel activation quantization for FP8 operations to enable higher throughput on the TensorCore MXUs.