MoEBlaze: Breaking the Memory Wall for Efficient MoE Training on Modern GPUs

Abstract: The pervasive "memory wall" bottleneck is significantly amplified in modern large-scale Mixture-of-Experts (MoE) architectures. MoE's inherent architectural sparsity leads to sparse arithmetic compute and also introduces substantial activation memory overheads -- driven by large token routing buffers and the need to materialize and buffer intermediate tensors. This memory pressure limits the maximum batch size and sequence length that can fit on GPUs, and also results in excessive data movements that hinders performance and efficient model scaling. We present MoEBlaze, a memory-efficient MoE training framework that addresses these issues through a co-designed system approach: (i) an end-to-end token dispatch and MoE training method with optimized data structures to eliminate intermediate buffers and activation materializing, and (ii) co-designed kernels with smart activation checkpoint to mitigate memory footprint while simultaneously achieving better performance. We demonstrate that MoEBlaze can achieve over 4x speedups and over 50% memory savings compared to existing MoE frameworks.

• The paper introduces a memory-efficient MoE training method that fuses token routing and expert computation to eliminate expensive activation buffers.
• It leverages GPU-specific features to achieve up to 4× memory reduction and 6.2× speedup in large-scale MoE configurations.
• The approach addresses the memory wall by using metadata-driven indexing and kernel fusion, enabling scalable training for large-context models.

MoEBlaze: Efficient Mixture-of-Experts Training via Activation and Routing Optimization

Introduction

Mixture-of-Experts (MoE) architectures are instrumental for scaling deep neural networks to trillion-parameter regimes, underpinning the conditional computation paradigm exploited in contemporary LLMs. MoE layers demand that only a subset of experts is activated per token, yielding compute savings but introducing pronounced memory overheads, especially during token routing and intermediate activation buffering. The divergent scaling of processor throughput versus memory bandwidth—the so-called “memory wall” (2601.05296)—is particularly acute in MoE training, severely throttling achievable batch sizes, sequence lengths, and device efficiency in distributed settings.

MoEBlaze directly addresses these bottlenecks with an architecture-co-designed training framework that eliminates expensive per-expert activation buffers and fuses routing, expert compute, and activation pipelines. The system is designed for modern high-bandwidth GPUs such as NVIDIA H100, leveraging hardware features (e.g., warp-group matrix multiplication, Tensor Memory Accelerator) to maximize memory bandwidth utilization and shorten activation lifetimes. Support is provided for large MoE layers in LLaMA-, Mixtral-, and DeepSeek-style models, including regimes with top-k routing and long context windows.

Memory Bottlenecks in MoE Training

Two principal sources of memory overhead in standard MoE systems are explicitly identified:

1. Token Routing Buffers: Conventional MoE implementations materialize per-expert activation buffers proportional to $L \times k \times d$, where $L$ is the total tokens per step, $k$ is the number of activated experts per token, and $d$ is the model dimension. For large batch/sequence regimes (e.g., DeepSeek), these buffers may reach ~94GB per MoE layer, which is prohibitive on production-scale GPUs.
2. Intermediate Activations in Nonlinear FFNs: Modern MoE FFNs employ complex activation functions (SiLU, SwiGLU) that further exacerbate memory pressure by requiring storage for multiple intermediate tensors during forward/backward passes. For instance, the activation footprint can approach ~98GB per FFN layer in large-scale configurations.

These memory footprints exceed device HBM and directly limit effective scaling, especially at longer context lengths and larger batch sizes.

The MoEBlaze Algorithm: Memory-Efficient Routing and Training

MoEBlaze proposes a meticulous re-architecting of the routing and expert computation pipeline, replacing conventional buffer materialization with lightweight, purely metadata-driven indexing structures:

Forward Pass

• Token Dispatch: No materialization of routed-token activation buffers. Instead, auxiliary indexing structures (per-expert token lists and per-token expert assignments) track routing decisions. These structures occupy negligible memory.
• Expert Computation: Expert MLPs operate via on-the-fly gathers from the original activation tensor, guided by token-expert indices. Only the intermediate result between two back-to-back MLPs (where necessary for backward) is temporarily buffered.
• Aggregation: Output aggregation is fused with the final MLP computation, applying on-the-fly reductions (based on token-expert maps) to directly produce the $(L, d)$ output, removing the need for large, materialized activation buffers.

Backward Pass

Backpropagation leverages reverse-mapping indices to perform “scatter” operations, allowing gradient expansion and reduction without ever materializing routed gradient tensors. Only necessary intermediate states are checkpointed, retaining the memory-efficiency of the forward pass.

Data Structures and Dispatch

Efficient construction of dispatch indices on GPU is obtained via parallel bitmap encoding, column-wise reductions, and atomic-free scatters. This enables deterministic, scalable, and highly parallel token-to-expert assignment without the sort-intensive global operations (multi-pass radix sort) used in previous work such as Megablocks.

Kernel Co-Design and Activation Checkpointing

Heavy memory traffic from intermediate activation storage is substantially mitigated by fusing expert GEMMs with non-linearity computation and activation checkpointing:

• SwiGLU Fusion: Both projections and activation epilogue are executed in a single kernel, with intermediate results resident in registers/shared memory, directly emitting only the final outputs to global memory.
• Activation Checkpointing: For computationally light (but memory-heavy) activations (e.g., SiLU in SwiGLU), intermediates are not buffered in forward; instead, recomputation is performed during backward. This reduces activation traffic from $O(Lkh)$ GB to negligible for practical sequence/expert sizes, with minimal recomputation overhead, especially on bandwidth-constrained hardware.

Backward pass gradients are aggregated in-place with tiling and warp specialization, eliminating temporary global buffers and maximizing arithmetic intensity.

Experimental Evaluation

Memory and Speed Results

Across seven LLM-style MoE configurations (varied in input dimension, expert count, routing k, batch, sequence), MoEBlaze demonstrates:

• Memory savings: Up to 3.6× reduction in peak activation memory for SiLU activations and 4× reduction for SwiGLU, e.g., from 40GB to 10GB in high-expert/high-dimension settings.
• Training speedups: Up to 6.2× end-to-end speedup versus the Megablocks baseline. Even in smaller configurations, speedup factors of 1.4×–3.7× are observed.

These improvements are attributed to complete elimination of per-expert token buffers, highly parallel index builds, and aggressive kernel fusion, particularly pronounced for complex activations.

Hardware Optimization

MoEBlaze kernels exploit H100-specific features—warp-group matrix multiplication, TMA, cluster tiling—to maximize memory bandwidth utilization and occupancy. Mixed precision (FP8/FP16 with FP32 accums) is used where beneficial, with router-weighted summation (for output aggregation) retaining higher precision for numerical stability.

Theoretical and Practical Implications

The memory wall fundamentally constrains scaling of sparse deep learning models. MoEBlaze’s elimination of routed activation buffers and its fusion of nonlinearity and computation offer a pathway to efficient trillion-parameter MoE model training with long contexts and large batches. The approach is complementary to prior efforts on sparse execution (Megablocks), routing throughput optimization (TurboMoE), and adaptive parallelism (Tutel, DeepSpeed-MoE), focusing specifically on routed activation memory.

The system’s indexing primitives and kernel fusion strategies are applicable in both single-device and distributed training regimes, paving the way for scalable multi-node MoE deployments. Future improvements will likely center on extending these dispatch and fusion paradigms to inter-node communication, adaptive load balancing, and multi-modal architectures.

Conclusion

MoEBlaze realizes substantial gains in memory and compute efficiency for MoE training on GPUs, driven by architecture-aware routing and activation management combined with hardware-tailored kernel fusion. The approach yields up to 4× reduction in activation memory and 6.2× training speedups, substantiating its suitability for scaling next-generation sparse LLMs under stringent hardware constraints. The design principles underlying MoEBlaze—metadata-driven routing, fused expert computation, activation checkpointing—are expected to inform future distributed and large-context model training strategies in the broader AI systems domain (2601.05296).