MUSA 上の MAGI-2 Preview が vLLM-Omni で 1 ステップあたり 565 ミリ秒を達成

TL;DR

Sand AI の 114B オーディオ・ビデオ生成モデル MAGI-2 Preview が、torchada 互換レイヤーを介して Moore Threads MUSA GPU 上で動作し、1 デノイジングステップあたり 565 ミリ秒(8×MTT S5000、272p、静的コンパイル)を達成しました。これは最初の動作バージョン(1 ステップ 15.6 秒)から 27.6 倍の改善です。この高速化は、ホスト同期の除去、MoE カーネルの融合、mHC 残差の書き換え、エキスパート並列シャーディングの最適化を行う一連のターゲット PR によるものです。


1. MAGI-2 Preview をユニークな推論ワークロードにするもの

  • シングルストリーム DiT バックボーンは、テキスト、ビデオ潜在、オーディオ潜在トークンを 1 つの自己注意シーケンスにパックします。モダリティ固有の重みを持つのは最初と最後の 4 層のみです。
  • adaLN なし – タイムステップは各トークンの最初の 64 チャンネルに直接エンコードされ、ビデオとオーディオのタイムステップを独立して処理できます。
  • マルチヘッド MoE(層 2〜37)は、各 3072 次元トークンを 12 個の 256 次元ヘッドに分割します。各ヘッドは 256 個のエキスパート(SwiGLU7)の上位 6 個にルーティングされ、さらに 1 つの共有エキスパートと 3 つのモダリティ別共有エキスパートがあります。ルーティングされるエキスパートはパラメータの約 95% を占めます。
  • mHC ハイパーコネクションは、20 回の Sinkhorn 反復で計算される二重確率 4×4 混合行列を持つ 4 ストリーム残差(合計幅 12288)を提供します。
  • サンプリングは、100 の未蒸留ステップを実行し、各ステップで 2 つの CFG ブランチを持ち、10 秒のビデオ全体を 1 パスでデノイズします。
メトリック 値
総パラメータ数 114B(トークンあたり約 6B アクティブ)
バックボーン 40 層 DiT、隠れ 3072、24 ヘッド×128
MoE 12 ヘッド×256 エキスパート、ヘッドあたり top-6
残差 mHC、4 ストリーム、20 Sinkhorn 反復
出力 10 秒、125 フレーム、12.5 fps、44.1 kHz ステレオ
対応解像度 272p(448×256)および 540p(896×512)

他のオープンビデオモデルと比較して、MAGI-2 Preview は、トークンごとの MoE ルーティング、mHC 残差、およびジョイントオーディオ・ビデオ生成を単一の自己注意ストリームで組み合わせた唯一のモデルです。


2. torchada で MUSA 上で CUDA ファーストモデルを実行する

  • torchada は、CUDA スタイルのインポート(torch.cuda、デバイス文字列、NCCL、CUDA Graph)をインターセプトし、MUSA 相当(torch.musa、MCCL、MUSAGraph)にマッピングします。このマッピングは意図的に真実(torch.cuda.is_available() は False のまま)であるため、プラットフォームチェックは torch.version.musa を使用する必要があります。
  • コード再利用: eager PyTorch、torch.compile/Inductor、およびアップストリームの Triton カーネル(融合 BF16 MoE、SwiGLU7、mHC)は MUSA 上で変更なしで実行されます。起動構成と一部のプラットフォーム固有のパスのみが異なります。
  • 主な適応:
    • torch.isfinite は MUSA では abs()<inf 式に置き換えられました(torchada #120)。
    • Int32 ルーティングカウント(scatter_add_)は MUSA では int64 を置き換え、呼び出しあたり約 39 ミリ秒のゲインをもたらしました(PR #8496)。
    • GEMM オペランドレイアウトはストリームメジャーに並べ替えられ、ステップあたり約 80 コピーが排除されました(PR #8510)。
    • Triton-MUSA には BF16 アトミックがないため、融合 BF16 MoE パスが常に選択されます。

「torchada は MUSA テンソルに対して is_cuda を true にするため、CUDA スタイルのコードは通常どおり実行されます。」 – torchada ドキュメント。


3. 最適化のタイムライン(ステップごとのレイテンシ)

ステージ ステップあたりの時間(ミリ秒) 主な変更
初期立ち上げ(9 月) 15,600 8×MTT S5000 でのベースライン
ホスト同期の除去 8,800 .item() → .tolist()、オールギャザーメタデータの削減
MoE 融合 GEMM(MATE) 2,390 グループ化 BF16 GEMM がエキスパートごとの Python ループを置き換え
融合 W13 + SwiGLU7 Triton 954 ゲート/アップ + アクティベーション用の単一カーネル
Int32 ルーティングカウント 756 PR #8496
K 次元タイリング(BLOCK_SIZE_K=64) 734 PR #8506
mHC 要素ごとの書き換え 734 → 734(ベースライン)→ 733.6 PR #8510(FP32 縮約)
Sinkhorn シングルカーネル –6.6 PR #8594
W13 アキュムレータ分割 590–591 PR #8855(カーネル 38〜40% 高速化)
フル PR スタック(10 月 19 日) 565 上記すべての PR 適用、静的コンパイル

27.6 倍の全体的な高速化(15.6 秒 → 565 ミリ秒)は、同じハードウェアとワークロード(272p、100 ステップ、T2VA プロンプト、SP8×CFG1+EP4)で測定されています。各項目は独自のベースラインに対して評価されています。数値は行間で累積されません。


4. 詳細な技術的利点

4.1 ホスト側同期

  • エキスパートごとの .item(): 一括 .tolist() に置き換えられ、ステップあたり約 3.4 秒節約。
  • オールギャザートークンカウントメタデータ: ステップあたり 39 回から 3 回に削減され、約 61 ミリ秒節約。
  • torch.isfinite: MUSA で abs()<inf として再実装され、ステップあたり約 1.8 秒削減。

4.2 MoE カーネル融合

  1. MATE グループ化 GEMM は、12×256 次元のエキスパート GEMM を単一の BF16 GEMM にまとめました(21.9 ミリ秒 → 0.69 ミリ秒)。
  2. 融合 W13 + SwiGLU7 Triton カーネルは、ゲート、アッププロジェクション、アクティベーションを組み合わせ、MoE ボトルネックを約 1.5 秒から 1 秒未満に削減しました。
  3. 重みレイアウトの事前インターリーブにより、呼び出しごとの転置がなくなり、ステップ時間が約 7.5% 短縮されました。
  4. Int32 ルーティングヒストグラムは、ルーティングフェーズから約 1.2 秒削減しました。
  5. K 次元タイリング(BLOCK_SIZE_K=64)は、約 2% の小幅なゲインをもたらしました。
  6. 決定論的ルート構築(PR #8592)は約 12 ミリ秒追加しました。

4.3 mHC 残差書き換え

  • 小さな K の bmm を FP32 要素ごとの縮約とストリームメジャーレイアウトに置き換え、ステップあたり約 250 ミリ秒(約 25%)削減。
  • 20 回の Sinkhorn 反復を単一の Triton カーネルに統合(-6.6 ミリ秒)。

4.4 アテンション

  • Torch 実装から MATE FlashAttention-3(FA3)に切り替え、学習済み LSE シンクを使用し、SP4×CFG2 構成で 16.8% 削減。
  • RoPE、モダリティグループ化、シンク LSE メタデータをキャッシュ(-6 ミリ秒)。

4.5 エキスパート並列シャーディング

  • ヘッドシャーディング EP(EP4)はシーケンス並列から分離され、最速の構成(565 ミリ秒)を実現。EP8 はシャードが多いほど通信オーバーヘッドが増えるため遅くなりました(633 ミリ秒)。
  • Ulysses 送受信周辺のバッファコピーを削除すると、さらに 10.6 ミリ秒節約(PR #8595)。

4.6 コンパイル

  • リージョナル torch.compile は、eager 実行と比較してステップあたりの時間を約 16% 削減。
  • 動的コンパイルで非トークン次元を固定すると、1.4 ミリ秒のカーネルが 55 マイクロ秒に短縮(PR #8627)。
  • MUSA での動的コンパイルでは、正しい出力を生成するために Sinkhorn 修正(#8594)が必要でした。

5. 測定されたパフォーマンス結果

5.1 PR ツリーのエンドツーエンドタイミング(静的 vs 動的)

並列構成 最終 PR ラウンド前(静的) 静的コンパイル(ミリ秒/ステップ) 動的コンパイル(ミリ秒/ステップ)
SP8×CFG1+EP4 714/715 665.7/669.9 678.3/678.9/683.7
SP4×CFG2+EP4 732.6 690.8 –
SP4×CFG2 783.1 729.9 –

565 ミリ秒の数値(セクション 4)は、2026 年 10 月 19 日に測定された最新のツリーからのものです。異なるコードベースとセッションを使用しているため、上記の行と直接比較することはできません。

5.2 W13 アキュムレータ分割の A/B テスト(PR #8855)

ツリー 静的コンパイル(ミリ秒/ステップ) 出力ハッシュ(オーディオ / ビデオ)
コントロール(分割なし) 690.8/691.1 32f1b4fca4 / 44b9747e9e
+ #8855(分割) 591.3/590.4 同一

14.5% の削減はビット同一です。同じ SHA-256 ハッシュが数値的等価性を確認しています。

5.3 最新ツリーでのエキスパート並列度スイープ

構成 平均ステップあたり時間(ミリ秒) EP4 に対する相対
SP8×CFG1、EP オフ(レガシー) 697.8 +23.5%
SP8×CFG1+EP4 565.0 ベースライン
SP8×CFG1+EP8 633.4 +12.1%

3 つの構成すべてで同一のオーディオおよびビデオハッシュが生成され、速度の違いが純粋に通信/計算のトレードオフであることを確認しています。


6. サンプル出力と品質に関するメモ

  • ベンチマークプロンプト(ゴールデンレトリバー、272p、SP8×CFG1+EP4)は、ステップあたり 565 ミリ秒で 10 秒のビデオを生成しました。クリップはブログアセットで入手できます。
  • 公式 Sand AI プロンプト(コーヒーショップ、アニメ、クレイアニメーション)は、T2VA と I2VA の両モードで実行されました。すべてのサンプルはテキスト記述と一致し、非ゼロのオーディオを含み、コーヒーショップのクリップには可聴の英語のダイアログが含まれています。
  • 540p プレビューは、静的で約 3.6 秒/ステップ、動的で約 3.9 秒で実行されます。これらの数値は、異なるコードツリーとコンパイルモードを使用しているため、272p のタイミングと直接比較することはできません。
  • 定量的なマルチシード品質評価はまだ実行されていません。視覚的および聴覚的忠実度は手動で検証されました。

7. アップストリームのステータスと互換性

7.1 PR の状況(2026 年 10 月 28 日現在)

PR 機能 ステータス 測定されたステップごとの影響
#7206 融合 BF16 MoE(W13) マージ済み –
#7243 設定可能な Dynamo 再コンパイル制限 マージ済み –
#7252 融合 SwiGLU7 マージ済み –
#7261 融合 mHC ポスト残差 マージ済み –
#8496 int32 ルーティングカウント マージ済み 1,952 → 756 ミリ秒
#8497 W13 の単一コピー マージ済み –
#8498 MUSA サポート(FA3 varlen + シンク) マージ済み –
#8506 MUSA での MoE GEMM 用 BLOCK_SIZE_K=64 マージ済み 754.5 → 733.9 ミリ秒
#8510 mHC FP32 縮約とストリームメジャーレイアウト マージ済み 984.7 → 733.6 ミリ秒
#8511 SP から独立したヘッドシャーディング EP オープン 864.6 → 717.5 ミリ秒
#8592 ビット正確なルート構築 マージ済み –12 ミリ秒
#8593 RoPE とシンク LSE メタデータキャッシュ オープン –6 ミリ秒
#8594 シングルカーネル Sinkhorn オープン –6.6 ミリ秒
#8595 Ulysses バッファコピーの削除 オープン –10.6 ミリ秒
#8627 アテンションヘッドレイアウトの固定(動的) マージ済み 1,429 マイクロ秒 → 55 マイクロ秒
#8855 W13 アキュムレータ分割(MUSA) オープン 690.8 → 590.4 ミリ秒(静的)
#8915 1080p リファイナー(ブロックウィンドウアテンション) ドラフト –

リストされたすべての PR はマージ済みまたはレビュー中です。この投稿のパフォーマンス数値は、各 PR が測定された時点のコードの状態を反映しています。

7.2 互換性の境界

  • 完全にベンチマーク済み: MUSA + SP8×CFG1+EP4 の 272p(セクション 5.1〜5.3)。
  • 部分サポート: 540p プレビュー(サンプルのみ)、I2VA モード(サンプルのみ)、1080p リファイナー(ドラフト実装、タイミング未測定)。
  • 未検証: レイヤーワイズオフロード、量子化、Cache-DiT、その他のメモリ節約技術。

8. 結論と今後の作業

  • 高速化の概要: ホスト同期の体系的排除、MoE カーネルの融合、mHC 残差の書き換え、エキスパート並列シャーディングの改善により、15.6 秒のベースラインから 565 ミリ秒/ステップ(27.6 倍)へ。
  • 残りのボトルネック:
    1. MoE W13 GEMM は依然としてメモリ帯域幅に制限されています。さらなるタイリング(128 幅の K)またはカスタムグループ化カーネルにより、さらに約 2〜3 ミリ秒削減できる可能性があります。
    2. 通信と計算のオーバーラップ: EP4 はステップあたり 60〜90 ミリ秒の集団通信を発生させますが、計算とオーバーラップしていません。
    3. アイドルホスト時間: FlowUniPC ホストロジック(約 19 ミリ秒)とカーネル内起動ギャップ(約 14 ミリ秒)が残っています。
    4. 動的コンパイルオーバーヘッド: いくつかのカーネルで int64 サイズ演算が使用されています。int32 に変換すると、失われたパフォーマンスの約 85% が回復します。
  • 次のステップ:
    • MUSA でグラフリプレイを有効にする(現在はリージョナルコンパイルより遅い)。
    • 1080p リファイナー統合を完了する(ブロックウィンドウアテンション、メモリプランニング)。
    • CI を拡張して、540p ベンチマークとマージ済み PR のみのメインブランチをカバーする。
    • 量子化とキャッシュ DiT 技術を調査して、さらなるレイテンシ削減を図る。

9. 謝辞

著者は、MAGI-2 Preview をオープンソース化し公式プロンプトを提供してくれた Sand AI、GPU スタックを提供してくれた Moore Threads MUSA、torchada、MATE チーム、および元のブログ投稿にリストされている PR の実装とレビューを行った多くの貢献者に感謝します。


10. 再現性

  • コンテナイメージ: registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview(ダイジェスト sha256:2785b4626d1da1e0765a4cb515b9afc5349a22c6417e14c81091add0c1a9b902)。
  • コマンド(8×MTT S5000 で実行):
    docker pull registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview
    docker run -d --name magi2 -v /ipfs:/ipfs -v "$PWD/magi2-out":/work \
      --entrypoint sleep registry.mthreads.com/mcconline/inference/vllm/vllm-omni:magi-2-preview infinity
    docker exec magi2 bash -lc '
      cd /opt/vllm-omni-src &&
      python /opt/magi2/instrument-steps.py &&
      VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONUNBUFFERED=1 \
      python /opt/magi2/magi2_e2e_gate.py \
        --ulysses-degree 8 --cfg-parallel-size 1 \
        --enable-expert-parallel --expert-parallel-size 4 \
        --resolution 272p --height 256 --width 448 \
        --seconds 10 --fps 12.5 --num-frames 125 --num-inference-steps 100 \
        --no-pin-cpu-memory --metrics-dir /work/steps --save-outputs /work/outputs'
    
  • 決定論的 Inductor 動作のための環境変数は、ソース投稿の付録 A.5 で説明されているように、コンテナ内で設定されます。

すべての数値、表、引用は、2026 年 10 月 12 日付けの vLLM ブログ投稿から直接取得したものです。追加のデータは作成されていません。

Sources