← Back to feed

After FlashAttention eliminated the memory inefficiency in attention, the next bottleneck in transformer training became visible: every other operator around the GEMM. Normalization, activations, residual additions, and reductions are executed as separate GPU kernels, each requiring a full round trip through global memory. They do little arithmetic relative to how much data they move.

A paper submitted to arXiv on May 19 from Han Guo, Jack Zhang, Arjun Menon, Driss Guessous, Vijay Thakkar, Yoon Kim, and Tri Dao attacks this directly. CODA — Compiler-Optimized Data-Aware (GPU) kernel Abstraction — rewrites these operators as GEMM-epilogue programs, executing while the GEMM output tile is still on-chip before it is written to global memory.

The Core Observation

Transformer training pipelines are built around dense linear algebra. GEMMs (general matrix multiplications) — the operations in attention, feed-forward layers, and projections — already run close to hardware peak on modern accelerators. The surrounding memory-bound operators are not.

The CODA paper’s central observation: many of these operators, currently exposed as separate framework kernels, can be algebraically reparameterized to execute as GEMM epilogues. LayerNorm, RMSNorm, SiLU, GELU, residual updates, and accumulation operations can all be rewritten to run while the GEMM output remains in on-chip register or shared memory — before the tensor is serialized back to DRAM.

The framework fixes the GEMM mainloop and exposes composable epilogue primitives for scaling, reductions, pairwise transformations, and accumulation. The constrained interface keeps the performance properties of expert-written GEMMs while covering “nearly all non-attention computations” in a standard transformer block.

Why Now

Three trends have made this problem increasingly relevant:

Attention is solved (mostly). FlashAttention and its successors have already reduced attention’s memory overhead to the point that the surrounding ops are now a larger fraction of total time. The law of diminishing returns redirects optimization effort.

Models are bigger. As parameter counts increase, the absolute time spent in normalization and activation layers grows proportionally. The overhead that was acceptable at 7B parameters is significant at 400B.

Training budgets are compressing. Efficiency improvements that reduce compute per training step have a direct commercial value as frontier training runs push into tens of billions of dollars. A reduction in memory-bound overhead compounds across millions of gradient steps.

The Author Signal

Tri Dao’s involvement carries weight. His FlashAttention papers in 2022 and 2023 moved from research to production deployment across essentially every major training framework. The technique now ships by default in PyTorch, JAX, and most commercial training stacks without user intervention.

CODA follows the same design philosophy: identify a structural inefficiency in how hardware resources are used, find the algebraic reparameterization that fixes it, build the abstraction that makes it composable. If the adoption trajectory resembles FlashAttention’s, the implementation path runs through PyTorch CUDA extensions and FlashAttention-3’s epilogue interface.

The paper does not publish benchmark numbers in the arXiv abstract, but the framing — “data movement an increasingly important bottleneck in otherwise highly optimized training stacks” — positions CODA as a training throughput improvement across hardware generations, not a narrow micro-optimization.

What to Watch

The two near-term signals worth tracking: whether the paper appears in ML frameworks’ roadmaps (PyTorch 2.x, XLA), and whether the authors release production-ready CUDA code. FlashAttention shipped usable kernels alongside the paper. If CODA follows the same pattern, framework integration timelines compress significantly.

For labs running training at scale, even a 5-10% reduction in time-to-train per step has real dollar value. At $50-100 billion annual training spend across the frontier, the addressable improvement is substantial.