Fused MoE Kernels

In a naive implementation, a Mixture of Experts forward pass requires launching multiple separate GPU kernels: one for the gating network, one to permute/sort tokens by expert, and several for the sparse matrix multiplications of each expert. Each kernel launch incurs significant latency from driver scheduling and setup. A fused kernel executes this entire sequence as one monolithic operation, eliminating this serial overhead and drastically reducing scheduler contention on the GPU.

The most significant performance gain comes from eliminating intermediate writes and reads to high-bandwidth memory (HBM). Without fusion, the outputs of each step (router scores, token permutation indices, intermediate expert outputs) must be written to global memory and then read by the next kernel. This memory-bound process is a major bottleneck. A fused kernel keeps these intermediate tensors in fast on-chip registers or shared memory, performing the entire computation in a single pass and dramatically cutting HBM traffic, which is often the limiting factor for throughput.

Fused kernels allow for sophisticated, warp-level load balancing that is impossible with separate kernels. They can:

• Dynamically distribute the variable workload of sparse expert matrices across streaming multiprocessors (SMs).
• Implement fine-grained synchronization (e.g., warp-level primitives) to handle the irregular computation graph where different experts process different numbers of tokens.
• This leads to near-ideal GPU utilization, hiding latency and ensuring all compute units are actively processing the sparse computation, not waiting on memory or other warps.

A core step in MoE is gathering tokens assigned to the same expert from scattered memory locations for efficient batched computation. A separate all-to-all communication or a sort/permute kernel is typically required. In a fused design, this gather/scatter logic is embedded within the kernel's indexing arithmetic. Tokens are logically routed and processed in-place or with minimal, coordinated data movement, avoiding the launch of a dedicated, memory-intensive permutation kernel. This is critical for performance in expert parallelism scenarios across multiple GPUs.

Modern inference serving uses continuous batching, where requests of varying sequence lengths are combined into a single batch. In MoE, this results in a highly irregular workload per expert. Fused kernels are uniquely suited to this dynamism because their internal logic (e.g., using parallel prefix sums for workload assignment) can adapt to variable token counts per expert at runtime without recompilation or suboptimal padding. This provides consistent low latency in production serving systems like vLLM with MoE support.

From a systems perspective, a single fused kernel call simplifies the execution graph, reducing pressure on the framework's (e.g., PyTorch's) dispatcher and Python interpreter. Instead of managing dependencies and launches for many small ops, the runtime schedules one heavyweight, efficient kernel. This reduces CPU overhead and context switching, which is vital for achieving high queries per second (QPS) in server environments. It also simplifies integration into compiler stacks like TorchInductor or OpenAI's Triton.

The primary goal of a fused MoE kernel is to eliminate the performance cost of launching multiple separate GPU kernels and moving intermediate data between global memory and registers. A naive implementation executes distinct steps:

• Compute router scores and perform top-k selection.
• Permute and gather tokens based on routing decisions.
• Launch separate GEMM (General Matrix Multiply) kernels for each expert's activated weight block.
• Scatter the results back to the original token order.

A fused kernel combines these steps into one monolithic CUDA kernel. This reduces kernel launch latency, minimizes global memory traffic by keeping data in faster cache hierarchies (L1/L2, shared memory), and allows for more aggressive compiler optimizations across the entire operation.

At the heart of the fused kernel is an optimized sparse matrix multiplication. Unlike a dense GEMM that computes over an entire matrix, an MoE layer only multiplies by the weight blocks of the top-k selected experts.

Key implementation techniques include:

• Block-Sparse Representation: Expert weights are stored in contiguous blocks. The kernel uses a lookup table of pointers to these blocks, indexed by the router's output.
• Warp-Level Specialization: Different warps (groups of 32 threads) within a CUDA block are assigned to compute on different experts or token groups, enabling fine-grained parallelism.
• Hierarchical Reduction: Partial results from processing different token-expert assignments are efficiently reduced within shared memory before being written to the final output tensor.

This approach avoids the overhead of a generic sparse linear algebra library, which is not optimized for the specific, structured sparsity pattern of MoE routing.

A major bottleneck in MoE is the data movement required to group tokens by their assigned expert for efficient computation (the "all-to-all" step in distributed settings). A fused kernel optimizes this by performing the permutation in-place within GPU registers and shared memory.

Instead of writing intermediate tensors to global memory, the kernel:

1. Loads a tile of input tokens and their routing assignments.
2. Uses warp shuffle instructions and shared memory buffers to locally reorder the tile so that tokens destined for the same expert are contiguous.
3. Feeds this reordered data directly into the subsequent block-sparse GEMM steps.

This eliminates separate gather/scatter kernels and their associated memory read/write cycles, significantly reducing latency, especially for smaller batch sizes.

Fused MoE kernels are implemented as custom operators within leading inference and training systems:

• vLLM: Integrates fused MoE kernels (e.g., for Mixtral models) within its PagedAttention scheduler, allowing for efficient memory management and continuous batching of MoE requests.
• NVIDIA TensorRT-LLM: Provides highly optimized fused MoE plugins for NVIDIA GPUs, leveraging Tensor Cores and advanced CUDA features for maximum throughput.
• PyTorch: Can utilize custom torch.autograd.Function implementations or leverage the Triton compiler (e.g., via triton.language) to write high-performance fused MoE kernels in Python-like syntax.
• DeepSpeed: Its MoE implementation includes optimized kernels for training, handling the complexities of expert parallelism and gradient synchronization.

The use of fused kernels directly targets the inference latency and throughput pillars. Benchmarks on models like Mixtral 8x7B show substantial improvements:

• Kernel Launch Overhead: Reduction from launching 5-10+ kernels per MoE layer to just 1 or 2.
• Memory Bandwidth: Can reduce global memory transactions by 30-50% for the MoE operation by fusing load/compute/store phases.
• End-to-End Latency: Can improve latency per token by 1.5x to 2x for the MoE layers, which are often the computational bottleneck in these models.

These optimizations are critical for achieving low-latency, high-throughput serving, making large MoE models viable for production use cases.

Building a robust fused MoE kernel involves navigating several engineering challenges:

• Variable Workload: The number of tokens assigned to each expert varies dynamically per request and batch. The kernel must handle this load imbalance without thread divergence causing performance cliffs.
• Expert Capacity: Must efficiently implement the capacity factor logic, handling tokens that exceed an expert's capacity (e.g., via buffering or a fallback path).
• Numerical Precision: Support for mixed precision inference (FP16, BF16, INT8) requires separate kernel variants or dynamic precision selection.
• Hardware Generality: Writing a kernel that performs well across different GPU architectures (e.g., Ampere, Hopper) and memory hierarchies is complex. Often, multiple kernel variants are maintained for different scenarios (e.g., small vs. large batch size).

A neural network architecture where the model is composed of multiple sub-networks (experts), and a routing mechanism dynamically selects a sparse subset of these experts to process each input. This enables massive parameter counts (e.g., over a trillion) with a conditional, much lower computational cost per token, as only a few experts are active for any given input.

The fundamental property of conditional computation models like MoE. Instead of a dense activation where all model parameters are used for every input, sparse activation means only a small, dynamically chosen subset of the model's total parameters is computed. This is the core mechanism that makes large MoE models computationally feasible during inference.

The gating network (or router) is a small, trainable component that decides expert assignment. Top-k gating is the most common strategy:

• The router outputs a score for each expert per token.
• It selects the k experts with the highest scores (e.g., top-2 in Mixtral).
• Only these k experts are activated for that token, enforcing sparsity. Variants like Noise Top-k add noise during training to improve load balancing.

A model parallelism strategy designed for MoE models at scale. Different experts are placed on different devices (e.g., GPUs). After the router decides assignments, an All-to-All communication collective is performed:

• Tokens are scattered from all devices to the devices hosting their assigned experts.
• Experts process their assigned tokens in parallel.
• Results are gathered back via another All-to-All. This is distinct from tensor or pipeline parallelism and is essential for distributing ultra-large MoE models.

A critical challenge in training MoE models. Without intervention, the router can collapse, always selecting the same few experts (expert starvation). Techniques to prevent this include:

• Auxiliary Load Balancing Loss: An added loss term that penalizes uneven routing distributions.
• Capacity Factor: A buffer hyperparameter that limits tokens per expert to prevent overflows and encourage spreading.
• Proper load balancing is necessary for training stability and to utilize all model parameters effectively.

The foundational mathematical operation that fused MoE kernels optimize. In a standard MoE layer (e.g., an MoE-FFN), the weight matrix is a large, block-sparse matrix where each block corresponds to an expert. A naive implementation would:

1. Select tokens for each expert.
2. Perform many small, independent dense matrix multiplications for each expert batch. Fused kernels combine steps 1 & 2 and the subsequent permutation of results into a single, highly optimized sparse GEMM operation, minimizing memory movement and kernel launch overhead.