vLLM v0.20.0 はハイブリッド SSM モデル向けに分散サービングを追加

はじめに

vLLM は NIXL ベースの KV コネクタを拡張し、ハイブリッド SSM-FA モデル向けに分散型プレフィル/デコードをサポートしました。標準的なトランスフォーマーの既存ワークフローは変更しません。 NVIDIA Nemotron-H のようなハイブリッドアーキテクチャは、Mamba スタイルの SSM レイヤーとフルアテンションレイヤーを交互に配置し、線形時間の SSM 効率とアテンションの表現力を組み合わせます。既存の NIXL 分散 P/D 設計は均一な KV キャッシュ形式を前提としていましたが、FA と SSM のレイヤーは状態を異なるレイアウトとサイズで保持するため、ハイブリッドモデルには当てはまりません。

背景: NIXL KV 転送ワークフロー

標準的なトランスフォーマーモデルでは、NIXL の分散 P/D はメモリ領域を登録し、ブロック単位のディスクリプタを作成し、ハンドシェイクを行い、RDMA を介してブロックを転送することで機能します。 各ワーカーは自分の KV キャッシュテンソルを NIXL に登録し、各ブロックに対して (address, length, device_id) を指定するディスクリプタを作成します。ハンドシェイクはプレフェッチ‑デコードのペアごとにメタデータを一度交換します。スケジューラはデコードワーカーに取得すべきブロック ID を指示し、デコードワーカーはブロック ID をディスクリプタ ID にマッピングし、RDMA READ を発行して完了をポーリングします。M 個の登録領域と各領域あたり N 個のブロックがある場合、ディスクリプタインデックスは r * N + b となります。

課題: FA と SSM の状態は根本的に異なる

ハイブリッドモデルは、FA レイヤーがトークンごとの K/V ペアを保持し、SSM レイヤーが固定サイズの conv ステートと SSM ステートを保持するため、均一なディスクリプタ前提を破ります。 FA レイヤーの KV キャッシュ形状は [num_blocks, 2, block_size, num_kv_heads, head_dim]（またはその変種）です。SSM レイヤーは conv ステート (conv_dim, state_len) と SSM ステート (num_heads, head_dim, state_size) を保持します。これらのステートはトークンベースではないため、ブロックサイズの概念が異なります。各 SSM ブロックは完全なステートスナップショットです。HMA は FA と SSM のグループ間でメモリをプールし、SSM 行にパディングを入れて両ビューが同じ物理テンソルとバイト単位のページサイズを共有できるようにします。その結果、均一な (address, length) を持つ単一のディスクリプタリストでは、FA と SSM の両ビューを正しくインデックス付けできません。

デュアルディスクリプタビュー

この解決策では、同一の物理メモリ上に 2 つの別々のディスクリプタリストを登録します。1 つは FA ディスクリプタ用、もう 1 つは SSM ディスクリプタ用で、単一の NIXL 転送ハンドルに連結されます。 FA ディスクリプタは最初の num_descs = M * N_phys スロットを占有し、各領域ごとに K と V を別々にインデックス付けします。SSM ディスクリプタはその後に続き、各 Mamba レイヤーはブロックごとに 4 つのサブディスクリプタ (x, B, C, SSM) で表現されます。ブロック ID からディスクリプタ ID へのマッピングは、FA グループでは region * N_phys + block_id、SSM グループでは mamba_region_id * N_log + block_id + num_descs を使用します。

物理ブロックサイズと論理ブロックサイズ

FA レイヤーはアテンションカーネル用に異なる物理ブロックサイズを必要とする場合がありますが、SSM レイヤーは論理ブロックを直接使用するため、FA と SSM のディスクリプタセクションでブロック数が異なります。 論理ブロックサイズ / カーネルブロックサイズ の比率を用いて、FA レイヤーのみ物理ブロック数 = 論理ブロック数 * 比率 を計算します。SSM レイヤーは常に論理ブロックを使用します。この情報は _physical_blocks_per_logical で管理され、テンソル並列サイズが異なる場合、プレフェッチとデコードのインスタンス間で異なることがあります。ブロック ID からディスクリプタ ID へのマッピングは、対象が FA か SSM かに応じて適切なストライドを使用します。

3 ディスクリプタ Conv 転送

異種テンソルパラレル構成では、conv ステートを DS 形式で配置し、各デコードランクが 3 つのディスクリプタ領域を介して連続した x、B、C スライスを読み取れるようにします。これにより、余分なバッファやリシェッフルなしでゼロコピー RDMA が可能になります。 DS レイアウト (dim, state_len) により、各サブプロジェクションはメモリ上で連続し、[x][B][C][SSM] の順になります。デコードワーカーは単一の NIXL READ 操作内で 3 つの連続した読み取りを発行し、conv ステートの自分のシャードを取得できます。これにより、デコード側でステージングバッファを確保する必要がなくなり、転送後のリシェッフルも不要です。また、所有する 1/TP の conv ステートのみを転送し、HMA のパディングバイトをスキップします。この手法は均質 TP と異種 TP の両方で機能し、後者の場合は Mamba レイヤーごとに 4 つのディスクリプタ領域 (x, B, C, SSM) を使用します。

実装例: Nemotron-H

具体例として、nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 を TP=2 の分散 P/D で提供する際のディスクリプタ登録と転送の流れを示します。 このモデルは Mamba と FA が交互に配置された 52 層を持ち、HMA によって 5 グループ（Mamba が 4、FA が 1）に分けられ、合計 6 つの共有 KV キャッシュテンソルが生成されます。FA レイヤーは [num_blocks, 2, block_size=400, 4, 128] を使用し、SSM レイヤーは [num_blocks, 3, 3072]（conv）と [num_blocks, 48, 64, 128]（ssm）を使用します。HMA パディング後、両ビューは同じバイト単位のページサイズを共有します。プレフェッチインスタンスは 6 つのテンソルを登録し、すべての領域 × N_phys ブロックに対して FA ディスクリプタ（K と V を別々に）を作成し、ブロックごとに 4 つのサブ領域を持つ Mamba ディスクリプタを追加します。デコードインスタンスは適切なストライドを用いてブロック ID をディスクリプタ ID にマッピングし、FA と Mamba ディスクリプタの両方を含む単一の事前設定された READ を発行し、完了をポーリングします。中間バッファやデータのリシェッフルは不要です。

パフォーマンス

ハイブリッド SSM モデル向けの分散 P/D は、デコードをプレフィルの干渉から分離することで、高い同時実行数において同居サービングのスループットと同等またはそれ以上の性能を示します。 NVLink を使用した 8 台の H200 GPU 上でのベンチマークでは、nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 に対し、同居ベースライン（単一インスタンス、TP=8）と分散構成（プレフェッチインスタンス TP=4 + デコードインスタンス TP=4、GPU 総数は同じ）を比較しました。並列ユーザー数を 8 から 256 に変化させた結果、分散構成のパレート曲線は高バッチサイズで同居曲線を上回り、デコードがプレフェッチから分離されたときに GPU あたりの秒間出力トークン数が増加することが示されました。

はじめに

ハイブリッド SSM モデルを分散 P/D で実行するには、VLLM_SSM_CONV_STATE_LAYOUT=DS を設定し、適切な引数でプレフェッチとデコードのインスタンスを起動します。 プレフェッチインスタンスの例コマンドは以下です:

VLLM_SSM_CONV_STATE_LAYOUT=DS vllm serve nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.85 \
    --trust-remote-code \
    --max-model-len 8192 \
    --block-size 128 \
    --no-disable-hybrid-kv-cache-manager \
    --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'

デコードインスタンスは kv_role を適切に設定した類似のコマンドを使用します（ソースには示されていません）。DS レイアウトは異種 TP の場合に必須ですが、その他の場合はオプションです。

制限事項と今後の課題

現在のサポートは Mamba2 のみです。Mamba1 モデルや GDN レイヤーはまだ未対応で、推測デコードとの連携は十分に検証されておらず、HMA が有効な場合の混在ブロックサイズはまだサポートされていません。 Mamba1 の SSM 時系列形状は conv 分解に必要な intermediate_size の再構築を妨げます。GDN のサポートは分散ロードマップに記載されています。推測デコードおよび HMA を伴うブロックサイズ比率の取り扱いは未解決の課題です。

謝辞

Thomas Parnell (IBM Research) と Roi Koren (NVIDIA) の貢献に感謝します。