Mixture-of-Experts kernel engineering on Amazon SageMaker AI

Amazon SageMaker AI has introduced kernel engineering improvements for Mixture-of-Experts (MoE) models, focusing on fusing token routing, data shuffling, and expert calculations into a single batched GPU operation. The team addressed hardware inefficiencies by transposing matrices to match NVIDIA Hopper’s Warp Group Matrix Multiply Accumulate (WGMMA) instruction layout, ensuring the large expert dimension occupies the primary operand slot and saturating the 64‑row hardware requirement. This change also activates asynchronous prefetching from shared memory, reducing idle cycles. The optimizations produce a 2–4 percent end‑to‑end throughput improvement on H200 GPUs and are upstreamed to vLLM version 0.24.0. Additionally, the choice of backend influences performance; profiling shows that for FP8 tensor‑quantized models on Hopper, the Triton backend consistently outperforms the default FlashInfer+CUTLASS path. Triton achieves this by consolidating the entire pipeline—FP8 scaling, matrix multiplication, and expert aggregation—into a single hardware pass, pulling data once from high‑bandwidth memory and processing it entirely within local SRAM registers, thereby eliminating costly global VRAM round‑trips. These kernel and backend optimizations collectively enhance MoE model efficiency, lowering request latency and operational cost for large‑scale language model deployments.