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.
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:
Key Mathematical Elements
The model's hybrid nature integrates three distinct mathematical formulations:
Standard self-attention mechanisms scale quadratically with sequence length (O(S2)), 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:
Where qt, kt, and 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 recurrent formulation allows the context window to scale linearly (O(S)) in memory, keeping the recurrent state footprint constant.
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.
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 soft-max probability distribution. Crucially, the model also incorporates one shared expert path that is always executed, functioning as a common representation layer:
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.
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).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. In designing our Mixture-of-Experts (MoE) routing layer, we evaluated two primary structural approaches to handle this cross-device transition:
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:
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:
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.
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:
"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:
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.
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.
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.
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 eliminate VREG spills and memory load stalls, doubling vector arithmetic throughput on the VPU.
Offloads the token-unpermutation and local reduction operations entirely to the SparseCore. By performing indirect gather and local reduction directly on the SC, we bypassed the materialization of intermediate, padded activation tensors in HBM, reducing the HBM read requirements from 20 down to 10 and writes from 15 down to 5, slashing MoE overhead.
To maximize hardware efficiency, our implementation leverages a chunk-level pipelined architecture rather than performing local reduction and the 8-device Reduce-Scatter sequentially on the full [81920,4096] tensor. The workload is partitioned into 4 distinct chunks. As soon as Chunk 1 completes its local unpermutation and gather-reduce on the SparseCore, it asynchronously kicks off its Reduce-Scatter collective across physical ICI links. Simultaneously, the SparseCore begins the local gather-reduce for Chunk 2. This strict chunk-level pipelining effectively overlaps and hides the cross-device network latency of the Reduce-Scatter behind the local compute of subsequent chunks.
The recurrent state updates in the 75% Gated DeltaNet (GDN) layers are highly susceptible to memory bandwidth bottlenecks due to constant recurrent state updates
To optimize this track, the performance team implemented a series of algorithmic fusions and precision co-designs:
The GDN recurrent update is preceded by a causal 1D convolution (K=4). Initially, this was compiled as an independent operation, forcing the intermediate convolution outputs to be written to and read from HBM. We designed a register-level sliding window algorithm that caches historical token states directly within the TPU's VPU registers. Fusing the 1D convolution and the GDN recurrent state update into a single execution block eliminated 6 redundant HBM round-trips (see PR #2823).
We restructured the linear attention update equations to exploit algebraic identities. By mathematically rearranging the operations, we completely skipped the expensive post-rank-1 matrix multiplication in the fused GDN kernels, reducing the computational footprint (see PR #2498).
Additionally, to further saturate the Vector Processing Unit (VPU), we transitioned the recurrent State Space Model (SSM) state variables from Float32 to BFloat16 precision. This doubled the vector arithmetic throughput on the VPU without compromising numerical convergence or output quality.
To prevent padding overhead from wasting MXU FLOPs during batched prefill execution, we optimized JAX-native chunked layouts in and introduced specialized sequence-handling routines that natively process ragged inputs in PR #2218, ensuring that variable sequence lengths do not introduce processing stragglers.
Rather than relying on separate execution stages, we designed and merged a fully-fused Pallas kernel in PR #3016 that compiles the causal 1D convolution and the entire GDN recurrent linear attention block into a single, unified execution unit on the VPU. By caching intermediate sequence and recurrent states directly within the local registers, this kernel completely bypasses the need to read and write intermediate activation tensors to VMEM or HBM.
This register-level fusion eliminates register-to-memory synchronization latency and provides a critical performance boost for both serving phases:
Serving Qwen 3.5 requires managing two heterogeneous attention state structures: the fixed-size recurrent linear attention state of Gated DeltaNet (GDN) and the dynamically growing standard attention Key-Value (KV) cache of Grouped Query Attention (GQA). Because the TPU v7 features 192GB of HBM capacity per chip (e.g. compared to the 288GB available on Blackwell GB300 GPUs - a ~50% capacity difference), HBM footprint optimization under high concurrency is a severe systems constraint. In PR #2416, we introduced a custom memory layout designed to align and store these hybrid attention states together in HBM. This layout minimizes padding and prevents memory fragmentation, directly reclaiming critical HBM headroom. This optimization increases the maximum supportable batch sizes, allowing the TPU to scale smoothly and sustain high serving throughput under heavy client concurrencies.
The systems engineering optimizations were validated under rigorous, empirical benchmarking scenarios. Below, we present the raw throughput results with our optimized JAX/Pallas stack on Ironwood TPU across 4 concurrency tiers.
Closing the Headroom Gap: Actual Throughput vs. Roofline Limits
Placing our empirical serving results side-by-side with our first-principles roofline limits at the baseline Concurrency 64 tier demonstrates the real-world efficiency of our custom Pallas kernels and DP+EP sharding topology:
This close alignment demonstrates that our low-level compiler and kernel fusions push the TPU hardware close to its physical execution limits, leaving minimal remaining headroom and proving the extreme efficiency of the open-source software stack.
Running large-scale Mixture-of-Experts models at scale requires not just raw throughput, but strict mathematical correctness. Under high concurrency, gating and routing matrices are highly sensitive to low-precision accumulation errors. In designing our custom JAX/Pallas gating kernels, the systems engineering team incorporated a dedicated Numerical Verification Layer to audit accumulation precision across our FP8 scaling blocks. By continuously monitoring the softmax distribution ranges and expert load balances, we verified that our Pallas-lowered gating weights maintain zero deviation from the high-precision Float32 reference path (see PR #2328 and PR #2674), guaranteeing high throughput alongside strict output quality.
To eliminate the remaining bottlenecks, our engineering team has structured an active optimization roadmap divided into two primary technical tracks:
top_k selection kernel directly on the VPU. Currently, routing logits are computed on the TensorCore and transferred to the VPU for top_k selection, introducing a serialization bottleneck. Fusing these operations will keep the routing pipeline local to the VPU.Optimizing massive open-weights models like Qwen 3.5 on modern accelerators is a challenge that cannot be solved by brute-force empirical changes. It requires a disciplined, first-principles systems engineering approach: mapping theoretical hardware limits via Roofline Modeling, systematically isolating bottlenecks through profile traces, and bridging the remaining gaps with hand-scheduled custom Pallas kernels and hardware-aware sharding topologies.
The optimization playbook developed for this deployment does more than just accelerate Qwen 3.5; it establishes a hardened, reusable open-source software stack that makes Google Cloud TPUs a highly general-purpose and competitive engine for the next generation of sparse Mixture-of-Experts architectures.
To explore our complete technical guides, access ready-to-run code templates, and learn more about optimizing frontier-class models on Google Cloud hardware, check out our new Google Cloud TPU Developer Hub.
This optimization work was a massive collaborative effort across multiple engineering, compiler, and management teams. We would like to express our gratitude to the following contributors who drove this initiative to success: