Mixture of experts Inference and serving·14 Oct 2026

AWS details MoE kernel fixes for SageMaker, upstreams them to vLLM

Three upstreamed vLLM fixes squeeze more throughput from Hopper GPUs

The facts

Swap-AB reorders matrix operands to fit Hopper's WGMMA 64-row layout, giving 2-4% throughput gains on H200; upstreamed to vLLM 0.24.0+.
Routing FP8 tensor-quantized MoE models to vLLM's Triton backend instead of the default FlashInfer+CUTLASS path yields a 10% throughput gain on H200, available in AWS LMI v28.
Adding do_not_specialize annotations stops Triton recompiling the _fused_inv_rope_fp8_quant_per_head kernel on every new token count, cutting Time to First Token 9% on H200 at 8-way tensor parallelism (vLLM PR #42849, available from vLLM 0.22.0).
AWS says the combined changes, built into the LMI Deep Learning Container on SageMaker AI, give up to 10% higher throughput and up to 9% faster TTFT.

AWS has detailed three GPU kernel optimizations for serving Mixture-of-Experts (MoE) models on its SageMaker AI platform. Each targets a specific inefficiency in how fused MoE layers run on Nvidia's Hopper architecture.

The first, called Swap-AB, reorders matrix multiplication operands. The large expert-output dimension now fills Hopper's fixed 64-row WGMMA instruction layout, instead of the tiny single-token batch typical of decoding. AWS says this produced a 2–4% end-to-end throughput improvement on H200 GPUs. The change has been upstreamed to vLLM and is available from version 0.24.0.

The second fix is a routing change rather than new kernel code. AWS found that for FP8 tensor-quantized models on Hopper, vLLM's default FlashInfer+CUTLASS backend underperforms the Triton backend. Triton fuses scaling, matrix multiplication and expert aggregation into a single pass through GPU SRAM, instead of repeated round-trips to high-bandwidth memory. Switching the default routing policy — a configuration change with no kernel modification — delivered a 10% end-to-end throughput gain for FP8 MoE models on H200. It is available now in AWS's Large Model Inference (LMI) v28 container.

The third optimization addresses Triton's just-in-time compiler, which recompiles kernels whenever it sees a new integer argument value, such as a changed token count. AWS identified one kernel, _fused_inv_rope_fp8_quant_per_head, recompiling on every new num_tokens value during MoE serving. Adding do_not_specialize annotations to its variable-length parameters eliminated the recompilations. That cut Time to First Token by 9% on H200 at eight-way tensor parallelism. The fix was merged upstream in vLLM pull request #42849 and is available from vLLM version 0.22.0.

AWS says the combined optimizations, now built into the LMI Deep Learning Container on SageMaker AI, deliver up to 10% higher throughput and up to 9% faster first-token latency.

Why it matters

These are upstream vLLM and LMI container changes, not SageMaker-exclusive tricks. Anyone self-hosting MoE models on Hopper GPUs can capture the same gains by updating vLLM or switching backend routing. The compiler-recompilation fix is the more interesting lesson here. Treating the JIT compiler itself as an optimization target can be worth more than writing new kernel code.

Primary source

AWS Machine Learning — aws.amazon.com

Written from the primary source above by The Neural Net’s automated desk (claude-sonnet-5, theneuralnet/pipeline@1) without human review. How this works. Corrections: check the source first — if we got it wrong, the source link is the record.