MoE Kernel Engineering Optimization Practices on Amazon SageMaker AI

As Mixture-of-Experts (MoE) models gradually become the mainstream architecture for large language models, the bottleneck in inference performance often lies not in the model structure itself, but in the implementation quality of the underlying kernels. The AWS engineering team recently shared a series of kernel-level optimizations they made for fused MoE models on Amazon SageMaker AI, including Swap-AB to improve WGMMA utilization, Triton backend routing to enable FP8 quantization, and eliminating redundant Triton JIT compilation. These optimizations delivered up to 10% throughput improvement and up to 9% improvement in time-to-first-token (TTFT) on H200 GPUs.

Why MoE Inference Is So Sensitive to Kernel Engineering

The core idea behind MoE architecture is to replace a single feed-forward layer with multiple expert networks, activating only a few experts per token. This sparse activation allows model parameter counts to scale dramatically while keeping compute growth relatively manageable. The trade-off, however, is that inference involves a large number of small-scale, irregular computational tasks: token-to-expert routing is dynamic, the number of tokens each expert receives is variable, and the matrix shapes of GEMM operations change with the batch.

These characteristics make it difficult for general-purpose kernels to fully utilize GPU compute. This is especially true on Hopper architecture (such as H200), where WGMMA (Warpgroup Matrix Multiply-Accumulate) instructions have strict requirements on matrix layout and pipelining. If the kernel's memory access pattern or data layout doesn't match, tensor core utilization drops noticeably. As a result, MoE inference performance optimization is largely a kernel engineering battle, not simply model compression or framework tuning.

WGMMA (Warpgroup Matrix Multiply-Accumulate) is an asynchronous matrix multiply-accumulate instruction introduced with NVIDIA's Hopper architecture. It is executed cooperatively by an entire warpgroup (typically 128 threads), can read operands directly from shared memory, and writes to tensor core accumulators. Compared to the previous-generation Ampere mma instructions, WGMMA supports larger matrix tiles and higher throughput, and can be pipelined with asynchronous copy mechanisms such as TMA (Tensor Memory Accelerator). However, its efficient operation depends on strict matrix layouts: operands must reside in shared memory in specific swizzle formats, and the division of labor among threads within a warpgroup must align with the instruction semantics. Once the kernel's memory access pattern, data layout, or synchronization strategy doesn't match, the tensor cores sit idle, and actual compute utilization can fall far below the theoretical peak. This is one of the fundamental reasons why irregular workloads like MoE struggle to run at full utilization on Hopper.

Three Key Kernel Optimizations

Swap-AB: Making WGMMA Run at Full Utilization

In standard GEMM, the A matrix typically represents activations and the B matrix represents weights. But in MoE fused kernels, the scale and layout of the weight matrix relative to activations may not map directly to WGMMA's optimal execution pattern. The Swap-AB strategy adopted by the AWS team essentially swaps the roles of the two operands, allowing the kernel to achieve better data reuse and fewer synchronization overheads in Hopper's WGMMA pipeline.

The value of this change is that it doesn't alter the mathematical equivalence of the model, yet significantly improves tensor core occupancy. For the grouped GEMMs and batched GEMMs that are abundant in MoE, this operand-level adjustment is often more directly effective than adjusting tiling strategies at a higher level.

The mathematical form of GEMM is C = A × B, where A typically represents activations in row-major order and B represents weights. But since matrix multiplication satisfies (A×B)ᵀ = Bᵀ×Aᵀ, swapping the two operands and transposing the result doesn't change the computational semantics—it only changes how data is arranged in memory and registers. On Hopper, WGMMA has different constraints on the source of A and B operands (registers vs. shared memory), major-order direction, and tile granularity. Under certain shapes, A is better suited for activations and B for weights, but in MoE fused kernels it's often exactly the opposite. Swap-AB leverages this equivalence transformation to place the heavier or more regularly laid-out side in the position WGMMA handles better, thereby reducing shared memory transfers and synchronization counts. It is a classic "zero mathematical cost, pure engineering gain" optimization.

Triton Backend Routing: A Practical Path to FP8 Quantization

FP8 quantization is currently an important means of reducing inference costs, but in MoE models, different experts may have different numerical distributions, and a uniform quantization strategy easily leads to precision loss. The AWS team leveraged Triton's backend routing mechanism to let different computational paths reach their own suitable backend implementations, thereby maintaining numerical stability under FP8 quantization.

This approach essentially pushes quantization granularity control down to the kernel scheduling layer. Compared to making global quantization configurations at the model level, backend routing allows finer-grained decisions: which experts or which computation stages are suitable for FP8, and which need to retain higher precision. For production environments pursuing extreme throughput, this flexibility is necessary.

FP8 is an 8-bit floating-point format, with common variants being E4M3 (4-bit exponent, 3-bit mantissa—smaller dynamic range but higher precision) and E5M2 (5-bit exponent, 2-bit mantissa—larger range but lower precision). Compared to FP16/BF16, FP8 can halve memory footprint and bandwidth requirements, and achieve higher compute throughput on GPUs that support FP8 tensor cores, such as Hopper. The cost, however, is that mantissa bits are extremely limited, making it sensitive to numerical distributions: if the dynamic range of a set of weights exceeds the representable range of the chosen format, overflow or underflow occurs. In MoE, different experts have varying training data and activation distributions, so uniform quantization can easily cause individual experts to suffer precision collapse. Triton's backend routing allows different branches of the same kernel to compile to different target implementations (such as different quantization formats or precision paths), thereby selecting the most appropriate numerical strategy per expert or per computation stage.

Eliminating Redundant Triton JIT Compilation

Triton's JIT compilation mechanism is very convenient during development, but in production inference, if a new input shape triggers recompilation every time, it introduces unpredictable latency jitter. MoE's dynamic routing is precisely prone to this problem: different batch routing results can lead to changes in kernel parameters, which in turn trigger repeated compilation.

The AWS team optimized this issue, reducing the frequency of JIT recompilation. Although the original article doesn't disclose specific implementation details, judging from the results, this optimization contributed significantly to TTFT improvement—TTFT is critical to user experience in online serving, and any compilation-level jitter gets amplified.

Triton is a Python-like kernel DSL for GPUs that uses just-in-time (JIT) compilation: developers describe kernel logic in Python, and Triton generates and compiles PTX/CUBIN on first invocation based on tensor shapes, data types, and other "specialization keys." This mechanism enables fast development iteration, but each new shape combination may trigger a full compilation, taking anywhere from tens of milliseconds to several seconds. MoE's dynamic routing causes the number of tokens received by each expert and each batch to constantly change, making shape combinations highly divergent and leading to low compilation cache hit rates. In production inference, this sporadic compilation directly manifests as long-tail latency and first-token jitter. Common mitigation approaches include: pre-compiling and caching common shapes, bucketing and merging shapes, fixing maximum shapes with masking, and warming up all possible kernel variants at deployment time.

Performance Gains and Engineering Insights

With the three optimizations combined, up to 10% throughput improvement and up to 9% TTFT improvement were achieved on H200 GPUs. This number isn't startling on its own, but considering that MoE inference kernel optimization space is often obscured by framework-level abstractions, the value of these low-level changes lies in the fact that they can be orthogonally stacked with upper-layer optimizations.

From an engineering perspective, this set of optimizations sends several signals. First, the tensor core capabilities of Hopper and subsequent architectures need to be unlocked through careful kernel layout—there's still a gap between "good enough" and "efficient" for general-purpose kernels. Second, quantization is not simply trading precision for speed; landing FP8 in MoE requires sufficient scheduling flexibility at the kernel level. Third, latency stability in production environments is as important as average throughput—seemingly marginal issues like JIT compilation get amplified at scale.

For teams deploying MoE models on SageMaker AI, these optimizations will likely appear in service performance in managed form, but understanding the underlying principles helps make more informed decisions when customizing kernels or choosing instance types. Kernel engineering is once again becoming the core battleground for large model inference optimization, and MoE has simply pushed this trend to the forefront.

One-line Summary

AWS made three kernel optimizations for fused MoE models on SageMaker AI, achieving up to 10% throughput improvement on H200.

Paragraph Summary

This article introduces kernel-level optimizations made by the AWS engineering team for fused MoE models on Amazon SageMaker AI. Due to dynamic routing and irregular shapes, MoE inference struggles to fully utilize GPU compute with general-purpose kernels—especially on Hopper architecture, where WGMMA has strict layout and pipeline requirements. The team adopted three optimizations: Swap-AB swaps GEMM operand roles to improve tensor core utilization; Triton backend routing provides finer-grained precision control for FP8 quantization; and eliminating redundant Triton JIT compilation reduces latency jitter in production. Combined, these three optimizations achieve up to 10% throughput improvement and up to 9% TTFT improvement on H200, demonstrating that kernel engineering is once again becoming the core battleground for large model inference optimization.