Mixture‑of‑Kittens (MoK)

概要 – MoK は、NVIDIA Blackwell GPU (NVL72) で大規模モデルをトレーニングするための決定論的混合エキスパート (MoE) メガカーネルです。MoE レイヤーで通常必要とされる計算と GPU 間通信を単一の CUDA カーネルに融合させることで、CPU‑GPU 同期の必要性を排除し、ネットワーク通信と計算をオーバーラップさせます。このライブラリは BF16 および新しい MXFP8 フォーマットをサポートし、順伝播および逆伝播の両方を提供します。

重要性

速度 – リポジトリ内のベンチマークでは、MXFP8 の順伝播において最速の公開ベースラインより最大 2.37 倍高速であり、BF16 でも同様の向上が見られます。
スケーラビリティ – カーネル内でディスパッチ/結合を処理することで、MoK は通常の CPU 側のボトルネックなしに、ラック内の多数の GPU にスケールします。
プロダクション対応 – Cursor の "Composer" モデルのトレーニングに内部使用されており、フル機能の Python API が付属しています。

はじめに (概要)

環境 – NVIDIA Blackwell SM100/SM103 GPU, Python 3.12+, PyTorch 2.10+, CUDA 13.0+。
インストール
```
pip install . --no-build-isolation   # or python setup.py install
```
古い SM100 シリコンを使用している場合は MOK_ARCH=SM100 を設定してください。
クイックチェックの実行
```
import mok
print(mok.__version__)
```

関数型 API の使用 (ほとんどのユーザーに推奨)

from mok import functional
config = functional.MoKConfig()          # ハイパーパラメータは後で調整
workspace = functional.get_workspace(
    config, torch.distributed.group.WORLD,
    device=x.device,
    num_local_tokens=..., hidden_size=..., topk=...)
schedule = functional.build_schedule(workspace, config, topk_experts, ...)
out, fwd_ctx = functional.forward(config, workspace, schedule, ...)
grads = functional.backward(config, workspace, schedule, fwd_ctx, ...)

カスタム制御が必要な場合は、操作レイヤー (mok.ops) を使用して生の CUDA カーネルを直接呼び出すことができます。

コアコンセプト

コンセプト	機能
操作レイヤー (`mok/ops.py`)	CUDA メガカーネルの薄いラッパー。すべてのテンソルを自分で管理します。
関数レイヤー (`mok/functional.py`)	ワークスペースの割り当て、スケジュールの作成、順伝播/逆伝播のオーケストレーションを処理します。
設定 (`MoKConfig`)	オーバーラップとメモリ使用量を制御する 5 つの調整可能なノブ（通信 SM 数、ミニバッチ粒度、マクロバッチサイズ、スケジュール容量倍率）。
ワークスペース	エキスパート並列グループ内のすべての GPU に存在する対称メモリバッファ。モデルごとに 1 回作成され、再利用されます。
MXFP8	低精度ウェイトフォーマット。ウェイトはカーネルに渡される前に `ops.mxfp8_quantize` で事前量子化されます。

使用シーン

ディスパッチ/結合ステップが実行時間を支配する大規模 MoE モデルのトレーニング。
Blackwell GPU (NVL72) にアクセスでき、新しい SM 機能を活用したい場合。
決定論的なトレーニングが必要な場合（非同期の CPU‑GPU 調整によるランダム性を排除）。
PyTorch に慣れており、torchrun を介してマルチ GPU ジョブを実行できる場合。

制限事項 / 考慮事項

ハードウェア固有 – NVIDIA Blackwell SM100/SM103 のみをサポート。古い GPU では動作しません。
チューニングが必要 – パフォーマンスは 5 つの設定パラメータに大きく依存します。ワークロードに合わせてスイープする必要があります。
対称メモリのオーバーヘッド – ワークスペースは各 GPU 上に大きなバッファを割り当てます。GPU メモリを適切に計画してください。
MXFP8 サポート は実験的であり、ウェイトの事前量子化が必要です。

コミュニティと貢献

ライセンス：Apache 2.0。
貢献を歓迎します。プロセスについては CONTRIBUTING.md を参照してください。
研究やプロダクションで MoK を使用する場合は、付属の arXiv 論文 (arXiv:2609.36070) を引用してください。

TL;DR – Mixture‑of‑Kittens は、最新の NVIDIA GPU 向けの高性能で決定論的な MoE カーネルスタックであり、計算と通信を単一のカーネルに集約することで大幅な高速化を実現します。Blackwell マシンにインストールし、少数の設定ノブを調整し、関数型 API を使用して PyTorch MoE トレーニングパイプラインに組み込んでください。

cursor/mixture-of-kittens