Mixture of Experts (MoE) architectures are a major trend for large‑scale model training: implementations such as DeepSeek, Qwen and Mixtral demonstrate that MoE can match or exceed dense model performance while using less training compute. Dropless MoE—where every selected token is processed by its assigned expert rather than being dropped or padded—preserves model quality but requires specialized runtime and kernel support to run efficiently at scale.
Challenges in MoE training
Production‑scale MoE training introduces bottlenecks not present in dense models: dynamic token routing, expert dispatch and gather, all‑to‑all communication, and ragged (irregular) GEMMs. Because the router is learned, expert loads can become heavily skewed and vary across batches; token counts per expert are data dependent and non‑uniform, producing ragged tensors that typical tensor libraries are not optimized for.
If dispatch and combine paths are not optimized, communication dominates execution time and GPUs are underutilized. Naive all‑to‑all implementations force GPUs to stall while data arrives.
Dropless vs capacity‑based MoE
- Dropless MoE: every token assigned to an expert is processed regardless of load imbalance. This preserves training data and model quality but demands kernels that accept variable token counts without padding or dropping.
- Capacity‑based MoE: assigns a fixed token budget to each expert and drops or pads overflow. This keeps computation regular and hardware friendly but trades off model quality or efficiency.
MegaBlocks and similar approaches reformulate expert computation as block‑sparse matrix multiplication, enabling experts to operate on varying token counts without dropping or padding, but requiring new block‑sparse kernels and dispatch/combine primitives.
Transformer Engine building blocks for dropless MoE in JAX
To make dropless MoE practical in JAX, Transformer Engine provides several building blocks:
- group‑aware MXFP8 quantization
- MXFP8 grouped GEMM for expert matmuls
- optimized Expert‑Parallel (EP) operations for dispatch and combine
Figure: an expert‑parallel MoE layer across two GPUs — the router assigns tokens to experts, dispatch moves tokens to the experts’ GPUs, grouped MLP runs two grouped GEMMs on variable‑length groups, and combine restores token order.
Optimization 1: Grouped GEMM
In a dense FFN every token uses the same weight matrix. In MoE, the router produces variable token counts per expert, breaking the regular GEMM shape. Previous methods used GEMM loops (which required Device→Host token counts and broke CUDA graphs) or batched GEMMs with worst‑case padding (extra compute). A grouped GEMM handles all expert matmuls in a single kernel with each expert’s actual token count, computing only valid regions.
Transformer Engine’s grouped_gemm / ragged_dot leverages cuBLAS and cuBLASLt to map onto high‑performance NVIDIA GEMM libraries and keep full Tensor Core utilization even with irregular expert shapes. On NVIDIA Blackwell GPUs this path also enables MXFP8 block scaling for expert matmuls using Transformer Engine grouped quantization kernels.
Optimization 2: Expert parallelism to fuse Dispatch and Combine
After the fused router kernels assign tokens, tokens must be permuted and sent across GPUs (Dispatch), and after processing the results must be routed back and accumulated (Combine). A naive serial chain causes GPUs to stall and data to be repeatedly read/written.
Transformer Engine’s EP implementation tightly fuses Dispatch and Combine into an integrated kernel path and uses NCCL EP, a communication backend tuned for the irregular, imbalanced traffic patterns of expert parallel routing. NCCL EP also implements token deduplication so tokens that must be sent to multiple experts traverse the network only once and are replicated on the receiving node, saving bandwidth.
Grouped GEMM addresses internal expert computation; EP handles everything around it.
Additional optimizations
- JAX host offloading: rematerialization APIs allow intermediate activations to be offloaded to host memory, saving high‑bandwidth device memory for DSv3 training (for example, offloading query and value projection results).
- XLA multistreaming collectives: default single‑stream collectives can serialize parallel operations. Multistream collectives let the compiler schedule independent collectives concurrently across separate CUDA streams, overlapping InfiniBand cross‑node transfers with intra‑node NVLink communication. The Latency Hiding Scheduler (LHS) analyzes replica groups to decide safe overlaps and reduce exposed collectives on the critical path.
Measured training impact with JAX + Transformer Engine
- On NVIDIA GB200, an unoptimized JAX baseline for DeepSeek‑V3 training achieved 103 TFLOPS/GPU with inter‑GPU communication consuming 84% of accumulated kernel time.
- Applying Transformer Engine targeted kernel optimizations in JAX raised throughput to 1,068 TFLOPS/GPU — a 10.4× improvement on that baseline.
- The reported end‑to‑end throughput improvement on DeepSeek‑V3 671B was approximately 10×.
- At large scale the system achieved 97% efficiency at 1,024 GPUs on NVIDIA GB300 NVL72 hardware, demonstrating that the communication and kernel optimizations preserve throughput across multirack deployments.
The authors plan further additions such as NVFP4, quantization fused with GEMM, and A2A overlap in future Transformer Engine JAX bindings.
How to get started and reproduce results
The Transformer Engine optimizations are available in the NVIDIA NGC MaxText container with Transformer Engine included. The container from September 9, 2026 (ghcr.io/nvidia/jax:maxtext-2026-09-09) or newer is recommended.
Recommended workflow:
- Start from the MaxText reference configuration and validate correctness on a small MoE model.
- Scale up while tracking step time, TFLOPS/GPU, MFU, grouped GEMM latency, and MoE dispatch/combine latency.
The article provides example MaxText configuration snippets for running DeepSeek‑V3 671B (model_name: "deepseek3-671b", max_target_length: 4096), training settings, TE MoEBlock flags (te_moe_block: true, te_gmm_quantization: "te_mxfp8", ragged_buffer_factor: 2.0, te_ep_overflow_check_every_n_steps: 20, prefuse_moe_weights: true), and cluster parallelism layout for different GPU counts. It also lists XLA GPU flags and environment variables used in their reproductions (for example XLA_PYTHON_CLIENT_MEM_FRACTION: 0.88, CUDA_DEVICE_MAX_CONNECTIONS: 16, and XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB: 180).
Note: the exact DeepSeek‑V3 reproducer configuration is model‑specific and requires tuning; using the TE MoEBlock is not mandatory to benefit from some optimizations.
Takeaways
Dropless MoE maintains model quality while Transformer Engine grouped GEMM and EP kernels make such training efficient at scale. The combined stack produced roughly 10× throughput improvements on DeepSeek‑V3 671B and sustained high multirack efficiency (97% at 1,024 GPUs). These optimizations are packaged in the NVIDIA NGC MaxText container with Transformer Engine for reproduction.
Acknowledgements
The work credits contributions by Abhinav Goel, MD Fahim Faysal Khan, Jane Liu, Terry Sun, Tj Xu, Ming Huang, Chase Roberts and Oleg Goncharov for MoE enablement and optimization in JAX, XLA and Transformer Engine; Artem Polyakov, Ke Wen and Subhadeep Bhattacharya for NCCL EP; and Igor Safanov for cuBLASLt contributions.
