CODA reparameterizes memory-bound operations such as norms, RoPE, residuals, and SwiGLU to execute as part of GEMM epilogues in transformer blocks

LLM training is built on fast MatMuls. But many surrounding ops still run as memory-bound kernels.

CODA reparameterizes them to hide in the matmul’s shadow, fused into its epilogue before results leave the chip.

Bonus: LLMs can write fast CODA kernels too (approaching SoLs).

After some mathematical rewrite, turns out all of transformer is a series of gemm + epilogue. Given a few optimized primitives, LLMs (and novice humans) can write speed-of-light kernels for all transformer ops!

Finally, huge thanks to the incredible team: @jcz42, Arjun, Driss, @tensorcore, @yoonrkim, and @tri_dao!

PDF: https://arxiv.org/abs/2605.19269 Code: https://github.com/HanGuo97/coda-kernels

the transformer reparameterizations lab has released even more reparameterizations :D

Faster training with CODA 🚀 most ops re-written as GEMM+epilogues achieving SoL

A significant portion of transformer training time is not spend doing matmuls, but for norms, residuals, RoPE, etc. This work shows that the entire transformer computation can be rewritten into matmul+epilogue, which is an extremely powerful abstraction for writing gpu kernels

The computational abstractions humans developed are great for building architectures, however they’re not necessarily the right abstractions for kernels. Han shows why 🔥

Pattern 1: GEMM + residual + RMSNorm + GEMM.

RMSNorm looks like a standalone memory-bound op. But CODA reparameterizes it to compute partial norm statistics in the first GEMM epilogue, run a small reduction, then apply scaling in the second GEMM epilogue.

Speaking of performance: a MatMul/GEMM has a mainloop and an epilogue, the final step before results are written to memory.

Epilogues are IO-efficient because data stays on-chip, and with the right design, like PingPong GEMM, they can overlap with tensor-core work.

The catch: locality.

ML frameworks make training code easy to write, but they hide the cost of many small, memory-bound ops.

That is why optimized training/inference stacks often rely on bespoke kernels and manual autograd. But these are hard to write and maintain.

Can we get the best of both?

Surprisingly, with the right algebraic reparameterization, almost the entire standard Transformer forward + backward pass can be expressed as:

(1) GEMM + epilogues (2) attention (3) a few cheap auxiliary reductions

Similar patterns apply to the backward pass. If the forward pass consists of a sequence of GEMM-Epilogue blocks (and ends with a GEMM), the backward pass does as well.

In practice, RMSNorm backward needs one extra algebraic trick.

Because epilogues are highly structured, CODA can start from an optimized GEMM template (QuACK PingPong) and compose a small set of fast primitives.

That turns out to be just enough structure for LLMs to write high-performance CuTeDSL kernels.

Pattern 2: GEMM + pairwise transforms.

Think RoPE, SwiGLU, and their backward passes. With the right layout, CODA processes them directly in epilogues.

Pattern 3: LM head + cross entropy.

Compute logits, then extract target logits and LSE in epilogues, as in Cut Cross Entropy.

Very cool work from an all-star team!

Turns out with the magic of hiding stuff in epilogue, GEMM-plus-epilogue programming can be surprisingly effective 🚀 Congrats to the great work @HanGuo97 @jcz42!

make GPUs go brrrrr

The headline trick is the commutation: r((xW0+z) ⊙ γ)W1 = r(((xW0+z) ⊙ γ)W1), so the row-wise scale slides past the second GEMM into its epilogue. That's why RMSNorm fuses and softmax doesn't. The LLM-written kernels are inside that envelope. What's the next op where you think a similar algebraic move exists?

@tri_dao Unrelated, but similar to how hedge funds try to get the most intelligent non-LLM model running with the least latency possible: do you see inference speed (in units of intelligence per seconds) competitions that might be winner-takes-all?