Efficient Mixture-of-Experts LLM Inference with Apple Silicon NPUs

1. Problem Statement

The paper addresses the challenge of running Mixture-of-Experts (MoE) Large Language Models (LLMs) efficiently on Apple Silicon NPUs (specifically the Apple Neural Engine, or ANE). While MoE architectures improve model capacity through sparse activation, they are fundamentally misaligned with the hardware constraints of mobile NPUs in three key ways:

1. Dynamic vs. Static Constraints: NPUs require static, pre-compiled compute graphs with fixed tensor shapes. MoE inference involves dynamic expert routing where the number of tokens assigned to each expert varies unpredictably per prompt, creating dynamic tensor shapes that NPUs cannot natively handle.
2. Unsupported Operators: MoE relies on irregular operations (e.g., top-k routing, scatter/gather, dynamic indexing) that are not optimized or supported by NPUs, forcing a fallback to the CPU.
3. Dispatch and Synchronization Overhead: Launching many small, independent expert kernels on an NPU incurs significant dispatch latency and synchronization costs. If the workload is split into too many small subgraphs, the overhead of CPU-NPU coordination can dominate execution time (accounting for >60% of runtime in worst cases), negating the performance benefits of offloading.

The authors focus specifically on the prefill phase of long-context workloads, which dominates end-to-end inference latency (80%) and CPU cycle usage (70%).

2. Methodology: NPUMoE

The authors propose NPUMoE, a runtime inference engine designed to offload dense, static computations to the NPU while maintaining a CPU/GPU fallback path for dynamic operations. The system relies on offline calibration to estimate expert capacity and popularity, driving three core techniques:

A. Static Tiers for Expert Capacity (§4.1)

To resolve the conflict between dynamic token routing and static NPU shapes:

• Concept: Instead of allocating a uniform capacity to all experts (which causes massive padding waste) or dynamic shapes (which NPU cannot support), the system assigns experts to static capacity tiers (e.g., Tier 1, 2, 3) based on their expected load.
• Mechanism: High-frequency ("hot") experts are mapped to larger capacity tiers, while low-frequency ("cold") experts get smaller tiers.
• Overflow Handling: If a hot expert receives more tokens than its tier capacity, the system prunes the overflow tokens based on an activation-based saliency score (discarding low-importance tokens) rather than padding, balancing accuracy and compute efficiency.

B. Grouped Expert Execution (§4.2)

To mitigate NPU concurrency limits and dispatch overhead:

• Concept: Instead of launching a separate compute graph for each expert, NPUMoE groups multiple experts into a single, fused, static dense FFN compute graph.
• Mechanism: Tokens are gathered and packed into a contiguous buffer where each expert processes a fixed slice. This allows a single NPU invocation to execute multiple experts simultaneously, amortizing the launch overhead and maximizing NPU utilization.
• Trade-off: The system carefully balances group size and per-expert capacity to avoid over-provisioning (wasted compute) or under-provisioning (token dropping).

C. Load-Aware Compute Graph Residency (§4.3)

To minimize CPU-NPU synchronization latency:

• Concept: Not all expert groups should run on the NPU. Small, bursty workloads (e.g., cold experts or small groups) suffer more from dispatch overhead than they gain from NPU speed.
• Mechanism: The system uses offline calibration to identify hot experts (highly active) and cold experts.
  • Hot Groups: Grouped experts with high demand are kept resident on the NPU to amortize launch costs.
  • Cold Groups: Experts with low demand or very small groups are relegated to the CPU/GPU fallback path, avoiding unnecessary NPU dispatch and synchronization.

3. Key Contributions

1. System Analysis: An in-depth analysis of the specific constraints of Apple Neural Engine (ANE) when applied to dynamic MoE workloads, highlighting the mismatch between sparse routing and static graph execution.
2. Novel Techniques: The introduction of three co-designed techniques:
   • Static capacity tiers to handle dynamic routing.
   • Grouped expert execution to reduce dispatch overhead.
   • Load-aware residency policies to optimize CPU-NPU coordination.
3. NPUMoE Engine: The first runtime inference engine specifically optimized for MoE models on mobile NPUs, utilizing offline calibration to drive runtime decisions.
4. Comprehensive Evaluation: A rigorous evaluation across multiple models and workloads demonstrating significant gains in latency, energy, and CPU usage without significant accuracy loss.

4. Experimental Results

The system was evaluated on Apple M2 Max and M2 Ultra devices using three MoE models (Phi-3.5-MoE, Phi-tiny-MoE, Qwen3-30B-A3B) and long-context workloads (HellaSwag, BoolQ, RULER).

• Latency: NPUMoE reduced prefill latency by 1.32x to 5.55x compared to baselines (Core ML CPU-only, naive Core ML, and ANEMLL).
• Energy Efficiency: It improved energy efficiency by 1.81x to 7.37x (measured in energy per token).
• CPU Usage: It reduced CPU cycle usage by 1.78x to 5.54x, effectively offloading compute to the NPU.
• Accuracy: The system maintained high accuracy with negligible degradation (<1.1% compared to FP16 baselines), even with token dropping rates of ~10-20%.
• Comparison to GPU: While GPUs were faster (1.26x–5.4x), they consumed significantly more energy (1.87x–4.48x) and power. NPUMoE prioritizes sustained on-device efficiency and reduced CPU contention over raw peak throughput.

5. Significance

This work demonstrates that mobile NPUs can be effectively utilized for sparse, dynamic MoE inference, a task previously thought to be too complex for static hardware accelerators. By transforming the MoE execution pipeline to fit NPU constraints (static shapes, grouped execution) while intelligently managing the CPU-NPU boundary, NPUMoE enables:

• Privacy-Preserving On-Device AI: Allowing large, complex MoE models to run locally without offloading to the cloud.
• System Responsiveness: By offloading heavy compute to the NPU, the CPU and GPU are freed for UI, I/O, and graphics tasks, preventing system lag during heavy LLM usage.
• Energy Efficiency: Extending battery life for on-device AI applications by leveraging the high-efficiency NPU for the bulk of the computation.

The paper concludes that careful systems co-design can bridge the gap between the dynamic nature of modern LLMs and the static constraints of mobile NPUs, making them viable targets for next-generation on-device AI.