OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

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. In this paper, we propose OmniMoE, a system-algorithm co-designed MoE framework that pushes granularity to the extreme with vector-level Atomic Experts, orchestrating their routing and execution at scale within a single MoE layer, while retaining a shared dense MLP for general-purpose processing. While 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$\times$ speedup) compared to PEER, demonstrating that massive-scale fine-grained MoE can be fast and accurate.

Large language models are becoming increasingly powerful, but they are also expensive to run because they use all of their internal parameters for every input. Mixture-of-Experts (MoE) models address this by activating only a small subset of parameters for each input, reducing computational cost while maintaining model capacity. However, making these expert subsets very small and numerous—which improves specialization—creates engineering challenges: deciding which experts to activate becomes expensive, and accessing many tiny experts scattered in memory is slow. Our work, OmniMoE, tackles both problems simultaneously through a co-design of the algorithm and system. We introduce a routing method that efficiently selects from a massive pool of fine-grained experts using a mathematical decomposition trick, and a scheduling strategy that reorganizes computation to avoid slow memory access patterns. Together, these innovations allow OmniMoE to use extremely small experts at scale without sacrificing speed. In experiments, OmniMoE matches or outperforms larger models while using significantly fewer active parameters, demonstrating that smarter expert orchestration can make AI models both more capable and more efficient.