AppleシリコンNPUによる効率的なMixture-of-Experts（MoE）LLM推論

要旨: Apple Neural Engine (ANE) は、すべての Apple Silicon チップに搭載されている専用のニューラル処理ユニット（NPU）です。Mixture-of-Experts（MoE）LLM は疎な活性化によって推論効率を改善しますが、NPU には次の 3 つの点で難しさがあります。専門家のルーティングが予測不能であり、NPU の形状に依存した制約と衝突する動的なテンソル形状を導入すること、top-k や scatter/gather などのいくつかの不規則な演算子が NPU に適していないこと、そして多数の小さな専門家カーネルを起動すると、ディスパッチおよび同期のオーバーヘッドが大きくなることです。NPU は AI 計算を CPU と GPU からオフロードするために設計されています。本研究の目的は、MoE 推論、特に長いコンテキストのワークロードがシステム資源を大量に消費する prefill 時に、こうしたオフロードを可能にすることです。
本論文では NPUMoE を提示します。これは、密で静的な計算を NPU にオフロードすることで Apple Silicon 上での MoE 実行を加速しつつ、動的な演算のための CPU/GPU のフォールバック経路を維持するランタイム推論エンジンです。NPUMoE はオフラインのキャリブレーションを用いて専門家の容量（capacity）と人気度（popularity）を推定し、これを駆動力として 3 つの主要技術を用います。(1) 動的な専門家ルーティングに対処するための専門家容量の静的ティア（Static tiers）、(2) NPU の並行性（concurrency）制限を緩和するためのグループ化された専門家実行（Grouped expert execution）、(3) CPU-NPU 間の同期オーバーヘッドを削減するための、負荷を考慮した専門家計算グラフのレジデンシ（residency）です。3 種類の代表的な MoE LLM と 4 つの長コンテキスト・ワークロードを用いた、Apple M シリーズデバイスでの実験により、NPUMoE は一貫してベースラインを上回り、レイテンシを 1.32x〜5.55x 削減し、エネルギー効率を 1.81x〜7.37x 改善し、効果的な NPU オフロードによって CPU サイクル使用量を 1.78x〜5.54x 削減することが示されています。