Mixture of experts (MoE) has become one of the defining architectural trends in large-scale AI model training. DeepSeek, Qwen, and Mixtral are examples of MoE models that match or exceed the performance of dense model counterparts at a fraction of the training compute.

MoE models provide efficient training through conditional computation. Instead of one dense feed-forward network (FFN) shared by all tokens, MoE replaces it with many smaller expert networks and a learned router that decides which Top-K experts to activate.

However, making MoE training efficient at scale is challenging. In DeepSeek-V3 training on NVIDIA GB300, an unoptimized baseline achieved just 103 TFLOPS/GPU with inter-GPU communication consuming 84% of accumulated kernel time. With the JAX Python library and NVIDIA Transformer Engine targeted kernel optimizations, that number rose to 1,068 TFLOPS/GPU, a 10.4x improvement.…