mKernel：マルチGPU・マルチノード融合カーネルライブラリ、GPU駆動通信を実現

GPU通信オーバーヘッドは、本番AIワークロードにおいて測定可能なボトルネックです。mKernelプロジェクトが引用するデータによると、通信はフォワードパスの43.6%、エンドツーエンドのトレーニング時間の32%を消費する可能性があります。一般的な混合エキスパート（MoE）モデルでは、デバイス間通信が総実行時間の最大47%を占めることがあります。カリフォルニア大学バークレー校のUCCLプロジェクトの研究者らは、ノード内NVLink通信、ノード間RDMA、および計算を単一のカーネルに融合する持続的CUDAカーネルのライブラリであるmKernelをリリースしました。

問題：ホスト駆動通信

マルチGPU通信の標準モデルはホスト駆動です。CPUが制御パスを実行し、NCCLやNVSHMEMなどのライブラリを呼び出します。ライブラリはAllReduce、AllGatherなどの集合操作をGPU間で発行します。計算と通信は別々のCUDAストリームで実行され、カーネル境界でオーバーラップします。

研究チームはこのアプローチに2つの問題を指摘しています：

（1）CPUはGPUの計算能力に追いついていません。GB300 NVL72ラックは72基のBlackwell Ultra GPUと36基のGrace CPUを統合し、720 PFLOP/s FP8/FP6、1.44 EFLOP/s FP4 Tensor Core性能、130 TB/sの全対全ラック内NVLink帯域幅を提供します。この速度では、マイクロ秒単位のホストオーケストレーションオーバーヘッド（cudaLaunchKernel呼び出し、CPU側の「全書き込み完了」チェック、ストリーム間イベント）がパイプラインバブルとして直接現れます。

（2）ホスト駆動システムは、粗いカーネル境界で計算と通信をオーバーラップします。ホスト側からはタイルやチャンクレベルでの細かいオーバーラップは不可能です。

代替案はGPU駆動通信です。GPU自身が転送をトリガーし、通信は計算と同じカーネルに融合されます。既存の融合カーネルライブラリのほとんどは、単一ノードまたは単一GPU内で動作します。mKernelはマルチノードケースを対象としています。

mKernelの機能

mKernelは持続的CUDAカーネルのライブラリです。各カーネルは、ノード内NVLink通信、ノード間RDMA、高密度計算を単一のカーネルに融合します。

• マルチGPU + マルチノード、1つのカーネル内：ノード内NVLinkとノード間RDMAの両方が同じ持続カーネル内に存在します。
• 細粒度のカーネル内オーバーラップ：計算と通信はタイル/チャンク粒度でオーバーラップし、ノード内およびノード間のGPU通信をカバーします。
• 持続カーネルとSMの専門化：CTAは自身の役割（計算、内部通信、送信、削減）を割り当てます。各役割に専用のSM数は形状ごとに調整可能です。
• libibverbsに基づくGPU駆動ネットワーキング：mKernelはGPU起動のRDMA書き込みを使用し、NCCLやNVSHMEMに依存しません。通信バックエンドは、パフォーマンスを最大化し、異種ネットワークデバイスをサポートするためにゼロから作成されています。

5つの融合カーネル

| カーネル | 融合内容 | 説明 | |----------|----------|------| | AllGather + GEMM | AllGather → GEMM | 各ランクはAのシャードを保持。ランクがNVLink/RDMAを介して他のランクのシャードを収集する間、ローカルGEMMはタイルが到着次第消費する。 | | GEMM + AllReduce | GEMM → AllReduce | 1回の起動でC = A @ Bを計算し、部分出力を全ランクで削減。出力タイルは生成されるとすぐに削減ツリーにプッシュされる。 | | MoE Dispatch + GEMM | All-to-Allディスパッチ → グループ化GEMM | MoEトークンをエキスパートランクにルーティングし（ノード内NVLink + ノード間all-to-all）、同じカーネル内でエキスパートごとにグループ化されたGEMMを実行。トークンは到着次第処理され、ステージングバッファのラウンドトリップは不要。 | | Ring Attention | Ring KV交換 → FlashAttention | ランク間のシーケンス並列注意。各ステップでKVチャンクをリング上で回転させ、ローカルFlashAttentionは以前に受信したチャンクを消費。計算とリングの送受信は単一の持続カーネル内で並行して実行。 | | GEMM + ReduceScatter | GEMM → ReduceScatter | C = A @ Bを計算し、出力をreduce-scatter。各出力タイルは生成されるとすぐに削減され、所有するランクに転送される。 |

評価設定

研究チームは、ノード間ファブリックのみが異なる2つの2ノード×8-H200クラスタでmKernelを評価しました：

• AWS EFA：2ノード×8 H200、ノード内NVLink、ノード間AWS EFA/SRD、各ノード16×200 Gb/s EFA。
• ConnectX-7：2ノード×8 H200、ノード内NVLink、ノード間InfiniBand、各ノード8×400 Gb/s NVIDIA ConnectX-7。

mKernelは、NCCL、Triton-distributed、Flux、Mercury、MagiAttention、Transformer-Engine、ring-flash-attentionとベンチマークされました。より大規模なベンチマークは現在進行中です。

バックエンドと要件

mKernelは2つのネットワーキングバックエンドをサポートします：

| バックエンド | マクロ | トランスポート | 実行環境 | |------------|--------|--------------|----------| | CX7 | INTERNODE_BACKEND_IBVERBS | libibverbs RC | ConnectX-7 / InfiniBand / RoCE | | EFA | INTERNODE_BACKEND_EFA | libibverbs + efadv (SRD) | AWS p5/p5e (H200, EFA) |

両バックエンドは同じホスト側APIと同一のGPUカーネルを共有します。プロキシ/セッション実装のみが異なります（CX7はsession.h、EFAはsession_efa.h）。要件：NVIDIA Hopper GPU（デフォルトビルドターゲットsm_90a）、CUDA 12.9、PyTorchを備えたPython。CX7バックエンドにはlibibverbs開発ヘッダとライブラリが必要です。EFAバックエンドには、libfabric、libibverbs、efadvを含むAWS EFAインストールが必要で、デフォルトのEFA_HOMEは/opt/amazon/efaです。

ロードマップと重要ポイント

mKernelは現在5つの融合カーネルを提供し、ConnectX-7とAWS EFAバックエンドをサポートしています。将来の計画には以下が含まれます：

• トポロジ認識ディスカバリ、配置、ルーティングを備えた完全な異種アクセラレータ/NICサポート
• ノード間メガカーネル：複数の融合ステップを1つのメガカーネルに統合し、トランスフォーマー層全体をカバー
• Blackwell GPUサポート

重要ポイント：

• mKernelはノード内NVLink、ノード間RDMA、計算を単一の持続CUDAカーネルに融合します。
• 引用データによると、MoEモデルでは通信オーバーヘッドが実行時間の最大47%を占めます。
• 5つのカーネルが含まれます：AllGather+GEMM、GEMM+AllReduce、MoE Dispatch+GEMM、Ring Attention、GEMM+ReduceScatter。
• GPU起動RDMAはlibibverbsを介して直接実装され、NCCLやNVSHMEMに依存しません。
• 現在Hopper GPU（sm_90a）とConnectX-7またはAWS EFAネットワーキングが必要です。Blackwellサポートはロードマップにあります。

詳細はリポジトリと技術詳細を参照してください。