NPUMoE: Efficiently Running Mixture-of-Experts Large Models on Apple Silicon NPU

The research team proposes the NPUMoE inference engine, which successfully offloads Mixture-of-Experts (MoE) large model inference to the Apple Neural Engine (ANE) through techniques like static layering, grouped execution, and load awareness. This achieves significant performance and energy efficiency improvements, providing a solution for edge devices to run large models efficiently.

The Computing Power Dilemma of Edge AI

As the capabilities of large language models improve, edge devices (laptops, tablets, phones) have an urgent need for local AI assistants, but they are limited by resources. Mixture-of-Experts (MoE) models, with their sparse activation mechanism, reduce computational requirements while maintaining model capacity, providing a path for edge deployment.

Apple Neural Engine: An Underrated AI Accelerator

The ANE integrated into Apple Silicon is a dedicated NPU that excels at core AI operations like matrix multiplication, with low energy consumption and high throughput. Offloading AI computations to the NPU can improve battery life and reduce heat, but the NPU is more suitable for regular static tasks, posing challenges for MoE deployment.

The dynamic sparsity of MoE models conflicts with the static regularity of NPUs. The core challenges include:

1. Unpredictable Expert Routing: The data dependency of the routing network leads to dynamic changes in tensor shapes, conflicting with the NPU's requirement for shape determination at compile time;
2. Presence of Irregular Operators: Operations like Top-k selection and scatter/gather have irregular index patterns, making it difficult to map to NPU's parallel matrix operations;
3. Fine-Grained Kernel Launch Overhead: Frequent launches of small expert kernels bring significant scheduling and synchronization overhead, offsetting the benefits of sparse computing.

To address the NPU adaptation challenges, NPUMoE adopts the following strategies:

1. Static-Layered Expert Capacity Management: Offline calibration of expert capacity and popularity, pre-dividing static layers to allow the NPU to determine tensor shapes at compile time;
2. Grouped Expert Execution: Combining multiple experts into batches and submitting them to the NPU, reducing kernel launch overhead and improving parallel utilization;
3. Load-Aware Computational Graph Residency: Intelligently deciding whether to keep computations on the NPU or fall back to CPU/GPU, minimizing cross-device synchronization overhead.

Evaluations on Apple M-series devices using 3 MoE models and 4 long-context workloads show:

• Latency Reduction: 1.32x to 5.55x, with significant acceleration in the long-context prefill phase;
• Energy Efficiency Improvement: 1.81x to 7.37x, extending mobile device battery life;
• CPU Usage Reduction: 1.78x to 5.54x, improving system responsiveness and heat dissipation; All optimizations maintain full model accuracy without affecting mathematical behavior.

Long-context capabilities (document analysis, code understanding, multi-turn dialogue, etc.) are an important trend in large model applications. NPUMoE fully leverages the NPU's advantages in the long-sequence prefill phase (computationally intensive), making it possible for edge devices to handle long contexts.

The success of NPUMoE reveals the importance of algorithm-hardware co-design in the edge AI era. The Apple NPU was not designed specifically for MoE, but through targeted scheduling optimizations, it becomes an ideal platform. This "hardware-software co-design" approach is worth learning for edge AI developers.

Future research directions for NPUMoE include:

• Expand to other NPU architectures like Qualcomm and MediaTek;
• Explore dynamic shape compilation techniques to reduce static layering constraints;
• Combine quantization and pruning techniques for more aggressive efficiency improvements;
• Develop adaptive load balancing strategies to handle complex scenarios. Once MoE architectures become mainstream, dedicated inference engines will be key components of edge AI infrastructure.