AWS boosts MoE kernel performance on SageMaker AI
AWS engineers have outlined kernel and backend changes for fused Mixture-of-Experts models running on Amazon SageMaker AI. The company says the work improves throughput and speeds first-token delivery on NVIDIA H200 GPUs.
Radwan AbdulhadiEditor and technology writer
AWS targets a costly MoE bottleneck
AWS engineers say kernel engineering is one of the largest opportunities for improving Mixture-of-Experts, or MoE, models because their operations account for a substantial share of total compute time. Inefficiencies in those operations can increase both request latency and serving costs. In a post on the AWS Machine Learning Blog, the company describes changes intended to improve fused MoE execution on Amazon SageMaker AI.
MoE is a sparsely activated architecture that expands the capacity of a large language model while limiting the computation performed for each token. A fused MoE system combines token routing, data movement and expert-network calculations into a batched GPU operation rather than handling them through separate sequential loops. That approach can reduce latency, but it also makes kernel implementation and debugging more complex.
Why decode workloads underuse Hopper GPUs
The optimization challenge is especially visible during decoding, when a system processes only one token per user sequence at a time. AWS notes that the resulting token dimension can become very small, often below 64 rows and sometimes equal to a single row. This shape does not align well with the Warp Group Matrix Multiply Accumulate, or WGMMA, instruction protocol used by NVIDIA Hopper GPUs.
WGMMA operates across a cooperative group of 128 threads and uses an asymmetric operand layout. The first operand has a fixed matrix-row dimension of 64, while the other operand uses a more variable column dimension. When the token matrix is thin, it cannot fill that fixed layout. As a result, part of the Tensor Core capacity remains idle while the instruction processes empty spatial slots.
The problem is amplified in fused MoE workloads because the kernel handles the active experts together. With a top-k value of eight, the effective output dimension can grow to tens or hundreds of thousands of parameters, while the decode token dimension remains extremely small. AWS says this contrast creates a layout bottleneck that is more severe than in a conventional dense layer.
Swap-AB reshapes the matrix operation
AWS addressed the mismatch with a Swap-AB transformation that transposes the matrix multiplication before it reaches the PTX assembly layer. The change places the large combined expert dimension in the leading operand position, allowing it to meet the Hopper hardware's 64-row requirement more effectively. The company says this lets Tensor Cores operate with greater local density and fewer unused execution lanes.
The transposed layout also supports more effective asynchronous prefetching. AWS says Hopper's copy engines can move data from shared memory into registers ahead of computation, and that placing the large, continuous expert-weight blocks in the leading slot helps keep the compute pipeline supplied. The optimization produced a reported 2% to 4% end-to-end output-throughput improvement across large-scale MoE models on H200 GPUs.
AWS says the change has been upstreamed to vLLM and is available in version 0.24.0 or later. That makes the optimization relevant to deployments using compatible vLLM releases rather than requiring a separate, private kernel implementation.
Backend routing favors Triton for FP8
Kernel shape is not the only performance variable. AWS also examined how vLLM's MoE oracle selects the backend responsible for a particular quantization scheme. The available paths include Triton and FlashInfer with CUTLASS, among others, and the best choice can vary with the model and workload.
For FP8 tensor-quantized models running on Hopper GPUs, AWS profiling found that the Triton backend consistently outperformed the default FLASHINFER_CUTLASS path. The advantage was most apparent in low-concurrency MoE workloads, where a custom operator-level implementation can avoid some of the limitations of more rigid framework-level matrix-multiplication configurations.
The finding reinforces the importance of workload-specific routing in model serving. Rather than treating one backend as universally optimal, the system can direct different quantization and execution patterns to the implementation that better matches their hardware and concurrency characteristics.
Reported gains for SageMaker AI hosting
AWS combines the Swap-AB layout change, Triton backend routing for FP8 and the elimination of Triton just-in-time recompilation in its performance work. According to the company's summary, the optimizations deliver up to 10% higher throughput and up to 9% faster Time to First Token on H200 GPUs for fused MoE models.
Time to First Token measures how quickly a serving system begins returning generated output, while throughput reflects the amount of work completed over time. Improvements in both measures can affect interactive responsiveness and infrastructure efficiency, although the reported results are tied to AWS's representative MoE model and testing conditions.
The work is presented as guidance for teams hosting models through Amazon SageMaker AI. AWS provides related deployment information in its Amazon SageMaker AI hosting documentation, while the detailed engineering discussion remains in the company's technical blog. The results illustrate how software changes at the kernel and backend layers can complement the hardware capabilities of Hopper GPUs.
Sources
- AWS Machine Learning BlogMixture-of-Experts kernel engineering on Amazon SageMaker AI
Continue chronologically
Related entity coverage
Related stories
AWS Announces SOCI-Indexed Images for Faster SageMaker AI
Amazon Web Services introduces Seekable OCI indexing for Amazon SageMaker AI jobs, cutting container image download times from minutes to seconds.
Fine-Tune Search Agents with Multi-Turn RL on SageMaker AI
Learn how to fine-tune LLM-powered search agents using multi-turn reinforcement learning on Amazon SageMaker AI for improved retrieval quality.
Generate Images and Video with vLLM-Omni on SageMaker AI
Deploy FLUX.2-klein and Wan2.1-VACE models using AWS vLLM-Omni Deep Learning Containers on Amazon SageMaker AI for text-to-image and video generation.
Build Efficient Agents Using Sequence-Level Knowledge Distillation
Learn how to transfer frontier model behavior to lower-cost student models using sequence-level knowledge distillation on Amazon SageMaker AI.
