vLLM AFD Plugin: MoE サービングのための Attention と FFN の分離

vLLM AFD Plugin: MoE サービングのための Attention と FFN の分離

vLLM は、Mixture-of-Experts (MoE) モデル向けに Attention-FFN Disaggregation (AFD) を実装する実験的な外部プラグインである vLLM AFD Plugin をリリースしました。このアーキテクチャは、Attention と Feed-Forward Network (FFN) コンポーネントを独立してデプロイされるサービスに分離し、vLLM のリクエストライフサイクルや OpenAI 互換のサービングインターフェースを変更することなく、それぞれの特定の計算リソース要件に基づいてスケーリングできるようにします。

Attention-FFN Disaggregation の根拠

MoE 推論において、Attention パスと FFN パスは根本的に異なる運用要件を持っています。Attention はステートフルであり、リクエストのスケジューリングと KV キャッシュ管理に結びついていますが、FFN (expert) パスはルーティング計算と all-to-all 通信が支配的です。

AFD は、主に 4 つのシステム設計上の課題に対処します:

  • 独立したスケーリング: Attention の容量はシーケンス長と KV キャッシュの負荷に依存し、エキスパートの容量はトークンのルーティングと負荷に依存します。AFD により、これらのパスで異なるランク・トポロジーを使用できます。
  • 実行時の責任: サービスを分割することで、FFN 側は軽量なコネクタ駆動のデーモンとして動作でき、エキスパート側でのスケジューリングや KV キャッシュの調整が不要になります。
  • バックエンド固有の通信: 共通のコネクタ・コントラクトにより、異なるハードウェア・バックエンド(例:CUDA および Ascend)が、独自の最適化されたデータパスとコレクティブ・ライブラリを実装できます。
  • 計算と通信のオーバーラップ: 非同期ディスパッチと MoE ubatching により、独立したステージをオーバーラップさせることができ、Attention パスがすべてのエキスパート作業を直列化してしまうのを防ぎます。

システムアーキテクチャ

vLLM AFD Plugin は、vllm.general_plugins エントリポイントと --additional-config チャネルを介して統合されるため、vLLM のソースツリーを修正する必要はありません。ランタイムは 3 つのコアコンポーネントで構成されます:

1. Attention Service

vLLM のスケジューラ、KV キャッシュ、バッチング、モデルのライフサイクル、およびサンプリングパスを保持します。プラグイン所有のモデルランナーが AFD メタデータを forward コンテキストにインストールし、状態(data-parallel, ubatch, layer, および graph)を FFN サービスに公開します。

2. FFN Service

リクエストトラフィック、スケジューラ、または KV キャッシュなしで動作します。メタデータとアクティベーションを受け取り、プラグイン所有のモデルラッパーを介して compute_ffn_output() を実行し、結果を Attention サービスに返すバックグラウンドループを実行します。

3. Connector Layer

各分離レイヤーにおけるブリッジとして機能し、隠れ状態(hidden states)と実行メタデータを Attention から FFN へ転送し、計算された出力を返します。

Connector とバックエンドのサポートマトリックス

Connector Backend Execution Recommended Stage Graph Support
P2pNcclAFDConnector GPU Synchronous P2P Decode FULL_DECODE_ONLY CUDA graph
CAMP2pAFDConnector NPU Synchronous CAMP2P/HCCL Decode FULL_DECODE_ONLY ACL graph
CAMAsyncAFDConnector NPU Asynchronous CAM Prefill Not supported

パフォーマンス・ベンチマーク

同期デコード・スループット (Ascend 910C 上の DeepSeek-V3.2 W8A8)

ベンチマークでは、従来の EP64 デプロイメントと CAMP2pAFDConnector を使用した AFD デプロイメントを比較しました。結果は、Attention と FFN のランク比が正規化スループット (tokens/s/die) に大きく影響することを示しています。

  • 16K 固定入力: 64A16F 構成 (64 Attention, 16 FFN ランク) は 258.9 tokens/s/die を達成し、EP64 ベースライン (232.6 tokens/s/die) に対して 11.3% の増加 を記録しました。48A16F 構成は 220.3 tokens/s/die でベースラインより低成績でした (-5.3%)。
  • 32K 固定入力: 64A16F 構成は 183.3 tokens/s/die を達成し、EP64 ベースライン (168.2 tokens/s/die) に対して 9.0% の増加 を記録しました。48A16F 構成はベースラインを 10.0% 下回りました。

これらの結果は、分離(disaggregation)だけで利得が保証されるわけではなく、Attention と FFN の間の具体的なランク割り当てが極めて重要であることを示しています。

非同期 Prefill パフォーマンス (Ascend 910C 上の DeepSeek V3.2 W8A8)

10 レイヤーの削減モデル上で CAMAsyncAFDConnector を使用した初期実験では、DP4PCP8 TP1 ベースラインと AFD レイアウト (Attention DP3PCP8 TP1 + FFN EP8) を比較しました。

リクエストレートが毎秒 12 リクエストの場合、中央値の Time to First Token (TTFT) は 15.1 秒から 8.0 秒 に減少し、約 47% の削減を実現しました。

実装とロードマップ

現在のサポート

  • モデル: DeepSeek V2/V3 ファミリー (DeepSeek V3.2 を含む) および GLM MoE DSA 用のラッパー。
  • ハードウェア: NVIDIA GPU および Ascend NPU。
  • 実行パス: eager、graph、および dual-batch 実行のサポート。同期コネクタ用の decode-only graph キャプチャ。
  • 要件: Python 3.10–3.13 および vLLM 0.19.1

今後の開発

  • アップストリームとの整合性: 新しい vLLM リリースの追跡と model runner v2 の評価。
  • 実行の柔軟性: graph モード、ubatch 数、および非同期ステージの拡張。
  • プロダクション検証: フルモデルおよび現実的なワークロードにおける安定性、精度、およびレイテンシの結果の公開。
  • カバレッジの拡大: より多くの MoE アーキテクチャとバックエンド転送の追加。
  • マルチモーダル統合: 自己回帰 (AR) および Diffusion Transformer (DiT) ステージにおける vLLM-Omni 内での AFD の適用の探索。
  • ヘテロジニアス・ハードウェア: TTFT とトークン間レイテンシを削減するために、異なるアクセラレータ・タイプにわたって Attention と FFN の役割をデプロイする方法の調査。

Sources