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에 존재하는 대칭 메모리 버퍼; 모델당 한 번 생성되고 재사용됩니다.
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