CODA: A GPU Kernel Abstraction for Efficient Transformer Training via GEMM-Epilogue Programming

We introduce CODA, a GPU kernel abstraction that expresses these computations as GEMM-plus-epilogue programs.

CODA is based on the observation that many Transformer operators exposed as separate framework kernels can be algebraically reparameterized to execute while a GEMM output tile remains on chip, before it is written to memory.

This constrained interface preserves the performance structure of expert-written GEMMs while remaining expressive enough to cover nearly all non-attention computation in the forward and backward pass of a standard Transformer block.

Across representative Transformer workloads, both human- and LLM-authored CODA kernels achieve high performance, suggesting that GEMM-plus-epilogue programming offers a practical path toward combining framework-level productivity with hardware-level efficiency.