Prime Intellect ha pubblicato Prime Mega MoE, un’estensione per l’addestramento del kernel GPU fuso Mega MoE di DeepSeek. Mentre il kernel di DeepSeek copriva solo il forward pass, Prime Intellect vi aggiunge un backward pass veloce, che si sovrappone agli altri calcoli, in modo che il kernel diventi utilizzabile per l’addestramento e non solo per l’inferenza. Secondo l’azienda, nel suo addestratore prime-rl, i passi di addestramento completi su modelli Qwen3 MoE girano fino al 34% più velocemente, e fino al 55% più velocemente rispetto a TorchTitan.
In un modello MoE (mixture of experts, o « miscela di esperti »), ogni token viene inviato solo a un piccolo sottoinsieme di esperti, distribuiti su diverse GPU. Ogni livello deve quindi inviare i token, far calcolare gli esperti, poi rinviare i risultati, durante il forward pass come il backward pass, la fase di retropropagazione utilizzata per l’apprendimento. In una configurazione standard, queste operazioni sono separate: la GPU ordina i token, attende il trasporto, calcola, poi attende di nuovo. Questi tempi di attesa sono altrettanto tempo di calcolo inattivo.
Un kernel fuso, completato per l’apprendimento
Il Mega MoE di DeepSeek, aggiunto alla libreria DeepGEMM nell’aprile 2026, sostituisce questa catena con una singola chiamata di kernel. Un kernel (kernel) è un programma ottimizzato che viene eseguito direttamente sulla GPU. Questo fonde il dispatch parallelo per esperti, i due strati lineari degli esperti, SwiGLU e la fase di combinazione. Mentre un gruppo di token è in transito, il kernel calcola su un gruppo già arrivato, ed evita i tensori intermedi grazie a buffer di lavoro situati unicamente sulla GPU. Questo kernel includeva tuttavia solo un’implementazione del forward pass.
Prime Intellect ha scritto il backward pass mancante, anch’esso in sovrapposizione. I tipi di dati supportati sono BF16, MXFP8, riservato alle GPU Blackwell, e FP8 a blocchi, riservato alle GPU Hopper. La configurazione di prime-rl propone anche una modalità fp8_bwd: forward pass in BF16, backward pass in FP8.
Cifre auto-dichiarate
Tutti i risultati sono pubblicati da Prime Intellect stessa. I passi di addestramento utilizzano nodi di 8 H200, in FP8, con sequenze di 16.384 token e un checkpointing di attivazione completo. Su Qwen3-30B-A3B, distribuito su 32 H200, un passo completo dura 2,78 secondi con la base prime-rl, 3,07 secondi con TorchTitan (FP8) e 2,08 secondi con Prime Mega MoE, ovvero guadagni del 34% e del 48%. Su Qwen3-235B-A22B, distribuito su 128 H200, le durate sono di 12,72 secondi, 14,98 secondi e 9,65 secondi, per guadagni del 32% e del 55%. La velocità passa, sul modello 30B, da 5.337 a 7.886 token al secondo e per GPU rispetto a TorchTitan, e sul modello 235B, da 1.094 a 1.697.
Misurato livello per livello su un singolo nodo, il guadagno è più marcato nel forward che nel backward. Il forward pass del modello Qwen3-30B-A3B passa da 12,0 ms a 3,9 ms, ovvero un fattore di 3,11. Il backward pass del modello Qwen3-235B-A22B passa da 30,5 ms a 24,6 ms, ovvero 1,24. La cifra « fino a 3,11x » riguarda quindi solo il forward pass. TorchTitan rivendica dal canto suo un guadagno fino al 41%, ma in una configurazione diversa (B200, modello da 671 miliardi di parametri, base BF16), che la fonte ritiene non direttamente confrontabile.
Prime Mega MoE è diffuso con licenza Apache 2.0. In prime-rl, si attiva come backend di dispatch: basta definire type = « mega_moe » sotto [trainer.model.moe.dispatch], poi scegliere la precisione tra bf16, mxfp8, fp8 o fp8_bwd. I kernel di addestramento sono forniti in una ruota DeepGEMM precompilata, disponibile sul repository prime-kernels, poiché la versione standard di DeepGEMM non li contiene.
