AWS described three optimizations for serving fused Mixture-of-Experts (MoE) models on Amazon SageMaker AI, reporting up to 10 percent higher throughput and up to 9 percent faster Time to First Token (TTFT, time until the first generated token). The work was developed with the vLLM community, open sourced, and applied in the LMI Deep Learning Container on SageMaker AI.

The first optimization swaps the operands in matrix multiplication to better fit NVIDIA Hopper GPUs’ WGMMA instruction layout. WGMMA uses a fixed 64-row M dimension, while MoE decoding often processes one token per user sequence, leaving M at 1 and many Tensor Core execution lanes idle. In fused MoE, a topk=8 setting combines selected experts’ intermediate dimensions into a very large N dimension. Transposing the multiplication puts that larger dimension in the first operand, meeting the 64-row layout and allowing asynchronous prefetching. AWS reports a 2–4 percent end-to-end output-throughput improvement across large-scale MoE models on H200 GPUs. The change is upstream in vLLM version 0.24.0 or later.

The second change routes FP8 tensor-quantized MoE computations on Hopper GPUs to Triton instead of the default FlashInfer+CUTLASS path. AWS profiling found Triton performed better in low-concurrency workloads because it fuses scaling, matrix multiplication, and expert aggregation into one GPU pass, avoiding intermediate global-memory read and write steps. Its dynamic scheduling also handles thin, variable token batches such as M < 64 or M = 1 more effectively than CUTLASS’s default tiling. AWS says changing the routing configuration, without modifying kernel code, raised end-to-end throughput by 10 percent for FP8 quantized fused MoE models on H200. The change is available in the AWS Large Model Inference (LMI) v28 container.

The third optimization prevents Triton from recompiling a kernel whenever it encounters a new token count. The _fused_inv_rope_fp8_quant_per_head kernel had been recompiling for each new num_tokens value; adding do_not_specialize annotations to variable-length parameters eliminated those repeated just-in-time compilations. AWS reports a 9 percent TTFT improvement on H200 at 8-way tensor parallelism, and says the change also avoids other GPUs waiting while one GPU recompiles. This optimization is upstream in vLLM version 0.22.0 or later. Together, the changes deliver up to 10 percent higher throughput and up to 9 percent faster TTFT in the LMI container on SageMaker AI.

AWS described three optimizations for serving fused Mixture-of-Experts (MoE) models on Amazon SageMaker AI, reporting up to 10 percent higher throughput and up to 9 percent faster Time to First Token (TTFT, time until the first generated token). The work was developed with the vLLM community, open sourced, and applied in the LMI Deep Learning Container on SageMaker AI.

The first optimization swaps the operands in matrix multiplication to better fit NVIDIA Hopper GPUs’ WGMMA instruction layout. WGMMA uses a fixed 64-row M dimension, while MoE decoding often processes one token per user sequence, leaving M at 1 and many Tensor Core execution lanes idle. In fused MoE, a topk=8 setting combines selected experts’ intermediate dimensions into a very large N dimension. Transposing the multiplication puts that larger dimension in the first operand, meeting the 64-row layout and allowing asynchronous prefetching. AWS reports a 2–4 percent end-to-end output-throughput improvement across large-scale MoE models on H200 GPUs. The change is upstream in vLLM version 0.24.0 or later.

The second change routes FP8 tensor-quantized MoE computations on Hopper GPUs to Triton instead of the default FlashInfer+CUTLASS path. AWS profiling found Triton performed better in low-concurrency workloads because it fuses scaling, matrix multiplication, and expert aggregation into one GPU pass, avoiding intermediate global-memory read and write steps. Its dynamic scheduling also handles thin, variable token batches such as M < 64 or M = 1 more effectively than CUTLASS’s default tiling. AWS says changing the routing configuration, without modifying kernel code, raised end-to-end throughput by 10 percent for FP8 quantized fused MoE models on H200. The change is available in the AWS Large Model Inference (LMI) v28 container.

The third optimization prevents Triton from recompiling a kernel whenever it encounters a new token count. The _fused_inv_rope_fp8_quant_per_head kernel had been recompiling for each new num_tokens value; adding do_not_specialize annotations to variable-length parameters eliminated those repeated just-in-time compilations. AWS reports a 9 percent TTFT improvement on H200 at 8-way tensor parallelism, and says the change also avoids other GPUs waiting while one GPU recompiles. This optimization is upstream in vLLM version 0.22.0 or later. Together, the changes deliver up to 10 percent higher throughput and up to 9 percent faster TTFT in the LMI container on SageMaker AI.