Mixture-of-Experts kernel engineering on Amazon SageMaker AI

This article describes kernel engineering optimizations for Mixture-of-Experts (MoE) models on Amazon SageMaker AI that improve performance and reduce latency.

• Swap-AB optimization transposes matrices to align with NVIDIA Hopper architecture WGMMA constraints, improving throughput by 2-4%
• Backend routing directs FP8 tensor-quantized models to Triton backend instead of CUTLASS for 10% throughput improvement
• Triton JIT recompilation elimination removes unnecessary kernel recompilations on variable parameters, reducing time-to-first-token by 9%
• Fused MoE merges token routing, data shuffling, and expert computations into single GPU operations
• Optimizations upstreamed to vLLM and available in LMI Deep Learning Container

These kernel optimizations collectively deliver significant performance gains for MoE model inference on Amazon SageMaker AI.