AWS details fused MoE kernel optimizations on SageMaker, up to 10% throughput gain
An AWS Machine Learning Blog post describes several fused Mixture-of-Experts (MoE) kernel optimizations that AWS developed with the vLLM community and open sourced. A Swap-AB transposition that moves the large expert dimension into the leading WGMMA operand slot delivered a 2–4% end-to-end output throughput improvement on H200 and was upstreamed to vLLM 0.24.0 and later. Routing FP8 tensor-quantized MoE workloads to the Triton backend instead of FlashInfer+CUTLASS added about 10% throughput (available in the AWS LMI v28 container), while adding do_not_specialize annotations to variable-length parameters removed repeated Triton JIT recompilations and improved time to first token by 9% on H200 at 8-way tensor parallelism (vLLM 0.22.0 and later). AWS says the optimizations together yield up to 10% higher throughput and up to 9% faster TTFT.
Teams self-hosting MoE inference can adopt mostly configuration-level changes to cut decode latency and cost.
Mixture-of-Experts (MoE) is a sparsely-activated architecture that scales large language model (LLM) capacity while keeping per-token compute low. Because its operations account for a large share of total compute time, inefficiencies there directly raise request latency and cost. Therefore, optimizing MoE operations through kernel engineering offers the largest opportunity for performance gains.
Fused MoE is a key optimization technique for the MoE architecture. A fused MoE system merges token routing, data shuffling, and specialized network calculations into a single batched operation on the GPU. This replaces executing those routing and expert computations as separate, sequential loop operations. While this provides latency advantages, it also brings complexity in kernel implementations and debugging capability. High GPU utilization for this architecture requires hardware-specific optimizations.
In this post, we describe the different kernel optimizations we implemented for fused MoE models and show you how you can benefit from this work. We used a representative MoE model as our primary example.
Swap-AB for WGMMA utilization
To maximize tensor processing throughput, the NVIDIA Hopper architecture relies on the Warp Group Matrix Multiply Accumulate (WGMMA) instruction protocol, which executes collectively across a cooperative warpgroup of 128 threads. This operand asymmetry is fundamentally an instruction set architecture (ISA)-level property. In the underlying wgmma.mma_async assembly definition, the hardware fixes the matrix row dimension (M) at a constant scale of 64 to define its parallel thread-mapping topology. The column dimension (N) remains highly variable. Operand A strictly maps across this fixed M dimension (using registers or shared memory), and Operand B maps across the variable N dimension (through a shared memory descriptor). These operand slots are therefore structurally non-interchangeable. Consequently, peak hardware tile utilization is governed by these rigid shape encodings, requiring that the M dimension equals or exceeds 64 to prevent compute execution lanes from sitting idle.
#Amazon Web Services#Amazon SageMaker AI#vLLM#NVIDIA H200#NVIDIA Hopper#Triton#FlashInfer#CUTLASS