Prime Intellect ha publicado Prime Mega MoE, una extensión de entrenamiento del kernel GPU fusionado Mega MoE de DeepSeek. Mientras que el kernel de DeepSeek solo cubría el pase forward, Prime Intellect le añade un pase backward rápido, que se solapa con los otros cálculos, de modo que el kernel se vuelve utilizable para el entrenamiento y no solo para la inferencia. Según la empresa, en su entrenador prime-rl, los pasos de entrenamiento completos en modelos Qwen3 MoE funcionan hasta un 34 % más rápido, y hasta un 55 % más rápido frente a TorchTitan.
En un modelo MoE (mixture of experts, o «mezcla de expertos»), cada token solo se envía a un pequeño subconjunto de expertos, distribuidos en diferentes GPU. Cada capa debe por tanto enviar los tokens, hacer calcular a los expertos, y luego reenviar los resultados, durante el pase forward como del pase backward, la fase de retropropagación utilizada para el aprendizaje. En una configuración estándar, estas operaciones están separadas: la GPU ordena los tokens, espera el transporte, calcula, y luego espera de nuevo. Estos tiempos de espera son tiempo de cálculo inactivo.
Un kernel fusionado, completado para el aprendizaje
El Mega MoE de DeepSeek, añadido a la biblioteca DeepGEMM en abril de 2026, reemplaza esta cadena por una sola llamada de kernel. Un kernel (kernel) es un programa optimizado que se ejecuta directamente en la GPU. Este fusiona el dispatch experto-paralelo, las dos capas lineales de los expertos, SwiGLU y la etapa de combinación. Mientras un grupo de tokens está en tránsito, el kernel calcula sobre un grupo ya llegado, y evita los tensores intermedios gracias a buffers de trabajo situados únicamente en la GPU. Este kernel solo incluía, sin embargo, una implementación del pase forward.
Prime Intellect ha escrito el pase backward faltante, también en solapamiento. Los tipos de datos soportados son BF16, MXFP8, reservado para GPU Blackwell, y FP8 por bloques, reservado para GPU Hopper. La configuración de prime-rl propone también un modo fp8_bwd: pase forward en BF16, pase backward en FP8.
Cifras autodeclaradas
Todos los resultados son publicados por Prime Intellect misma. Los pasos de entrenamiento utilizan nodos de 8 H200, en FP8, con secuencias de 16 384 tokens y un checkpointing de activación completo. En Qwen3-30B-A3B, distribuido en 32 H200, un paso completo dura 2,78 segundos con la base prime-rl, 3,07 segundos con TorchTitan (FP8) y 2,08 segundos con Prime Mega MoE, es decir, ganancias del 34 % y del 48 %. En Qwen3-235B-A22B, distribuido en 128 H200, las duraciones son de 12,72 segundos, 14,98 segundos y 9,65 segundos, para ganancias del 32 % y del 55 %. El rendimiento pasa, en el modelo 30B, de 5 337 a 7 886 tokens por segundo y por GPU frente a TorchTitan, y en el modelo 235B, de 1 094 a 1 697.
Medido capa por capa en un solo nodo, la ganancia es más marcada en el forward que en el backward. El pase forward del modelo Qwen3-30B-A3B pasa de 12,0 ms a 3,9 ms, es decir, un factor de 3,11. El pase backward del modelo Qwen3-235B-A22B pasa de 30,5 ms a 24,6 ms, es decir, 1,24. La cifra «hasta 3,11x» solo concierne, por tanto, al pase forward. TorchTitan reclama por su parte una ganancia de hasta el 41 %, pero en una configuración diferente (B200, modelo de 671 mil millones de parámetros, base BF16), que la fuente considera no directamente comparable.
Prime Mega MoE se distribuye bajo licencia Apache 2.0. En prime-rl, se activa como backend de dispatch: basta con definir type = «mega_moe» bajo [trainer.model.moe.dispatch], y luego elegir la precisión entre bf16, mxfp8, fp8 o fp8_bwd. Los kernels de entrenamiento se proporcionan en una rueda DeepGEMM precompilada, disponible en el repositorio prime-kernels, ya que la versión estándar de DeepGEMM no los contiene.
