还在手写CUDA内核？CODA来了！LLM和新手也能让Transformer跑出光速
5 月 22 日，Tri Dao 在社交媒体上转发了 Han Guo 的一条推文。他还写道：「经过一些数学重写，结果发现 Transformer 的所有内容都是一系列 GEMM + epilogue（矩阵乘法加尾声）。给定一些优化的原语，LLM（以及新手）就可以为所有 Transformer 操作编写光速内核！」
Tri Dao 是 FlashAttention 系列的核心作者之一，而这条推文则指向了他们当天发布的一篇论文： C OD A 。
论文标题：CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs
论文地址：https://arxiv.org/abs/2605.19269
代码地址：https://github.com/HanGuo97/coda-kernels
这个名字，读起来像「终曲」，念起来像「CUDA」。来自 MIT、普林斯顿、Together AI 和 Meta 的研究者， 试图用 一套新的 编程抽象，把 Transformer 训练里那些鲜少被人关注、却持续消耗时间的「散碎计算」，系统性地消化掉。
背景：训练大模型的「偷懒税」
要理解 CODA 在解决什么问题，先要明白大模型训练的时间都去哪了。
在一块英伟达 H100 上训练一个 LLaMA-3 风格的 1B 参数模型，大部分人会直觉地认为：时间都花在矩阵乘法和注意力计算上，毕竟那才是「真正的计算」。这个直觉大体上没错： 矩阵乘法（GEMM）和注意力确实占据了主要算力 。
但如果你打开性能分析器仔细看，会发现还有一批「小算子」在安静地消耗着时间：归一化（RMSNorm）、激活函数（SwiGLU、RoPE）、残差加法、跨层规约…… 它们单个计算量不大，却频繁地把大 型中间 张量从显存里搬进搬出。
这就是所谓的「内存带宽瓶颈」：好比一个厨艺绝顶的厨师，但每做一道菜都要把食材从远处的仓库搬来、用完再送回去，而不是放在手边的台面上。厨师的手速再快，等待搬运的时间也是真实的浪费。
更糟糕的是，随着英伟达的 FP8、FP4 等低精度格式让矩阵计算越来越快，这些「搬运」操作的相对成本反而在上升：矩阵乘法加速了，但张量搬进搬出的成本并没有同比缩短。
论文中有一组数据很直观：在 H100 上用 TorchTitan 训练 1B 参数模型时，非矩阵乘法操作占据了相当一部分的端到端运行时间，且随着 FP8 精度的引入，这一比例还会进一步凸显。
现有的编程框架对此几乎无能为力。PyTorch 把 Transformer 的计算表达成一串算子序列，算子之间有清晰的边界。这种边界对于自动微分（autograd）非常友好，却恰好阻止了跨算子的融合优化：每一个算子边界，往往就是一次不必要的显存写回。
CODA：「尾声」里藏着宝藏
CODA 的出发点是一个朴素的观察。
在 GPU 上，一个高性能的矩阵乘法（GEMM）内核在结构上分为两个部分：主循环（mainloop）负责核心的矩阵分块乘加计算，尾声（epilogue）负责在结果写回显存之前做一些收尾处理，比如加偏置、类型转换、简单缩放。
尾声存在的意义，在于此时矩阵乘法的输出还「活在」片上寄存器里，还没有落地到全局显存。这是一个短暂的黄金窗口：如果能在这个时刻多做一些计算，就可以完全省掉一次显存写入再读出的往返。
CODA 的核心洞察是： Transformer 里那些内存密集型操作，其实很多可以被代数地重新参数化，塞进这个「尾声」窗口里执行。
这需要一点数学技巧。以最常见的 GEMM-RMSNorm-GEMM 模式为例：一个矩阵乘法的结果，经过残差加法、RMS 归一化，然后再做另一个矩阵乘法。传统做法是三个独立算子串行执行，中间结果两次落地显存。
CODA 团队发现，RMS 归一化中的行缩放因子 r，因为是每行共享的标量，它和后面的矩阵乘法满足交换律：可以把 r 的应用从「第二个 GEMM 之前」推迟到「第二个 GEMM 的尾声」。推迟之后，第一个 GEMM 的尾声只需要计算局部的「分块均方根」（partial RMS），由一个极轻量的辅助规约内核合并，而完整的 RMSNorm 计算消失了。
类似的重新参数化，对 SwiGLU、RoPE（旋转位置编码）、交叉熵损失等操作同样适用，甚至对反向传播也成立。论文中有一个定理证明：只要前向尾声是「分块局部」的，反向传播就自动继承相同的结构。具体请访问原论文查看。
五种「积木」和一套「乐高语言」
CODA 不是一个具体的融合内核，而是一套编程抽象。
它固定住经过专家优化的 GEMM 主循环，然后在尾声位置暴露五类可组合的基本原语：
逐元素变换（residual 加法、激活函数、RoPE）
向量加载与存储（广播 RMSNorm 权重）
矩阵分块加载与存储（保存中间激活供反向传播使用）
分块规约（局部均方根、分块 log-sum-exp）
有状态变换（在线归一化所需的 max 和 sum-exp 统计）
用这五类积木，一个标准 Transformer 的前向和反向传播中、除注意力之外的几乎全部操作都可以被覆盖。
更有意思的是这套抽象对「谁来写代码」的宽容度。论文在实验中评估了两种实现模式：一种是人工程序员撰写，另一种是用 Claude Code 来生成 —— 给定 CODA 的原语说明、若干示例和实现日志，由 AI 完成大部分内核代码，人工轻度监督。
两种模式的性能表现均达到了较高水平。Tri Dao 在推文中说「LLM 以及新手就可以编写光速内核」，这正是论文实验结果在现实层面的映射。
实验结果
CODA 的基准测试选择的是较为苛刻的对手：cuBLAS 加上 torch.compile，以及专为 LLM 优化的 Liger Kernel 和 FlashInfer。
论文对每个内核评估了两种实现： CODA (LLM) 由 Claude Code 生成，研究者提供原语说明、若干示例和一份持续更新的实现技巧日志，AI 完成主体代码，人工做轻度监督； CODA (Human) 由人工程序员独立编写，使用同样的高层重参数化思路，但不依赖 CODA 原语集本身。两组结果都与 cuBLAS + torch.compile、Liger Kernel、FlashInfer 等优化库进行对比。
在单算子层面，以 GEMM-RMSNorm-GEMM 这一典型模式为例，CODA 在对应 1B、7B、70B 三个模型规模的隐藏维度下均实现了对 cuBLAS + PyTorch 基线的超越。SwiGLU、RoPE、交叉熵等尾声组合也有类似表现。
LLM 生成的内核在大多数基准上与人工手写版本不相上下，个别配置下甚至略有超越。这在 GPU 内核优化这个历来门槛极高的领域，是一个颇为罕见的结论。