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

Pitch¶

MoEBlaze shatters the critical memory wall bottleneck in Mixture-of-Experts (MoE) training by eliminating the substantial activation overheads caused by traditional token routing mechanisms. By introducing a fused kernel design that avoids materializing intermediate buffers alongside a smart activation checkpoint scheme, the framework achieves over 4× speedups and 50% memory savings compared to state-of-the-art systems, enabling more efficient scaling of trillion-parameter models.

1. Executive Summary¶

MoEBlaze presents a memory-efficient training framework for Mixture-of-Experts (MoE) models that eliminates the substantial activation memory overhead caused by traditional token routing mechanisms. The system achieves over 4× speedups (up to 6.2× on SwiGLU configurations) and over 50% memory savings (up to 3.6× reduction in peak activation memory) compared to MegaBlocks, the state-of-the-art MoE training system. The key innovations are a fused kernel design that avoids materializing intermediate token routing buffers, and a smart activation checkpoint scheme specifically optimized for modern activation functions like SwiGLU on NVIDIA H100 GPUs.

2. Context and Motivation¶

The Core Problem: The "Memory Wall" in MoE Training¶

The fundamental problem this paper addresses is the memory wall bottleneck in Mixture-of-Experts training—a situation where memory bandwidth and capacity, rather than raw compute throughput, limit system performance. This problem has become acute as language models scale to trillion-parameter sizes with longer sequence lengths and larger batch sizes.

The memory wall manifests specifically in MoE architectures through two mechanisms:

Architectural sparsity reduces compute density. In MoE models, each token activates only a small subset of experts (typically \(k\) experts out of \(E\) total), meaning the arithmetic intensity (FLOPs per byte of memory access) is inherently lower than dense models. This makes memory bandwidth the limiting factor more often.
Activation memory overhead from token routing. Before experts can process tokens, those tokens must be routed to the correct experts. Conventional implementations allocate large per-expert buffers to compact and store dispatched tokens. These buffers consume memory proportional to \(L \times k \times d\) (number of tokens × experts per token × model dimension), which can reach ~94 GB for a single MoE layer in production models like DeepSeek (Section 2.1).

Why This Matters: Real-World Impact¶

This problem has direct practical consequences for the AI industry:

Batch size and sequence length limits. Memory pressure caps the maximum batch size and sequence length that can fit on GPUs, directly limiting training efficiency and model capability. Longer contexts are increasingly important for modern LLMs.
Distributed training overhead. As models exceed single-device HBM capacity, training must be distributed across more GPUs and nodes, increasing pressure on interconnect bandwidth and making communication a bottleneck.
Model scaling constraints. The memory wall hinders efficient scaling to larger models, which are empirically known to achieve better performance (scaling laws).

Prior Approaches and Their Limitations¶

The paper identifies two broad categories of prior work:

Token dropping and capacity-limited routing. Early MoE systems like Switch Transformers and GShard used capacity-limited routing, where each expert has a fixed capacity \(C \approx \gamma \cdot \frac{Lk}{E}\). Tokens exceeding this capacity are dropped or routed to a residual path. While this simplifies implementation by guaranteeing fixed-size buffers, it degrades model quality and requires tuning the capacity factor \(\gamma\).

Dropless routing with optimized computation. More recent systems like MegaBlocks and TurboMoE focus on dropless routing (every token is processed) with efficient handling of dynamic, variable-length expert workloads. MegaBlocks reformulates MoE computation as block-sparse operations to avoid padding, mapping well to GPU architectures. However, these systems still require storing indices and compacted token data, maintaining a memory footprint proportional to \(L \times k \times d\). The routing buffers remain a significant memory consumer.

A Quantitative Example of the Problem¶

The paper provides a concrete illustration using DeepSeek's MoE configuration (Section 2.1): - \(L \approx 2\) million tokens - \(k = 4\) active experts per token - \(d = 6144\) model dimension - 2 bytes per element (bfloat16)

Memory for routing buffer: \(L \times d \times k \times 2 \text{ bytes} \approx 94 \text{ GB}\)

For the FFN intermediate activations with hidden dimension \(h = 24576\): Memory: \(2L \times h \approx 98 \text{ GB}\)

A single MoE layer can consume nearly 200 GB—far exceeding the ~80 GB HBM on an H100 GPU.

How This Paper Positions Itself¶

MoEBlaze takes a fundamentally different approach: instead of optimizing how to manage and store intermediate buffers, it eliminates them entirely through a fused kernel design. The paper positions itself against MegaBlocks specifically, which the authors identify as the current state-of-the-art for high-performance sparse training. The paper explicitly targets two bottlenecks:

Token routing buffers (Section 3): Replace materialized per-expert activation buffers with lightweight index lists, performing on-the-fly gathers during computation.
SwiGLU activation intermediates (Section 5): Fuse kernel operations and apply smart recomputation to reduce memory traffic for modern activation functions.

The approach is co-designed for modern GPU architectures (specifically NVIDIA H100), leveraging features like warp-group matrix multiplication and tensor memory acceleration.

3. Technical Approach¶

3.1 Reader Orientation¶

MoEBlaze is a systems/implementation paper that redesigns the token routing and expert computation kernels in MoE training to eliminate intermediate memory buffers while improving throughput.

3.2 Big-Picture Architecture (Diagram in Words)¶

The MoEBlaze framework has three major phases that execute within a single fused kernel:

Memory-Efficient Token Dispatch. Instead of physically moving tokens into per-expert buffers, this phase constructs lightweight index data structures (expert_token_indices, expert_token_offsets, token_expert_indices, token_index_map) that track routing decisions.
On-the-Fly Expert Computation. The expert MLP computations access tokens directly from the original, unpermuted activation tensor using the index structures, rather than from materialized buffers. Only the intermediate result between the two MLP layers is buffered.
Direct Output Aggregation. Expert outputs are summed and reduced directly into the final output tensor, fused with the second MLP computation, eliminating intermediate storage.

The backward pass uses the same index structures to perform inverse operations without intermediate expansion buffers.

3.3 Roadmap for the Deep Dive¶

First, the memory-efficient token routing algorithm (Section 3 of paper), explaining how the forward and backward passes work without materializing buffers.
Second, the data structures and their efficient construction (Section 4 of paper), which enable the routing algorithm.
Third, the kernel co-design for SwiGLU (Section 5 of paper), addressing the secondary memory bottleneck from complex activation functions.
Fourth, how these components integrate into end-to-end training.

3.4 Detailed, Sentence-Based Technical Breakdown¶

Notation and Problem Setup¶

The paper focuses on token-choice MoE, where each token selects which experts to route to (as opposed to expert-choice where experts select tokens). The notation is: - \(L\): Number of routed token instances (batch size × sequence length) - \(K\) (or \(k\)): Number of selected experts per token - \(E\): Total number of experts - \(d\): Model dimension (input hidden size) - \(h\): FFN hidden dimension

The gating network produces routing decisions through: \(\(\text{topk\_experts} = \text{TopK}(\text{softmax}(W_g x))\)\)

where \(W_g \in \mathbb{R}^{E \times d}\) are the gating parameters.

The Conventional Approach (What MoEBlaze Replaces)¶

In conventional implementations (Figure 1, left), token routing follows three stages:

1. Token Dispatch. Tokens are physically compacted into per-expert buffers. For each token, the system identifies its top-\(k\) experts and copies the token's activation vector into each expert's buffer. This requires allocating an \((L \times k, d)\) "routed token buffer."

2. Expert Computation. Each expert processes its assigned tokens from its buffer. The FFN for expert \(E_i\) computes: \(\(E_i(x) = W_{2,i} \cdot \sigma(W_{1,i} x)\)\)

where \(\sigma\) is the activation function (ReLU, GELU, or SwiGLU). Intermediate activations (the output of the first linear layer) are stored for the backward pass.

3. Output Aggregation. Expert outputs are combined via weighted summation: \(\(y = \sum_{i=1}^{E} g_i(x) \cdot E_i(x)\)\)

where \(g_i(x)\) is the gating score for expert \(i\).

The critical inefficiency is that the routed token buffer of size \(L \times k \times d\) must be materialized and stored throughout computation, creating massive memory pressure.

MoEBlaze's Memory-Efficient Forward Pass¶

MoEBlaze eliminates the routed token buffer by performing on-the-fly gathers from the original activation tensor (Section 3.1).

Token Dispatch Phase. Instead of creating buffers for routed tokens, MoEBlaze generates four lightweight indexing data structures (detailed in Section 4.1 below):

expert_token_indices: Which tokens are assigned to each expert
expert_token_offsets: Where each expert's token list begins
token_expert_indices: Which experts each token is routed to
token_index_map: Positions of each token's outputs in the intermediate buffer

These structures store only indices (integers), not activations, so their memory footprint is negligible compared to the \(O(L \times k \times d)\) activation buffer they replace.

Expert Computation Phase. When processing expert \(E_i\), the kernel uses expert_token_indices to gather the relevant tokens directly from the original \((L, d)\) input tensor. This is an on-the-fly gather operation—tokens are fetched when needed rather than pre-copied. Critically, only the intermediate result between the two back-to-back MLP layers (the output of the first linear transformation) is buffered for the backward pass.

Output Aggregation Phase. The summation is tightly fused with the second MLP computation. Using token_expert_indices, the kernel directly reduces the expert outputs into the final \((L, d)\) output tensor. There is no intermediate \((L \times k, d)\) buffer holding expert outputs before aggregation.

MoEBlaze's Memory-Efficient Backward Pass¶

The backward pass (Section 3.2) must propagate gradients without access to the materialized routed token buffer that conventional approaches use.

Challenge. In conventional implementations, the gradient of the \((L, d)\) output must be "expanded" to an \((L \times k, d)\) "routed gradient tensor" before backpropagating through the MLP experts. This expansion requires the stored routed token activations.

MoEBlaze's Solution. The system uses the same index structures in reverse:

Expert Summation Backward. The \((L, d)\) gradient tensor is scattered to the materialized intermediate MLP result tensor using the token-mapping structure. This maps each output gradient back to its corresponding expert contributions.
Expert Computation Backward. Gradients flow backward through the MLPs. The previously checkpointed intermediate result between the two MLP layers is used for computing weight gradients.
Token Gradient Accumulation. Gradients from all \(k\) experts each token was routed to are accumulated using the token index data structure, performing on-the-fly reductions rather than materializing an intermediate expansion buffer.

The Index Data Structures (Section 4.1)¶

The four data structures are the foundation of MoEBlaze's approach. Figure 2 provides a concrete example with \(L=6\) tokens, \(E=4\) experts, and \(k=2\).

expert_token_indices (size \(L \times k\)): A compact tensor storing token IDs assigned to each expert, concatenated across all experts. For the example, given routing decisions: - Token 0 → experts {2, 3} - Token 1 → experts {0, 1} - Token 2 → experts {0, 3} - Token 3 → experts {1, 2} - Token 4 → experts {0, 3}

The expert assignments are: - Expert 0: tokens {1, 2, 4} - Expert 1: tokens {1, 3} - Expert 2: tokens {0, 3} - Expert 3: tokens {0, 2, 4}

Concatenated: expert_token_indices = [1, 2, 4, 1, 3, 0, 3, 0, 2, 4]

expert_token_offsets (size \(E + 1\)): An array storing exclusive prefix sums of token counts per expert. This allows finding each expert's token range via expert_token_indices[offsets[i]:offsets[i+1]]: expert_token_offsets = [0, 3, 5, 7, 10]

token_expert_indices (size \(L \times k\)): The inverse mapping storing which experts each token is routed to, ordered by token ID: token_expert_indices = [2, 3, 0, 1, 1, 2, 0, 3]

token_index_map (size \(L \times k\)): Stores the position of each token within the concatenated expert_token_indices list, allowing efficient gathering of a token's \(k\) expert outputs. For token 0: token_index_map[0] = {5, 7} since token 0 appears at positions 5 and 7 in expert_token_indices.

Efficient Data Structure Construction (Section 4.2)¶

Building these structures naively would require sorting \((token\_id, expert\_id)\) tuples by expert ID, which involves expensive multi-pass radix sort on GPUs with \(O(Lk)\) data movement across multiple global memory passes.

MoEBlaze replaces sorting with a three-step parallel algorithm that is atomic-free and minimizes global memory access:

Step 1: Build Dense Token-Expert Map. Construct an \(L \times E\) dense bitmap where entry \((i, e)\) stores token ID \(i\) if token \(i\) is routed to expert \(e\). Each warp processes a disjoint tile of tokens, writing each \((i, e)\) pair at most once (guaranteeing no intra-warp collisions). This is fully parallelizable.

Step 2: Compute Expert Lengths. Launch one CTA (Cooperative Thread Array) per expert, counting non-zero entries in that expert's column of the dense map. Warp-level reductions aggregate counts, producing expert_lengths. A prefix sum then generates expert_offsets.

Step 3: Route Indices to Gates. Generate expert_token_indices by computing each token's destination position in the concatenated array. This uses a two-phase strategy: - Tile-level scan: Each CTA processes one expert's assigned tokens, computing local exclusive scans in shared memory. - Global offset addition: CTA-local counts are added to the pre-computed global expert_offsets to yield final positions.

A final parallel kernel reads from the dense map and writes to expert_token_indices at computed positions—fully parallel, no atomics required.

Kernel Co-Design for SwiGLU (Section 5)¶

The second major memory bottleneck comes from modern activation functions like SwiGLU (Equation in Section 5.1): \(\(\text{SwiGLU}(x; W_1, W_2) = \text{SiLU}(xW_1) \cdot (xW_2)\)\)

where \(\text{SiLU}(u) = u \cdot \sigma(u)\) and \(\sigma\) is the sigmoid function.

SwiGLU requires two first-layer projections (\(a = xW_1\) and \(b = xW_2\)), the SiLU activation, and an element-wise product. Conventional kernels materialize all intermediates (\(a\), \(b\), \(\sigma(a)\), \(\text{SiLU}(a)\), and the product) to global memory for the backward pass.

Key Observations (Section 5.2): 1. Activation function computation is memory-bandwidth bound on modern GPUs (point-wise operations on "tall-and-skinny" matrices where \(L \gg d\)). 2. Activation footprints scale linearly with batch size and sequence length, becoming prohibitive at scale.

MoEBlaze's Solution: Fused Kernel + Activation Checkpoint.

Kernel Fusion. The two first-layer projections and SwiGLU epilogue are fused into a single kernel: - Input \(x\) is loaded once - Both GEMMs (\(xW_1\) and \(xW_2\)) are computed simultaneously, streaming through registers/shared memory - SiLU(\(a\)) is computed in-register and immediately multiplied with \(b\) - Only the final output is written to global memory

This eliminates global writes/reads of \(a\), \(b\), and intermediate activations. Input reads of \(x\) are halved compared to separate kernels.

Activation Checkpoint (Recomputation). During the forward pass, MoEBlaze does not save SiLU intermediates. During the backward pass: \(\(S_{\text{recomp}} = \text{SiLU}(A)\)\)

The SiLU is recomputed from the stored \(A\) (which is needed anyway for weight gradients). Since SiLU is element-wise and memory-bandwidth bound, recomputation costs little compared to the memory savings.

Backward Pass Gradients. (Algorithm 1, lines 26-28): \(\(\nabla A = \nabla Y_{\text{swi}} \odot B \odot \nabla \text{SiLU}(A)\)\) \(\(\nabla B = \nabla Y_{\text{swi}} \odot S_{\text{recomp}}\)\)

Gradients from both branches are aggregated in-place via tiled reductions, eliminating temporary global buffers.

End-to-End Training Flow¶

Algorithm 1 summarizes the complete training process:

Forward: 1. Load input tokens \(X\) once 2. Compute fused SwiGLU: \((A, B), Y_{\text{swi}} \leftarrow \text{FusedSwiGLU}(X, W_1, W_2)\), with SiLU computed transiently 3. Final projection: \(Y_{\text{out}} \leftarrow Y_{\text{swi}} W_3\) 4. Store only \(A\), \(B\), \(Y_{\text{swi}}\) for backward (not SiLU intermediates)

Backward: 1. Compute weight gradient: \(\nabla W_3 \leftarrow Y_{\text{swi}}^T \nabla Y_{\text{out}}\) 2. Backpropagate to SwiGLU output: \(\nabla Y_{\text{swi}} \leftarrow \nabla Y_{\text{out}} W_3^T\) 3. Load stored \(A\), \(B\); recompute \(S_{\text{recomp}} \leftarrow \text{SiLU}(A)\) 4. Compute activation gradients: \(\nabla A\), \(\nabla B\) using the recomputed SiLU 5. Fused backward through projections: \(\nabla W_1\), \(\nabla W_2\), \(\nabla X\)

Design Choice Justifications¶

Why on-the-fly gathers over pre-materialized buffers? Pre-materialized buffers require \(O(L \times k \times d)\) memory per MoE layer. On-the-fly gathers require only \(O(L \times k)\) indices (integers), which is orders of magnitude smaller. The gather overhead is amortized by the reduced memory traffic.

Why fusion over separate kernels? Separate kernels for each operation incur repeated global memory reads/writes and kernel launch overhead. Fusion keeps data in registers/shared memory, converting memory-bound operations to compute-bound where possible.

Why recomputation over storing all intermediates? SwiGLU's intermediates scale as \(O(L \times h)\), which is substantial at scale. Since SiLU is element-wise (cheap to compute) and memory-bandwidth bound, recomputation trades compute for memory—the right tradeoff when memory is the bottleneck.

Why the three-step parallel construction over sorting? GPU radix sort requires multiple global memory passes proportional to key width, creating a multi-kernel pipeline with high launch latency. The three-step approach is atomic-free, minimizes global memory access, and maps efficiently to GPU parallelism.

4. Key Insights and Innovations¶

Innovation 1: Eliminating Materialized Routing Buffers Entirely¶

The most fundamental contribution is the conceptual shift from "how to efficiently manage routing buffers" (the approach of MegaBlocks and TurboMoE) to "how to avoid needing routing buffers at all." This is achieved by maintaining only lightweight index data structures and performing on-the-fly gathers during computation.

This is not an incremental optimization of existing approaches—it changes the fundamental algorithmic structure of MoE computation. The memory savings are proportional to the size of the original routing buffer (\(O(L \times k \times d)\)), which is substantial at scale (94 GB in the DeepSeek example). The index structures require only \(O(L \times k)\) integers.

Innovation 2: Atomic-Free Parallel Data Structure Construction¶

Replacing expensive GPU radix sort with a three-step parallel algorithm (dense map construction → expert length computation → position-aware writing) is a significant systems innovation. The key insight is that the routing problem can be decomposed into operations that map naturally to GPU parallelism without requiring atomic operations or global sorting.

This matters because sorting-based approaches create a multi-kernel pipeline with high kernel launch latencies and multiple passes over \(O(Lk)\) data. MoEBlaze's approach minimizes global memory traffic and kernel launches, which is critical when \(L\) is in the millions.

Innovation 3: Specialized SwiGLU Kernel Fusion with Smart Recomputation¶

While kernel fusion is a known optimization technique, the paper's specific fusion strategy for SwiGLU—simultaneously computing two projections and fusing the activation epilogue, combined with selective recomputation during backward—addresses the specific memory characteristics of modern activation functions.

The insight is recognizing that SwiGLU intermediates are memory-bandwidth bound (point-wise operations on tall-and-skinny matrices) rather than compute-bound, making recomputation cheap relative to memory storage. This is a non-obvious design choice: naively, one might assume recomputation hurts performance, but the paper demonstrates the opposite by exploiting the operation's memory-bound nature.

Innovation 4: End-to-End Co-Design Across Dispatch, Compute, and Checkpoint¶

The paper's fourth innovation is treating MoE training as an integrated system rather than optimizing components in isolation. The memory-efficient dispatch algorithm, the index data structure construction, and the SwiGLU kernel fusion are co-designed to work together:

The dispatch algorithm produces index structures in a format directly consumable by the fused compute kernel.
The checkpoint strategy only stores activations that are needed by the backward pass given the fused kernel design.
The fused kernel's memory access patterns are designed assuming on-the-fly gathers rather than sequential buffer reads.

This integrated approach yields gains that would not be possible from optimizing each component independently.

5. Experimental Analysis¶

Evaluation Methodology¶

Hardware and Software. All experiments conducted on a single NVIDIA H100 Tensor Core GPU with PyTorch 2.0.1 and CUDA 12.1.

Benchmark Configurations. The paper evaluates seven MoE configurations (Table 1) varying input dimension (\(d \in \{512, 1024, 2048\}\)), number of experts (\(E \in \{4, 8, 16\}\)), top-k (\(k \in \{1, 2, 4\}\)), batch size (\(B \in \{16, 32\}\)), and sequence length (\(L \in \{512, 1024, 2048\}\)). FFN hidden dimension is fixed at \(4 \times\) input dimension.

Baselines. Primary baseline is MegaBlocks (Gale et al., 2023), identified as the state-of-the-art sparse training system using custom kernels and block-sparse operations.

Metrics. 1. Training Speed: Speedup factor relative to MegaBlocks in end-to-end single training pass (forward + backward), excluding optimizer updates. 2. Activation Memory Consumption: Total memory allocated for intermediate activation tensors, measured using PyTorch's saved tensor hooks.

Activation Functions. Both ReLU/SiLU and SwiGLU are evaluated.

Main Quantitative Results¶

Memory Efficiency with SiLU Activation (Figure 3)¶

MoEBlaze consistently reduces activation memory across all seven configurations:

Maximum savings at conf4 (D=2048, E=16, k=4, L=1024, B=32): MoEBlaze requires 6,100 MB vs. MegaBlocks' 22,000 MB—a 3.6× reduction.
Smaller savings at conf1 (D=512, E=4, k=1): Less pronounced savings expected, since savings scale with sequence length \(L\) and \(k\), both small in conf1.
Across all configurations, MoEBlaze maintains a clear and consistent memory advantage.

Training Speed with SiLU (Figure 4)¶

Speedups range from 1.4× to 3.7× across configurations:

Maximum speedup at conf4: 3.7× improvement.
The paper attributes speedups to three factors: (1) optimized token dispatch reducing latency, (2) lightweight data structure construction avoiding expensive multi-pass sorting, (3) fused kernels leveraging H100 features (warp-group MMA, tensor memory accelerator).

Memory Efficiency with SwiGLU (Figure 5)¶

SwiGLU inherently requires more memory due to gating and element-wise operations. MoEBlaze's advantage is even more pronounced:

conf3 (D=1024, E=16, k=4): MegaBlocks requires over 40,000 MB, while MoEBlaze uses approximately 10,000 MB—a 4× reduction.
Peak activation memory is often less than half of the baseline across configurations.

Training Speed with SwiGLU (Figure 6)¶

Speedups are higher and more consistent than SiLU, ranging from 2× to 6.2×:

The paper attributes this to: (1) more complex SwiGLU computation exposing greater optimization opportunities for fused kernels, (2) memory-bandwidth savings being more critical when intermediate sizes are larger.

Assessment: Do the Experiments Support the Claims?¶

Claim 1: Over 50% memory savings. Strongly supported. Across both SiLU and SwiGLU, MoEBlaze consistently uses less than half of MegaBlocks' activation memory in most configurations. The 3.6× reduction at conf4 (SiLU) and 4× reduction at conf3 (SwiGLU) exceed the claimed threshold.

Claim 2: Over 4× speedups. Supported, with qualification. The 6.2× speedup at conf4 (SwiGLU) and 3.7× at conf4 (SiLU) bracket this claim. Some SiLU configurations show lower speedups (1.4× at conf1), but the paper accurately reports the range.

Claim 3: Memory savings scale with problem size. Supported. Smaller configurations (conf1 with k=1) show smaller memory savings, while larger configurations (conf3, conf4 with larger d, E, k) show substantial savings—consistent with the theoretical analysis that savings are proportional to \(L \times k \times d\).

Limitations and Missing Experiments¶

Single-device evaluation. The paper acknowledges this limitation (Section 8) and notes that extending to distributed training is future work. This is a significant gap: distributed MoE training introduces additional bottlenecks (inter-expert communication, all-to-all operations) that single-device experiments cannot capture.

No ablation studies. The paper does not isolate the contribution of each optimization (dispatch algorithm, data structure construction, kernel fusion, recomputation). It's unclear whether the gains come primarily from one technique or require all three together.

No accuracy evaluation. The paper claims "without compromising accuracy" but does not report model convergence metrics, loss curves, or downstream task performance. The focus is purely on system efficiency.

Limited baseline comparison. MegaBlocks is the only baseline. No comparison to TurboMoE, DeepSpeed-MoE, or other recent systems. The paper cites TurboMoE in related work but does not benchmark against it.

Fixed hyperparameters. FFN hidden dimension is always \(4 \times\) input dimension; no exploration of how different FFN expansion ratios affect the memory-speed tradeoff.

Summary of Key Numbers¶

Metric	Configuration	MoEBlaze	MegaBlocks	Improvement
SiLU Memory	conf4	6,100 MB	22,000 MB	3.6× reduction
SwiGLU Memory	conf3	~10,000 MB	>40,000 MB	~4× reduction
SiLU Speedup	conf4	—	—	3.7× faster
SwiGLU Speedup	conf4	—	—	6.2× faster

6. Limitations and Trade-offs¶

Assumption: Single-Device Execution with Contiguous Memory¶

All experiments in the paper are conducted on a single NVIDIA H100 GPU (Section 6.1), and the algorithm design assumes contiguous activation tensors that can be randomly accessed via gather operations. This assumption creates several constraints:

Distributed settings introduce communication overhead not accounted for. In production MoE training at scale, experts are typically sharded across multiple GPUs and nodes. The on-the-fly gather approach requires random access to the original \((L, d)\) activation tensor, which may not be locally available if tokens and experts are distributed. The paper acknowledges this limitation in Section 8:

"While this paper primarily focuses on single-device performance, we note that the core mechanisms of MoEBlaze are also applicable to distributed settings."

However, the claim of applicability is not substantiated with experiments. In distributed MoE, the primary bottleneck becomes inter-expert all-to-all communication for token dispatch and result aggregation—precisely the operations MoEBlaze redesigns. Whether the memory-efficient approach transfers to distributed settings, where network bandwidth may dominate memory bandwidth, is an open question.

Random access patterns may not suit all memory hierarchies. The on-the-fly gather operation assumes that randomly accessing the original activation tensor is efficient. On GPUs with deep memory hierarchies (registers → shared memory → L1/L2 cache → HBM), random gathers can cause cache thrashing if the index patterns are not locality-friendly. The paper does not analyze the cache behavior of its gather patterns or compare against sequential buffer reads on the memory hierarchy.

Assumption: SiLU/SwiGLU Recomputation Is Always Beneficial¶

The activation checkpoint strategy for SwiGLU relies on the assumption that recomputing \(\text{SiLU}(A)\) during the backward pass is cheaper than storing it. This assumption is justified by the observation that activation functions are memory-bandwidth bound on GPUs (Section 5.2):

"Computation of activation functions is generally memory bandwidth bound on modern GPUs due to two primary reasons: 1) activation function's computation is mostly point-wise operations and modern GPU is highly capable of such operations, 2) in LLM training, we are usually handling the case where the number of tokens is far larger than the embedding dimension \(L \gg d\)."

However, this assumption may not hold in all regimes:

Small batch sizes or short sequences. When \(L\) is small, the memory traffic reduction from not storing SiLU intermediates may be marginal, and the recomputation cost may not be amortized.
Different hardware generations. The experiments are on H100 with its specific memory bandwidth (3.35 TB/s) and compute capabilities. On older GPUs with lower memory bandwidth or different cache hierarchies, the recomputation tradeoff may shift.
Alternative activation functions. The paper analyzes SiLU and SwiGLU but does not consider other modern activations (GeGLU, SquaredReLU) that may have different compute-to-memory ratios.

The paper does not provide a theoretical or empirical characterization of when recomputation becomes unfavorable.

No Ablation Studies: Contribution Attribution Is Unclear¶

A significant limitation is the absence of ablation experiments to isolate the contribution of each optimization. MoEBlaze comprises three major components:

Memory-efficient token dispatch (eliminating routing buffers)
Atomic-free parallel data structure construction (replacing sorting)
Fused SwiGLU kernels with smart recomputation

The reported speedups (up to 6.2×) and memory savings (up to 4×) are for the combined system. It is unclear:

What portion of memory savings comes from eliminating routing buffers vs. SwiGLU fusion?
What portion of speedup comes from dispatch optimization vs. kernel fusion vs. reduced memory traffic?
Are all three components necessary, or could a simpler combination achieve similar results?

The lack of ablation makes it difficult for practitioners to determine which techniques to adopt if they cannot implement the full system, or to understand which component is most critical for different MoE configurations.

No Accuracy or Convergence Evaluation¶

The paper claims to improve efficiency "without compromising accuracy" (Abstract, Section 1) but provides no experimental validation of model quality. The evaluation focuses entirely on memory consumption and training speed for a single MoE layer.

Missing evidence includes: - Training loss curves comparing MoEBlaze to MegaBlocks - Downstream task performance (e.g., perplexity on language modeling benchmarks) - Convergence rate (steps to reach a target loss) - Numerical stability analysis (floating-point precision differences between materialized buffers and on-the-fly gathers)

This is a notable omission because on-the-fly gathers with on-the-fly reductions could theoretically introduce numerical differences from different operation ordering, particularly for large models with many experts. The claim of no accuracy compromise should be verified empirically.

Limited Baseline Comparison¶

MegaBlocks is the only system baseline evaluated. While MegaBlocks is a reasonable choice as a state-of-the-art MoE training system, the paper cites but does not benchmark against other relevant systems:

TurboMoE (Aminabadi et al., 2025): Cited in related work as introducing "fused, metadata-driven kernels and data-layout transformations." The paper does not explain why TurboMoE is not included as a baseline, despite being directly comparable in its optimization target (fused kernels for MoE).
DeepSpeed-MoE (Rajbhandari et al., 2022): A widely used production system for distributed MoE training.
Tutel (Hwang et al., 2023): Another MoE system with adaptive parallelism.

Without comparison to these systems, it is unclear whether MoEBlaze's gains are specific to the MegaBlocks baseline or represent genuine improvements over the broader state-of-the-art.

Fixed Configuration Parameters¶

The experimental configurations (Table 1) explore a range of input dimensions, expert counts, and top-k values, but several parameters are held fixed without justification:

FFN hidden dimension: Always \(4 \times\) input dimension. Different expansion ratios (e.g., \(8\times\) as in some LLMs) would change the activation memory footprint and may affect the relative benefit of MoEBlaze.
Token-choice routing only: The paper explicitly focuses on token-choice MoE (Section 2) and does not evaluate expert-choice routing, where experts select tokens rather than tokens selecting experts. Expert-choice has different memory and computation patterns.
Gating function: All experiments use standard Top-K gating. No evaluation of alternative routing mechanisms (e.g., hash-based routing, learned routing) that may produce different index distributions.

Edge Cases Not Addressed¶

Several edge cases and operational scenarios are not discussed:

Load imbalance. When some experts receive many more tokens than others (a common issue in MoE training), the on-the-fly gather pattern may create workload imbalance across GPU threads. The paper does not analyze how MoEBlaze handles skewed expert assignments or whether the parallel construction algorithm degrades under imbalance.

Variable sequence lengths within a batch. The paper assumes all tokens in a batch have the same sequence length \(L\). In practice, batches often contain sequences of different lengths, requiring padding or special handling. Whether MoEBlaze's data structures and kernels generalize to variable-length sequences is not addressed.

Gradient accumulation. Many production training runs use gradient accumulation to simulate larger batch sizes. The memory-efficient backward pass may interact differently with gradient accumulation, but this is not discussed.

Scalability to Extreme Configurations Not Tested¶

The largest configuration tested (conf4) has \(L = 32 \times 1024 = 32,768\) tokens, \(d = 2048\), \(E = 16\) experts. While this is reasonable, it does not approach the scale of frontier models:

DeepSeek-V3 (cited in related work): 671B parameters with 37B active per token—orders of magnitude larger.
Frontier LLMs: Often have \(E = 128\) or more experts, \(d = 8192+\) dimensions, and \(L\) in the millions for long-context training.

Whether MoEBlaze's three-step parallel construction and on-the-fly gathers remain efficient when \(L \times k\) reaches millions and \(E\) exceeds 64 is unknown. The dense bitmap in Step 1 has size \(L \times E\), which could become a memory concern itself at extreme scales.

7. Implications and Future Directions¶

How This Work Changes the Landscape¶

MoEBlaze demonstrates that the conventional wisdom—that MoE training must materialize large routing buffers—is fundamentally unnecessary. This shifts the design space for MoE systems from "how to manage memory pressure" to "how to restructure computation to avoid memory pressure in the first place."

This conceptual shift has several implications for the field:

Memory efficiency and throughput are not zero-sum. Prior approaches often traded memory for speed (e.g., storing more activations to reduce recomputation) or speed for memory (e.g., aggressive checkpointing). MoEBlaze achieves both simultaneously—up to 6.2× speedup with 4× memory reduction—demonstrating that the two objectives can be aligned through careful algorithm-hardware co-design.

The "memory wall" is addressable at the algorithm level. The paper's title—"Breaking the Memory Wall"—is not hyperbole. By restructuring the MoE computation to eliminate intermediate buffers, the system achieves efficiency gains that would not be possible through hardware improvements alone. This reinforces a growing recognition in systems research: algorithmic innovations can often deliver larger efficiency gains than hardware advances.

Index structures as a first-class design primitive. The paper elevates index data structures (expert_token_indices, expert_token_offsets, etc.) from auxiliary bookkeeping to central components of the computation model. This pattern—maintaining lightweight metadata to avoid materializing heavy intermediates—could generalize to other sparse computation domains beyond MoE.

Follow-Up Research This Work Enables or Suggests¶

1. Distributed MoE training with memory-efficient dispatch. The most immediate extension is applying MoEBlaze's techniques to multi-GPU, multi-node MoE training. The key challenges are:

Communication optimization. In distributed settings, the token dispatch phase involves all-to-all communication between GPUs holding different experts. MoEBlaze's index-based approach could potentially reduce communication volume by sending only indices rather than full activations, but this requires careful coordination with the expert compute kernels.
Hierarchical data structures. The flat index structures may need to be extended to hierarchical forms that account for expert placement across GPUs and nodes.
Overlap of communication and computation. The on-the-fly gather pattern could be adapted to hide communication latency by overlapping data transfers with expert computation on local experts.

2. Generalization to other sparse computation patterns. The principle of avoiding materialized intermediates through on-the-fly gathers applies beyond MoE:

Sparse attention mechanisms. Long-context models use sparse attention patterns (e.g., sliding windows, random sparse attention). Similar index-based approaches could reduce memory for attention intermediates.
Mixture-of-Depths. Recent architectures conditionally skip layers for certain tokens. The dispatch/aggregation patterns are similar to MoE and could benefit from analogous optimizations.
Conditional computation in other modalities. Vision transformers and multimodal models increasingly use sparse routing; the techniques may transfer.

3. Hardware-aware kernel synthesis. MoEBlaze's fused kernels are manually designed for H100. Future work could explore:

Automatic kernel fusion. Compiler techniques to automatically determine which operations to fuse based on hardware characteristics (memory bandwidth, shared memory size, tensor core capabilities).
Portable implementations. Adapting the kernels to other hardware (AMD MI300, Intel Gaudi, Apple Silicon) with different memory hierarchies and compute capabilities.
Different precision regimes. Extending the fusion and recomputation strategies to FP8 training or mixed-precision inference.

4. Theoretical analysis of the memory-compute tradeoff. The paper empirically demonstrates that recomputation is beneficial for SwiGLU, but does not provide a theoretical framework. A rigorous analysis of:

When recomputation is optimal (operation arithmetic intensity vs. memory bandwidth)
The breakeven point for different activation functions and hardware
The interaction of recomputation with gradient checkpointing policies

would provide principled guidance for future systems.

5. Integration with load balancing and routing research. MoEBlaze assumes a fixed routing policy. Future work could investigate:

Co-design of routing policies with dispatch efficiency. Certain routing patterns (e.g., routing nearby tokens to the same expert) may be more amenable to MoEBlaze's gather patterns.
Dynamic capacity management. Combining MoEBlaze's memory-efficient dispatch with load balancing losses that prevent expert collapse.

6. End-to-end training evaluation. The most critical missing validation is full training runs:

Train an LLM from scratch or fine-tune a pretrained model using MoEBlaze
Report loss curves, convergence speed, and downstream task performance
Compare to baseline systems on equal training time or equal compute budget

This would validate the "without compromising accuracy" claim and establish practical adoption pathways.

Practical Applications and Downstream Use Cases¶

Production MoE training infrastructure. MoEBlaze directly addresses the memory constraints that limit batch size and sequence length in large-scale MoE training. Organizations training trillion-parameter models could:

Fit larger batch sizes in the same GPU memory, improving throughput
Train with longer contexts without increasing GPU count
Reduce the number of GPUs required for a given model configuration, lowering costs

On-device MoE inference. The memory-efficient techniques could extend to inference, where memory constraints are even more stringent:

Mobile or edge deployment of MoE models
Inference on GPUs with limited memory (e.g., consumer GPUs, older datacenter hardware)
Multi-model serving where memory must be shared across models

Research prototyping. The framework enables researchers to experiment with MoE architectures at larger scales than previously feasible:

Testing new routing mechanisms without hitting memory walls
Exploring expert specialization patterns at higher expert counts
Rapid iteration on MoE configurations during architecture search

When to Prefer This Method Over Alternatives¶

Prefer MoEBlaze when:

You are training MoE models with SwiGLU or SiLU activations (where the fused kernel benefits are largest)
You are memory-limited rather than compute-limited (batch size or sequence length capped by GPU memory)
You are training on H100 or similar high-bandwidth GPUs where the fusion and recomputation strategies are tuned
Your MoE configuration has moderate to large \(L \times k \times d\) (where routing buffer elimination provides significant savings)
You can adopt the entire framework (index structures + fused kernels) rather than individual components

Prefer alternatives when:

You need distributed multi-node training immediately (MoEBlaze is single-device only)
You require expert-choice routing (MoEBlaze is designed for token-choice)
You are using activation functions other than SiLU/SwiGLU (ReLU, GeGLU) where fusion benefits may differ
Your MoE has very small routing buffers relative to other memory consumers (e.g., very small \(k\) or \(L\))
You need battle-tested, production-grade code (MoEBlaze is a research artifact; MegaBlocks, DeepSpeed-MoE, and Tutel have more deployment history)

For reproduction or integration, the critical choices are:

Implement the three-step parallel construction (dense map → expert lengths → route indices) rather than sorting-based dispatch, as this is where much of the speedup originates.
Use the four index data structures (expert_token_indices, expert_token_offsets, token_expert_indices, token_index_map) as the dispatch format; they are designed for on-the-fly gathers.
Fuse SwiGLU projections and activation epilogue into a single kernel if using SwiGLU; load input once, compute both GEMMs simultaneously, and perform SiLU in-register.
Apply recomputation for SiLU intermediates during backward; store only \(A\) and \(B\) (the first-layer projections), not SiLU(\(A\)).
Ensure index structures persist across forward and backward passes; they are needed for both on-the-fly gathers and on-the-fly reductions.