Custom Kernels for Mixture-of-Experts

The padding trap

A standard transformer feed-forward layer receives a fixed-shape tensor and calls one big matrix multiplication. Mixture-of-Experts (MoE) breaks this assumption entirely. A router sends each token to one or two of E experts; the number of tokens landing on any given expert varies every forward pass. If expert 3 receives 47 tokens this step and expert 7 receives 312, a naive implementation pads every expert's batch to capacity = C tokens and then calls E separate GEMMs, one per expert.

That approach wastes compute in two distinct ways. First, padded zeros still burn FLOPs and memory bandwidth. Second, launching E small GEMMs is catastrophically inefficient: each GEMM needs to reach a tile-occupancy sweet spot before it becomes bandwidth-bound in the right way; tiny GEMMs never get there. Switch Transformer (Fedus et al., 2021) reported that naive padding with a capacity factor of 1.25 drops about 3-10% of tokens on overflow, yet the kernel utilisation is still poor because the per-expert batches are too small to hide memory latency.

The right answer is a kernel purpose-built for variable-length expert batches.

Grouped GEMM: the first-principles solution

The grouped GEMM abstraction computes a set of independent matrix multiplications in a single kernel launch:

for i in 0..E:
    C_i = A_i @ W_i        # A_i has shape (n_i, d_model), W_i has shape (d_model, d_ff)

Here n_i differs per expert. A grouped GEMM kernel tiles across all expert problems simultaneously, scheduling warps to whichever sub-problem has the best occupancy at that moment. NVIDIA's cuBLAS exposes cublasGemmBatchedEx and the more efficient cublasGemmGroupedBatchedEx, but both require host-side metadata arrays describing each problem's size and data pointer.

The tricky part is that those metadata arrays must be assembled on the GPU after routing (because routing is dynamic), then read back by the kernel. Naive implementations add a CPU round-trip between routing and the GEMM, breaking the computation graph and preventing CUDA graph capture.

High-performance MoE stacks (Tutel, FasterMoE, Megablocks) solve this by fusing the sort/dispatch step with the grouped GEMM launch: the routing kernel writes its output metadata directly into GPU-resident descriptor buffers, and the GEMM kernel reads from those buffers without ever touching the CPU.

Block-sparse kernels: the MegaBlocks approach

MegaBlocks (Gale et al., 2022) reframes MoE computation as a block-sparse matrix multiplication. Instead of E separate dense sub-problems, all expert weight matrices are conceptually stacked into one large block-sparse matrix W of shape (E * d_ff, d_model), and the token-to-expert assignments define the sparsity pattern.

# Conceptual block-sparse view
W_sparse = block_diag(W_0, W_1, ..., W_{E-1})   # shape: (E*d_ff, d_model)
X_routed = scatter(X, routing_indices)            # permuted token matrix
Y = X_routed @ W_sparse                          # block-sparse GEMM

This formulation never pads and never drops tokens. The kernel traverses only non-zero blocks, so its memory footprint scales with actual traffic, not maximum capacity. Gale et al. report up to 40% throughput improvement over Tutel and 2.4x over Megatron-LM on certain configurations, precisely because tiles map cleanly to warp-level parallelism without empty work.

The key hardware insight is that modern GPUs have efficient sparse-tensor-core support (Ampere introduced 2:4 fine-grained sparsity), but MoE sparsity is coarser and block-structured, so custom block-sparse kernels outperform the fine-grained hardware path for typical expert sizes (d_ff = 1024-8192).

Writing MoE kernels in Triton

Triton is particularly well suited to MoE kernels because it lets you express tiled block-sparse loops in Python-like syntax while the compiler handles register allocation, shared-memory staging, and auto-tuning. A simplified Triton grouped GEMM kernel skeleton looks like:

@triton.jit
def grouped_gemm_kernel(
    A_ptr, B_ptr, C_ptr,
    expert_offsets_ptr,   # precomputed cumulative token counts per expert
    N, K, M,
    stride_ak, stride_bk, stride_cm,
    BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, BLOCK_M: tl.constexpr,
):
    expert_id = tl.program_id(axis=0)
    tile_n    = tl.program_id(axis=1)
    tile_m    = tl.program_id(axis=2)

    start = tl.load(expert_offsets_ptr + expert_id)
    end   = tl.load(expert_offsets_ptr + expert_id + 1)
    n_tokens = end - start          # variable per expert

    # Standard tiled GEMM accumulation below this point ...

The expert_offsets_ptr array is the critical piece: it is produced by the routing kernel on-device, which means the full forward pass (route -> sort -> GEMM) can be captured in a CUDA graph. Triton's @triton.autotune decorator then searches BLOCK_N/BLOCK_K/BLOCK_M combinations at first run and caches the best config per (n_tokens, d_model, d_ff) triple.

For Triton tutorials on fused GPU kernels, see Fused Attention tutorial which illustrates the same tiling discipline applied to attention.

Interaction with XLA and torch.compile

GShard (Lepikhin et al., 2020) was implemented on TPUs via XLA, which has its own grouped-matmul primitive (jax.lax.batch_matmul combined with scatter/gather ops). XLA's strength is whole-graph fusion: the compiler can fuse the routing softmax, top-k selection, scatter, and matmul into a single HLO graph and then lower it to efficient TPU systolic-array instructions. On GPUs via JAX, the same XLA path applies but generally underperforms hand-written CUDA/Triton kernels because XLA's GPU backend was historically less mature for sparse patterns.

torch.compile with Inductor can fuse simple MoE forward passes, but dynamic shapes (variable n_i) force graph breaks unless you annotate the dynamic dimensions explicitly with torch.export constraints. As of PyTorch 2.x, the recommended approach for production MoE is to use a custom Triton kernel via torch.library.custom_op and register it as a backend-agnostic op so torch.compile treats it as a black box and does not try to retrace through it.

When it falls down

Expert imbalance destroys utilisation. If routing collapses to a few popular experts (a known failure mode of learned routers), grouped GEMM tile counts become highly uneven. One expert's sub-problem fills hundreds of tiles while most others have single-digit tiles. The CUDA scheduler cannot dynamically reassign warps across sub-problems once the kernel is launched, so you end up with most SMs idle. Auxiliary load-balancing losses (z-loss, auxiliary balance loss in Switch) are not optional; they are prerequisites for the kernel to be efficient.

Capacity-factor vs. throughput trade-off. Eliminating the capacity factor (as MegaBlocks does) solves padding waste but makes the GEMM shapes fully dynamic. Dynamic shapes prevent static tile-count computation, which complicates grid-launch sizing and can degrade auto-tuner cache hit rates.

Communication overhead in multi-GPU MoE. On a single node, expert dispatch requires All-to-All communication across GPUs. The kernel efficiency gains described above are measured on-GPU; the All-to-All can dominate total latency when the model spans many devices. FasterMoE and DeepSpeed-MoE overlap All-to-All with local compute using double-buffered pipelines, but this requires careful pinning of CUDA streams and is brittle across driver versions.

Quantised MoE kernels are immature. INT8 grouped GEMMs require per-expert scale factors, which adds another dynamic metadata array. Most production quantisation toolkits (bitsandbytes, GPTQ) were designed for dense layers and do not support grouped GEMM natively; custom kernels are needed and the ecosystem is still fragmentary.

torch.compile and dynamic dispatch. Even with registered custom ops, torch.compile may recompile on every new (n_tokens_per_expert) shape seen during warmup. Always use torch.compile(..., dynamic=True) and pad bucket sizes for inference to bound recompilation.

Further reading

• Switch Transformers: Scaling to Trillion Parameter Models (Fedus et al., 2021)
• MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (Gale et al., 2022)
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (Lepikhin et al., 2020)
• Triton Fused Attention Tutorial (block-sparse tiling discipline)