vLLMコミュニティ、H20のMoEスケジューリングの非効率を調査　専門家負荷に応じたカーネル選択の修正を提案

9月25日、vLLMコミュニティはH20の推論性能に関する報告を提出し、FlashInferのSM90 FP8混合専門家モデル（MoE）経路では、負荷の推定が適切でないために非効率な行列演算カーネルが選ばれる可能性があると指摘した。報告者はFlashInferに修正を提出したが、9月27日時点の確認ではまだマージされておらず、特定のvLLM正式版によって性能が低下したとも確認されていない。[コミュニティの報告](https://github.com/vllm-project/vllm/issues/58799)

FlashInferの`cutlass_fused_moe`は、CUTLASSを使って専門家演算を融合して実行する。APIは、エキスパート並列の規模とランク、およびFP8ブロックスケーリングに対応している。今回の問題は既存の実行経路に関するもので、複数GPUで疎なモデルをデプロイするチームにとって特に関係がある。[公式APIドキュメント](https://docs.flashinfer.ai/generated/flashinfer.fused_moe.cutlass_fused_moe.html)

報告によると、スケジューラは当初、入力トークン数を作業量の推定値としていた。しかし、カーネルの分岐と行列のタイル選択に必要なのは、専門家ごとのデータ行数だ。テストでは288トークンが各トークンあたり6人の専門家を選び、400個の物理的な専門家スロットに分散した結果、平均はわずか4.32行だった。修正では推定値を「トークン数×top-k÷物理的な専門家の総数」に変更し、切り上げて最低1とする。この推定は、最悪ケースのワークスペース容量を計算した後に適用される。[修正案](https://github.com/flashinfer-ai/flashinfer/pull/5560)

この事例は、バッチ全体のサイズだけでは、専門家が実際に処理する行列の形状を表せないことを示している。平均値もルーティングの偏りを十分には捉えられない。少数の専門家にデータの大半が集中し、他のランクに処理対象がない場合は、最適なカーネル選択が異なる可能性がある。そのため、均等分散時と負荷集中時のテスト結果は分けて解釈する必要がある。

報告者は8基のH20で均等なルーティングを行い、288トークンのローカルMoEレイテンシが1.808ミリ秒から0.564ミリ秒に短縮し、約3.21倍高速化したと測定した。テストでは5回のウォームアップ後に24回サンプリングしたが、計測にGPU間通信は含まれておらず、GPUクロックも固定されていない。この数値をオンラインサービスのスループットに直接換算することはできない。報告者は、モデルの重みに依存しない合成再現スクリプトも添付しており、他の開発者が演算子の挙動を検証できる。[テスト方法と結果](https://github.com/vllm-project/vllm/issues/58799)

検証にはまだ不足がある。ベンチマークは変更を限定したソースコードのバックポート版で実施されており、上流の完全なビルドと、変更を加えない状態でのテスト実行は未完了だ。また、独立した参照実装との比較テストがベースライン版と修正版の両方で失敗しており、原因は分かっていない。ルーティングが高度に集中する条件では、単一の行列演算段階が約7.6%遅くなることもあった。エンジニアリングチームは上流でのレビューを追い、モデル、量子化方式、並行実行条件を固定して、エンドツーエンドのレイテンシと出力の正確性を比較する必要がある。現時点の証拠では、他のGPUやエキスパート並列の規模にも一般化できるとは言えない。[検証上の制約](https://github.com/flashinfer-ai/flashinfer/pull/5560)