Cursor abre kernel para treino de MoE que chega a 2.37x mais rápido

O kernel tira a CPU do caminho

O Mixture-of-Kittens implementa um megakernel determinístico de mixture-of-experts (MoE) criado para sistemas NVL72. Ele funde todas as etapas de cálculo e comunicação do MoE, sobrepõe processamento e rede entre GPUs e elimina a sincronização entre CPU e GPU.

O projeto suporta os formatos BF16 e MXFP8. Também cobre tanto o passe forward, que calcula a saída, quanto o backward, usado para atualizar o modelo durante o treino.

O código oferece uma API de baixo nível em mok/ops.py. Essa camada chama diretamente os kernels CUDA e deixa a equipe controlar os dados, os tensores simétricos e a coordenação dos lançamentos.

A API funcional fica em mok/functional.py. Ela cria a memória de trabalho, coordena os lançamentos e organiza o fluxo de execução para quem não precisa controlar cada etapa.

A Cursor usa a camada funcional no treino de produção do Composer. O README descreve o uso em uma infraestrutura interna distribuída por vários racks NVL72.

A instalação exige hardware específico

Para instalar o projeto, a equipe precisa de GPUs NVIDIA Blackwell SM100 ou SM103, como GB200 NVL72 ou GB300 NVL72. O ambiente também exige Python 3.12 ou posterior, PyTorch 2.10 ou posterior e CUDA 13.0 ou posterior.

A compilação depende de uma correspondência exata entre o CUDA do PyTorch e as versões principal e secundária do toolkit instalado no sistema. O MoK compila para SM103 por padrão; quem usa SM100 precisa definir a variável de ambiente MOK_ARCH.

O repositório não descreve um serviço externo para executar o kernel. O processamento acontece na infraestrutura de GPUs do usuário, enquanto a configuração e os lançamentos ficam sob controle da própria equipe.

Na adoção pela API funcional, o usuário cria uma configuração e um workspace. Depois, chama schedule(...) uma vez para montar o dispatch e o combine e reutiliza a mesma agenda em forward(...) e backward(...).

O desempenho depende de cinco parâmetros ajustáveis. A equipe precisa calibrar o número de SMs de comunicação em cada passe, o tamanho do minibatch, o buffer do macrobatch e o multiplicador de capacidade da agenda.

O README recomenda entre 4 e 52 SMs de comunicação para forward e backward. Também indica minibatches entre 2048 e 16384 e sugere elevar o macrobatch para preencher a memória disponível da GPU.

O parâmetro schedule_capacity_multiplier controla a capacidade reservada para tokens enviados a um único rank. O valor padrão é 0.5; a Cursor recomenda usar uma margem maior no começo do treino, quando o desequilíbrio entre especialistas costuma ser pior, e reduzir o valor depois.

Nos benchmarks de camadas MoE, o MoK chegou a 2.37x no forward MXFP8, 1.78x no backward MXFP8, 1.92x no forward BF16 e 1.58x no backward BF16, sempre contra o baseline mais rápido.

O repositório cursor/mixture-of-kittens tinha 195 estrelas, 14 forks e 4 issues abertas. A Cursor criou o projeto em 2026-07-29, publicou o último commit em 2026-08-03 e identificou a versão mais recente como v0.1.0.

O projeto tinha 6 dias de vida na coleta dos dados. O volume inicial de estrelas chama atenção, mas a pouca idade e as quatro issues abertas recomendam testar compatibilidade e desempenho antes de colocar o kernel no treino principal.

A licença Apache-2.0 permite usar, alterar e redistribuir o código conforme seus termos. O MoK não serve para equipes sem GPUs NVIDIA Blackwell SM100 ou SM103, sem CUDA 13.0 ou posterior ou sem uma operação de treino distribuído entre GPUs.