Prime Intellect lanza Prime Mega MoE: kernel MoE fusionado entrenable, con un pase hacia adelante 3.11x más rápido

Prime Intellect ha lanzado Prime Mega MoE, una extensión de entrenamiento del kernel Mega MoE fusionado de DeepSeek. El kernel de DeepSeek solo ejecutaba el pase hacia adelante. Prime Intellect añadió un pase hacia atrás rápido y superpuesto. Eso hace que el kernel fusionado sea utilizable para entrenamiento, no solo para inferencia. En su prime-rl trainer, los pasos completos de entrenamiento en modelos Qwen3 MoE se ejecutan hasta un 34% más rápido. Frente a TorchTitan, la ganancia alcanza el 55%. Para los equipos que entrenan modelos de mezcla de expertos (MoE), eso es un recorte directo de horas de GPU.

TL;DR

• Tamaño: Un kernel de GPU, no un modelo. Probado en Qwen3-30B-A3B y Qwen3-235B-A22B (128 expertos, enrutamiento top-8).
• Funciona en: GPUs NVIDIA Hopper (FP8 por bloques) y Blackwell (MXFP8), BF16 en ambas. Las cifras publicadas usan nodos de 8×H200, de 32 a 128 GPUs.
• Rendimiento: Pasos completos de entrenamiento entre un 32% y un 34% más rápidos que la línea base de prime-rl. Entre un 48% y un 55% más rápidos que TorchTitan.
• Mejor caso: pase hacia adelante 3.11x más rápido en una capa MoE de Qwen3-30B-A3B (de 12.0 ms a 3.9 ms).
• Peor caso: pase hacia atrás 1.24x más rápido en una capa MoE de Qwen3-235B-A22B (de 30.5 ms a 24.6 ms).

¿Qué es Prime Mega MoE?

Prime Mega MoE es un kernel MoE fusionado y superpuesto para entrenamiento de Prime Intellect. Se basa en el kernel Mega MoE de la biblioteca de código abierto DeepGEMM de DeepSeek.

¿Por qué el entrenamiento de MoE está limitado por la comunicación?

En un modelo MoE, cada token va a un pequeño subconjunto de expertos. Esos expertos están en GPUs diferentes. Así que cada capa envía tokens, ejecuta el cálculo de los expertos y devuelve los resultados. Esto ocurre en el pase hacia adelante y de nuevo en el pase hacia atrás.

En una configuración estándar, estas son operaciones separadas. La GPU ordena los tokens y espera en la capa de transporte. Luego calcula, devuelve los resultados y vuelve a esperar. Esas esperas son tiempo de GPU inactivo.

¿Cómo funciona el kernel fusionado?

Mega MoE de DeepSeek, añadido a DeepGEMM en abril de 2026, reemplaza esa cadena con una sola llamada al kernel. Fusiona el despacho en paralelo de expertos, las 2 capas lineales de expertos, SwiGLU y la combinación. Mientras un grupo de tokens está en tránsito, el kernel calcula sobre un grupo que ya llegó. También evita los tensores intermedios mediante buffers de scratch solo de GPU. Pero solo incluía una implementación del pase hacia adelante.

Prime Intellect escribió el kernel del pase hacia atrás que faltaba, también superpuesto. Los tipos soportados son BF16, MXFP8 (solo Blackwell) y FP8 por bloques (solo Hopper). La configuración de prime-rl también expone un fp8_bwd modo: forward en BF16, backward en FP8.

¿Qué tan rápido es Prime Mega MoE?

Todos los números siguientes son autoreportados por Prime Intellect. Los pasos de entrenamiento usan nodos de 8×H200, FP8, secuencias de 16,384 tokens y full activation checkpointing. Cada resultado es la media de los pasos 6 a 20.

Paso completo de entrenamiento, segundos (menor es mejor):

Frente a TorchTitan, el throughput subió de 5,337 a 7,886 tokens/s/GPU en el modelo de 30B. En el modelo de 235B, subió de 1,094 a 1,697. La línea base prime-rl usa torch all-to-all más expertos FP8 de DeepGEMM.

Una sola capa MoE, milisegundos (1 nodo, 8×H200, EP=8, 128 expertos, top-8, enrutamiento aleatorio):

El patrón es claro. Las ganancias hacia adelante son grandes, y las ganancias hacia atrás son menores. El titular «hasta 3.11x» se refiere solo a la pasada hacia adelante.

¿Cómo se compara con otras opciones MoE fusionadas?

La ganancia del 41% de TorchTitan proviene de una configuración diferente (B200, modelo de 671B, línea base BF16). No es directamente comparable.

¿Cómo lo usan los desarrolladores?

Prime Mega MoE es un backend de dispatch en prime-rl. Configure type = "mega_moe" bajo [trainer.model.moe.dispatch]. Elija precision de bf16, mxfp8, fp8 o fp8_bwd. Los kernels de entrenamiento se incluyen en una rueda precompilada de DeepGEMM de prime-kernels. La versión estándar de DeepGEMM no los incluye.

Conclusiones clave

• Prime Mega MoE añade una pasada hacia atrás fusionada y superpuesta al Mega MoE de DeepSeek.
• Los pasos de entrenamiento completos se ejecutan un 32% a 34% más rápido que la línea base de prime-rl.
• Frente a TorchTitan, los pasos se ejecutan un 48% a 55% más rápido en los modelos MoE de Qwen3.
• Las ganancias hacia adelante alcanzan 3.11x, pero las ganancias hacia atrás caen a 1.24x en las capas grandes.
• Se incluye en prime-rl como el mega_moe backend de dispatch.

Consulta el artículo del blog de Prime Mega MoE, el anuncio en Xy el repositorio prime-rl. Además, no dudes en seguirnos en Twitter y no olvides unirte a nuestro 150k+ML SubReddit y suscribirte a nuestro boletín. ¡Espera! ¿estás en telegram? ahora también puedes unirte a nosotros en telegram.

¿Necesitas asociarte con nosotros para promocionar tu repositorio de GitHub O página de Hugging Face O lanzamiento de producto O webinar, etc.? Conéctate con nosotros

La entrada Prime Intellect lanza Prime Mega MoE: kernel MoE fusionado entrenable, pase hacia adelante 3.11 veces más rápido apareció primero en MarkTechPost.