Mixture‑of‑Kittens (MoK)

What it is – MoK is a deterministic mixture‑of‑experts (MoE) megakernel for training large models on NVIDIA Blackwell GPUs (NVL72). It fuses the compute and the inter‑GPU communication that MoE layers normally require into a single CUDA kernel, removing the need for CPU‑GPU synchronisation and overlapping work with networking. The library supports BF16 and a new MXFP8 format, and provides both forward and backward passes.

Why it matters

Speed – Benchmarks in the repo show up to 2.37× faster than the fastest public baseline for MXFP8 forward passes and similar gains for BF16.
Scalability – By handling dispatch/combine inside the kernel, MoK scales across many GPUs in a rack without the usual CPU‑side bottlenecks.
Production‑ready – Used internally for training the "Composer" model at Cursor, and shipped with a full‑featured Python API.

Getting started (high‑level)

Environment – NVIDIA Blackwell SM100/SM103 GPUs, Python 3.12+, PyTorch 2.10+, CUDA 13.0+.

Install

pip install . --no-build-isolation   # or python setup.py install

Set MOK_ARCH=SM100 if you are on the older SM100 silicon.

Run a quick sanity check
```
import mok
print(mok.__version__)
```

Use the functional API (recommended for most users)

from mok import functional
config = functional.MoKConfig()          # tune hyper‑params later
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, ...)

The ops layer (mok.ops) lets you call the raw CUDA kernels directly if you need custom control.

Core concepts

Concept	What it does
Ops layer (`mok/ops.py`)	Thin wrappers around the CUDA megakernel; you manage all tensors yourself.
Functional layer (`mok/functional.py`)	Handles workspace allocation, schedule creation, and forward/backward orchestration.
Config (`MoKConfig`)	Five tunable knobs (communication SM count, minibatch granularity, macrobatch size, schedule capacity multiplier) that control overlap and memory usage.
Workspace	Symmetric memory buffers that exist on every GPU in the expert‑parallel group; created once per model and reused.
MXFP8	A low‑precision weight format; weights are pre‑quantized with `ops.mxfp8_quantize` before being passed to the kernel.

When to use it

Training large MoE models where the dispatch/combine step dominates runtime.
You have access to Blackwell GPUs (NVL72) and want to exploit their new SM capabilities.
You need deterministic training (no randomness from asynchronous CPU‑GPU coordination).
You are comfortable with PyTorch and can run multi‑GPU jobs via torchrun.

Limitations / considerations

Hardware specific – Only supports NVIDIA Blackwell SM100/SM103; will not run on older GPUs.
Tuning required – Performance depends heavily on the five config parameters; you’ll need to sweep them for your workload.
Symmetric memory overhead – Workspaces allocate large buffers on every GPU; plan GPU memory accordingly.
MXFP8 support is experimental and requires you to pre‑quantize weights.

Community & contribution

License: Apache 2.0.
Contributions are welcomed; see CONTRIBUTING.md for the process.
Cite the accompanying arXiv paper (arXiv:2609.36070) if you use MoK in research or production.

TL;DR – Mixture‑of‑Kittens is a high‑performance, deterministic MoE kernel stack for the newest NVIDIA GPUs, offering substantial speedups by collapsing compute and communication into a single kernel. Install it on a Blackwell machine, tune the small set of config knobs, and use the functional API to plug it into your PyTorch MoE training pipeline.

cursor/mixture-of-kittens