OmniMoE: Un MoE Efficiente attraverso l'Orchestrazione di Esperti Atomici su Larga Scala

OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

February 5, 2026

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

cs.AI

Abstract

Le architetture Mixture-of-Experts (MoE) si stanno evolvendo verso una granularità più fine per migliorare l'efficienza dei parametri. Tuttavia, i progetti MoE esistenti affrontano un compromesso intrinseco tra la granularità della specializzazione degli esperti e l'efficienza di esecuzione hardware. Proponiamo OmniMoE, un framework co-progettato a livello di sistema e algoritmo che spinge la granularità degli esperti al suo estremo logico. OmniMoE introduce Atomic Experts a livello vettoriale, consentendo un instradamento e un'esecuzione scalabili all'interno di un singolo layer MoE, mantenendo al contempo un ramo MLP denso condiviso per l'elaborazione generica. Sebbene questo design atomico massimizzi la capacità, pone serie sfide per la complessità di instradamento e l'accesso alla memoria. Per affrontare queste problematiche, OmniMoE adotta una co-progettazione sistema-algoritmo: (i) un Cartesian Product Router che scompone lo spazio massiccio degli indici per ridurre la complessità di instradamento da O(N) a O(√N); e (ii) un Expert-Centric Scheduling che inverte l'ordine di esecuzione per trasformare ricerche sparse e vincolate dalla memoria in efficienti operazioni dense su matrici. Validato su sette benchmark, OmniMoE (con 1.7B di parametri attivi) raggiunge un'accuratezza zero-shot del 50.9% su sette benchmark, superando i baseline a granularità grossolana (es. DeepSeekMoE) e fine (es. PEER). Fondamentalmente, OmniMoE riduce la latenza di inferenza da 73ms a 6.7ms (un'accelerazione di 10.9 volte) rispetto a PEER, dimostrando che un MoE su larga scala a granularità fine può essere veloce e accurato. Il nostro codice è open-source all'indirizzo https://github.com/flash-algo/omni-moe.

English

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.