Title: SNaX: sparse narrow accelerated mixture of experts
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: Mixture of Experts, GPU, kernel

Abstract: Mixture of Experts (MoE) models have emerged as the de-facto architecture for scaling up language models without significantly increasing the computational cost. Existing MoE methods optimize system efficiency or model architecture independently. We show that as MoE models get more granular and sparser, they become more memory-bound, and jointly optimizing the algorithms and the kernel design leads to a major improvement in MoE training throughput. We first propose a memory-efficient algorithm to compute the forward and backward of MoE with minimal activation saved. We then design GPU kernels that overlap memory IO latency with compute, benefiting all MoE architectures. Finally, we propose a novel "token rounding" method that minimizes the wasted compute brought by tile quantization. As a result, our method SNaX reduces 45% activation memory and has 1.87x compute throughput improvement on Hopper GPUs compared to state-of-the-art BF16 MoE kernel for a fine-grained 7B MoE. Concretely, SNaX on 64 H100s achieves almost the same total throughput as ScatterMoE on 96 H100s for a 7B MoE model training with token-choice rounding while training with FSDP-2. Under high MoE sparsity settings, our tile-aware token rounding algorithm yields an additional 1.18x speedup on kernel execution time compared to vanilla top-K routing while maintaining similar downstream performance.

TL;DR: We propose an algorithm and system codesign for fine-grained and sparse MoE.