OmniMoE: Эффективная смесь экспертов за счет масштабируемого управления атомарными экспертами

OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

February 5, 2026

Авторы: Jingze Shi, Zhangyang Peng, Yizhang Zhu, Yifan Wu, Guang Liu, Yuyu Luo

cs.AI

Аннотация

Архитектуры смеси экспертов (MoE) эволюционируют в сторону более мелкой гранулярности для повышения эффективности использования параметров. Однако существующие проекты MoE сталкиваются с неизбежным компромиссом между степенью специализации экспертов и эффективностью аппаратного выполнения. Мы предлагаем OmniMoE — совместно спроектированный на уровне системы и алгоритмов фреймворк, который доводит гранулярность экспертов до логического предела. OmniMoE вводит векторные атомарные эксперты, обеспечивая масштабируемую маршрутизацию и выполнение в рамках одного слоя MoE, сохраняя при этом общую плотную ветвь MLP для обработки общего назначения. Хотя такой атомарный дизайн максимизирует ёмкость, он создаёт серьёзные проблемы для сложности маршрутизации и доступа к памяти. Для их решения OmniMoE использует совместное проектирование системы и алгоритма: (i) маршрутизатор на основе декартова произведения, который разлагает огромное пространство индексов, снижая сложность маршрутизации с O(N) до O(√N); и (ii) ориентированное на экспертов планирование, которое инвертирует порядок выполнения, превращая разрозненные, ограниченные памятью поиски в эффективные плотные матричные операции. Проверка на семи бенчмарках показала, что OmniMoE (с 1.7B активными параметрами) достигает 50.9% точности в режиме zero-shot, превосходя грубозернистые (например, DeepSeekMoE) и мелкозернистые (например, PEER) базовые линии. Что важно, OmniMoE сокращает задержку вывода с 73 мс до 6.7 мс (ускорение в 10.9 раз) по сравнению с PEER, демонстрируя, что крупномасштабные мелкозернистые MoE могут быть быстрыми и точными. Наш код открыт по адресу 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.