MoEBlaze Cuts the Routing Buffers Behind the MoE Memory Wall

TL;DR for operators

Sparse expert activation does not guarantee a small GPU-memory footprint. Conventional token routing can physically duplicate token representations into expert-specific storage, while expert FFNs preserve additional intermediates for backward computation. In the paper’s illustrative DeepSeek-like case, one materialized routing buffer alone is estimated at roughly 94 GB.

MoEBlaze¹ attacks that data movement rather than changing the router or MoE architecture. It keeps compact token-to-expert location information, gathers original activations when expert computation needs them, reduces results directly back into token order, builds the required routing metadata without global sorting, and combines fused SwiGLU kernels with selective recomputation.

On seven single-MoE-layer configurations running on one NVIDIA H100, MoEBlaze reports lower saved activation memory than MegaBlocks throughout the tested configurations. Reported forward-plus-backward speedups range from roughly 1.4× to 3.7× in the section described as SiLU and roughly 2× to 6.2× under the SwiGLU evaluation.

For teams training sparse models, the practical implication is to benchmark activation footprint and memory traffic alongside active parameters and FLOPs. The evidence supports that decision at the MoE-layer level on one H100. It does not yet establish equivalent savings for complete LLM training, optimizer-inclusive runtimes, expert-parallel communication, or distributed clusters.

Sparse arithmetic can still produce a large memory problem

MoE looks economical at first glance because each token activates only a subset of the available experts. That reduces arithmetic relative to evaluating every expert for every token.

The difficulty appears in what happens between those arithmetic operations. A conventional implementation may copy or reorder each token representation into storage organized by expert. If a token selects several experts, its representation can appear several times. The paper models this routing-buffer footprint as growing with token volume, embedding dimension, and top-$k$ routing. Its DeepSeek-like illustration reports roughly 94 GB for one routing buffer.

The expert network adds another burden. Forward computation produces intermediate activations that may need to survive until backward computation. The paper gives a separate illustrative FFN-intermediate footprint of roughly 98 GB under its large-model example.

Those numbers are illustrations rather than measurements from the seven H100 benchmark configurations, but they explain the systems problem MoEBlaze is designed around: sparsity reduces which experts calculate, yet routing and backward-state management can still move and retain very large tensors.

For infrastructure planning, this changes what should be measured. A model can look inexpensive from its active-parameter count while still hitting GPU-memory capacity or bandwidth limits because of activation movement.

MoEBlaze keeps the routing map instead of the routed copies

MoEBlaze’s primary contribution is a different representation of the routing process.

Rather than preserving large expert-ordered copies of token activations, it stores compact structures describing which tokens belong to which experts and where those relationships are located. Expert computation then gathers the required inputs directly from the original, unpermuted activation tensor.

The same idea extends to the output side. Instead of keeping another large routed output tensor before restoring token order, MoEBlaze fuses aggregation with the second expert MLP and reduces expert results directly into the final token-ordered output. During backpropagation, the routing mappings are reused so expanded routed-gradient buffers can also be avoided.

This is a systems redesign, not a new routing policy. The paper does not claim a different expert-selection objective or a new MoE model architecture. Its intervention is in how selected-token information is represented, moved, and retained on the GPU.

That distinction matters for adoption. A training-infrastructure team evaluating the approach is primarily considering a change to execution and memory management rather than a change to model semantics.

Compact routing only helps if constructing it is cheap

Replacing activation copies with indices creates another requirement: those indices must be built quickly enough that dispatch preparation does not become the new bottleneck.

MoEBlaze therefore constructs expert-centric routing metadata using a dense token-expert map, per-expert counting, prefix sums, and a location-map procedure designed to avoid both global sorting and atomic write contention.

This part of the paper is best understood as an implementation mechanism supporting the main architecture, rather than a separate empirical thesis. The metadata-driven training path depends on efficient routing-index construction; an expensive sort or highly contended dispatch stage could erase part of the benefit obtained by eliminating routed buffers.

For systems teams, the broader design criterion is useful: replacing a large tensor with metadata is economically attractive only when generating and consuming that metadata costs less than the movement it removes.

SwiGLU makes activation lifetime a second optimization target

Routing buffers are not the only tensors consuming bandwidth and capacity. SwiGLU introduces two projected branches:

A conventional training path can materialize several values associated with those projections and the elementwise activation because backward computation will need them later.

MoEBlaze fuses the two first-layer projections with the activation epilogue so the shared input can be loaded once and transient values can remain in registers or shared memory where possible. It also avoids storing selected inexpensive intermediates, recomputing the SiLU-related values during backward instead.

That recomputation is selective. It is not presented as generic whole-layer checkpointing. The trade is specifically to repeat relatively cheap operations when doing so avoids larger global-memory reads, writes, and saved tensors.

This extends the paper’s central idea from routing representation to activation lifetime: reduce what must travel through global memory and shorten how long temporary values need to remain materialized.

The benchmark shows memory savings can coincide with higher throughput

The paper’s main empirical evidence compares MoEBlaze with MegaBlocks across seven representative single-layer MoE configurations on one NVIDIA H100, using PyTorch 2.0.1 and CUDA 12.1. Configurations span input dimensions from 512 to 2048, 4 to 16 experts, top-$k$ values from 1 to 4, batch sizes of 16 or 32, and sequence lengths from 512 to 2048.

Saved activation memory is lower with MoEBlaze across the tested configurations.

One reported example in the section described as SiLU is conf4, where MoEBlaze uses about 6.1 GB of saved activation memory versus about 22 GB for MegaBlocks—roughly a 3.6× reduction. In that evaluation, reported training-speed improvements range from approximately 1.4× to 3.7×.

Under the SwiGLU evaluation, conf3 uses approximately 10 GB with MoEBlaze versus more than 40 GB with MegaBlocks. Reported speedups in that evaluation range from approximately 2× to 6.2×.

These comparisons are the paper’s main performance evidence. The routing-data structures and fused-kernel design explain the proposed mechanism; the source package does not provide a separate causal ablation establishing exactly how much of each speedup comes from each component.

The manuscript also contains inconsistent ReLU, SiLU, and SwiGLU labeling in parts of Section 6 and its figure captions. The numerical results above therefore retain the terminology of the sections in which they are reported rather than attempting to reconcile those labels.

The business value is additional room inside a fixed GPU budget

For a training-infrastructure team already constrained by HBM capacity, lower activation memory can change which configurations fit on existing hardware. The relevant decisions include increasing batch size, extending sequence length, or accommodating a larger MoE configuration without adding GPUs solely to absorb routing and intermediate activations.

Higher layer throughput can also reduce GPU time when the workload resembles the memory-intensive regimes tested here. The paper’s results are particularly relevant when routing volume, model dimension, top-$k$, or activation complexity make data movement a substantial part of execution cost.

A practical evaluation should therefore track at least three quantities together:

This is a Cognaptus inference from the systems evidence, not a demonstrated cluster-level ROI calculation. Actual economics depend on whether these layer-level bottlenecks remain dominant inside the complete training stack.

What still needs distributed validation

The experimental boundary is substantial enough to affect deployment decisions.

All reported benchmarks use one H100 and one MoE layer. Optimizer updates are excluded from the timing comparison. MegaBlocks is the primary baseline. The paper does not report complete-model training, downstream accuracy or convergence comparisons, multi-GPU expert parallelism, network communication effects, multi-node scaling, repeated-run uncertainty, or statistical significance tests.

Those omissions do not negate the reported memory reductions. They determine what can be concluded from them.

The evidence supports a narrower statement: under the seven tested single-H100 MoE-layer configurations, eliminating routed activation materialization and reducing intermediate storage coincides with substantially lower saved activation memory and higher forward-plus-backward throughput than MegaBlocks.

Before an infrastructure team treats 6.2× as a training-cluster expectation, it needs to verify whether communication, optimizer work, dense layers, synchronization, and distributed expert placement become the dominant costs once the local MoE data path is cheaper.

Sparse-model economics should include the cost of moving activations

MoEBlaze’s most transferable contribution is not the largest reported speedup. It is the systems framing behind the result.

Sparse computation determines how much expert arithmetic is executed. It does not determine how efficiently the surrounding activations are represented, transported, preserved, and reconstructed for backward computation.

For teams deciding whether MoE will fit a fixed GPU fleet or lower training cost, those quantities belong in the capacity model alongside FLOPs and active parameters. MoEBlaze provides evidence that, at least on a single H100 MoE layer, redesigning that data path can improve memory footprint and throughput at the same time.

The next test is whether those gains survive contact with the rest of the training system.

Cognaptus: Automate the Present, Incubate the Future.