Mixture‑of‑Kittens (MoK)

它是什么 – MoK 是一个用于在 NVIDIA Blackwell GPU (NVL72) 上训练大型模型的确定性混合专家 (MoE) 超级内核。它将 MoE 层通常需要的计算和 GPU 间通信融合到一个 CUDA 内核中，消除了 CPU‑GPU 同步的需求，并实现了计算与网络通信的重叠。该库支持 BF16 和新的 MXFP8 格式，并提供前向和反向传播功能。

为什么它很重要

速度 – 仓库中的基准测试显示，在 MXFP8 前向传播中比最快的公开基准测试快 2.37 倍，BF16 也具有类似的增益。
可扩展性 – 通过在内核内部处理分发/合并，MoK 可以在机架内的多个 GPU 上扩展，而不会出现通常的 CPU 端瓶颈。
生产就绪 – 在内部用于训练 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`)	五个可调旋钮（通信 SM 数量、小批量粒度、宏批量大小、调度容量乘数），用于控制重叠和内存使用。
工作空间	存在于专家并行组中每个 GPU 上的对称内存缓冲区；每个模型创建一次并重复使用。
MXFP8	一种低精度权重格式；权重在传递给内核之前使用 `ops.mxfp8_quantize` 进行预量化。

何时使用

训练大型 MoE 模型，其中分发/合并步骤占据了运行时间的主要部分。
您可以访问 Blackwell GPU (NVL72) 并希望利用其新的 SM 功能。
您需要确定性训练（没有来自异步 CPU‑GPU 协调的随机性）。
您熟悉 PyTorch 并可以通过 torchrun 运行多 GPU 作业。

限制 / 注意事项

硬件特定 – 仅支持 NVIDIA Blackwell SM100/SM103；无法在旧版 GPU 上运行。
需要调优 – 性能在很大程度上取决于五个配置参数；您需要针对您的工作负载进行扫描。
对称内存开销 – 工作空间会在每个 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