高度な融合カーネルがMoE学習スループットを加速
NVIDIAは、大規模AIの基盤であるMoEモデルの学習効率を向上させるため、CuTe DSL対応の統合MLPカーネルを公開した。本技術はメモリ競合、同期オーバーヘッド、GLU活性化関数の処理遅延を解消し、ハードウェアとソフトウェアの協調設計でTensor Coreの稼働効率を最大化する。 従来手法では、動的なカーネル起動がCPU依存を招き、GLU演算の融合に中間テンソルの読み書きが必要となっていた。NVIDIAは入力とゲートテンソルを再配置する技術でGEMMエピローグ内にGLU演算を完全統合した。また、GroupGEMM内でグループごとのトークン数をGPUメモリ上で管理し、ホストデバイスの同期を排除。1イテレーション全体をCUDA Graphで実行するSync-free処理を可能にした。活性化・量子化・転置演算をGEMMに直接融合し、BF16テンソルの余分な読み書きを撤廃している。 マイクロベンチマークでは順伝播で最大1.3倍、逆伝播で最大2.1倍の高速化を実現。実環境のプリトレーニング試験では、DeepSeek-V3でエンドツーエンド性能が8%、GPT-OSSで93%の向上が確認された。通信カーネルとのオーバーラップにより、フルスタック学習パイプラインのスループットが大幅に強化された。 本機能はcuDNN Frontend、Transformer Engine、Megatron-Coreから利用可能である。GitHubに導入手順が公開されており、各抽象化レイヤーに応じて組み込める。今後はJAX対応の追加、Ahead-of-Timeコンパイルの実装、より複雑な融合パターンのサポートを進める予定である。本技術により、計算リソース制約下におけるMoEモデルのスケーラビリティがさらに向上する。
