AWS releases MoE inference changes for SageMaker LMI and vLLM

AWS has released three changes for serving mixture-of-experts models on NVIDIA Hopper GPUs: a backend-routing policy in its SageMaker Large Model Inference v28 container and two code changes contributed to vLLM. Operators now have new ways to address underused tensor-core capacity and repeated compilation, though the accompanying performance figures come from AWS and an upstream contributor, not independent tests.

A fused MoE operation rolls token routing, data movement and expert-network calculations into one batched GPU operation. The AWS engineering post identifies three sources of serving overhead: matrix layout during token-by-token decoding, backend selection for FP8 models and Triton recompilation when request shapes change. Each change has its own availability floor; they are not a single software release. DataPhoenix has separately covered AWS’s multi-node vLLM inference stack for Amazon EKS.

The first vLLM change, swap-AB, applies when a fused-MoE matrix multiplication has a token dimension below 64. Hopper’s WGMMA instructions map the leading matrix dimension to a fixed 64-row layout, meaning a small decode batch can leave execution capacity idle.

The merged swap-AB change transposes the operation from A[M,K] × B[K,N] to B[N,K] × A[K,M]. That puts the larger expert-output dimension first without altering the mathematical result. AWS reports a 2% to 4% increase in end-to-end output-token throughput across large MoE models on H200 GPUs, but the opened evidence contains no independent reproduction. The company says swap-AB is available in vLLM 0.24.0 or later; the pull request was merged before the v0.24.0 release.

The LMI-specific change modifies routing policy, not kernel code. For FP8 tensor-quantized fused-MoE workloads on Hopper, LMI v28 chooses the Triton backend instead of the default FlashInfer+CUTLASS path. AWS attributes a 10% end-to-end throughput improvement on H200 to the configuration change. It did not publish a complete model-and-workload matrix or a before-and-after results table for the figure, and the opened sources provide no independent reproduction.

The third change tackles just-in-time compilation in vLLM’s fused inverse-RoPE and FP8-quantization kernel. Because Triton can specialize integer arguments, a previously unseen token count can trigger another compilation. The merged JIT change adds do_not_specialize handling for variable-length values, preventing them from creating redundant compiled variants. The vLLM 0.22.0 release notes include the change, making v0.22.0 its shipped version floor.

AWS says the JIT change made time to first token 9% faster on an eight-H200 setup with eight-way tensor parallelism. The contributor-posted run used DeepSeek-V4-Flash, a ShareGPT workload, maximum concurrency of eight, 480 prompts and 300 output tokens. Mean time to first token dropped from 370.14 milliseconds to 338.37 milliseconds, or about 8.6%, while p99 time to first token climbed from 632.74 milliseconds to 790.74 milliseconds. This single run supports the rounded mean improvement for that setup, but it does not show an improvement in tail latency.