論文「Efficient Mixture-of-Experts LLM Inference with Apple Silicon NPUs」の技術的サマリー

この論文は、Apple Silicon チップに搭載された専用ニューラルプロセッシングユニット（NPU: Apple Neural Engine, ANE）を用いて、大規模言語モデル（LLM）の「Mixture-of-Experts（MoE）」アーキテクチャを効率的に推論するためのランタイム推論エンジン**「NPUMoE」**を提案するものです。MoE モデルの動的な特性と NPU の静的な制約の間のミスマッチを解決し、特に長いコンテキストを持つプリフィル（prefill）フェーズにおける遅延、エネルギー効率、CPU 負荷を大幅に改善することを目的としています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

---

1. 問題定義 (Problem)

MoE アーキテクチャは、スパースな活性化によりモデル容量を効率的にスケーリングしますが、Apple Silicon のようなモバイル NPU での推論には以下の 3 つの重大な課題があります。

1. 予測不可能なエキスパートルーティングと動的なテンソル形状:
   • MoE では、トークンごとにどのエキスパート（専門家の FFN ブロック）を活性化するかを動的に決定します（Top-k ルーティング）。
   • これにより、各エキスパートへの入力トークン数がランダムに変化し、動的なテンソル形状が生じます。
   • NPU は事前コンパイルされた静的な形状の計算グラフに最適化されているため、この動的な形状は NPU の制約と衝突し、CPU へのフォールバックを招きます。
2. 不規則な演算子のサポート不足:
   • Top-k 選択、Scatter/Gather（分散・集約）、動的インデックス付けなどの MoE 固有の演算子は、NPU が苦手とする不規則なメモリアクセスや条件分岐を含みます。
   • これらの演算は NPU ではなく CPU で実行されざるを得ず、データ転送のオーバーヘッドが発生します。
3. ディスパッチと同期のオーバーヘッド:
   • 多数の小さなエキスパートカーネルを個別に NPU に起動すると、ディスパッチオーバーヘッドと CPU-NPU 間の同期コストが実行時間自体を上回る可能性があります。
   • NPU の並列処理能力が十分に活用されず、スループットが低下します。

特に、長いコンテキストを扱うタスクでは、推論時間の 80% と CPU サイクルの 70% が「プリフィル（入力トークンの処理）」フェーズで消費されるため、このフェーズの最適化が重要です。

2. 手法とシステム設計 (Methodology)

NPUMoE は、NPU を主要な計算エンジンとして活用しつつ、動的な操作は CPU/GPU にフォールバックさせるハイブリッドなアプローチを採用しています。その核心となる 3 つの技術は以下の通りです。

(1) エキスパート容量のための静的ティア (Static Tiers for Expert Capacity)

• 課題: 各エキスパートへのトークン数のばらつき（不均衡）を NPU の静的形状制約に適合させる。
• 解決策: オフラインでの較正（Calibration）を用いて、各エキスパートの期待されるワークロード（人気度）を推定します。
  • エキスパートを「人気度」に基づいて複数の「ティア（階層）」に分類し、各ティアに固定された容量（例：高頻度エキスパートは大きな容量、低頻度は小さな容量）を割り当てます。
  • 容量を超えたトークンは、精度低下を最小限にするための「活性化に基づく重要度スコア」を用いて剪定（プルーニング）するか、パディングで処理します。
  • これにより、動的な入力サイズを少数の静的な形状に変換し、NPU での実行を可能にします。

(2) グループ化されたエキスパート実行 (Grouped Expert Execution)

• 課題: 個々のエキスパートごとに計算グラフを起動すると、ディスパッチオーバーヘッドとキューの直列化が発生する。
• 解決策: 複数のエキスパートを単一の静的な高密度 FFN 計算グラフにグループ化して実行します。
  • 1 つの起動呼び出しで複数のエキスパートを処理し、起動コストを共有します。
  • グループ内のエキスパートは同じ容量ティアに属するように配置され、入力バッファを連続したスライスとしてパッキングすることで、動的なインデックス付けを排除します。
  • これにより、NPU の並列性を最大化し、ディスパッチオーバーヘッドを大幅に削減します。

(3) 負荷感知型計算グラフ常駐 (Load-Aware Compute Graph Residency)

• 課題: 小さなバースト的なワークロード（冷たいエキスパート）を NPU に送ると、同期オーバーヘッドが計算コストを上回る。
• 解決策: エキスパートの「人気度（ホット/コールド）」に基づいて、計算グラフをどのデバイス（NPU または CPU）に配置するかを動的に決定します。
  • ホットエキスパート: 頻繁に使用されるグループは NPU に常駐させ、起動コストを償却します。
  • コールドエキスパート: 頻繁に使用されないグループや、非常に小さなグループは CPU 上で実行し、不必要な CPU-NPU 同期を回避します。
  • これにより、NPU の利用率を最大化しつつ、システム全体のレイテンシを最小化します。

3. 主要な貢献 (Key Contributions)

1. NPU 上での MoE 推論の課題分析: Apple Neural Engine における MoE 推論の動的特性と静的制約の間のミスマッチを詳細に分析し、ボトルネックを特定しました。
2. 3 つの革新的な技術の提案:
   • 静的ティアによるエキスパート容量の最適化。
   • グループ化によるエキスパート実行の効率化。
   • 負荷感知による計算グラフの適切な配置（NPU/CPU 間）。
3. NPUMoE の実装と評価: 初のオンデバイス MoE 推論エンジンとして実装し、Apple M シリーズデバイス（M2 Max, M2 Ultra）上で包括的な評価を行いました。

4. 実験結果 (Results)

Apple M2 Ultra/M2 Max 上で、Phi-3.5-MoE、Phi-tiny-MoE、Qwen3-30B-A3B の 3 つのモデルと、HellaSwag、BoolQ、RULER などの長コンテキストワークロードを用いて評価されました。

• レイテンシの改善: ベースライン（CPU 実行、Naïve な CoreML、ANEMLL）と比較して、プリフィル遅延が 1.32 倍〜5.55 倍 短縮されました。
• エネルギー効率の向上: トークンあたりのエネルギー消費が 1.81 倍〜7.37 倍 改善されました。
• CPU 負荷の削減: CPU サイクル使用量が 1.78 倍〜5.54 倍 削減され、NPU へのオフロードが効果的に行われていることが確認されました。
• 精度への影響: トークンの剪定やパディングを行っても、推論精度の低下は 1.1% 未満 と極めて小さく、実用性を損ないませんでした。
• 詳細な分析:
  • 専門家 FFN の実行時間がプリフィル遅延の 86% 以上を占めており、NPUMoE はこの部分を最適化することで大幅な高速化を実現しています。
  • グループサイズと容量の組み合わせを最適化することで、ディスパッチオーバーヘッドとパディングのトレードオフをバランスさせています。

5. 意義と結論 (Significance)

• モバイル NPU の有効活用: 従来の研究が GPU や CPU に依存していた MoE 推論において、Apple Silicon の NPU を主要な計算リソースとして活用できることを実証しました。
• システム共設計の重要性: モデルの動的な挙動とハードウェアの静的な制約のミスマッチを、システムレベルの設計（オフライン較正、グラフ構造の最適化、動的配置）によって解決できることを示しました。
• オンデバイス AI の未来: プライバシーが保たれ、バッテリー駆動時間が長いオンデバイス LLM 推論において、MoE モデルを長コンテキストタスクでも実用的に動作させるための基盤技術を提供しています。

この研究は、動的なスパースモデルを静的な NPU 上で効率的に実行するための新しいパラダイムを示しており、将来的なモバイル AI 推論の標準的なアプローチとなる可能性があります。