10月14日、AWSが「Mixture-of-Experts kernel engineering on Amazon SageMaker AI」と題した記事を公開した。この記事では、Amazon SageMaker AI上でMixture-of-Experts（MoE）モデルのカーネル最適化を実装し、スループット向上とレイテンシ削減を実現した技術的手法について詳しく紹介されている。3つの最適化のうち1つは「設定変更だけでスループット10%改善」という実用的な知見であり、いずれもvLLMにアップストリームされてオープンソースとして利用可能だ。
MoEカーネル最適化が重要な理由
Mixture-of-Experts（MoE）は、トークンごとの計算量を抑えながらLLMの容量をスケールさせるスパース活性化アーキテクチャだ。Mistral系やDeepSeek、Qwenなど近年の主要モデルに採用が広がっている。MoE演算は総計算時間の大きな割合を占めるため、ここの非効率がそのままリクエストレイテンシとコストに直結する。
AWSのチームが今回実施した最適化は3つ。H200上での実測値として、スループット最大10%向上・TTFT（Time to First Token）最大9%短縮という結果が得られている。いずれもvLLMにアップストリームされており、オープンソースとして利用可能だ。
最大のポイント：Swap-ABによるWGMMA活用率の改善
最も技術的に興味深いのが「Swap-AB」最適化だ。
NVIDIAのHopperアーキテクチャが採用するWGMMA（Warp Group Matrix Multiply Accumulate）命令には構造的な制約がある。行次元（M）が64以上でないとTensor Coreの実行レーンが無駄になるという非対称性だ。
MoEのデコードフェーズでは、1ユーザーあたり1トークンずつ処理するため、バッチサイズが極端に小さくなる（M=1が典型的）。通常の行列積の形式：
A[M,K] x B[K,N] → C[M,N]
この配置では、Aが「1×4096」のような極薄い水平ベクトルになり、Tensor Coreの64行スロットをほぼ埋められない。シリコンが空の計算スロットをサイクルとして消費する状態になる。
さらにFused MoEカーネルでは、N次元が複数のエキスパートをまとめた形（例：topk=8の場合、数万〜数十万パラメータ規模）に膨れ上がる。この非対称性が問題を悪化させる。
解法は行列転置の代数的性質を利用することだ：
BT[N,K] x AT[K,M] → CT[N,M]
転置によってエキスパート次元（N=8192など）を第1オペランドに移すことで、2つの効果が得られる：
- Tensor Coreの稼働率最大化：大きなNが64行の閾値を容易に超え、無駄な実行サイクルがなくなる
- 非同期プリフェッチの有効化：Hopperの専用コピーエンジンが連続した大きな重みブロック（BT）を先読みし、パイプラインを飽和させ続ける
数式の正確性は維持したまま、PTXアセンブリ層に到達する前に行列乗算をプログラム的にフリップするという実装だ。結果、H200上のMoEモデルでエンドツーエンドのスループットが2〜4%向上した。
FP8テンソル量子化のバックエンド切替：設定変更だけで10%改善
次に大きな効果を得られたのが、FP8量子化モデルのバックエンド選択だ。
vLLMは量子化方式によってTriton、FlashInfer+CUTLASSなど複数のバックエンドにルーティングする仕組みを持つ。プロファイリングの結果、HopperでのFP8テンソル量子化モデルではTritonバックエンドがデフォルトのFLASHINFER_CUTLASSを一貫して上回ることが判明した。
理由はレジスタフュージョンにある。CUTLASSはFP8のテンソルスケーリング、行列乗算、エキスパート集約を別々の処理として扱う。対してTritonバックエンドはこのパイプライン全体を単一のハードウェアパスに統合し、グローバルHBMからのデータ読み込みを1回で済ませ、GPU内部のSRAMレジスタ上で完結させる。
また、CUTLASSはバッチサイズが大きく予測可能な行列に最適化されており、デコードフェーズのM<64やM=1という可変バッチプロファイルでは非効率なブロックタイリングが生じる。Tritonの動的スケジューリングはこの点を補う。
バックエンドルーティングの変更はカーネルコードを一切変えずに純粋な設定変更で完結し、FP8量子化MoEモデルのスループットを10%向上させた。AWS Large Model Inference（LMI）v28コンテナで利用可能だ。
Triton JIT再コンパイルの排除：TTFTを9%改善
3つ目の最適化は「見えないオーバーヘッド」の排除だ。
Tritonはデフォルトで整数引数を特殊化（specialize）する。カーネルが新しい値で呼び出されるたびにJITコンパイルが走る。可変バッチサイズのサービング環境では、これが継続的なコンパイルオーバーヘッドとなりTTFTを圧迫する。
MoEサービングでは_fused_inv_rope_fp8_quant_per_headカーネル（回転位置エンベディングとFP8量子化を融合したカーネル）がnum_tokensの新しい値ごとに再コンパイルしていた。可変長パラメータにdo_not_specializeアノテーションを追加するだけで再コンパイルが完全になくなった（PR #42849）。
特に8-wayテンソル並列のマルチGPU環境では、1台のGPUが再コンパイルで停止すると残りの7台がアイドル待ちになる。全クラスタで再コンパイルを排除することで、8-wayテンソル並列のH200でTTFTが9%短縮した。
3つの最適化まとめ
| 最適化 | 効果 | 適用条件 |
|---|---|---|
| Swap-AB（WGMMA活用） | スループット+2〜4% | MoEデコードフェーズ全般 |
| FP8バックエンドをTritonに切替 | スループット+10% | FP8量子化MoEモデル |
| JIT再コンパイル排除 | TTFT-9% | 8-wayテンソル並列 |
3つの最適化はいずれもvLLMにアップストリームされており、SageMaker AIのLMI Deep Learning Containerに適用済みだ。Swap-ABのような「数式の正確性を保ちながらハードウェアのISA制約に合わせてオペランドを入れ替える」アプローチは、他のアーキテクチャへの応用も効く考え方だ。
詳細はMixture-of-Experts kernel engineering on Amazon SageMaker AIを参照していただきたい。
