MoE

Foundation Model Kernel Optimization in 2026: A Field Guide Across Dense, MoE, Multimodal, and Diffusion Most “kernel optimization” conversations still start with FLOPs. In 2026 that is usually the wrong first question. Foundation-model runtime is dominated by HBM traffic, KV cache, temporary tensors, collectives, dynamic permutation, and launch overhead. FlashAttention, fused linear–cross-entropy, paged KV, MoE grouped GEMM, and DeepEP-style dispatch all share one essence: do not materialize intermediates, or make each byte travel once. ...

Large MoE Performance: The Three Walls After Sparsity Sparsity made MoE cheap on paper. At production scale it made training harder than dense: total parameters grow with E, per-token FLOPs grow with k, and the gap between those two numbers is exactly where systems break. The useful framing is not “optimize the MoE kernel.” It is the one NVIDIA’s Megatron-Core MoE report uses (arXiv:2603.07685): Memory, Communication, and Compute Efficiency are three coupled walls. Push on one and pressure shows up in another. ByteDance’s MegaScale-MoE (arXiv:2505.11432) proves the same thesis from the other direction — on 1,440 Hoppers, communication was ~44% of forward time before their redesign, and fixing parallelism + overlap delivered 1.88× over Megatron-LM. ...

MoE Internals: DeepSpeed-Megatron Expert Parallel Implementation DeepSpeed-Megatron MoE is easy to misread if you only inspect one initialization function. Megatron builds the usual TP, PP, and DP topology. DeepSpeed then wraps the model, finds MoE modules, and gives those modules expert-parallel groups. The topology is split across the training framework and the MoE layer implementation. This post connects the pieces: deepspeed.initialize(), EP groups, expert-DP groups, MoELayer, and the difference between a DeepSpeed MoE wrapper and a Megatron-integrated SwitchMLP-style module. For routing, capacity, and All-to-All fundamentals, start with MoE Parallelism Principles. ...

MoE Parallelism Principles: GShard, Expert Parallel, and All-to-All Mixture-of-Experts is sparse compute wrapped in a communication problem. The model idea is simple: replace a dense FFN with many expert FFNs and route each token to a few of them. The distributed-training problem starts when those experts live on different GPUs. Shazeer et al. introduced the sparsely-gated MoE layer in Outrageously Large Neural Networks. GShard made the pattern practical at scale with top-2 routing, capacity, load-balancing loss, and all-to-all dispatch in GShard. Switch simplified the router to top-1 in Switch Transformers. This post explains the shared systems model underneath them. ...