Prime Intellect publicou o Prime Mega MoE, uma extensão de treino para o kernel GPU fundido Mega MoE da DeepSeek. Enquanto o kernel da DeepSeek apenas cobria a passagem direta, a Prime Intellect adicionou-lhe uma passagem inversa rápida, que se sobrepõe com os outros cálculos, de modo que o kernel se torna utilizável para treino e não apenas para inferência. Segundo a empresa, no seu treinador prime-rl, as etapas completas de treino em modelos Qwen3 MoE rodam até 34% mais rápido, e até 55% mais rápido face ao TorchTitan.
Num modelo MoE (mixture of experts, ou « mistura de peritos »), cada token é enviado apenas para um pequeno subconjunto de peritos, distribuídos por diferentes GPUs. Cada camada deve, portanto, enviar os tokens, fazer calcular os peritos e depois reenviar os resultados, durante a passagem direta e a passagem inversa, a fase de retropropagação usada para aprendizagem. Numa configuração padrão, estas operações estão separadas: o GPU ordena os tokens, aguarda o transporte, calcula e depois aguarda novamente. Estes tempos de espera são tempos de cálculo inativos.
Um kernel fundido, completado para aprendizagem
O Mega MoE da DeepSeek, adicionado à biblioteca DeepGEMM em abril de 2026, substitui esta cadeia por uma única chamada de kernel. Um kernel (kernel) é um programa otimizado que executa diretamente no GPU. Este funde o despacho paralelo por perito, as duas camadas lineares dos peritos, SwiGLU e a etapa de combinação. Enquanto um grupo de tokens está em trânsito, o kernel calcula sobre um grupo já recebido, e evita tensores intermédios graças a buffers de trabalho localizados apenas no GPU. Contudo, este kernel incluía apenas uma implementação da passagem direta.
A Prime Intellect escreveu a passagem inversa em falta, também em sobreposição. Os tipos de dados suportados são BF16, MXFP8 (reservado a GPUs Blackwell) e FP8 por blocos (reservado a GPUs Hopper). A configuração do prime-rl propõe igualmente um modo fp8_bwd: passagem direta em BF16, passagem inversa em FP8.
Números auto-declarados
Todos os resultados são publicados pela própria Prime Intellect. As etapas de treino usam nós de 8 H200, em FP8, com sequências de 16 384 tokens e checkpointing de ativação completo. No Qwen3-30B-A3B, distribuído por 32 H200, uma etapa completa demora 2,78 segundos com a base prime-rl, 3,07 segundos com TorchTitan (FP8) e 2,08 segundos com Prime Mega MoE, ou seja, ganhos de 34% e 48%. No Qwen3-235B-A22B, distribuído por 128 H200, as durações são de 12,72 segundos, 14,98 segundos e 9,65 segundos, para ganhos de 32% e 55%. O débito passa, no modelo 30B, de 5 337 para 7 886 tokens por segundo e por GPU face ao TorchTitan, e no modelo 235B, de 1 094 para 1 697.
Medido camada a camada num único nó, o ganho é mais marcado na passagem direta do que na inversa. A passagem direta do modelo Qwen3-30B-A3B passa de 12,0 ms para 3,9 ms, ou seja, um fator de 3,11. A passagem inversa do modelo Qwen3-235B-A22B passa de 30,5 ms para 24,6 ms, ou seja, 1,24. O número « até 3,11x » diz respeito, portanto, apenas à passagem direta. O TorchTitan reivindica, por seu lado, um ganho de até 41%, mas numa configuração diferente (B200, modelo de 671 mil milhões de parâmetros, base BF16), que a fonte considera não diretamente comparável.
O Prime Mega MoE é distribuído sob licença Apache 2.0. No prime-rl, ativa-se como backend de despacho: basta definir type = « mega_moe » em [trainer.model.moe.dispatch], e depois escolher a precisão entre bf16, mxfp8, fp8 ou fp8_bwd. Os kernels de treino são fornecidos numa roda DeepGEMM pré-compilada, disponível no repositório prime-kernels, pois a versão standard da DeepGEMM não os contém.
