Cursor ha pubblicato con licenza Apache 2.0 MoK, acronimo di Mixture-of-Kittens, una tecnologia open source progettata per accelerare l’addestramento dei modelli Mixture-of-Experts sui sistemi Nvidia GB300 NVL72 e GB200 NVL72. Il kernel viene già utilizzato per addestrare Composer, il modello di programmazione sviluppato internamente dall’azienda, e integra in un unico componente deterministico tutte le operazioni di comunicazione e calcolo eseguite all’interno dei livelli MoE.
Nei modelli Mixture-of-Experts, i token devono essere distribuiti tra diversi esperti collocati su GPU differenti. Il trasferimento dei dati può occupare più della metà del tempo complessivo di addestramento, trasformando la comunicazione tra acceleratori nel principale collo di bottiglia. Cursor ha quindi progettato MoK partendo dalla comunicazione, anziché concentrarsi soltanto sull’ottimizzazione delle singole operazioni matematiche.
MoK è stato sviluppato specificamente per le configurazioni rack Blackwell nelle quali 72 GPU sono collegate all’interno di un unico dominio NVLink. La pianificazione delle operazioni viene eseguita direttamente sulle GPU, riducendo gli interventi della CPU e le sincronizzazioni tra processore e acceleratori. Comunicazione e calcolo vengono inoltre sovrapposti, in modo che il trasferimento dei dati possa proseguire mentre le GPU eseguono le operazioni sui token già disponibili.
Il sistema combina due modalità di trasferimento. Accanto al tradizionale schema push, nel quale una GPU invia i dati verso la destinazione, MoK utilizza un meccanismo pull che permette alla GPU ricevente di recuperare direttamente le informazioni necessarie. Questa configurazione ha aumentato fino al 29% l’utilizzo della larghezza di banda NVLink ed elimina parte dei segnali di conferma normalmente scambiati tra le GPU al termine delle operazioni.
I token vengono gestiti attraverso un Ring Token Buffer di dimensione fissa, riutilizzato ciclicamente senza richiedere l’intervento della CPU. MoK sfrutta inoltre la funzione Cluster Launch Control dell’architettura Blackwell per eseguire contemporaneamente le comunicazioni NVLink all’interno del rack e i trasferimenti RDMA tra rack differenti.
Il kernel supporta sia la precisione BF16 sia il formato MXFP8 e comprende nello stesso megakernel il calcolo dei pesi del router e le operazioni SwiGLU. L’ordine delle operazioni in virgola mobile è mantenuto costante, così che uno stesso input produca sempre lo stesso risultato. Questa caratteristica consente di impiegare MoK anche nelle attività di post-addestramento con apprendimento per rinforzo e nelle procedure di verifica dei modelli.
Nei test condotti su un singolo rack GB300 NVL72, MoK ha superato implementazioni pubbliche come DeepEP e la combinazione NCCL-PyTorch. Il miglioramento massimo della capacità di elaborazione ha raggiunto 2,37 volte nella propagazione in avanti con MXFP8 e 1,78 volte nella fase di retropropagazione. Con BF16, l’aumento è stato pari a 1,92 volte nel passaggio forward e a 1,58 volte nel backward.
In un ambiente reale utilizzato per l’addestramento di Composer con 512 GPU, il numero di token elaborati ogni secondo da ciascun acceleratore è passato da 760,9 a 1.070,2. La capacità complessiva di addestramento è quindi aumentata di circa il 41%, pari a un miglioramento di 1,41 volte.
MoK funziona esclusivamente con GPU Nvidia Blackwell SM100 e SM103 installate nei sistemi GB300 NVL72 o GB200 NVL72. L’esecuzione richiede almeno Python 3.12, PyTorch 2.10 e CUDA 13.0, limitando l’utilizzo alle organizzazioni che dispongono di grandi infrastrutture Blackwell, come laboratori di ricerca, operatori di GPU cloud e centri nazionali di supercalcolo.
Cursor ha collegato lo sviluppo del kernel ai progressi compiuti dagli strumenti di programmazione basati sull’intelligenza artificiale. Secondo l’azienda, questi sistemi sono ormai in grado di automatizzare gran parte dell’ottimizzazione dei singoli kernel e consentono anche a gruppi di dimensioni ridotte di affrontare componenti più complessi, come i megakernel distribuiti destinati ai modelli Mixture-of-Experts.