高度な融合カーネルによる MoE 学習スループットの向上

Mixture-of-experts (MoE) モデルは、現代の大規模 AI システムの基盤コンポーネントとして急速に普及しました。各トークンに対してパラメータの一部のみを活性化しながらモデル容量を実質的に大幅に拡大できるため、実用的な計算リソース予算内でパフォーマンスをスケーリングする究極のアプローチを提供し、広く採用されています。モデル規模がさらに拡大するにつれ、これらのブロックの最適化はトレーニングのスループット最大化にとって極めて重要になります。 これらの限界を押し広げるために、私たちは CuTe DSL（Domain Specific Language）でカスタムビルドされた、稠密モデルおよび MoE モデル向けの高度に融合された MLP カーネルを導入します。本来的なメモリおよび同期のボトルネックに対処することで、これらの新しいカーネルは非融合パスと比較して 1.3 倍から 2 倍のカーネルレベルでの高速化を実現し、フルイテレーション CUDA Graphs（Graph）における同期フリーの MoE 実行を可能にします。 NVIDIA のフルスタック DeepSeek-V3 プリートレーニング設定において、この最適化はエンドツーエンドのパフォーマンスで 8% の向上をもたらします。同様に GPT-OSS プリートレーニング設定では、93% のエンドツーエンドパフォーマンス向上が実現されます。トレーニング時間の短縮やハードウェア利用率の最適化をいずれ目指す場合でも、これらのカーネルは今日から cuDNN Frontend で利用可能であり、Transformer Engine や Megatron-Core を通じてシームレスにアクセスできます。 その仕組みを理解するには、現代の MoE ブロックを悩ませる最大のボトルネック 3 つを体系的に分解し、ハードウェアを意識したソフトウェアコードデザインを通じて Tensor Cores（テンソルコア）が絶えず稼働し続けるようにスタックを再設計した方法を明らかにする必要があります。 ## MoE ブロックにおけるトレーニングのボトルネック克服 MoE モデルのスループットを最大化するためには、まず計算サイクルがどこで消費されているかを正確にマッピングする必要がありました。MoE ブロック内の標準的なトレーニング反復の実行タイムラインをプロファイルした際、3 つのシステムレベルのボトルネックが浮き彫りになりました。 - アクティベーションのボトルネック：アクティベーション関数は通常、メモリーバウンド型のカーネルや大規模なテンソルの読み書き操作を引き起こし、これらの間隔中に Tensor Cores が未利用のままになります。 - CPU による制約/オーバーヘッド：ルーティングされたエキスパートでは、エキスパートあたりのトークン数が実行時に計算され、通常は CPU で処理されます。CPU が GPU の速度に追いつけない場合、CPU 側の操作がボトルネックとして顕在化します。これに対応するためには、CPU との同期や介入を必要としないカーネルを構築する必要があります。 - 量子化のコスト：アクティベーション関数と同様に、高精度なテンソルを低精度へ量子化するとメモリーバウンド型のカーネルとなり、Tensor Cores がアイドル状態になります。 これらの課題に対し、MoE ブロックの再設計において cuTE DSL（DSL）で記述されたカスタムカーネルを採用し、同期不要の MoE 向けに特別に設計された 3 つのカーネルファミリーを導入しました。 - GroupGemm + Quantize - GroupGemm + Activation + Quantize/Transpose - GroupGemm + dActivation + Quantize/Transpose サポートされている活性化関数は、SwiGLU、GeGLU、sReLU であり、クランプ処理とスケーリングの追加オプションも用意されています。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/06/image3-4-625x370.png)*Figure 1. Fusing operations into a single custom kernel in the forward and backward pass with CuTe DSL* ## 融合された GEMM エピローグによる GLU 活性化関数の最適化 ゲート付き線形関数は最近非常に人気を集めており、現代のモデルのほとんどは SwiGLU や GeGLU など、ゲート付き線形ユニット (Gated Linear Unit: GLU) 活性化関数の変種を使用しています。これらの活性化関数は FC1 レイヤーの出力をチャンクに分割し、それらを結合して最終的な GLU 出力を作成します。私たちは、順伝播と逆伝播の両方において GEMM と対応する GLU 演算をシームレスに統合した融合カーネルを実装しました。 GLU 活性化関数は、GEMM のエピローグ内で融合させることが容易ではありません。なぜなら、GLU はテンソルの2つの異なるチャンク（入力とゲート）へのアクセスを必要とするからです。通常、これら2つのチャンクは異なるスレッドブロックによって計算され、2 つの出力を結合するためには、カーネルが両方の出力をグローバルメモリに書き込む必要があります。この融合を実現するために、重みを入力の列とゲートの列として再パックします。これにより、同じスレッドブロックが入力テンソルの半分のタイル幅とゲートテンソルの半分のタイル幅の両方にアクセスできるようになります。その結果、入力とゲートをグローバルメモリを経由することなくエピローグ内で結合することが可能になります。この再パックは、トレーニング開始前、チェックポイント読み込み中に実行できます。 同様に、逆伝播（バックプロパゲーション）においても、エピローグは GEMM の出力を読み取り、dSwiGLU を計算し、それを量子化してグローバルメモリに書き戻します。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/06/image5-1-625x634.png)*図 2. 入力とゲートの重みがパックされることで、スレッドブロックは CUDA Core 内で SwiGLU 出力を計算するために両方の重みにアクセスできるようになります*** 特筆すべきは、これらの融合パターンが中間テンソルの読み書きを単に排除するだけでなく、残存するメモリ操作を GEMM そのものと直接オーバーラップさせることで、利用効率を最大化することです。 SwiGLU、GeGLU、sReLU のようなコア活性化関数を超えて、これらのカーネルはネイティブに特徴量のスケーリング、テンソルのクランプ、バイアスベクトルの加算を含む融合されたエピローグ操作を処理します。 ホストとデバイスの同期および CPU 起動オーバーヘッドの排除 従来、カーネルが実行する作業量は、起動時のブロック数によって定義され、形状情報はホスト側で利用可能である必要があります。例えば、マルチストリームグループ化 GEMM では、異なる GEMM が別々のストリーム上で実行されます（ここで N はグループ数です）。各グループあたりのトークン数はランタイムで決定されるため、リソース利用率を最大化するために CPU はこれらの動的サイズを持つ GEMM を別々のストリーム上で起動する必要があります。 これにより 2 つの主要な問題が生じます。第一に、起動するカーネルの数はローカルエクスパートの数に応じてスケールします。第二に、カーネル起動前にホスト側で形状情報を取得するために同期ポイントが必須となります。これらの課題に対処するため、CuTe DSL GroupGEMM カーネルは GPU メモリ内でグループあたりのトークンを追跡します。これにより、反復処理中の CPU 依存性が排除され、反復全体にわたって CUDA Graphs が可能となり、CPU のボトルネックが効果的に解消されます。 MXFP8 および NVFP4 量子化の融合による露出メモリオーバーヘッドの削減 事前学習における MXFP8 や NVFP4 といった低精度レシピの人気は高まっており、これらの精度は精度への影響を最小限に抑えつつ、大幅な速度向上をもたらします。これらの低精度レシピでは、狭い精度の GEMM（行列乗算）演算のために、活性化関数の後に量子化と転置が行われます。 MXFP8 の場合、量子化カーネルは活性化関数の出力（BF16）を読み取り、MXFP8 形式の出力および逆伝播用の転置されたバージョンの出力を書き出します。新たに設計したカーネルはこの量子化ステップを GEMM カーネル自体に融合させることで、BF16 テンソルの追加読み書きを不要にしています。同様に NVFP4 の場合、カーネルは順伝播用に BF16 形式の出力とテンソルごとの amax（配列最大値）を生成し、逆伝播では出力の転置されたアダマール回転に対する amax を計算します。これにより、テンソルごとの amax 計算のための追加メモリパスが不要になります。 ## カーネルレベルでの向上から事前学習全体の速度向上へ ユニットレベルのマイクロベンチマーク全体を通じて、これらの融合カーネルは大幅な速度向上をもたらします。従来の非融合実行パスと比較して、順伝播は最大 **1.3 倍**、逆伝播は最大 **2.1 倍** 加速されます。 これらの速度向上をエンドツーエンドのトレーニングスループット向上に結びつけるために、以下の機能もサポートしています: - ダイナミックスケジューリング: エキパート並列性やデータ並列性からの通信など、他のカーネルとの効率的なオーバーラップをサポートします。 - 設定可能なクラスタマージンにより、カーネルを実行する SM（ストリーミングマルチプロセッサ）の数を制限することで、ユーザーが SM リソースの一部を予約できる機能を提供し、他のカーネルが GPU 上で同時に起動・実行するための余裕を残します。 これらの同期不要なカーネルによりエンドツーエンドの CUDA グラフの利用や通信カーネルとの効率的なオーバーラップが可能となるため、個別のカーネルごとの高速化に加え、アプリケーション全体レベルでははるかに大きな速度向上が期待できます。内部テストでは、Deepseekv3 において最大**8% のエンドツーエンド速度向上**、GPT-OSS 事前学習ランにおいて最大**93% のエンドツーエンド速度向上**を確認しています。 私たちは引き続き新しいカーネルの追加や、これらのカーネルに対する新機能のサポートを継続しています。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/06/image1-6.webp)*図 3. GB200 における異なる活性化関数パターンでの速度向上。ベースラインは Transformer Engine の最適化カーネルを使用* ## CuTe DSL 統合カーネルをどのように活用するか これらのカーネルは、異なる抽象度レベルで使用可能です。 - cuDNN Front-end (v1.23.0+): これらのカーネルは cuDNN Frontend ライブラリに格納されています。ユーザーはこのライブラリをソフトウェアスタックにインストールし、そこから直接これらのカーネルを呼び出すことができます。Cudnn-Frontend はまた、これらのカーネル用のラッパーも提供しており、初回の呼び出し時にカーネルをコンパイルし、その後の呼び出しではキャッシュされたオブジェクトを再利用します。ユーザーは、カーネルを直接呼び出すか、またはラッパー API を通じてカーネルにアクセスするかを選択できます。また、これらのカーネルに対して AOT（Ahead of time）コンパイルサポートをライブラリに導入する作業も積極的に進めており、これによりカーネルを cubins としてコンパイルし、ディスク上にキャッシュできるようにします。 - Transformer Engine (v2.15+): ユーザーはまた、Transformer Engine を通じてこれらのカーネルを使用することもできます。Transformer Engine は、transformer_engine.pytorch.ops 構造体を通じてこれらの操作を公開しています。これらの操作は、transformer_engine.pytorch.ops.Sequential ブロックを使用して組み合わせることができ、内部でパターンマッチングを行い、cuDNN フロントエンドライブラリから融合されたカーネルを呼び出します。 - Megatron Core (26.04-alpha.rc2+): ユーザーはまた、Megatron Core を通じてこれらのカーネルを使用することもできます。この場合、適切な一連のノブ（設定項目）を使用するだけで機能を呼び出すことができます。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/06/image4-3.webp)*図 4****.**** ユーザーは、CUDA スタック内の異なる抽象化レイヤー（CuDNN Frontend、Transformer Engine、または Megatron Core）のいずれからでも、これらの融合カーネルをシームレスに統合して選択することができます* ## 次のステップ 私たちは現在、より多くの融合パターンや JAX などのより多くのフレームワークのサポートなど、複数の新機能の開発を積極的に進めています。 アクティベーションの再計算、最適なカーネルを選択するためのヒューリスティック、コンパイルコストを削減するための事前コンパイル（Ahead of Time (AOT) Compilation）、CPU オーバーヘッドの削減など、複数のカーネル最適化が進行中です。 特定の活性化関数をご希望の場合は、CuDNN カーネルを自分で調整し、PR を通じて貢献していただくことを推奨します。あるいは、cuDNN フロントエンドでその機能を追跡できるよう、Issue を作成してください。 コミュニティからのフィードバックを大歓迎いたします！ ## 始め方 これらのカーネルの実行方法については、[GitHub](https://github.com/NVIDIA/cudnn-frontend/tree/develop/benchmark/cutedsl_fusion_kernels) の手順に従ってください。

高度な融合カーネルによる MoE 学習スループットの向上

背景や根拠まで確認しますか？

調べる

選ぶ

サイト