AWSは、Amazon SageMaker AIでMixture-of-Experts（複数の専門家モデルを使い分ける方式）モデルを提供するため、融合型MoEカーネルの最適化を3つ説明しました。最大10％のスループット向上と最大9％のTime to First Token（TTFT、最初の生成トークンが出るまでの時間）の短縮を報告しています。vLLMコミュニティと共同で開発し、オープンソース化したうえで、SageMaker AIのLMI Deep Learning Containerに適用しました。

最初の最適化では、行列乗算のオペランドを入れ替え、NVIDIA Hopper GPUのWGMMA命令の配置に合わせます。WGMMAはM次元が固定の64行ですが、MoEのデコードではユーザー系列ごとに1トークンを処理することが多く、Mは1となり、Tensor Coreの実行レーンが多数使われません。融合型MoEのtopk=8設定では、選ばれた専門家の中間次元をまとめることでN次元が非常に大きくなります。乗算を転置して大きい次元を第1オペランドに置くと、64行の配置に合い、非同期プリフェッチが可能になります。AWSによると、H200 GPU上の大規模MoEモデルで、エンドツーエンドの出力スループットが2～4％向上しました。この変更はvLLMバージョン0.24.0以降に取り込まれています。

2つ目の変更では、Hopper GPU上のFP8テンソル量子化MoE計算を、既定のFlashInfer+CUTLASS経路ではなくTritonに振り分けます。AWSのプロファイリングでは、Tritonはスケーリング、行列乗算、専門家の集約を1回のGPU処理に融合し、中間データをグローバルメモリーで読み書きする工程を省くため、同時実行数の少ない処理で高い性能を示しました。動的スケジューリングも、M < 64やM = 1のような小さく変動するトークンバッチを、CUTLASSの既定タイル配置より効率よく処理します。AWSは、カーネルコードを変更せず経路の設定を切り替えることで、H200上のFP8量子化済み融合型MoEモデルのエンドツーエンドスループットが10％向上したとしています。この変更はAWS Large Model Inference（LMI）v28コンテナで利用できます。

3つ目の最適化では、新しいトークン数に遭遇するたびにTritonがカーネルを再コンパイルするのを防ぎます。_fused_inv_rope_fp8_quant_per_headカーネルは、それまで新しいnum_tokens値ごとに再コンパイルしていました。可変長パラメーターにdo_not_specialize注釈を加え、繰り返し発生する実行時コンパイルをなくしました。AWSは、H200でテンソル並列処理を8分割した場合にTTFTが9％改善し、1基のGPUが再コンパイルする間にほかのGPUが待つ状況も避けられると報告しています。この最適化はvLLMバージョン0.22.0以降に取り込まれています。3つの変更により、SageMaker AI上のLMIコンテナでスループットが最大10％向上し、TTFTが最大9％短縮するとAWSは述べています。

AWSは、Amazon SageMaker AIでMixture-of-Experts（複数の専門家モデルを使い分ける方式）モデルを提供するため、融合型MoEカーネルの最適化を3つ説明しました。最大10％のスループット向上と最大9％のTime to First Token（TTFT、最初の生成トークンが出るまでの時間）の短縮を報告しています。vLLMコミュニティと共同で開発し、オープンソース化したうえで、SageMaker AIのLMI Deep Learning Containerに適用しました。

最初の最適化では、行列乗算のオペランドを入れ替え、NVIDIA Hopper GPUのWGMMA命令の配置に合わせます。WGMMAはM次元が固定の64行ですが、MoEのデコードではユーザー系列ごとに1トークンを処理することが多く、Mは1となり、Tensor Coreの実行レーンが多数使われません。融合型MoEのtopk=8設定では、選ばれた専門家の中間次元をまとめることでN次元が非常に大きくなります。乗算を転置して大きい次元を第1オペランドに置くと、64行の配置に合い、非同期プリフェッチが可能になります。AWSによると、H200 GPU上の大規模MoEモデルで、エンドツーエンドの出力スループットが2～4％向上しました。この変更はvLLMバージョン0.24.0以降に取り込まれています。

2つ目の変更では、Hopper GPU上のFP8テンソル量子化MoE計算を、既定のFlashInfer+CUTLASS経路ではなくTritonに振り分けます。AWSのプロファイリングでは、Tritonはスケーリング、行列乗算、専門家の集約を1回のGPU処理に融合し、中間データをグローバルメモリーで読み書きする工程を省くため、同時実行数の少ない処理で高い性能を示しました。動的スケジューリングも、M < 64やM = 1のような小さく変動するトークンバッチを、CUTLASSの既定タイル配置より効率よく処理します。AWSは、カーネルコードを変更せず経路の設定を切り替えることで、H200上のFP8量子化済み融合型MoEモデルのエンドツーエンドスループットが10％向上したとしています。この変更はAWS Large Model Inference（LMI）v28コンテナで利用できます。

3つ目の最適化では、新しいトークン数に遭遇するたびにTritonがカーネルを再コンパイルするのを防ぎます。_fused_inv_rope_fp8_quant_per_headカーネルは、それまで新しいnum_tokens値ごとに再コンパイルしていました。可変長パラメーターにdo_not_specialize注釈を加え、繰り返し発生する実行時コンパイルをなくしました。AWSは、H200でテンソル並列処理を8分割した場合にTTFTが9％改善し、1基のGPUが再コンパイルする間にほかのGPUが待つ状況も避けられると報告しています。この最適化はvLLMバージョン0.22.0以降に取り込まれています。3つの変更により、SageMaker AI上のLMIコンテナでスループットが最大10％向上し、TTFTが最大9％短縮するとAWSは述べています。