Cursor abre código de megakernel MoE para GPUs NVL72

O que é o Mixture-of-Kittens (MoK)

A Cursor open-sourced o Mixture-of-Kittens (MoK), um megakernel de treinamento MoE (Mixture-of-Experts) para racks GB300 NVL72 que funde toda a comunicação e computação da camada MoE em um único kernel determinístico. O MoK já alimenta o treinamento do Composer (modelo de código agêntico da Cursor) em dezenas de milhares de GPUs.

O problema

Ao escalar o treinamento do Composer, a camada MoE tornou-se o principal gargalo - consumindo mais da metade do tempo total de treinamento. Abordagens anteriores otimizavam apenas a computação, tratando a comunicação inter-GPU separadamente. Com a migração para GB300 NVL72s, dois fatores mudaram o cenário: o NVLink permite sobreposição fina de comunicação e computação entre 72 GPUs, mas as CPUs Grace integradas são lentas, fazendo com que as GPUs ficassem ociosas esperando trabalho do lado da CPU.

Escolha da direção de comunicação

Uma contribuição central é a análise de quando usar comunicação pull vs push. Contrariando a prática comum (push), o MoK usa:

• Pull para dispatch (forward e backward reverse-combine): elimina sinalização cross-GPU (18 µs vs 103 µs do push) e entrega até 29% maior utilização de bandwidth NVLink sob desbalanceamento de experts.
• Push para combine: reutiliza a mesma tabela de schedule construída para o dispatch.

A tabela de schedule é construída uma única vez e reutilizada nas quatro operações de comunicação (forward dispatch/combine e backward reverse-dispatch/reverse-combine), sem necessidade de ordenação.

Granularidade do overlap computação-comunicação

O MoK usa um minibatch tunável - nem fino demais (tensor cores não saturam) nem grosso demais (espera longa pelos primeiros tokens). A heurística: escolher um minibatch que produza pelo menos duas ondas completas de SMs por GEMM agrupado. Para Kimi 2.5 (H=7168, I=2048, 148 SMs), isso dá T ≥ 2368 tokens, confirmado empiricamente.

O overlap é implementado com inter-SM overlapping: alguns SMs fazem comunicação (comms SMs) e outros fazem computação (comp SMs), sinalizando-se via contadores locais. Com TMA, menos de um terço dos SMs satura o NVLink.

Ring token buffers (macrobatch)

Para lidar com a dinamicidade do MoE sem sincronização CPU-GPU nem token dropping, o MoK usa um ring buffer de algumas centenas de MB. Dispatch e combine são intercalados para que slots sejam drenados e reutilizados imediatamente. Um ring reverso no forward minimiza o replay de ativações no backward.

Outras características

• Megakernel: fusão total via software partitioning (mais confiável que green contexts).
• Determinismo: operações em ordem fixa, output bitwise-identical.
• CLC (Cluster Launch Control): permite overlap com RDMA inter-rack.
• MXFP8: treinamento em MXFP8 com shared expert em BF16.
• Router weight gradients: cálculo estilo SonicMoE fundido no backward do SwiGLU.

Resultados

Em benchmarks de camada MoE individual (EP=64, 2048 tokens/GPU, NVL72):

• Até 2,37x mais throughput MXFP8 forward vs o melhor baseline público.
• Até 1,78x MXFP8 backward, 1,92x BF16 forward, 1,58x BF16 backward.

Em produção (512 GPUs, vários racks NVL72):

• 1,41x mais tokens/segundo end-to-end vs a stack anterior baseada em DeepEP.

Agentes e megakernels

A Cursor destaca que agentes de IA ajudaram a escrever o megakernel - automatizando kernels single-operator e acelerando o trabalho complexo com uma equipe pequena.