RaMP:Mixture-of-Experts向けのランタイム対応メガカーネル・ポリモーフィズム

arXiv cs.AI / 2026/4/30

💬 オピニオンDeveloper Stack & InfrastructureModels & Research

要点

  • この論文は、Mixture-of-Experts(MoE)推論の性能がバッチサイズだけでなくエキスパートのルーティング分布にも左右されるにもかかわらず、既存の本番システムのディスパッチはバッチサイズのみに依存しており、カーネルスループットの最大10〜70%を活かせていないと指摘しています。
  • ルーティングを考慮したディスパッチフレームワークRaMPを提案し、ハードウェア定数から最適化の効き方が分かる「パフォーマンス領域」分析で、複数のアーキテクチャ(未検証のものを含む)でも正しく予測できることを示します。
  • RaMPは、ランタイムのエキスパート・ヒストグラムから最速に近いカーネル構成を選ぶための4パラメータの「ウェーブコスト」モデルを備え、網羅探索に対して平均後悔度0.93%を達成し、各モデルあたり初回のプロファイリングは10〜24分で済むとしています。
  • この手法はカーネルに依存せず、CTAグリッド幾何に基づくだけなので、Alpha-MoEに対してソース変更なしで適用できることが示されます。
  • さらに、CuTe DSLカーネルで134〜268通りのポリモーフィック構成を引き出すことで、vLLM serving(Triton上)でカーネル1.22×、エンドツーエンド1.30×など、他のバックエンド(DeepGEMM、FlashInfer CUTLASS)でも向上が報告されています。

要旨: Mixture-of-Experts(MoE)推論における最適なカーネル構成は、バッチサイズとエキスパートのルーティング分布の両方に依存する。しかし実運用のシステムではバッチサイズだけに基づいてディスパッチされるため、カーネルのスループットの10〜70%が実現されない。私たちは、ルーティングを考慮したディスパッチのための枠組みRaMPを提案する。性能領域(performance-region)分析により、ハードウェア定数だけから、各最適化がいつ有効になるかを導出し、8つの検証済みアーキテクチャについて(さらに3つの未見アーキテクチャを含めて)正しく予測できることを示す。4パラメータのウェーブコストモデルが、実行時のエキスパートヒストグラムから最速の構成を選択し、網羅的探索に対して平均リグレット(後悔度)0.93%を達成する。これは、モデルごとに1回限りのプロファイリングを10〜24分行うだけで適合(フィット)できる。モデルはCTAグリッドの幾何学形状のみに依存するため、カーネルに非依存である。Alpha-MoEに適用すると、ソース修正なしで1.14xをもたらす。共同設計されたCuTe DSLカーネルにより134〜268の多相(ポリモーフィック)構成を公開することで、RaMPは静的ディスパッチに対して1.22xのカーネル速度向上を提供し、Triton上でのvLLMサービングでは1.30x、DeepGEMMでは1.41x、FlashInfer CUTLASSでは1.13xのエンドツーエンド速度向上を実現する。