Los modelos Mixture-of-Experts (MoE) se consolidaron como pieza fundamental de los sistemas de inteligencia artificial a gran escala. Su atractivo es claro: permiten escalar la capacidad total del modelo activando solo una fracción de los parámetros por cada token, lo que se traduce en mejor desempeño con presupuestos de cómputo más razonables. A medida que los modelos crecen, optimizar este bloque pasa a ser crítico para sostener el ritmo de entrenamiento.

Para empujar ese límite, NVIDIA presentó una familia de kernels MLP fusionados para modelos densos y MoE, construidos a medida con NVIDIA CuTe DSL. Al atacar de raíz los cuellos de botella de memoria y sincronización, los nuevos kernels entregan 1,3× a 2× de speedup a nivel de kernel respecto a los caminos no fusionados, y habilitan ejecución MoE sync-free para CUDA Graphs de iteración completa.

En el setup de pre-entrenamiento full-stack de NVIDIA para DeepSeek-V3, esta optimización aporta 8% de mejora end-to-end. En el setup equivalente para GPT-OSS, la mejora es del 93% end-to-end. Los kernels están disponibles desde hoy en NVIDIA cuDNN Frontend y se acceden sin fricción vía NVIDIA Transformer Engine y NVIDIA Megatron-Core.

¿Dónde se va el cómputo en un bloque MoE?

Para maximizar el throughput, primero hubo que mapear con detalle dónde se gastan los ciclos de cómputo. Al perfilar la línea de tiempo de una iteración estándar dentro del bloque MoE, NVIDIA identificó tres cuellos de botella sistémicos:

  • Activaciones: las funciones de activación generan kernels limitados por memoria con grandes operaciones de lectura y escritura de tensores, que dejan a los Tensor Cores subutilizados.
  • Dependencia de la CPU: con expertos enrutados, los tokens por experto se calculan en tiempo de ejecución y suelen procesarse en CPU. Si la CPU no le sigue el ritmo a la GPU, las operaciones quedan expuestas. La solución pasa por kernels que no necesiten sincronización ni intervención del host.
  • Costo de cuantización: igual que las activaciones, pasar tensores de alta a baja precisión produce kernels memory-bound que mantienen ociosos los Tensor Cores.

NVIDIA aborda estos desafíos rediseñando el bloque MoE con kernels custom escritos en cuTE DSL, e introduce una familia de tres kernels para el MoE sync-free:

  • GroupGemm + Quantize
  • GroupGemm + Activation + Quantize/Transpose
  • GroupGemm + dActivation + Quantize/Transpose

Las funciones de activación soportadas son SwiGLU, GeGLU y sReLU, con opción de añadir clamping y escalado.

Figura 1. Fusión de operaciones en un único kernel custom en forward y backward pass con CuTe DSL
Figura 1. Fusión de operaciones en un único kernel custom en forward y backward pass con CuTe DSL

Optimizando GLU vía epílogos GEMM fusionados

Las funciones Gated Linear Unit se volvieron muy populares y la mayoría de los modelos modernos usan alguna variante: SwiGLU, GeGLU y similares. Estas activaciones dividen la salida de la capa FC1 y la combinan para escribir el output GLU final. NVIDIA implementa un kernel fusionado que mezcla la GEMM con la operación GLU correspondiente en forward y back prop.

Las activaciones GLU no son triviales de fusionar dentro del epílogo de la GEMM porque la GLU necesita acceso a dos chunks distintos del tensor: input y gate. Normalmente esos dos chunks los calcularían thread blocks distintos y, para combinarlos, el kernel debería escribirlos a memoria global. Para lograr la fusión, NVIDIA reempaqueta los pesos en columnas de input y gate, garantizando que el mismo thread block tenga acceso al medio ancho de tile del tensor input y al medio ancho del tensor gate. Eso permite combinar ambos en el epílogo sin pasar por memoria global. El reempaquetado puede ocurrir antes de empezar el entrenamiento, durante la carga del checkpoint.

En el back prop, el epílogo lee la salida de la GEMM, calcula la dSwiGLU, la cuantiza y la escribe a memoria global.

Figura 2. Pesos de input y gate empaquetados para que el thread block tenga acceso a ambos y compute la salida SwiGLU dentro del CUDA core
Figura 2. Pesos de input y gate empaquetados para que el thread block tenga acceso a ambos y compute la salida SwiGLU dentro del CUDA core

Estos patrones de fusión no solo eliminan lecturas y escrituras de tensores intermedios, sino que maximizan utilización al superponer las operaciones de memoria restantes directamente con la GEMM. Más allá de SwiGLU, GeGLU y sReLU, los kernels manejan nativamente operaciones de epílogo fusionado como escalado de features, clamping de tensores y suma de vectores de bias.

Eliminando la sincronización host-device

Tradicionalmente, el trabajo que hace un kernel se define por el conteo de bloques al momento del launch, lo que exige que la información de shape esté disponible en el host. Por ejemplo, un grouped GEMM multi-stream lanza G GEMMs distintos en streams separados, donde G es la cantidad de grupos. Como el número de tokens por grupo se determina en tiempo de ejecución, la CPU debe lanzar esas GEMMs de tamaño dinámico en streams separados para maximizar el uso de recursos.

Esto trae dos problemas: la cantidad de kernels a lanzar escala con la cantidad de expertos locales, y se requiere un punto de sincronización obligatorio para recuperar la información de shape en el host antes de lanzar el kernel. Para resolverlo, los kernels GroupGEMM de CuTe DSL trackean los tokens por grupo dentro de la memoria de la GPU. Eso elimina la dependencia con la CPU durante la iteración y habilita CUDA graphs para toda la iteración completa, removiendo de raíz el cuello de botella del host.

Fusionando cuantización MXFP8 y NVFP4

La popularidad de recetas de precisión reducida como MXFP8 y NVFP4 para pre-entrenamiento crece sin parar: estas precisiones aportan speedups significativos con mínimo impacto en accuracy. En estas recetas, la activación es seguida de cuantización y transposición para la operación GEMM de precisión angosta.

Para MXFP8, el kernel de cuantización lee la salida de la activación (BF16) y escribe la salida MXFP8 más una versión transpuesta para el back prop. Los nuevos kernels fusionan esa cuantización dentro del propio kernel GEMM, eliminando lectura y escritura adicionales del tensor BF16. Para NVFP4, el kernel produce la salida BF16 y el amax por tensor para el forward prop; en el back prop calcula el amax para la rotación de Hadamard transpuesta de la salida. Eso elimina la pasada extra de memoria para el cálculo del amax por tensor.

De ganancias por kernel a speedups en pre-entrenamiento

En microbenchmarks de unidad, los kernels fusionados entregan un speedup considerable: aceleran el forward pass hasta 1,3× y el backward pass hasta 2,1× respecto a los caminos no fusionados.

Para que esas ganancias se traduzcan en throughput end-to-end, los kernels también soportan:

  • Scheduling dinámico para superponer eficientemente con otros kernels (comunicación de paralelismo de expertos, paralelismo de datos, etc.).
  • Cluster Margin configurable, que permite reservar un margen de recursos SM limitando el kernel a menos SMs, dejando espacio para que otros kernels se lancen y ejecuten en paralelo en la GPU.

Además del speedup por kernel, como los kernels sync-free habilitan CUDA graphs de iteración completa y se solapan con los kernels de comunicación, el speedup a nivel aplicación es bastante mayor. En testing interno, NVIDIA ve hasta 8% end-to-end en DeepSeek-V3 y hasta 93% end-to-end en pre-entrenamiento de GPT-OSS.

Figura 3. Speedup en distintos patrones de activación sobre GB200. El baseline usa los kernels optimizados del Transformer Engine
Figura 3. Speedup en distintos patrones de activación sobre GB200. El baseline usa los kernels optimizados del Transformer Engine

¿Cómo se usan los kernels CuTe DSL?

Los kernels están disponibles en distintos niveles de abstracción.

  • cuDNN Front-end (v1.23.0+): los kernels viven en la librería cuDNN Frontend. Se instala en el stack del usuario y se invocan directamente desde ahí. CuDNN-Frontend ofrece un wrapper que compila el kernel en la primera invocación y luego reutiliza el objeto cacheado. Hay opción de invocar el kernel directo o vía el wrapper API. NVIDIA está trabajando en soporte AOT (Ahead of time) para que los kernels se compilen como cubins y se cacheen en disco.
  • Transformer Engine (v2.15+): los usuarios pueden acceder a los kernels vía Transformer Engine, que los expone a través del constructor transformer_engine.pytorch.ops. Estas operaciones se combinan con el bloque transformer_engine.pytorch.ops.Sequential, que internamente hace pattern matching para invocar el kernel fusionado desde la librería cuDNN Frontend.
  • Megatron Core (26.04-alpha.rc2+): los kernels también se invocan vía Megatron Core, donde las features se activan con los flags correspondientes.
Figura 4. Los usuarios pueden integrar los kernels de fusión desde cualquiera de las capas de abstracción del stack CUDA: cuDNN Frontend, Transformer Engine o Megatron Core
Figura 4. Los usuarios pueden integrar los kernels de fusión desde cualquiera de las capas de abstracción del stack CUDA: cuDNN Frontend, Transformer Engine o Megatron Core

¿Qué viene?

NVIDIA trabaja activamente en nuevas features: más patrones de fusión y soporte para frameworks adicionales como JAX. La pieza original aclara que el contenido generado por IA puede resumir información de forma incompleta, por lo que recomienda verificar los datos relevantes en la documentación de referencia.