cursor/mixture-of-kittens

何を解決するか

Mixture-of-Kittens (MoK) は、ハイエンドのNVIDIA Blackwell NVL72システム上でMixture-of-Experts (MoE) モデルを訓練する際の効率的なボトルネックを解決します。特に、CPU-GPU同期によるオーバーヘッドと、MoE訓練中に計算と通信が分離されていることによる遅延をターゲットとしており、これにより訓練プロセスが遅延することがあります。

仕組み

MoKは、すべてのMoE計算と通信を1つのCUDAカーネルに統合する決定論的「メガカーネル」です。これにより、コンピューティングとGPU間ネットワーキングを設定可能な粒度で重ね合わせ、CPU-GPU同期の完全な不要化を実現します。BF16およびMXFP8の精度形式をサポートし、順伝播と逆伝播の両方を処理できます。

対象ユーザー

このプロジェクトは、NVIDIA Blackwell GPU（SM100またはSM103）と大規模なMoE訓練に取り組む研究者およびエンジニア向けに設計されています。たとえば、プロダクションレベルのLLMを訓練している人などです。

特徴

• 高いパフォーマンス: 最速のベースラインと比較して、MXFP8順伝播で最大2.37倍高速、BF16順伝播で最大1.92倍高速。
• 統合カーネル: 計算と通信を1つの操作に統合し、ハードウェア利用効率を最大化。
• 精度サポート: 最適化された訓練のためのBF16およびMXFP8をサポート。
• 二層API: CUDAカーネル呼び出しに直接対応する低レベル「Opsレイヤー」と、生産環境への統合を容易にする高レベル「Functionalレイヤー」を提供。