PaperlayerView exact revision on arXiv

Source-linked AI summary

OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

Jingze Shi, Zhangyang Peng, Yizhang Zhu, Yifan Wu, Guang Liu, Yuyu Luo

arXiv:2602.05711v2cs.CLcs.AI

TL;DR

Fine-grained MoEs offer parameter efficiency but challenge hardware execution efficiency. OmniMoE combines atomic experts with a Cartesian Product Router and Expert-Centric Scheduling, achieving higher accuracy and a 10.9× inference speedup over PEER.

  • Problem

    Fine-grained MoEs seek parameter efficiency, but the field lacks a way to reconcile their specialization with the hardware efficiency of coarse-grained designs.

  • Method

    OmniMoE combines a shared dense MLP with massive vector-level atomic experts, using Cartesian Product Routing and Expert-Centric Scheduling for scalable activation and execution.

  • Results

    10.9× inference speedup over PEER accompanies consistently improved average accuracy and stronger performance on most benchmarks under comparable activated-parameter budgets.

  • Takeaways & Limitations

    The results support holistic system-algorithm co-design as a way to make massive-scale fine-grained MoE activation practical at scale.

  • Takeaways & Limitations

    The evaluation varies only the FFN module and relies on controlled pre-training from scratch rather than off-the-shelf checkpoint comparisons.

Abstract

from arXiv · show

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.

1. Introduction

MoE expert granularity balances routing precision and system efficiency: coarse-grained designs waste computation, while fine-grained designs face expressivity, routing, and memory-access challenges. OmniMoE addresses this trade-off with atomic experts, a shared dense MLP, Cartesian Product Routing, and Expert-Centric Scheduling.

  • Motivation: Coarse-grained MoEs use hardware-efficient dense matmuls and contiguous VRAM access but activate large blocks that waste computation on irrelevant parameters.DeepSeek-V3 and KIMI-K2 exemplify this design with 256 and 384 experts, respectively.
  • Motivation: Fine-grained MoEs improve activation efficiency and flexibility through lightweight experts, but massive scaling restricts expressivity and increases routing and memory-access overhead.Existing designs such as PEER reduce experts to static parameter vectors, limiting computation to linear vector aggregation.
  • OmniMoE: OmniMoE combines a shared dense MLP for general semantic knowledge with routed fine-grained experts for specialized long-tail knowledge retrieval.This hybrid parallel design is intended to combine precise parameter activation with hardware efficiency.
  • OmniMoE: Atomic Experts provide a minimal routable unit parameterized by two vectors, while Dynamic Expert Assembly composes massive experts into token-specific parameter configurations.The formulation supports highly specialized, high-expressivity compositions across a massive expert pool.
  • OmniMoE: The Cartesian Product Router decomposes the one-dimensional expert index space into a two-dimensional grid and factorizes routing into independent low-dimensional projections.This reduces routing cost from linear in N to proportional to sqrt(N), making large-scale routing practical.
  • OmniMoE: Expert-Centric Scheduling inverts token-centric execution, groups requests targeting the same experts, and converts scattered accesses into contiguous, coalesced memory operations.The strategy targets poor locality and low GPU efficiency caused by fine-grained routing.

2. Methodology

OmniMoE combines vector-level atomic experts with Dynamic Expert Assembly, a Cartesian Product Router, and Expert-Centric Scheduling to scale sparse routing while improving execution efficiency. The methodology retrieves token-specific expert parameters, factorizes massive expert indexing, and executes grouped dense computations.

  • Overall Architecture: OmniMoE follows a standard MoE layer with sparse top-K routing and a universally activated shared dense MLP branch.For each token, routing scores select activated experts whose weighted outputs are combined with the shared MLP output.
  • Atomic Experts and DEA: Atomic Experts are minimal routable units whose sparse retrieval and composition through DEA form a token-conditioned assembled expert.DEA retrieves the most relevant atomic experts and assembles their parameters into a fused computation, giving each token a unique assembled expert.
  • Atomic Experts and DEA: DEA stores all N atomic-expert parameters in two global matrices, then gathers top-K indexed rows and combines them with their routing scores.This centralized repository enables compact, token-local parameter blocks before weighted assembly.
  • Cartesian Product Router: The Cartesian Product Router represents each expert as a coordinate on an Nr×Nc grid and predicts separate row and column distributions instead of an N-way classifier.Factorized projections Wr and Wc produce log-probabilities whose sums implicitly score grid coordinates without materializing the full score matrix.
  • Expert-Centric Scheduling: Expert-Centric Scheduling clusters tasks by active experts, sorts tokens within groups, and executes each group with fused high-throughput Grouped GEMM kernels.Dense parameter blocks, stacked token inputs, and gating weights improve parameter reuse, memory locality, input-read coalescing, and output scattering.

3. Experiments

Experiments show that OmniMoE combines strong accuracy, throughput, latency, scalability, and parameter efficiency through its hybrid dense-plus-atomic design, Cartesian Product Router, and Expert-Centric Scheduling. Ablations confirm that the shared dense branch supports quality while the two system components deliver efficient routing and execution.

  • Experimental Setup: OmniMoE evaluates controlled FFN comparisons, speed and memory sweeps, and scaling-law families from 280M-A80M to 6.4B-A1.7B against matched Dense counterparts.All models share the same Transformer backbone, varying only the FFN module; activated parameters count unique parameters used for one token.
  • Downstream Performance: 50.9 zero-shot accuracy is achieved by 6.4B-A1.7B OmniMoE, exceeding DeepSeekMoE by +0.7 and PEER by +2.0.OmniMoE also improves knowledge-intensive TriviaQA by +1.1 and OBQA by +1.4 versus coarse-grained models.
  • Reasoning Performance and Throughput: ∼14.2k tok/s throughput is achieved by OmniMoE, near Dense at ∼14.8k tok/s and above PEER at ∼11.6k tok/s and PKM at ∼7.9k tok/s.The hybrid architecture combines precise retrieval with stable reasoning capacity, whereas purely fine-grained designs without a shared dense backbone struggle with sustained multi-step logic.
  • End-to-End Efficiency and Scalability: 6.7 ms latency at 4,096 tokens is achieved with 28M activated parameters, versus 73 ms for PEER and 102 ms for DeepSeekMoE.This corresponds to 10.9× and 15.2× speedups, respectively, while maintaining memory comparable to coarse-grained MoEs; expert-centric batching enables the gain.
  • End-to-End Efficiency and Scalability: Constant communication cost is observed once the expert pool exceeds the active token count, supporting scaling to millions of experts.Under matched training FLOPs and activated-parameter budgets, OmniMoE consistently attains the lowest perplexity among baselines.
  • Ablation Studies: 30.6× latency and 337.5× memory result from replacing the Cartesian Product Router, while token-centric execution causes 24.8× latency and 417.7× memory.Removing the shared dense MLP slightly improves efficiency but worsens perplexity to 1.2×, knowledge performance to 0.91×, and reasoning performance to 0.79×.

4. Related Work

Related work spans coarse-grained MoE architectures with lightweight routing and system optimizations targeting kernel efficiency, communication, and padding overhead. OmniMoE differs by applying Expert-Centric Scheduling to fine-grained atomic experts, converting scattered accesses into contiguous batched operations.

  • MoE Architectures: MoE architectures scale capacity with bounded per-token cost through conditional computation.Early MoE Transformers primarily use coarse-grained FFN experts with lightweight routing, including Switch Transformers’ Top-1 gating and GShard.
  • MoE Architectures: Switch Transformers and GShard exemplify early coarse-grained FFN MoEs with lightweight routing.Switch Transformers uses Top-1 gating.
  • Efficient MoE Systems: DeepSpeed-MoE, Fast-MoE, and MegaBlocks optimize coarse-grained MoEs through GEMM kernels and variable-length sequence handling.These optimizations mitigate padding overheads and generally emphasize kernel fusion and communication scheduling.
  • Efficient MoE Systems: OmniMoE’s Expert-Centric Scheduling targets fine-grained atomic experts by converting scattered memory accesses into contiguous batched operations.SonicMoE instead uses memory-efficient algorithms, minimal activation caching, and tile-aware token rounding to reduce padding waste in Grouped GEMM kernels.

5. Conclusion

OmniMoE is a system-algorithm co-designed MoE framework combining a shared dense MLP with massive atomic experts for precise parameter activation. It makes large-scale expert activation practical through the Cartesian Product Router and Expert-Centric Scheduling.

  • 5. Conclusion: OmniMoE integrates a shared dense MLP for general-purpose reasoning with massive atomic experts for long-tail knowledge retrieval.This combination enables more precise parameter activation.
  • 5. Conclusion: The framework makes large-scale expert activation practical through two key innovations: the Cartesian Product Router and Expert-Centric Scheduling.
  • 5. Conclusion: Together, these components yield a dramatic 10.9× improvement.

A. Complexity Analysis

This section shows that OmniMoE reduces routing complexity through Cartesian factorization and substantially lowers expert-parameter I/O through expert-centric scheduling, while keeping scheduling overhead below 5% of total latency.

  • Routing Complexity: 2 · 10^9 FLOPs per token arise from standard routing when N = 10^6, making million-scale expert routing computationally prohibitive.Standard parameter storage is Mstd = d · N.
  • Top-K Selection Complexity: The fused Block-wise Merge Selection kernel avoids full score materialization and eliminates the dominant O(N) global-memory I/O bottleneck.Scores remain in registers/SRAM, making top-K selection compute-bound and effectively negligible in latency on modern GPUs.
  • Memory Traffic Analysis: η ≫ 1 indicates substantial parameter-I/O reduction when L · K ≫ |Eactive|, because each active expert is loaded exactly once and shared across assigned tokens.Expert-centric scheduling inverts execution order and groups tokens by unique activated expert.
  • Scheduling Overhead and Token Traffic: < 5% of total latency is occupied by scheduling, as radix-sorted routing tasks amortize preprocessing overhead against GEMM speedups.Hierarchical sorting also produces sequential, coalesced token reads rather than the random access patterns of token-centric baselines.

B. Experimental Setup

The experiments use controlled from-scratch pre-training and method-specific optimized kernels to isolate architectural effects. Speed, memory, and language-modeling evaluations vary activation budgets and tokens under specified configurations.

  • Experimental methodology: Baselines are trained from scratch on an identical corpus to isolate the impact of MoE architectural design from data and training-recipe differences.The comparison represents Gshard and DeepSeekMoE as architectural prototypes rather than using off-the-shelf checkpoints.
  • Implementation: Coarse-grained baselines use NVIDIA CuTile kernels, fine-grained baselines use optimized Triton fused kernels, and OmniMoE uses custom Triton Expert-Centric Scheduling.The implementations are chosen to maximize hardware utilization for each evaluated method.
  • Speed and memory benchmarking: Speed and memory tests sweep either activated parameters or activated tokens while fixing the other at its minimum configuration.Activated tokens are fixed at 1K when sweeping Act Params; Act Params is fixed at each method’s minimum when varying Act Tokens.
  • Language modeling: Language-modeling experiments span 80M, 200M, 680M, and 1.7B activation scales with specified training steps, batch tokens, learning rates, depths, model widths, and embedding tying.Total Params and Act Params denote total model size and per-token active parameters, respectively.

C. Expert Parallelism Communication Overhead

OmniMoE’s expert-parallel communication remains bounded by activated experts rather than total model capacity, while communication volume grows linearly with sequence length. Even at millions of experts or ultra-long sequences, reported latency remains low.

  • Scalability with Number of Experts: Communication latency and memory usage saturate as experts increase beyond the activation count.Figure A evaluates scaling from 1K to 2M experts under fixed sequence length.
  • Scalability with Sequence Length: Communication volume scales linearly with sequence length as token distribution and gradient aggregation dominate expert-parallel communication.Figure B studies sequence lengths from 1K to 128K under a fixed number of experts.
  • Scalability with Number of Experts: 16,384 activated experts mark the saturation point, after which backward communication stabilizes at approximately 80MB regardless of expert count.In an 8-GPU environment, the corresponding communication latency is 0.521 ms.
  • Scalability with Sequence Length: Approximately 6GB communication volume at ultra-long sequences corresponds to an estimated 15 ms latency on a 64-GPU cluster.The result indicates efficient communication for long-sequence training.
Loading 2602.05711v2…