El artículo presenta una nueva investigación llamada CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs, cuyo objetivo principal es optimizar la eficiencia del entrenamiento de modelos Transformer, especialmente resolviendo las operaciones «intensivas en memoria» que parecen dispersas pero que, acumulativamente, consumen mucho tiempo.

El 22 de mayo, Tri Dao retuiteó una publicación de Han Guo en las redes sociales. También escribió: «Después de algunas reescrituras matemáticas, resultó que todo en Transformer es una serie de GEMM + epílogo (multiplicación de matrices más coda). ¡Dados algunos primitivos optimizados, LLM (y los novatos) pueden escribir kernels de velocidad de la luz para todas las operaciones de Transformer!». Tri Dao es uno de los autores principales de la serie FlashAttention, y este tuit apuntaba a un artículo que publicaron ese día: CODA. Título del artículo: CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs. Dirección del artículo: https://arxiv.org/abs/2605.19269. Dirección del código: https://github.com/HanGuo97/coda-kernels.

Este nombre, que se lee como «coda», suena como «CUDA». Investigadores del MIT, Princeton, Together AI y Meta intentan utilizar un nuevo conjunto de abstracciones de programación para digerir sistemáticamente los «cálculos dispersos» que rara vez se notan en el entrenamiento de Transformer, pero que consumen tiempo continuamente.

Antecedentes: El «impuesto a la pereza» del entrenamiento de grandes modelos. Para entender qué problema está resolviendo CODA, primero hay que entender a dónde va el tiempo del entrenamiento de grandes modelos. En una NVIDIA H100, al entrenar un modelo de 1B parámetros al estilo LLaMA-3, la mayoría de la gente intuitivamente pensaría que el tiempo se gasta en la multiplicación de matrices y el cálculo de la atención, ya que ese es el «cálculo real». Esta intuición es en gran medida correcta: la multiplicación de matrices (GEMM) y la atención sí ocupan la mayor parte de la potencia de cálculo. Pero si se abre el analizador de rendimiento y se mira con atención, se encontrará un lote de «pequeños operadores» que consumen tiempo en silencio: normalización (RMSNorm), funciones de activación (SwiGLU, RoPE), suma residual, reducción entre capas… Su cálculo individual no es grande, pero frecuentemente mueven grandes tensores intermedios dentro y fuera de la memoria de la GPU.

Esto es lo que se conoce como el «cuello de botella del ancho de banda de la memoria»: es como un chef con una habilidad culinaria excelente, pero que tiene que traer los ingredientes desde un almacén lejano cada vez que cocina un plato, y luego devolverlos después de usarlos, en lugar de ponerlos en la encimera a su alcance. No importa lo rápido que sea el chef, el tiempo de espera para el transporte es una pérdida real. Lo que es peor, a medida que los formatos de baja precisión de NVIDIA, como FP8 y FP4, hacen que el cálculo de matrices sea cada vez más rápido, el coste relativo de estas operaciones de «transporte» está aumentando: la multiplicación de matrices se ha acelerado, pero el coste de mover los tensores dentro y fuera no se ha reducido en la misma proporción. En el artículo hay un conjunto de datos muy intuitivo: al entrenar un modelo de 1B parámetros con TorchTitan en H100, las operaciones que no son de multiplicación de matrices ocupan una parte considerable del tiempo de ejecución de extremo a extremo, y esta proporción se destacará aún más con la introducción de la precisión FP8. Los marcos de programación existentes son casi impotentes ante esto. PyTorch expresa el cálculo de Transformer como una serie de secuencias de operadores, con límites claros entre los operadores. Este límite es muy amigable para la diferenciación automática (autograd), pero precisamente impide la optimización de la fusión entre operadores: cada límite de operador es a menudo una escritura innecesaria en la memoria de la GPU.

CODA: El tesoro está escondido en el «epílogo». El punto de partida de CODA es una observación sencilla. En la GPU, un kernel de multiplicación de matrices (GEMM) de alto rendimiento se divide estructuralmente en dos partes: el bucle principal (mainloop) se encarga del cálculo principal de multiplicación y suma de bloques de matrices, y el epílogo se encarga de algunos procesos de acabado antes de que el resultado se escriba de nuevo en la memoria de la GPU, como la adición de sesgos, la conversión de tipos y el escalado simple. El significado de la existencia del epílogo es que la salida de la multiplicación de matrices todavía «vive» en los registros del chip en este momento, y aún no se ha asentado en la memoria global de la GPU. Esta es una ventana dorada fugaz: si se puede hacer más cálculo en este momento, se puede ahorrar por completo un viaje de ida y vuelta de escritura y lectura de la memoria de la GPU.

La principal idea de CODA es que muchas de las operaciones intensivas en memoria en Transformer pueden ser realmente reparametrizadas algebraicamente e insertadas en esta ventana de «epílogo» para su ejecución. Esto requiere un poco de habilidad matemática. Tomemos como ejemplo el patrón GEMM-RMSNorm-GEMM más común: el resultado de una multiplicación de matrices se somete a una suma residual, una normalización RMS y luego se realiza otra multiplicación de matrices. El enfoque tradicional es ejecutar tres operadores independientes en serie, con el resultado intermedio aterrizando en la memoria de la GPU dos veces. El equipo de CODA descubrió que el factor de escala de fila r en la normalización RMS, ya que es un escalar compartido por cada fila, satisface la ley conmutativa con la multiplicación de matrices posterior: la aplicación de r puede retrasarse desde «antes del segundo GEMM» hasta «el epílogo del segundo GEMM». Después del retraso, el epílogo del primer GEMM sólo necesita calcular la «raíz cuadrática media parcial» (RMS parcial) local, que se fusiona mediante un kernel de reducción auxiliar muy ligero, y el cálculo completo de RMSNorm desaparece.

Una reparametrización similar se aplica también a operaciones como SwiGLU, RoPE (codificación de posición rotatoria) y la pérdida de entropía cruzada, e incluso es válida para la retropropagación. El artículo contiene una demostración del teorema: siempre que el epílogo de avance sea «local por bloques», la retropropagación hereda automáticamente la misma estructura. Consulte el artículo original para obtener más detalles.

Cinco tipos de «bloques de construcción» y un «lenguaje LEGO». CODA no es un kernel de fusión específico, sino un conjunto de abstracciones de programación. Fija el bucle principal GEMM optimizado por expertos y luego expone cinco tipos de primitivas básicas combinables en la posición del epílogo: 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 de bloques de matrices (guardar la activación intermedia para su uso en la retropropagación), reducción de bloques (raíz cuadrática media local, log-sum-exp por bloques), transformación con estado (estadísticas max y sum-exp necesarias para la normalización en línea). Con estos cinco tipos de bloques de construcción, se puede cubrir casi todas las operaciones, excepto la atención, en la propagación hacia adelante y hacia atrás de un Transformer estándar.

Lo que es más interesante es la tolerancia de esta abstracción a «quién escribe el código». El artículo evalúa dos modos de implementación en los experimentos: uno es escrito por programadores humanos, y el otro es generado por Claude Code – dadas las instrucciones primitivas de CODA, varios ejemplos y registros de implementación, la IA completa la mayor parte del código del kernel, con supervisión humana ligera. El rendimiento de ambos modos alcanzó un alto nivel. Tri Dao dijo en un tuit que «LLM y los novatos pueden escribir kernels de velocidad de la luz», que es precisamente el reflejo de los resultados experimentales del artículo en el mundo real.

Resultados experimentales. La prueba de referencia de CODA elige un oponente relativamente duro: cuBLAS más torch.compile, así como Liger Kernel y FlashInfer, que están optimizados para LLM. El artículo evalúa dos implementaciones para cada kernel: CODA (LLM) es generado por Claude Code, y los investigadores proporcionan instrucciones primitivas, varios ejemplos y un registro de trucos de implementación actualizado continuamente, la IA completa el cuerpo principal del código y los humanos realizan una supervisión ligera; CODA (Human) es escrito independientemente por programadores humanos, utilizando las mismas ideas de reparametrización de alto nivel, pero sin depender del conjunto de primitivas CODA en sí. Ambos conjuntos de resultados se comparan con bibliotecas optimizadas como cuBLAS + torch.compile, Liger Kernel y FlashInfer.

A nivel de operador único, tomando como ejemplo el patrón típico GEMM-RMSNorm-GEMM, CODA ha logrado superar la línea de base cuBLAS + PyTorch en las dimensiones ocultas correspondientes a las escalas de modelo de 1B, 7B y 70B. Las combinaciones de epílogos como SwiGLU, RoPE y la entropía cruzada también tienen un rendimiento similar. El kernel generado por LLM es comparable a la versión escrita a mano por humanos en la mayoría de los puntos de referencia, e incluso la supera ligeramente en algunas configuraciones. Esta es una conclusión bastante rara en el campo de la optimización del kernel de la GPU, que históricamente ha tenido un umbral muy alto.

Los beneficios de la retropropagación son particularmente notables: el kernel inverso de GEMM-Residual-PartialRMS-GEMM tiene una aceleración de 1,6 a 1,8 veces en comparación con la línea de base, y el inverso de SwiGLU también tiene una mejora de aproximadamente 1,4 a 1,6 veces. En esta dirección, la diferencia entre LLM y la implementación humana también es pequeña. Esto no es sorprendente: la retropropagación implica naturalmente más acceso a tensores intermedios, y los beneficios de la fusión del epílogo son mayores; y el diseño primitivo de CODA es lo suficientemente claro como para que el modelo de IA pueda completar correctamente la combinación.

En la prueba de referencia de extremo a extremo de la capa Transformer completa, la aceleración de avance de CODA es de aproximadamente el 5% al 20% en diferentes escalas, y el efecto es más significativo en tamaños de modelo más grandes (correspondientes a la dimensión oculta de la escala 70B). En términos de precisión numérica, la reparametrización de CODA ajusta el tiempo de aplicación del factor de escala RMSNorm, pero los experimentos muestran que su error numérico es comparable a la implementación de referencia de PyTorch, e incluso menor en algunas configuraciones – gracias a que el bucle principal GEMM en sí tiene un acumulador de mayor precisión.

Qué puede hacer CODA: una hoja de referencia rápida. Antes de entrar en una perspectiva más amplia, aclaremos los límites de la capacidad de CODA. Cobertura: casi todos los cálculos, excepto la atención y la incrustación de palabras, en la propagación hacia adelante y hacia atrás de un Transformer estándar (como la arquitectura LLaMA), incluyendo RMSNorm, suma residual, activación SwiGLU, codificación de posición rotatoria RoPE, pérdida de entropía cruzada, y el cálculo del gradiente inverso de las operaciones anteriores. Efecto de aceleración: en la dimensión oculta correspondiente a la escala de 1B a 70B, hay diferentes grados de mejora en comparación con la línea de base cuBLAS + torch.compile a nivel de operador único, entre los cuales el beneficio de la retropropagación es el más significativo (algunos kernels pueden alcanzar más de 1,6 veces); la aceleración de avance de extremo a extremo de la capa Transformer completa es de aproximadamente el 5% al 20%, y el efecto es más pronunciado en tamaños de modelo más grandes. Quién puede usarlo: CODA se basa en CuTeDSL (Python DSL de NVIDIA CUTLASS) para su implementación, y soporta dos métodos de escritura de kernel, programador humano y modelo de IA, y ambos métodos pueden lograr un alto rendimiento. Limitaciones actuales: actualmente sólo soporta escenarios de una sola GPU, sin involucrar el entrenamiento distribuido; la reparametrización se dirige principalmente a la arquitectura Transformer estándar, y la aplicabilidad de otras arquitecturas está por verificar.

Conclusión. CODA no es un trabajo aislado. Es una implementación concreta de una clase de ideas: en la GPU, el verdadero espacio de optimización a menudo no está en «qué calcular», sino en «cómo mover». FlashAttention permite que el cálculo de la atención «viva» en la memoria del chip, y CODA intenta que la normalización y la función de activación también «vivan allí». Triton reduce el umbral para escribir kernels personalizados, y ThunderKittens, TileLang, etc. exploran aún más este espacio en diferentes niveles. Estos trabajos apuntan conjuntamente en la misma dirección: unificar realmente la conveniencia de la expresión del gráfico de operadores de PyTorch y la eficiencia de ejecución cercana a CUDA escrita a mano en un marco programable.

La última frase del tuit de Tri Dao merece ser saboreada de nuevo: «LLM y los novatos pueden escribir kernels de velocidad de la luz para todas las operaciones de Transformer». Detrás de esto hay una lógica más profunda: cuando la abstracción de la programación está suficientemente bien diseñada, el propio modelo de IA puede participar en la optimización de su propia infraestructura de entrenamiento. Este ciclo es el lugar más intrigante de CODA. Desde esta perspectiva, el nombre «CODA» puede tener otro significado. En la música clásica, la Coda es la sección al final de una pieza musical que resume toda la pieza. Aquí, es el «epílogo» del kernel GEMM – y escribir bien este epílogo puede ser el siguiente capítulo importante en la mejora de la eficiencia del sistema de entrenamiento de Transformer.

[Máquina del Corazón]