Deploying powerful Mixture-of-Experts (MoE) models like Qwen3.5-397B on high-performance hardware accelerators presents unique and complex system engineering challenges. Simply loading its massive 400 GB weight footprint into High Bandwidth Memory (HBM) and ensuring maximum hardware utilization demands a meticulous, first-principles engineering approach, rather than relying on trial-and-error.
As the landscape of open-weights models grows increasingly intricate, engineering teams can no longer afford to spend months optimizing each new model family individually. To address this scalability hurdle, our performance team has pioneered a modular, model-agnostic optimization strategy. We break down models into independent, self-contained building blocks, each accompanied by hardware-aware cost models.
These pre-optimized modules, such as Batched RPA, Grouped GEMMs, and SparseCore unpermutation, are designed for near-zero friction porting when a new architecture emerges. This innovative approach allows our engineers to achieve state-of-the-art serving performance well ahead of initial projections. Our focus has shifted from localized model optimization to global, platform-level scalability, ensuring rapid deployment and peak efficiency.
Revolutionizing Qwen 3.5 MoE on Ironwood (TPU v7x)
This report details how we systematically applied our global optimization playbook to the Qwen 3.5 MoE model on the cutting-edge Ironwood (TPU v7x) platform. By leveraging our extensive library of reusable JAX/Pallas kernels, we focused our efforts only on Qwen 3.5’s novel components. Specifically, we targeted its Gated DeltaNet (GDN) linear attention and Attention Data Parallelism.
Our strategic optimizations yielded significant performance uplifts for both decode-heavy and prefill-heavy workloads. Between April and June 2026, we improved inference performance by approximately 3.1x for decode-heavy workloads and an astounding 4.7x for prefill-heavy workloads (at the 512 Concurrency tier). Furthermore, by integrating these modular optimizations directly into open-source serving frameworks like vLLM and SGLang, we’ve removed legacy software barriers, offering a seamless and production-ready migration path for global enterprise workloads at scale.
Understanding Qwen 3.5’s Architecture
Qwen 3.5 boasts a total of 397 billion parameters, yet it uses a highly sparse routing scheme that activates only 17 billion parameters per token per forward pass. This sparse configuration means a mere 4.3% routing activation ratio. The result is a model that delivers the expressive capacity and intelligence of a 400B-class model, but with the inference footprint and execution speed of a much smaller 20B-class system.
The official model weights and configuration are readily available via the Qwen3.5-397B-A17B Hugging Face Repository. For a deeper dive into Qwen 3.5’s unique hybrid linear attention and gating components, explore the technical reports: Qwen3.5: Nobody Agrees on Attention Anymore and Gated DeltaNet for Linear Attention.
The network comprises 60 layers, a hidden dimension D=4096, and a padded vocabulary size of 248,320 tokens. Instead of a uniform Transformer layer stack, Qwen 3.5 employs a highly customized hybrid layout consisting of 15 repeating structural blocks. Each block follows a 3:1 ratio:
- 3x Gated DeltaNet (GDN) Linear Attention layers
- 1x Standard Grouped Query Attention (GQA) layer
- Each Attention layer is followed by a Mixture-of-Experts (MoE) Feed-Forward Network (FFN)
This repeating sequence means the model’s hybrid nature integrates three distinct mathematical formulations:
- Gated DeltaNet (GDN) Linear Attention: Scales linearly with sequence length, perfect for long contexts.
- Grouped Query Attention (GQA): Anchors linear attention retrieval and compresses the KV cache footprint.
- Mixture-of-Experts (MoE) Feed-Forward Networks: Provides massive capacity with sparse activation.
The Ingenious Gated DeltaNet (GDN)
Standard self-attention mechanisms suffer from quadratic scaling with sequence length, creating a significant bottleneck for generating long contexts. GDN elegantly solves this by implementing linear attention, using 64 linear attention heads for Values (V) and 16 heads for Queries and Keys (QK), each with a head dimension of 128. Instead of constructing a pairwise softmax attention matrix, GDN maintains a constant-sized hidden state matrix per head that acts as a recurrent memory.
At each token step t, the state matrix is updated using the delta rule: St = St-1 + ΔtqtvtT, where qt, kt, vt 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 formulation allows the context window to scale linearly (O(S)) in memory, keeping the recurrent state footprint constant.
To ground the linear attention retrieval, 25% of the model’s layers utilize standard GQA. This GQA configuration 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 significantly compresses the KV cache footprint during generation, though it imposes stringent hardware-level sharding constraints.
Optimizing the MoE FFN Layers
The feed-forward network (FFN) layers are intelligently sharded into 512 small experts, each 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 for all tokens.
This native multimodal MoE architecture can process text, image, and video inputs through an early-fusion training paradigm on trillions of multimodal tokens. The model supports a native context length of 262,144 tokens, extensible to over 1,010,000 tokens using YaRN RoPE scaling. This remarkable capability is a testament to its advanced design.
Benchmarking and Workload Analysis
To systematically identify, profile, and resolve compiler and kernel bottlenecks, our systems engineering team established a rigorous multi-dimensional evaluation matrix. This matrix is based on real-world, asymmetric workloads designed to stress various hardware execution subsystems. Our benchmarking sweeps included diverse scenarios:
- Prefill-heavy workloads: Large batch sizes and short decode steps (e.g., Concurrency 512, 8K input/16 decode tokens).
- Decode-heavy workloads: Small batch sizes and long decode steps (e.g., Concurrency 512, 128 input/1K decode tokens).
The benchmark was executed on an enterprise-grade single-host cluster featuring 8x TPU v7x (Ironwood) devices. This powerful setup allowed us to push the boundaries of performance and identify optimization opportunities.
We conducted a first-principles roofline analysis for the Qwen 3.5 workload under a standard 8K/1K configuration at 64 concurrency to pinpoint theoretical bounds and execution stalls. 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).
This comprehensive roofline model, applied to our standard evaluation workload (Concurrency 64, with an 8K/1K prefill-heavy and 1K/8K decode-heavy sequence length layout), established the absolute, application-level throughput bounds (in tokens-per-second per physical chip) for each serving phase. For a prefill-heavy workload, it revealed a throughput limit of 32,000 tokens/sec/chip, while for decode-heavy, the limit was 4,000 tokens/sec/chip.
Overcoming Sharding Challenges with Hybrid Parallelism
The specific architectural constraints of Qwen 3.5—specifically, having exactly 2 KV heads in the GQA layers and 512 experts in the MoE layers—rendered traditional uniform sharding approaches ineffective. 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) would necessitate fractional head sharding (2/8 = 0.25 heads per device), which is physically impossible on hardware.
Replicating the heads locally across 8 cores would duplicate the physical KV cache memory footprint on every device, nullifying the memory-saving benefits of GQA. This memory redundancy severely limits the HBM headroom available for active KV caches under high-load workloads. Consequently, the server engine would be forced 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 critical bottleneck, we co-designed a revolutionary hybrid sharding scheme (PR #2577). This scheme combines 8-way Attention Batch Sharding (Data Parallelism, DP=8) 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 prevents duplicating the immense 400 GB parameter footprint across all nodes while keeping collective payload sizes manageable. This intelligent hybrid approach allowed us to fully unlock the model’s potential.
Optimized MoE Routing and Communication
Transitioning between Attention DP and MoE EP necessitates efficient cross-device token routing. In designing our Mixture-of-Experts (MoE) routing layer, we evaluated two primary structural approaches to handle this cross-device transition:
- Option A (Pre-Routing): Tokens are pre-routed to their respective expert devices before expert computation begins.
- Option B (Post-Routing): Tokens are processed locally and then routed to experts, followed by collective communication.
Because deterministic latency is paramount for real-world serving, we opted for Option B and subsequently developed low-level communication fusions to optimize its collective pathways. In a naive Option B implementation, preparing for local MoE computation requires broadcasting three distinct pieces of data across the cluster to every device rank:
- The full input tokens for local expert execution.
- The top-k expert indices (routing metadata).
- The top-k expert weights (routing metadata).
Each collective communication call incurs 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. Since 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. This allows a single All-Gather across the data dimension (ShardingAxisName.MLP_DATA) for both routing metadata blocks, unpacking them locally and effectively 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:
- Phase 1 (intra-chip reduction): Locally sum partial outputs from experts on the same chip.
- Phase 2 (inter-chip scatter): Scatter the reduced results across chips to their final data-parallel destinations.
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 performs vector additions for micro-batch i-1, effectively hiding communication latency behind computation. This meticulous optimization ensures smooth and efficient data flow.
Custom Kernels and Memory Management
By authoring custom kernels using the JAX custom kernel language, Pallas, we bypassed the standard XLA lowering path. This allowed us to directly control VMEM layout, registers, and memory scheduling across the three primary execution tracks: Vector Processing Unit (VPU), TensorCore, and Inter-Chip Interconnect (ICI).
Managing the KV cache for the 25% GQA layers demands 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 resulted 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 significantly reduced the decode step latency under Concurrency-512 from 428µs to 283µs, achieving an impressive 33.8% kernel-level speedup. To further saturate the memory bus, we designed batched RPA kernels. This design groups multiple decode streams 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. These innovations were crucial for maximizing hardware efficiency and performance.
Source: Google Developers Blog
