OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

Abstract: Mixture-of-Experts (MoE) architectures are evolving towards finer granularity to improve parameter efficiency. However, existing MoE designs face an inherent trade-off between the granularity of expert specialization and hardware execution efficiency. We propose OmniMoE, a system-algorithm co-designed framework that pushes expert granularity to its logical extreme. OmniMoE introduces vector-level Atomic Experts, enabling scalable routing and execution within a single MoE layer, while retaining a shared dense MLP branch for general-purpose processing. Although this atomic design maximizes capacity, it poses severe challenges for routing complexity and memory access. To address these, OmniMoE adopts a system-algorithm co-design: (i) a Cartesian Product Router that decomposes the massive index space to reduce routing complexity from O(N) to O(sqrt(N)); and (ii) Expert-Centric Scheduling that inverts the execution order to turn scattered, memory-bound lookups into efficient dense matrix operations. Validated on seven benchmarks, OmniMoE (with 1.7B active parameters) achieves 50.9% zero-shot accuracy across seven benchmarks, outperforming coarse-grained (e.g., DeepSeekMoE) and fine-grained (e.g., PEER) baselines. Crucially, OmniMoE reduces inference latency from 73ms to 6.7ms (a 10.9-fold speedup) compared to PEER, demonstrating that massive-scale fine-grained MoE can be fast and accurate. Our code is open-sourced at https://github.com/flash-algo/omni-moe.

• The paper introduces a novel MoE architecture that leverages atomic experts with a dense MLP to balance fine-grained specialization and overall expressivity.
• It employs a Cartesian Product Router to factorize the routing index space, reducing computational overhead and significantly lowering latency.
• The expert-centric scheduling optimizes hardware utilization by grouping tasks and ensuring contiguous memory access, which boosts throughput during inference.

OmniMoE: System-Algorithm Co-Design for Efficient Atomic Mixture-of-Experts at Scale

Introduction

OmniMoE proposes a Mixture-of-Experts architecture that reconciles the long-standing tension between fine-grained parameter efficiency and hardware execution efficiency in large-scale neural networks. While prior works either activate large, monolithic experts (coarse-grained MoEs) or millions of lightweight units (fine-grained MoEs), both paradigms suffer from significant inefficiencies: coarse-grained approaches waste compute and memory by activating irrelevant parameters, and fine-grained designs suffer severe bandwidth bottlenecks due to scattered memory accesses and reduced expressivity. OmniMoE shifts expert design to its logical atomic extreme—parameterizing each expert by an input and output vector—while introducing a co-designed routing and execution system that renders such fine granularity scalable, expressive, and hardware-efficient.

Figure 1: Contrasting activation and optimization patterns for coarse-grained, fine-grained, and OmniMoE architectures, highlighting the system-level and algorithmic bottlenecks mitigated by the proposed approach.

OmniMoE Architecture

The OmniMoE architecture is a dual-pathway system that integrates:

• A routed branch with atomic experts: Each atomic expert is a minimal unit parameterized by a pair of vectors, assembled dynamically on a per-token basis, maximizing specialization and minimizing redundant computation.
• A universally active shared dense MLP: This branch provides general-purpose reasoning and semantic stability, compensating for the limited expressivity of atomic units.

The core routing mechanism incorporates a Cartesian Product Router, which factorizes the routing index space to reduce computational and storage cost from $O(N)$ to $O(\sqrt{N})$, making million-scale expert selection practical. Downstream, all selected atomic experts are dynamically assembled into token-dependent, dense parameter blocks.

Figure 2: The OmniMoE framework, with parallel routed and dense branches, the former employing a Cartesian Product Router and dynamic expert assembly for token-specific projection, the latter guaranteeing shared semantic coverage.

Routing and Systemic Scheduling

Atomic Experts and Dynamic Assembly: Each atomic expert $E_i$ is realized as

$E_i(x) = \sigma(x {w^{in}_i}^\top) w^{out}_i,$

where $w^{in}_i$, $w^{out}_i$ are vectors and $\sigma$ is the activation function (SwiGLU in their implementation). For each input token, top-$K$ relevant experts are dynamically retrieved using the routing mechanism, and their parameters are assembled on-the-fly, guaranteeing that only relevant—and all relevant—parameters are activated.

Cartesian Product Router: To circumvent the scaling collapse of naive routers in the high-expert regime, each expert is indexed as a pair $(i,j)$ in a $N_r \times N_c$ grid. Scoring is factorized via two projections, dramatically reducing computational complexity and parameter storage without sacrificing coverage of the expert pool.

Expert-Centric Scheduling: OmniMoE inverts conventional token-centric scheduling for MoEs, in which each token independently retrieves its routed experts—resulting in random, bandwidth-inefficient memory accesses. By grouping tasks by active expert and sorting by token ID within each expert, memory lookups become contiguous and highly reusable. This enables efficient Grouped GEMM operations, aligning computation with GPU architectural strengths and dramatically improving throughput.

Figure 3: Comparison between token-centric and expert-centric scheduling, showing how OmniMoE's strategy yields dense, fused batched compute aligned with GPU resources.

Empirical Results

OmniMoE is validated across seven language modeling and reasoning benchmarks, and against established FFN variants including Dense, GShard, DeepSeekMoE, PKM, and PEER. Evaluation focuses on accuracy, latency, and memory utilization under fixed activated parameter budgets and at matched scales.

• Accuracy: OmniMoE achieves a zero-shot average accuracy of 50.9%, outperforming DeepSeekMoE (+0.7) and PEER (+2.0). Gains are especially notable in knowledge-intensive (e.g., TriviaQA, OBQA) and reasoning-intensive (ARC, HellaSwag) tasks, demonstrating the complementarity of the hybrid design.
• Efficiency: OmniMoE reduces inference latency from 73ms (PEER) and 102ms (DeepSeekMoE) to 6.7ms at 4K tokens (10.9–15.2× speedup) and maintains memory usage consistent with coarse-grained MoEs—despite activating more parameters.

  Figure 4: End-to-end latency and memory versus active parameters and token count, showing that OmniMoE dominates both fine- and coarse-grained baselines on efficiency.

• Scaling Laws: Under matched computational budgets, OmniMoE achieves consistently lower perplexity, reflecting improved scaling characteristics in both compute and parameter efficiency.

  Figure 5: Scaling curves for validation perplexity, showing superior trade-off for quality and cost with OmniMoE.

• Communication and Parallelism: Communication cost for expert parallelism saturates once the number of unique activated experts matches activation demand, and grows only linearly with sequence length, not total model size—contrasting with traditional approaches where communication is a function of overall expert pool size.

  Figure 6: Communication overhead is constant beyond the activation count, decoupling parallel communication from total expert pool size.

  

  Figure 7: Communication overhead grows strictly linearly with sequence length, confirming scalability for long-sequence training.

Ablations and Analysis

Ablation studies isolate the contributions of the shared dense MLP, Cartesian Product Router, and expert-centric scheduling:

• Removing the dense MLP compromises knowledge and reasoning while yielding negligible efficiency gains.
• Substituting the Cartesian Product Router incurs catastrophic degradation in routing efficiency and expert utilization, collapsing into mode collapse (only 4% of experts active).
• Dropping expert-centric scheduling explodes latency (24.8×) and memory (417×) without improving quality, confirming that hardware alignment is the main bottleneck at fine expert granularity.

Implications and Future Perspectives

OmniMoE demonstrates that co-designing algorithmic and system aspects of MoE architectures enables practical deployment of extreme fine-grained expert pools with uncompromised expressivity, scalability, and resource utilization. The findings have several practical and theoretical implications:

• Practical Efficiency: Dramatic reductions in latency and bandwidth cost validate the deployment of atomic-level expert routing even at massive scale, paving the way for resource-flexible, energy-efficient deployment in production systems.
• Theoretical Capacity: OmniMoE empirically confirms scaling-law hypotheses that increasing activation granularity (when paired with efficient routing and scheduling) improves both compute and parameter efficiency.
• Composable Routing: The Cartesian Product Router suggests broader applicability for subspace-factorized attention and routing mechanisms in structured sparsity scenarios.
• Future Architectures: Holistic system-algorithm codesign should become the default paradigm when pushing the limits of expert capacity, sparsity, and model modularity.

Conclusion

OmniMoE establishes that hardware-efficient, highly expressive, and scalable atomic expert routing is achievable through tight integration of architectural innovation (dynamic expert assembly, hybrid pathway design), novel algorithmic routing (Cartesian Product Router), and hardware-aligned scheduling (expert-centric scheduling). The framework sets the benchmark for future research on scalable, sparse, and efficient model architectures, particularly as model capacities continue to scale beyond current regimes, and motivates new explorations into structured routing strategies, load-balancing, and expert specialization at scale.

Reference:

"OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale" (2602.05711)