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

这是一篇关于在 Apple Silicon 芯片的神经网络加速器（NPU，即 Apple Neural Engine, ANE）上高效运行混合专家模型（Mixture-of-Experts, MoE）大语言模型（LLM）推理的论文。论文提出了名为 NPUMoE 的运行时推理引擎。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

背景：

• MoE 架构的兴起： 现代大模型（如 PhiMoE, Qwen3 MoE）广泛采用 MoE 架构，通过稀疏激活（Sparse Activation）来扩展模型容量，提高推理效率。
• 设备端推理的需求： 在移动设备（如 iPhone, Mac）上进行 LLM 推理可以保护隐私并实现个性化，但面临资源竞争（CPU 需处理 UI/IO，GPU 需处理图形渲染）。
• NPU 的优势与局限： Apple Neural Engine (ANE) 是专为 AI 计算设计的 NPU，具有高吞吐量和低功耗。然而，现有的 NPU 设计主要针对静态、稠密的张量操作（如标准 Transformer 层），而 MoE 的推理过程具有高度的动态性和稀疏性，导致两者不匹配。

核心挑战：

1. 动态路由与形状不匹配： MoE 的专家路由（Expert Routing）是不可预测的，导致每个专家接收的 Token 数量动态变化。NPU 需要预先编译的静态计算图（Static Compute Graphs），无法直接处理动态张量形状。
2. 不规则算子支持不足： MoE 涉及 Top-k 选择、Scatter/Gather（分散/聚合）、动态索引等不规则操作，这些在 NPU 上不受支持或效率极低，迫使计算回退到 CPU。
3. 调度与同步开销大： 如果为每个小专家单独启动计算图，会导致大量的 NPU 启动（Dispatch）开销和 CPU-NPU 同步延迟。在细粒度工作负载下，同步开销可能占运行时间的 60% 以上，严重降低 NPU 利用率。
4. 长上下文预填充（Prefill）瓶颈： 对于长上下文任务，预填充阶段占用了 80% 的端到端延迟和 70% 的 CPU 周期，是优化的关键目标。

2. 方法论 (Methodology)

NPUMoE 的核心设计理念是：将 NPU 作为主要执行引擎处理稠密、静态计算，同时保留 CPU/GPU 作为动态操作和回退路径，通过离线校准和运行时策略最大化 NPU 利用率。

系统包含三个关键技术：

(1) 专家容量的静态层级 (Static Tiers for Expert Capacity)

• 问题： 专家接收的 Token 数量分布极度不平衡（长尾分布），直接为所有专家分配最大容量会导致大量填充（Padding）浪费，而统一小容量会导致 Token 丢弃。
• 方案：
  • 利用离线校准（Offline Calibration）统计每个专家在预填充阶段的预期负载和流行度。
  • 将专家划分为几个静态容量层级（Tiers）。高频使用的“热门专家”分配较大的容量层级，低频的“冷门专家”分配较小的层级。
  • 溢出处理： 当 Token 数量超过静态容量时，基于激活值的重要性（Saliency Score）进行剪枝（Pruning），丢弃低重要性 Token，而非简单丢弃或过度填充。

(2) 分组专家执行 (Grouped Expert Execution)

• 问题： 为每个专家单独启动 NPU 计算图会导致严重的序列化执行和启动开销。
• 方案：
  • 将多个专家打包到一个单一的、静态的稠密 FFN 计算图中。
  • 执行流程： 将 Token 按专家 ID 分组 -> 分配一个大的输入缓冲区 -> 将不同专家的 Token 填充到该缓冲区的不同切片中 -> 一次性在 NPU 上执行所有专家的 FFN 计算 -> 将结果聚合。
  • 优势： 将多次 NPU 启动合并为一次，显著降低了启动开销和同步延迟，提高了 NPU 的并行利用率。

(3) 负载感知的计算图驻留 (Load-Aware Compute Graph Residency)

• 问题： 并非所有计算都适合放在 NPU 上。冷专家（Cold Experts）或极小的分组任务，其计算量不足以抵消 CPU-NPU 的同步和调度开销。
• 方案：
  • 基于离线校准得出的专家流行度，动态决定计算图的驻留位置。
  • 热门专家组（Hot Groups）： 驻留在 NPU 上，因为其高负载足以摊销启动成本。
  • 冷门专家组（Cold Groups）： 直接回退到 CPU/GPU 执行，避免不必要的 NPU 调度开销。
  • 这种策略实现了计算单元的智能分配，平衡了延迟和能效。

3. 主要贡献 (Key Contributions)

1. 深入分析： 首次深入分析了 MoE 推理在 Apple Neural Engine 上的挑战，揭示了动态路由与 NPU 静态约束之间的根本冲突。
2. 三项创新技术： 提出了静态容量层级、分组专家执行和负载感知驻留策略，有效解决了动态性、算子支持和调度开销问题。
3. NPUMoE 系统实现： 构建了首个针对移动 NPU 优化的 MoE 推理引擎，实现了从离线校准到运行时调度的全流程优化。
4. 开源与评估： 在 Apple M2 系列芯片上进行了全面评估，代码将开源。

4. 实验结果 (Results)

实验在 Apple M2 Max 和 M2 Ultra 设备上，使用 Phi-3.5-MoE、Phi-tiny-MoE 和 Qwen3-30B-A3B 三个模型，针对长上下文预填充任务进行了测试。

• 延迟降低： 相比基线（CoreML CPU-only, CoreML Naïve, ANEMLL），NPUMoE 将预填充延迟降低了 1.32x – 5.55x。
• 能效提升： 每个 Token 的能耗降低了 1.81x – 7.37x。
• CPU 占用减少： CPU 周期使用量减少了 1.78x – 5.54x，有效释放了 CPU 资源用于其他系统任务。
• 精度损失极小： 在 Token 丢弃率约为 10-20% 的情况下，模型精度损失小于 1.1%（例如 BoolQ 数据集上无损失，HellaSwag 上仅损失 1.07%）。
• 消融实验： 证明了三项技术均对性能有显著贡献，其中分组执行（Grouped Execution）对降低延迟至关重要，而负载感知驻留对能效提升最大。

5. 意义与影响 (Significance)

• 解锁移动设备上的 MoE 推理： 证明了尽管 MoE 具有动态性，但通过系统级的协同设计（System Co-design），可以在资源受限的移动 NPU 上高效运行稀疏大模型。
• 优化设备端用户体验： 通过大幅降低 CPU 占用和能耗，使得在保持系统响应速度（UI 流畅、后台任务正常）的同时，能够运行更强大的本地大模型。
• 为未来 NPU 设计提供方向： 揭示了当前 NPU 在处理动态稀疏工作负载时的瓶颈，并为未来的硬件/软件栈优化（如更好的动态算子支持、更灵活的图调度）提供了实证依据。
• 长上下文处理： 特别针对长上下文预填充阶段进行了优化，这对于 RAG（检索增强生成）和长文档分析等实际应用至关重要。

总结： NPUMoE 通过巧妙的软件优化策略，成功弥合了 MoE 模型的动态特性与 Apple NPU 静态执行约束之间的鸿沟，实现了在移动设备上高效、低功耗且高精度的 MoE 大模型推理。