CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

Transformer training systems are built around dense linear algebra, but memory-bound operators like normalization, activation, residual update, and reduction account for a significant portion of total time. These repeatedly move large tensors to and from global memory with little computation, making data transfer the primary bottleneck.

CODA is a GPU kernel abstraction that expresses these operations as GEMM-plus-epilogue programs.

The core idea: many Transformer operators exposed as separate kernels in frameworks are algebraically reparameterized to run while the GEMM output tile still resides on-chip (before being written to memory).

By fixing the GEMM main loop and providing a set of composable epilogue primitives for scaling, reduction, pairwise transformation, and accumulation, the interface—though limited—covers nearly all operations in the forward/backward pass of a standard Transformer block, except for attention.

On representative Transformer workloads, both human-written and LLM-written CODA kernels achieved high performance, demonstrating that GEMM-plus-epilogue programming is a practical path to combining framework-level productivity with hardware-level efficiency.

The author list includes Tri Dao, known for FlashAttention, and Yoon Kim, among others.