Authors: @Amanda_Liang @KaranGoel @Jaideep_Singh1
Mistral 3 is a Mixture-of-Experts (MoE) model powering modern enterprise AI workloads. This blog post outlines how we optimized Mistral 3 large MoE model inference on Google’s Ironwood (TPU v7x), achieving a 1.5x performance gain.
By utilizing hybrid sharding, replacing linear VPU summations with tree reductions, optimizing GMM/MLA kernels, and adopting asynchronous scheduling, the team successfully boosted throughput by up to 48% while maintaining benchmark accuracy neutrality.
Below, we detail the key optimizations that made this acceleration possible and present a technical playbook for scaling MoE inference on Ironwood TPUs.
Key highlights
-
Optimized Serving Parameters: Defined key production variables, such as safety buffers for max-model-len, and optimized prefill throughput (max-num-batched-tokens).
-
Hybrid Sharding (DP Attention + EP/TP MoE): Decoupled sharding by assigning 16-way data parallelism (DP=16) to attention to maximize KV cache capacity for single-head MLA, while using expert parallelism (EP=16) in MoE.
-
SparseCore Gather-Reduce Kernel: Replaced linear O(N) VPU summation chains with parallel O(log N) binary tree reduction, reducing MoE unpermute stalls.
-
MLA / GMM Kernel Tuning: Retuned the MLA / GMM kernel for Mistral 3 specific serving setup.
1. Roofline & gap analysis
We conducted an empirical headroom analysis across both Prefill and Decode phases to isolate performance gaps, focusing on high concurrency, and based on the following sharding strategies.
-
Attention Data Parallelism (DP=16):
-
The Multi-Head Latent Attention (MLA) Challenge: MLA features compressed KV representations with a single KV head (num_kv_heads = 1). Standard tensor parallelism cannot shard 1 head across 16 chips without head replication.
-
DP: By distributing data batches across 16 attention replicas, KV cache pressure is reduced, enabling higher request concurrency and greater throughput.
-
-
MoE Expert Parallelism (EP=16):
-
EP minimizes collective operation overhead relative to Tensor Parallelism (TP). While TP relies on frequent, low-latency inter-core communications, EP reduces total aggregated transmission time by requiring fewer collective synchronization steps.
-
EP facilitates exhaustive activation sharding across the expert dimension.
-
Prefill bottlenecks
- Gather-Reduce Overhead: Under pure tensor parallelism, unoptimized JAX-native gather-reduce operations executing on tensor cores (TC) accounted for a significant portion of end-to-end prefill latency. Offloading these collective operations to TPU sparse core (SC) will be more efficient, whether running exposed or overlapped with compute.
- Dispatch Overhead: Token dispatch introduced additional latency overhead. Typically occurring post-attention and prior to the MLP (represented as the GMM kernel in execution traces), this operation permutes and reorders activation vectors in memory so that tokens assigned to the same expert are stored contiguously.
- GMM Efficiency: During the prefill phase, we identified that compute-bound GMM kernels were performing significantly below the hardware roofline.
Decode bottlenecks
- Attention: The attention kernel was only around a quarter of the roofline.
- Host-to-Device Stalls: Gathering token payloads across 16 DP ranks led to large step-to-step transmission gaps when CPU-TPU transfer was synchronous.
2. Production serving setup
-
Tokenizer Handling: The max-model-len could simply be 2k+0.5k=2560. However, custom ShareGPT prompt generators extract token ID arrays, detokenize them into raw text strings, and pass them to vLLM. Because detokenization/re-tokenization is non-bijectional, re-tokenizing raw prompt strings in vLLM often produces slightly more tokens than the nominal array length. Adding a 512-token safety buffer to max-model-len eliminated benchmark request rejection failures.
-
Prefill Token Chunking: Configured specifically for prefill optimization. Setting max-num-batched-tokens ~= seq_len * batch_size / att_dp maximizes prefill operational intensity and chunking efficiency across MXUs.
-
Adaptive Token Routing: Token permutation can be executed either via indirect memory indexing, dense matrix multiplication, or offloading to sparse core. Setting ONEHOT_MOE_PERMUTE_THRESHOLD enables one-hot matrix multiplication on TPU MXUs for small batch sizes, executing token permutation at peak hardware matrix speed. For larger batches, the system transitions to SparseCore Gather/Reduce to maximize overlapping.
3. Low-level optimizations
SparseCore gather reduce kernel & tree reduction
Under pure Tensor Parallelism, token routing and reduction operations in MoE layers previously defaulted to unoptimized JAX native execution. Profiling revealed that gather-reduce operations consumed a third of end-to-end execution time.
To resolve this bottleneck, we introduced a SparseCore dense gather-reduce kernel. Furthermore, we optimized the kernel’s accumulation path using a binary tree reduction:
-
Sequential Reduction (Previous): Summing N gathered rows was performed linearly: ((r_0 + r_1) + r_2) + r_3. Each step waited for the preceding addition, creating a critical path latency of O(N).
-
Tree Reduction (Optimized): Restructuring accumulation into parallel pairs: (r_0 + r_1) + (r_2 + r_3 reduced critical path latency from linear O(N) to logarithmic O(log N). This parallelization shortened VPU dependency chains.
GMM tiling optimization
During prefill benchmarking, we identified a gap between realized GMM time and theoretical roofline. The GMM v2 kernel relied on a hard-coded tiling heuristic (t_m = 128) optimal for standard shapes but suboptimal for Mistral 3 prefill shapes.
By running shape-aware tiling tuning, we identified optimal block shapes for Mistral 3 prefill tensors. This shape-aware tuning yielded ~75% roofline v.s ~50% before.. For a quick prototype, we could just hardcode the tiling in the specific code.
Asynchronous scheduling & pipeline overlap
In data parallelism (DP), token payloads must be gathered from all DP ranks across TPU devices. Transferring these payloads back to the host CPU introduced large latency gaps between decode steps.
Activating async_scheduling enabled pipelined TPU-to-CPU token transmission. Overlapping network transmission with ongoing TPU execution effectively hid transfer overhead.
MLA kernel tuning
To optimize kernel performance for specific model deployment configurations, users can utilize the open-source MLA tuning tool available here. In our implementation, MLA parameters were tailored specifically to match the Mistral 3 Large serving setup. By adjusting these parameters to fit your serving environment and running the tuner, you can achieve optimal results for your workload. A complete implementation example is available here.
4. Performance & quality benchmarks
Extended context workloads (20K Input / 0.5k Output)
-
Throughput Win: Achieved 587.07 output tok/s (68.80 req/min), representing a +46.4% throughput improvement over the TPU unoptimized baseline (401.07 tok/s / 47.00 req/min).
-
Latency Reduction: Decreased P50 Inter-Token Latency (ITL) by 28.6% (falling from 48.08 ms down to 34.33 ms, yielding 29.13 tok/s P50 speed).
Standard context workloads (2K Input / 0.5k Output)
-
Throughput Win: Achieved 2,910.16 output tok/s (341.03 req/min), delivering a +48.3% request throughput boost over the unoptimized baseline (1,932.00 tok/s / 230.00 req/min).
-
Latency Reduction: Lowered P50 ITL by 30.0% (falling from 44.94 ms down to 31.45 ms, yielding 31.79 tok/s P50 speed).
Model quality & accuracy neutrality
Across all optimizations, evaluation loss and GSM8K benchmark accuracy remained strictly neutral compared to FP16/BF16 baselines (achieving 91.7% to 92.0% accuracy), matching full-precision quality standards while operating at maximum throughput.
Conclusion & future roadmap
Through systematic roofline gap analysis and low level optimizations, Google’s Ironwood TPUs have proven to be a top-tier platform for serving large MoE models like Mistral 3.
Looking forward, our team is continuing to push Ironwood’s low-concurrency limits by optimizing lower batch size decode. In low-concurrency decode, tokens route sparsely to 128 experts, leaving many experts with 0 or 1 tokens. Operation shifts from compute-bound to a strictly latency- and memory-bound regime. We should consider:
-
Add multi-bucket pre-compilation (e.g., compilation_sizes = {16, 32}) to minimize padding overhead during decode
-
Expert imbalance mitigations: Leverage tensor parallelism or hybrid (TP+EP) in MoE, or duplicate hot experts.
-
Disaggregate serving: Decoupling the prefill and decode phases would enable the application of phase-specific optimization strategies, allowing for more granular tuning of the latency-bound decode regime. Under this configuration, decode performance becomes the primary bottleneck, and we could potentially leverage higher chip counts to further accelerate step-by-step token generation.
-
Disabling XLA Dot Strength Reduction: When adding low-concurrency compilation sizes (such as batch sizes 16 and 32), we noticed XLA compiled dot_general matrix multiplications into elementwise VPU loop fusions rather than convolution fusions. For memory-bound decode workloads at low batch sizes, this strength reduction severely degrades FLOPs utilization. Disabling xla_tpu_enable_dot_strength_reduction forced XLA to preserve standard matrix multiplication, enabling fusion with elementwise operations.