Transformer训练提速新思路：CODA将所有操作重写为矩阵乘法+尾声，告别内存带宽瓶颈

5月22日，Tri Dao在社交媒体上转发了Han Guo的一条动态，并附言：“通过巧妙的数学重写，Transformer的每个组成部分本质上都变成了一连串的GEMM加尾声（矩阵乘法+收尾处理）。只要拥有经过优化的基础原语，无论是大语言模型（LLM）还是新手开发者，都能为所有Transformer操作编写出速度极快的内核！”

Tri Dao是FlashAttention系列的核心作者之一，而这条推文指向了他们当天发布的一篇论文：CODA。

• 论文标题：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内核优化这个历来门槛极高的领域，是一个颇为罕见的结论。

反向传播的收益尤为突出：GEMM-Residual-PartialRMS-GEMM的反向内核相比基线加速幅度可达1.6至1.8倍，SwiGLU反向也有约1.4至1.6倍的提升。这个方向上，LLM与人工实现的差距同样微小。这并不奇怪：反向传播天然涉及更多中间张量的存取，尾声融合的收益就更大；而CODA的原语设计足够清晰，使得AI模型能够正确地完成组合。

深度重写与降重版本

在完整的 Transformer 层端到端基准测试中，CODA 的前向传播加速幅度在不同模型规模下介于 5% 到 20% 之间，且对于较大模型（对应 70B 规模的隐藏维度），其加速效果更为显著。

关于数值精度：CODA 通过重参数化调整了 RMSNorm 缩放因子的应用时机。然而，实验结果表明，其数值误差与 PyTorch 参考实现处于同一水平，在某些配置下甚至更小——这得益于 GEMM 主循环本身具备更高精度的累加器。

CODA 的能力速查单

在探讨更宏观的视角前，我们先明确 CODA 的功能边界：

• 覆盖范围：标准 Transformer（如 LLaMA 架构）的前向和反向传播中，除注意力和词嵌入之外的大部分计算，包括 RMSNorm、残差加法、SwiGLU 激活函数、RoPE 旋转位置编码、交叉熵损失及其对应的反向梯度计算。
• 加速效果：在对应 1B 至 70B 规模的隐藏维度下，单算子层面相比 cuBLAS + torch.compile 基线有不同程度的提升，其中反向传播的收益最为突出（部分内核可达 1.6 倍以上）；完整 Transformer 层的端到端前向加速约为 5% 至 20%，在较大模型尺寸下效果更佳。
• 适用人群：CODA 基于 CuTeDSL（NVIDIA CUTLASS 的 Python DSL）实现，支持人工程序员和 AI 模型两种内核编写方式，且两种方式均能达到高性能。
• 当前限制：目前仅支持单 GPU 场景，不涉及分布式训练；重参数化主要针对标准 Transformer 架构，其他架构的适用性有待进一步验证。

结语

CODA 并非孤立的工作。它体现了一类核心思想：在 GPU 上，真正的优化空间往往不在于“算什么”，而在于“怎么搬”。

FlashAttention 让注意力计算“驻留”在片上内存中，而 CODA 则试图让归一化和激活函数也“驻留”其中。Triton 降低了编写自定义内核的门槛，ThunderKittens、TileLang 等工具则在不同层面进一步探索这一领域。这些工作共同指向同一个目标：将 PyTorch 算子图的表达便利性与接近手写 CUDA 的执行效率，真正统一到一套可编程的框架中。

Tri Dao 推文中的最后一句话值得回味：“LLM 以及新手就可以为所有 Transformer 操作编写光速内核。”这背后蕴含着一个更深的逻辑：当编程抽象设计得足够好时，AI 模型本身就能参与到自身训练基础设施的优化中。这个循环，正是 CODA 最耐人寻味之处。

从这个角度看，“CODA”这个名字或许另有深意。在古典音乐中，Coda 是乐曲末尾收束全篇的段落。在这里，它是 GEMM 内核的“尾声”——而写好这段尾声，或许正是 Transformer 训练系统效率提升的下一个重要篇章。