Prime Intellect heeft Prime Mega MoE uitgebracht, een trainingsuitbreiding voor de gefuseerde GPU-kernel Mega MoE van DeepSeek. Waar de kernel van DeepSeek alleen de forward pass ondersteunde, voegt Prime Intellect er een snelle backward pass aan toe, die overlapt met andere berekeningen, waardoor de kernel bruikbaar wordt voor training en niet alleen voor inferentie. Volgens het bedrijf draaien volledige trainingsstappen op Qwen3 MoE-modellen in hun prime-rl trainer tot 34% sneller, en tot 55% sneller vergeleken met TorchTitan.
In een MoE-model (mixture of experts, of ‘mengsel van experts’) wordt elke token alleen naar een kleine subset van experts gestuurd, verdeeld over verschillende GPU’s. Elke laag moet daarom de tokens versturen, de experts laten berekenen en vervolgens de resultaten terugsturen, zowel tijdens de forward pass als de backward pass, de terugpropagatiefase die voor leren wordt gebruikt. In een standaardconfiguratie zijn deze operaties gescheiden: de GPU sorteert de tokens, wacht op transport, berekent en wacht dan weer. Deze wachttijden zijn verloren rekentijd.
Een gefuseerde kernel, compleet gemaakt voor training
De Mega MoE van DeepSeek, toegevoegd aan de DeepGEMM-bibliotheek in april 2026, vervangt deze keten door een enkele kernel-aanroep. Een kernel (kernel) is een geoptimaliseerd programma dat direct op de GPU draait. Deze fuseert de expert-parallelle dispatch, de twee lineaire lagen van de experts, SwiGLU en de combinatiestap. Terwijl een groep tokens onderweg is, berekent de kernel op een groep die al is aangekomen, en vermijdt het tussenliggende tensoren dankzij werkbuffers die alleen op de GPU staan. Deze kernel bevatte echter alleen een implementatie van de forward pass.
Prime Intellect heeft de ontbrekende backward pass geschreven, ook overlappend. De ondersteunde datatypen zijn BF16, MXFP8 (gereserveerd voor Blackwell-GPU’s) en blok-FP8 (gereserveerd voor Hopper-GPU’s). De prime-rl-configuratie biedt ook een fp8_bwd-modus: forward pass in BF16, backward pass in FP8.
Zelfgerapporteerde cijfers
Alle resultaten zijn gepubliceerd door Prime Intellect zelf. De trainingsstappen gebruiken nodes van 8 H200’s, in FP8, met sequenties van 16.384 tokens en volledige activatie-checkpointing. Op Qwen3-30B-A3B, verdeeld over 32 H200’s, duurt een volledige stap 2,78 seconden met de prime-rl-basis, 3,07 seconden met TorchTitan (FP8) en 2,08 seconden met Prime Mega MoE, wat winsten van 34% en 48% oplevert. Op Qwen3-235B-A22B, verdeeld over 128 H200’s, zijn de tijden 12,72 seconden, 14,98 seconden en 9,65 seconden, voor winsten van 32% en 55%. De doorvoer gaat, op het 30B-model, van 5.337 naar 7.886 tokens per seconde per GPU vergeleken met TorchTitan, en op het 235B-model van 1.094 naar 1.697.
Gemeten laag voor laag op een enkele node, is de winst groter aan de voorkant dan aan de achterkant. De forward pass van het Qwen3-30B-A3B-model gaat van 12,0 ms naar 3,9 ms, een factor 3,11. De backward pass van het Qwen3-235B-A22B-model gaat van 30,5 ms naar 24,6 ms, een factor 1,24. Het cijfer ‘tot 3,11x’ heeft dus alleen betrekking op de forward pass. TorchTitan claimt zelf een winst tot 41%, maar in een andere configuratie (B200, model van 671 miljard parameters, BF16-basis), die de bron niet direct vergelijkbaar vindt.
Prime Mega MoE wordt vrijgegeven onder de Apache 2.0-licentie. In prime-rl wordt het geactiveerd als dispatch-backend: het volstaat om type = « mega_moe » in te stellen onder [trainer.model.moe.dispatch], en vervolgens de precisie te kiezen uit bf16, mxfp8, fp8 of fp8_bwd. De trainingskernels worden geleverd in een voorgecompileerd DeepGEMM-wheel, beschikbaar in de prime-kernels-repository, omdat de standaardversie van DeepGEMM ze niet bevat.
