O artigo apresenta uma nova pesquisa intitulada "CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs", cujo objetivo principal é otimizar a eficiência do treinamento de modelos Transformer, especialmente abordando operações "intensivas em memória" que parecem dispersas, mas que, acumuladas, consomem muito tempo.

Autor do artigo, fonte: Machine Heart

Em 22 de maio, Tri Dao compartilhou em mídias sociais um tweet de Han Guo. Ele também escreveu: "Após algumas reescritas matemáticas, descobriu-se que tudo no Transformer é uma série de GEMM + epílogo (multiplicação de matrizes mais epílogo). Dadas algumas primitivas otimizadas, LLM (e iniciantes) podem escrever kernels de velocidade da luz para todas as operações do Transformer!"

Tri Dao é um dos autores principais da série FlashAttention, e este tweet aponta para um artigo publicado por eles naquele dia: CODA.

Título do artigo: CODA: Reescrevendo Blocos Transformer como Programas de Epílogo GEMM
Endereço do artigo: https://arxiv.org/abs/2605.19269
Endereço do código: https://github.com/HanGuo97/coda-kernels

Este nome soa como "finale" e é pronunciado como "CUDA". Pesquisadores do MIT, Princeton, Together AI e Meta tentam eliminar, de forma sistemática, os cálculos dispersos e pouco notados que consomem tempo continuamente no treinamento de Transformers, usando um novo conjunto de abstrações de programação.

O "imposto da preguiça" para treinar grandes modelos

Para entender o problema que a CODA está resolvendo, primeiro é preciso compreender onde o tempo é gasto no treinamento de grandes modelos.

Treinar um modelo de 1 bilhão de parâmetros no estilo LLaMA-3 em uma única NVIDIA H100, a maioria das pessoas teria a intuição de que todo o tempo é gasto em multiplicações de matrizes e cálculos de atenção, já que esses são os “verdadeiros cálculos”. Essa intuição está essencialmente correta: as multiplicações de matrizes (GEMM) e a atenção realmente consomem a maior parte da capacidade de processamento.

Mas se você abrir o analisador de desempenho e observar com atenção, perceberá que há um conjunto de "operadores menores" consumindo silenciosamente tempo: normalização (RMSNorm), funções de ativação (SwiGLU, RoPE), adição residual, redução entre camadas... Eles têm baixa carga de cálculo individual, mas frequentemente movem grandes tensores intermediários para dentro e para fora da memória da GPU.

Isso é o que se chama de "gargalo de largura de banda de memória": é como um chefe de cozinha extremamente habilidoso, mas que precisa trazer os ingredientes de um armazém distante e devolvê-los após cada prato, em vez de tê-los à mão sobre a bancada. Mesmo que as mãos do chefe sejam rápidas, o tempo esperando pelo transporte é um desperdício real.

Pior ainda, à medida que os formatos de baixa precisão da NVIDIA, como FP8 e FP4, tornam os cálculos matriciais cada vez mais rápidos, o custo relativo dessas operações de “mudança” aumenta: as multiplicações matriciais aceleram, mas o custo de transferir tensores para dentro e para fora não diminui na mesma proporção.

Um conjunto de dados no artigo é bastante intuitivo: ao treinar um modelo de 1 bilhão de parâmetros no H100 usando o TorchTitan, operações não de multiplicação de matrizes ocupam uma parcela significativa do tempo total de execução, e essa proporção se torna ainda mais pronunciada com a introdução da precisão FP8.

Os frameworks de programação atuais têm quase nenhuma capacidade de lidar com isso. O PyTorch expressa o cálculo do Transformer como uma sequência de operadores, com fronteiras claras entre eles. Essas fronteiras são muito amigáveis para a diferenciação automática (autograd), mas justamente impedem a otimização de fusão entre operadores: cada fronteira de operador frequentemente corresponde a uma gravação desnecessária na memória.

CODA: Um tesouro escondido no "Epílogo"

O ponto de partida da CODA é uma observação simples.

Na GPU, um kernel de multiplicação de matriz de alto desempenho (GEMM) é estruturado em duas partes: o loop principal (mainloop), responsável pelo cálculo central de multiplicação e adição de blocos de matriz, e o epílogo (epilogue), que realiza tratamentos finais antes de gravar o resultado na memória VRAM, como adição de viés, conversão de tipo e escala simples.

A importância do epílogo reside no fato de que a saída da multiplicação de matrizes ainda está "viva" nos registradores integrados e ainda não foi gravada na memória global da GPU. Este é um breve janela de ouro: se puder ser feito mais cálculo neste momento, poderá-se economizar completamente uma ida e volta de gravação e leitura da memória da GPU.

A principal descoberta da CODA é que muitas operações de alta demanda de memória no Transformer podem ser reparametrizadas algebricamente e executadas dentro desta janela de "epílogo".

Isso exige um pouco de habilidade matemática. Tomando como exemplo o padrão mais comum GEMM-RMSNorm-GEMM: o resultado de uma multiplicação de matriz passa por adição residual, normalização RMS e, em seguida, outra multiplicação de matriz. A abordagem tradicional executa três operadores independentes em sequência, com os resultados intermediários sendo gravados duas vezes na memória da GPU.

A equipe CODA descobriu que, na normalização RMS, o fator de escala de linha r, sendo um escalar compartilhado por linha, comuta com a multiplicação matricial subsequente: a aplicação de r pode ser adiada do período "antes do segundo GEMM" para o "final do segundo GEMM". Após essa adiamento, o final do primeiro GEMM precisa apenas calcular a "raiz média quadrada parcial" (partial RMS), que é combinada por um kernel de redução auxiliar extremamente leve, eliminando assim o cálculo completo do RMSNorm.

Essa mesma reparametrização aplica-se também a operações como SwiGLU, RoPE (Positional Encoding Rotacional) e função de perda de entropia cruzada, e até mesmo ao backpropagation. O artigo apresenta um teorema que prova: sempre que a fase final da frente for “blocos locais”, o backpropagation herda automaticamente a mesma estrutura. Consulte o artigo original para mais detalhes.

Cinco "peças" e um "idioma LEGO"

CODA não é um kernel de fusão específico, mas sim um conjunto de abstrações de programação.

Ele fixa o loop principal GEMM otimizado por especialistas e expõe, na parte final, cinco tipos de primitivas compostas:

Transformação elemento a elemento (adição residual, função de ativação, RoPE)
Carregamento e armazenamento de vetores (broadcast de pesos RMSNorm)
Carregamento e armazenamento (salvamento de ativações intermediárias para uso na retropropagação) em blocos de matriz
Blocagem por redução (RMS local, blocagem log-sum-exp)
Transformações de estado (estatísticas máximas e sum-exp necessárias para a normalização online)

Com estes cinco tipos de blocos, quase todas as operações de um Transformer padrão, exceto a atenção, podem ser cobertas.

Mais interessante ainda é a tolerância desse abstrato em relação a "quem escreve o código". No experimento, o artigo avaliou dois modos de implementação: um escrito por programadores humanos e outro gerado pelo Claude Code — com base nas instruções primitivas do CODA, alguns exemplos e logs de implementação, a IA realiza a maior parte do código do kernel, com supervisão leve humana.

O desempenho de ambos os modos atingiu níveis elevados. Tri Dao disse em um tweet: “LLMs e iniciantes podem escrever núcleos de velocidade da luz”, o que reflete diretamente os resultados dos experimentos do artigo na prática.

Resultados do experimento

A benchmark da CODA escolhe adversários desafiadores: cuBLAS com torch.compile, além dos Liger Kernel e FlashInfer otimizados para LLMs.

O artigo avalia duas implementações para cada núcleo: CODA (LLM), gerada pelo Claude Code, com os pesquisadores fornecendo instruções de primitivas, alguns exemplos e um log de dicas de implementação atualizado continuamente, enquanto a IA desenvolve o código principal e os humanos realizam supervisão leve; CODA (Human), escrita independentemente por programadores humanos, utilizando a mesma abordagem de reparametrização de alto nível, mas sem depender do próprio conjunto de primitivas CODA. Os resultados de ambos os grupos são comparados com bibliotecas otimizadas como cuBLAS + torch.compile, Liger Kernel e FlashInfer.

A nível de operador único, com o padrão típico GEMM-RMSNorm-GEMM como exemplo, o CODA superou a base cuBLAS + PyTorch em todas as dimensões ocultas correspondentes aos três tamanhos de modelo: 1B, 7B e 70B. Comportamentos semelhantes foram observados em combinações finais como SwiGLU, RoPE e entropia cruzada.

Os núcleos gerados por LLM são comparáveis às versões escritas manualmente na maioria dos benchmarks e, em algumas configurações, até ligeiramente superiores. Essa é uma conclusão bastante rara no campo historicamente de alto nível de dificuldade da otimização de núcleos GPU.

Os ganhos na propagação reversa são particularmente notáveis: o kernel de propagação reversa GEMM-Residual-PartialRMS-GEMM apresenta aceleração de 1,6 a 1,8 vezes em comparação com a linha de base, enquanto a propagação reversa do SwiGLU também melhora em cerca de 1,4 a 1,6 vezes. Nesta direção, a lacuna entre o LLM e a implementação manual é igualmente pequena. Isso não é surpreendente: a propagação reversa envolve naturalmente mais acessos a tensores intermediários, o que amplifica os benefícios da fusão no final; e o design dos primitivos do CODA é suficientemente claro para permitir que modelos de IA realizem corretamente a combinação.

Em benchmarks end-to-end de camadas Transformer completas, a aceleração forward da CODA varia de aproximadamente 5% a 20% em diferentes escalas, sendo mais significativa em tamanhos de modelo maiores (correspondentes à dimensão oculta de 70B).

Em termos de precisão numérica, a reparametrização da CODA ajusta o momento em que o fator de escala do RMSNorm é aplicado, mas experimentos mostram que seu erro numérico é comparável à implementação de referência do PyTorch e, em algumas configurações, até menor — graças ao acumulador de maior precisão presente no loop principal da GEMM.

O que o CODA pode fazer: um guia de referência rápida para esclarecer os limites das capacidades do CODA antes de adotar uma perspectiva mais ampla.

Cobertura: Quase todo o cálculo, exceto atenção e embeddings de palavras, nas propagações direta e reversa de Transformers padrão (como a arquitetura LLaMA), incluindo RMSNorm, adição residual, ativação SwiGLU, codificação de posição rotacional RoPE, perda de entropia cruzada e os cálculos de gradiente reverso das operações acima.
Efeito de aceleração: Em dimensões ocultas de escala de 1B a 70B, há melhorias variadas no nível de operador em comparação com a base cuBLAS + torch.compile, sendo o benefício na retropropagação mais significativo (alguns kernels alcançam mais de 1,6 vezes); a aceleração end-to-end da frente para camadas Transformer completas é de aproximadamente 5% a 20%, com efeitos mais pronunciados em modelos maiores.
Quem pode usar: CODA, implementado com base no CuTeDSL (DSL em Python do NVIDIA CUTLASS), suporta dois métodos de escrita de kernels: por programadores humanos e por modelos de IA, ambos alcançando desempenho de alta performance.
Limite atual: atualmente, apenas o cenário de GPU única é suportado, sem treinamento distribuído; a reparametrização é principalmente voltada para a arquitetura Transformer padrão, e a aplicabilidade a outras arquiteturas ainda precisa ser verificada.

Conclusão

CODA não é um trabalho isolado. É uma implementação concreta de uma classe de ideias: em GPUs, o verdadeiro espaço de otimização muitas vezes não está em "o que calcular", mas em "como mover".

FlashAttention trouxe o cálculo de atenção para a memória no chip; o CODA tenta fazer o mesmo com normalização e funções de ativação. O Triton reduziu a barreira para escrever kernels personalizados, e projetos como ThunderKittens e TileLang exploram ainda mais esse espaço em diferentes níveis. Esses trabalhos apontam todos na mesma direção: unificar, em um único framework programável, a facilidade de expressão do grafo de operadores do PyTorch com a eficiência de execução próxima ao CUDA escrito manualmente.

A última frase do tweet de Tri Dao merece ser refletida novamente: «LLMs e iniciantes podem escrever núcleos de velocidade da luz para todas as operações Transformer.» Por trás disso há uma lógica mais profunda: quando o design da abstração de programação é suficientemente bom, os modelos de IA podem participar diretamente da otimização da própria infraestrutura de treinamento. Esse ciclo é o aspecto mais intrigante do CODA.

Do ponto de vista deste ângulo, o nome "CODA" talvez tenha um significado mais profundo. Na música clássica, Coda é o trecho final que encerra uma peça. Aqui, é o "epílogo" do núcleo GEMM — e escrever bem este epílogo pode ser o próximo capítulo importante para o aumento da eficiência dos sistemas de treinamento Transformer.