El artículo presenta una nueva investigación titulada "CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs", cuyo objetivo principal es optimizar la eficiencia del entrenamiento de modelos Transformer, especialmente abordando operaciones "intensivas en memoria" que parecen dispersas pero que acumulan tiempos de ejecución significativos.

Autor del artículo, fuente: Machine Learning China

El 22 de mayo, Tri Dao compartió en redes sociales un tweet de Han Guo. También escribió: "Tras algunas reescrituras matemáticas, resulta que todo el Transformer es una serie de GEMM + epílogo (multiplicación de matrices más epílogo). Con algunos primitivos optimizados, los LLM (y los principiantes) pueden escribir kernels a velocidad de luz para todas las operaciones del Transformer."

Tri Dao es uno de los autores principales de la serie FlashAttention, y este tweet apunta a un artículo que publicaron ese día: CODA.

• Título del artículo: CODA: Reescribir bloques Transformer como programas GEMM-Epílogo
• Dirección del artículo: https://arxiv.org/abs/2605.19269
• Dirección del código: https://github.com/HanGuo97/coda-kernels

Este nombre suena como «finale» y se pronuncia como «CUDA». Investigadores del MIT, Princeton, Together AI y Meta intentan eliminar sistemáticamente los cálculos dispersos y poco atendidos, pero que consumen constantemente tiempo, en el entrenamiento de Transformers, mediante una nueva abstracción de programación.

El "impuesto a la pereza" para entrenar modelos grandes

Para entender qué problema resuelve CODA, primero debes saber dónde se gasta el tiempo en el entrenamiento de modelos grandes.

Entrenar un modelo de 1B parámetros con estilo LLaMA-3 en una NVIDIA H100, la mayoría de las personas asumirían intuitivamente que el tiempo se dedica principalmente a las multiplicaciones matriciales y los cálculos de atención, ya que estos son los «verdaderos cálculos». Esta intuición es en gran medida correcta: las multiplicaciones matriciales (GEMM) y la atención realmente consumen la mayor parte de la capacidad de cómputo.

Pero si abres el analizador de rendimiento y lo examinas detenidamente, descubrirás que hay un conjunto de «operadores pequeños» que consumen tiempo silenciosamente: normalización (RMSNorm), funciones de activación (SwiGLU, RoPE), suma residual, reducción entre capas... Cada uno tiene un bajo costo computacional, pero frecuentemente transfieren grandes tensores intermedios de y hacia la memoria GPU.

Esto es lo que se conoce como «cuello de botella de ancho de banda de memoria»: es como tener un chef de cocina excepcional, pero que debe trasladar los ingredientes desde un almacén lejano cada vez que prepara un plato, y devolverlos después, en lugar de tenerlos a mano sobre la encimera. Por muy rápido que sean sus movimientos, el tiempo esperando el transporte es una pérdida real.

Lo peor es que, a medida que los formatos de baja precisión de NVIDIA, como FP8 y FP4, aceleran los cálculos matriciales, el costo relativo de estas operaciones de «transferencia» aumenta: la multiplicación matricial se ha acelerado, pero el costo de mover tensores dentro y fuera no ha disminuido proporcionalmente.

En el artículo, un conjunto de datos es muy claro: al entrenar un modelo de 1B parámetros en H100 con TorchTitan, las operaciones no de multiplicación de matrices representan una parte significativa del tiempo total de ejecución, y esta proporción se vuelve aún más pronunciada con la introducción de la precisión FP8.

Los marcos de programación actuales tienen casi ninguna capacidad para manejar esto. PyTorch expresa el cálculo de Transformer como una secuencia de operadores con límites claros entre ellos. Estos límites son muy amigables para la diferenciación automática (autograd), pero precisamente impiden la optimización de fusión entre operadores: cada límite de operador suele representar un escritura innecesaria en memoria.

CODA: «El epílogo» esconde un tesoro

El punto de partida de CODA es una observación sencilla.

En una GPU, un núcleo de multiplicación matricial de alto rendimiento (GEMM) se divide estructuralmente en dos partes: el bucle principal (mainloop), que se encarga del cálculo central de multiplicación y suma por bloques matriciales, y el epílogo (epilogue), que realiza tareas finales antes de escribir los resultados de vuelta a la memoria de vídeo, como sumar un sesgo, convertir tipos y aplicar escalado simple.

La razón de ser del epílogo es que la salida de la multiplicación de matrices aún "vive" en los registros integrados y aún no se ha escrito en la memoria gráfica global. Este es un breve período dorado: si se puede realizar más cálculo en este momento, se puede eliminar por completo un ciclo de escritura y lectura de la memoria gráfica.

La insight central de CODA es que muchas de las operaciones intensivas en memoria en los Transformer pueden reparametrizarse algebraicamente y ejecutarse dentro de esta ventana de "colas".

Esto requiere un poco de habilidad matemática. Tomemos como ejemplo el patrón más común GEMM-RMSNorm-GEMM: el resultado de una multiplicación de matrices, seguido de una suma residual, normalización RMS y luego otra multiplicación de matrices. El enfoque tradicional ejecuta tres operadores independientes en serie, con los resultados intermedios escritos dos veces en la memoria de la GPU.

El equipo de CODA descubrió que, en la normalización RMS, el factor de escala por fila r, al ser un escalar compartido por cada fila, conmuta con la multiplicación matricial posterior: se puede posponer la aplicación de r desde «antes del segundo GEMM» hasta «al final del segundo GEMM». Tras este posponimiento, al final del primer GEMM solo es necesario calcular la «raíz media cuadrada parcial» (partial RMS), la cual se combina mediante un núcleo de reducción extremadamente ligero, eliminando así el cálculo completo de RMSNorm.

La misma reparametrización se aplica también a operaciones como SwiGLU, RoPE (positional encoding rotacional) y pérdida de entropía cruzada, e incluso se cumple para la propagación hacia atrás. El artículo presenta un teorema que demuestra: siempre que la etapa final del forward sea "bloqueadamente local", la propagación hacia atrás hereda automáticamente la misma estructura. Consulte el artículo original para más detalles.

Cinco «bloques» y un «lenguaje LEGO»

CODA no es un núcleo de fusión específico, sino un conjunto de abstracciones de programación.

Fija el bucle principal GEMM optimizado por expertos y expone en la sección final cinco tipos de primitivas componibles:

• Transformación elemento por elemento (suma residual, función de activación, RoPE)
• Carga y almacenamiento de vectores (difusión de pesos RMSNorm)
• Carga y almacenamiento por bloques de matriz (guardar activaciones intermedias para la propagación hacia atrás)
• Block reduction (local RMS, block log-sum-exp)
• Cambios de estado (estadísticas máximas y sum-exp necesarias para la normalización en línea)

Con estos cinco tipos de bloques, se pueden cubrir casi todas las operaciones de la propagación hacia adelante y hacia atrás de un Transformer estándar, excepto la atención.

Lo más interesante es la flexibilidad de este abstracción respecto a "quién escribe el código". En los experimentos, el artículo evaluó dos modos de implementación: uno en el que los programadores humanos escribieron el código, y otro en el que se utilizó Claude Code para generar el código: dadas las especificaciones primitivas de CODA, algunos ejemplos y los registros de implementación, la IA realizó la mayor parte del código del núcleo, con supervisión humana ligera.

El rendimiento de ambos modos alcanzó niveles elevados. Tri Dao dijo en un tweet: «Los LLM y los principiantes pueden escribir núcleos de velocidad de luz», lo cual es precisamente la aplicación práctica de los resultados experimentales del artículo.

Resultados del experimento

La prueba de referencia de CODA elige como oponentes más exigentes: cuBLAS con torch.compile, así como Liger Kernel y FlashInfer optimizados para LLM.

El artículo evalúa dos implementaciones para cada núcleo: CODA (LLM), generada por Claude Code, con instrucciones de primitivas, varios ejemplos y un registro actualizado continuamente de técnicas de implementación proporcionados por los investigadores; la IA genera el código principal, mientras que los humanos realizan una supervisión ligera; CODA (Human), escrita independientemente por programadores humanos, utilizando la misma idea de reparametrización de alto nivel, pero sin depender del conjunto de primitivas CODA. Los resultados de ambos grupos se comparan con bibliotecas optimizadas como cuBLAS + torch.compile, Liger Kernel y FlashInfer.

A nivel de operador único, con el patrón típico GEMM-RMSNorm-GEMM como ejemplo, CODA supera la línea de base de cuBLAS + PyTorch en las dimensiones ocultas correspondientes a tres tamaños de modelo: 1B, 7B y 70B. Combinaciones finales como SwiGLU, RoPE y entropía cruzada también muestran un rendimiento similar.

Los núcleos generados por LLM son comparables a las versiones escritas manualmente en la mayoría de las pruebas, y en algunas configuraciones incluso ligeramente superiores. Este es un resultado bastante raro en el campo de la optimización de núcleos GPU, tradicionalmente de alta barrera de entrada.

Los beneficios de la propagación hacia atrás son particularmente notables: el kernel de propagación hacia atrás de GEMM-Residual-PartialRMS-GEMM logra una aceleración de 1.6 a 1.8 veces en comparación con la línea base, y el SwiGLU hacia atrás también muestra una mejora de aproximadamente 1.4 a 1.6 veces. En esta dirección, la brecha entre los LLM y las implementaciones manuales es igualmente pequeña. Esto no es sorprendente: la propagación hacia atrás implica naturalmente un mayor acceso a tensores intermedios, lo que aumenta los beneficios de la fusión al final; y el diseño de los primitivos de CODA es lo suficientemente claro como para permitir que los modelos de IA completen correctamente la combinación.

En el benchmark end-to-end de la capa Transformer completa, la aceleración hacia adelante de CODA oscila entre aproximadamente un 5% y un 20% en diferentes escalas, siendo más notable en tamaños de modelo más grandes (correspondientes a dimensiones ocultas de 70B).

En cuanto a la precisión numérica, la reparametrización de CODA ajusta el momento en que se aplica el factor de escalado de RMSNorm, pero los experimentos muestran que su error numérico es comparable al de la implementación de referencia de PyTorch, e incluso menor en algunas configuraciones, gracias a que el bucle principal de GEMM posee un acumulador de mayor precisión.

¿Qué puede hacer CODA: una hoja de referencia para aclarar los límites de las capacidades de CODA antes de adoptar una perspectiva más amplia.

• Cobertura: casi todo el cálculo en la propagación hacia adelante y hacia atrás de Transformers estándar (como la arquitectura LLaMA), excepto la atención y los embeddings de palabras, incluyendo RMSNorm, adición residual, activación SwiGLU, codificación de posición rotacional RoPE, pérdida de entropía cruzada y los cálculos de gradientes inversos de las operaciones anteriores.
• Efecto de aceleración: En dimensiones ocultas de escala de 1B a 70B, se logra una mejora variable a nivel de operador en comparación con la línea de base cuBLAS + torch.compile, siendo el beneficio en la propagación hacia atrás el más significativo (algunos kernels alcanzan más de 1.6 veces); la aceleración end-to-end en la capa Transformer completa oscila entre un 5% y un 20%, con efectos más pronunciados en modelos de mayor tamaño.
• Quién puede usar: CODA, implementado sobre CuTeDSL (DSL de Python para NVIDIA CUTLASS), que admite dos métodos de escritura de kernels: por programadores humanos y por modelos de IA, ambos capaces de lograr alto rendimiento.
• Restricción actual: Actualmente solo se admite el escenario de una sola GPU, sin entrenamiento distribuido; la reparametrización se enfoca principalmente en la arquitectura Transformer estándar, y la aplicabilidad a otras arquitecturas aún debe verificarse.

Conclusión

CODA no es un trabajo aislado. Es una implementación concreta de una clase de ideas: en las GPU, el verdadero espacio de optimización rara vez está en "qué calcular", sino en "cómo mover los datos".

FlashAttention permite que los cálculos de atención "vivan" en la memoria integrada en el chip; CODA intenta lograr lo mismo con la normalización y las funciones de activación. Triton reduce la barrera para escribir kernels personalizados, y proyectos como ThunderKittens y TileLang exploran aún más este espacio en distintos niveles. Estos esfuerzos convergen hacia una misma dirección: unificar de manera real en un marco programable la facilidad de expresión del grafo de operadores de PyTorch con la eficiencia de ejecución cercana al CUDA escrito a mano.

La última frase del tuit de Tri Dao merece ser reflexionada nuevamente: «Los LLM y los principiantes pueden escribir núcleos a velocidad de luz para todas las operaciones de Transformer». Detrás de esto hay una lógica más profunda: cuando la abstracción de programación está diseñada lo suficientemente bien, los modelos de IA pueden participar directamente en la optimización de su propia infraestructura de entrenamiento. Este ciclo es lo más intrigante de CODA.

Desde este punto de vista, el nombre «CODA» podría tener un significado más profundo. En la música clásica, la coda es el pasaje final que cierra una pieza. Aquí, es el «epílogo» del núcleo GEMM — y redactar bien este epílogo podría ser el próximo capítulo importante para mejorar la eficiencia del sistema de entrenamiento de Transformer.