AWS Details Two MoE Kernel Tweaks That Lift SageMaker AI Throughput on H200

Why MoE kernels matter

Mixture-of-Experts is a sparsely activated architecture that scales the capacity of large language models while keeping per-token compute low. AWS says MoE operations account for a large share of total compute time, so inefficiencies there raise request latency and cost directly. Kernel engineering on these operations therefore offers the largest opportunity for performance gains, according to the company.

A key technique is fused MoE. It merges token routing, data shuffling and specialized network calculations into a single batched operation on the GPU, instead of running routing and expert computations as separate, sequential loops. AWS notes that this brings latency advantages but also makes kernel implementation and debugging more complex, and that high GPU utilization needs hardware-specific optimizations. The post uses a representative MoE model as its main example.

Swap-AB: fixing a mismatch with Hopper Tensor Cores

The first optimization targets the Warp Group Matrix Multiply Accumulate (WGMMA) instruction used by the NVIDIA Hopper architecture. AWS explains that the hardware fixes the matrix row dimension (M) at 64, while the column dimension (N) is variable. Peak utilization therefore requires M to equal or exceed 64.

That is a problem in the decode phase, where the system processes one token per user sequence at a time. The token row count M is then often below 64 and frequently just 1. In the standard layout, matrix A becomes a thin vector, for example 1 x 4096, and cannot fill the 64-row bucket. AWS says this leaves a large portion of Tensor Core capacity starved.

The effect is stronger in MoE layers. In a fused MoE kernel, N is the number of experts chosen multiplied by the expert intermediate size. For a model running a topk=8 routine, AWS says the effective N dimension grows to tens or hundreds of thousands of parameters.

The fix relies on matrix transposition: computing B^T x A^T instead of A x B, which leaves the mathematical result unchanged. This moves the large expert dimension into the first operand slot. According to AWS, it fills the 64-row threshold and also lets Hopper's hardware copy engines prefetch data asynchronously from shared memory into registers. The company reports a 2, 4 percent end-to-end output throughput improvement across large-scale MoE models on H200. The change has been upstreamed to vLLM and is available in version 0.24.0 or later.

Switching the backend for FP8 models

The second gain comes from choosing a different backend. The vLLM MoE oracle routes quantization schemes to backends such as Triton or FlashInfer+CUTLASS. AWS says profiling showed that for FP8 tensor-quantized models on Hopper, the Triton backend consistently outperforms the default FLASHINFER_CUTLASS path.

AWS attributes this to how the two handle the computation. CUTLASS treats FP8 tensor scaling, matrix multiplication and expert aggregation as discrete operations. Triton consolidates the pipeline into a single hardware pass: data is read from global high-bandwidth memory once, processed in local SRAM registers, then written to the final destination. AWS adds that CUTLASS is natively optimized for dense, large and predictable matrix dimensions, so its default tiling and thread mapping become inefficient for the ultra-thin batches (M < 64 or M = 1) typical of decode. Triton, by contrast, maps small, fragmented token arrays dynamically across the active expert weights.

Updating the routing policy is a pure configuration change with no kernel code modified. AWS still reports a 10 percent end-to-end throughput improvement for FP8 quantized fused MoE models on H200.

What it means for users

For teams serving MoE models on Hopper-class GPUs, the source points to two practical steps. The transposition optimization is available in vLLM 0.24.0 or later, and the backend change requires only a configuration update. All figures are AWS's own measurements on H200 and apply to the workloads it describes, so results on other models or hardware may differ.