Amazon SageMaker AI

融合MoEは、トークンのルーティング、データのシャッフル、エキスパートネットワークの計算を別々の逐次ステップとして実行するのではなく、単一のバッチ化されたGPU操作に統合する。この構造はレイテンシの改善に役立つが、カーネルのチューニングは難しくなり、NvidiaのHopperアーキテクチャ——H200 GPUを含む世代——の制約に特に敏感になる。

最初のボトルネックはデコード時に現れる。システムはユーザーシーケンスごとに一度に1トークンを処理する。バッチサイズは1行または64行未満に縮小する一方、エキスパートの重み行列は大きいままである。HopperのWGMMA命令は第1オペランドに固定の64行レイアウトを想定しているため、細いトークンベクトルではTensor Coreの容量の多くが未使用のままになる。AWSのSwap-AB転置は大きなエキスパート次元をその先頭スロットに昇格させ、レイアウトを飽和させるとともにHopperがメモリを非同期でプリフェッチできるようにする。

2つ目の効果は、FP8量子化されたMoE処理をデフォルトのFLASHINFER_CUTLASSパスではなくvLLMのTritonバックエンドにルーティングしたことから得られた。Tritonはスケーリング、行列積、エキスパート集約を高速なオンチップメモリを通る1パスに融合し、ライブ配信の小さく断片化したトークンバッチに適している。3つ目の修正は、トークン数の変化によって引き起こされるTriton JITの繰り返し再コンパイルを排除したもので、8基のGPUすべてが歩調を合わせる必要があるマルチGPUクラスタでは重要である。これらの変更により、H200 GPUで最大10%のスループット向上と最大9%の初回トークン時間短縮が実現した。