Mixture-of-Experts kernel engineering on Amazon SageMaker AI
aws ·
This article details kernel engineering optimizations for Mixture-of-Experts (MoE) models on Amazon SageMaker AI, focusing on performance gains for large language model inference. The authors describe three key techniques: Swap-AB transposition to align matrix dimensions with NVIDIA Hopper’s WGMMA hardware requirements, backend routing adjustments to use Triton instead of CUTLASS for FP8 quantization, and the elimination of unnecessary Triton JIT recompilations to reduce latency. These optimizations address specific bottlenecks in MoE architectures, such as thin batch sizes during decoding and variable token counts. Implementing these changes resulted in up to 10% higher throughput and a 9% improvement in Time to First Token (TTFT) on H200 GPUs. Several of these enhancements have been upstreamed to the vLLM open-source community. The work demonstrates how deep hardware-specific tuning can significantly enhance efficiency in sparse activation models.
aws