CODA：把 Transformer 周邊運算塞進 GEMM epilogue，少搬資料才是真加速

CODA 將 Transformer 的 normalization、activation、residual update、reduction 等 memory-bound 周邊運算改寫成 GEMM-plus-epilogue programs，趁 GEMM output tile 還在 chip 上時完成計算，減少 global memory round-trip。真正啟發在於用受限 composable API 讓 GPU kernel 最佳化更可被 LLM / 工程師組合。

2026年5月22日3 分鐘閱讀👁 5

ML Systems / GPU Kernels / Transformer Training

這則 Threads 介紹 arXiv 論文 CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs。核心不是發明新的矩陣乘法，而是把 normalization、activation、residual update、reduction 等 memory-bound operator 重新寫成 GEMM-plus-epilogue program，趁 GEMM output tile 還在 chip 上時完成後續計算，減少寫回 global memory 再讀回來的資料搬運。

核對資訊：arXiv:2605.19269v2，作者為 Han Guo、Jack Zhang、Arjun Menon、Driss Guessous、Vijay Thakkar、Yoon Kim、Tri Dao；作者標註機構包含 MIT、Princeton、Together AI、Meta。官方 code 在 HanGuo97/coda-kernels，README 說 CODA built on CUTLASS CuTeDSL、targets NVIDIA Hopper H100 GPUs。Threads 提到「Tri Dao（Anthropic）」這點與 arXiv 作者欄不一致，本文以 arXiv metadata 為準。

問題：GEMM 很快，周邊搬資料很慢

Transformer 訓練的 FLOPs 主要在 matrix multiplication 與 attention；但 normalization、activation、residual update、reduction 等算術密度低，會反覆把大型 intermediate tensors 在 global memory 與處理器間搬來搬去。當 FP8 / FP4 讓矩陣乘法更快時，資料搬運瓶頸反而更顯眼。

CODA 的觀察

很多 Transformer operators 雖然在 framework 裡是分開 kernel，但可以代數重參數化，在 GEMM output tile 還留在 chip 上、尚未寫回 memory 前執行。這就是 GEMM-plus-epilogue 的核心。

抽象層設計

CODA 固定 GEMM mainloop，只暴露一小組 composable epilogue primitives：scaling、reductions、pairwise transformations、accumulation。這讓開發者保留 expert-written GEMM 的效能骨架，同時用受限 API 組合更多 Transformer 周邊運算。

Agent / LLM codegen 的意義

Hacker News 高訊號留言指出：epilogue fusion 本身不是新概念，Triton / NVIDIA 內部工具早能做；真正有意思的是設計一個 restricted、composable API，讓 LLM 不必手刻底層 GPU 最佳化，而是組合專家寫好的 blocks。

層級	傳統做法	CODA 方向
Framework graph	PyTorch / autograd 把 Transformer block 表成多個 operator；operator boundary 常變成 materialization boundary。	把可融合的非 attention 計算重寫成 GEMM epilogue 組合，減少中間 tensor 寫回。