CODA性能实测：新手也能让Transformer跑出光速的代码优化方案

5月22日，Tri Dao在社交媒体上转发了一篇论文，并附上了核心洞察：“通过数学上的重参数化，Transformer的所有计算本质上都能被重构为一系列GEMM（通用矩阵乘法）加上一个尾声操作。只要提供了正确的优化原语，即便是LLM或编程新手，也能为所有Transformer操作编写出接近硬件极限的高性能内核。”

作为FlashAttention系列工作的核心作者之一，Tri Dao所指的正是当天发布的新研究：CODA。

• 论文标题：CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs
• 论文地址：https://arxiv.org/abs/2605.19269
• 代码地址：https://github.com/HanGuo97/coda-kernels

CODA这个名字，既像“终曲”，发音又接近“CUDA”。这项由MIT、普林斯顿、Together AI和Meta研究者共同完成的工作，旨在通过一套全新的编程抽象，系统性地解决Transformer训练中那些看似零碎、却持续消耗宝贵时间的“内存搬运”问题。

背景：大模型训练的隐性成本

理解CODA的价值，需要先剖析大模型训练的时间开销构成。

在英伟达H100上训练一个10亿参数规模的模型，直觉上矩阵乘法和注意力计算是算力消耗的主力。这个判断基本正确：矩阵乘法和注意力机制确实占据了绝大部分的浮点运算。

然而，性能剖析工具会揭示另一批“沉默的消耗者”：RMSNorm层归一化、SwiGLU激活函数、RoPE旋转位置编码、残差连接、跨层规约等操作。它们单个计算量不大，却频繁迫使大型中间张量在GPU全局显存和计算单元之间往返搬运。

这正是典型的内存带宽瓶颈。好比顶尖厨师被要求每处理完一道工序，就必须将半成品送回远处的仓库，而不是在案板上直接进行下一步。厨师的技艺再精湛，等待搬运的时间也无法避免。

随着英伟达FP8、FP4等低精度格式大幅提升矩阵计算速度，这些数据搬运操作的相对成本反而在增加——矩阵乘法变快了，但数据进出显存的时间并未同比缩短。

论文中的数据直观展示了这一点：在H100上使用TorchTitan训练10亿参数模型时，非矩阵乘法操作占据了可观的端到端运行时间。引入FP8精度后，这一比例会进一步扩大。

现有深度学习框架对此优化空间有限。以PyTorch为例，其将Transformer计算表达为清晰的算子序列，这种设计便于自动微分，却阻碍了跨算子的融合优化：每个算子边界通常对应一次不必要的中间结果显存写回。

CODA：挖掘“尾声”的黄金窗口

CODA的出发点直接而高效。

GPU上的高性能矩阵乘法内核通常分为两部分：主循环执行核心的分块乘加计算，“尾声”则在结果写回显存前进行收尾处理，如添加偏置、转换数据类型或执行缩放。

“尾声”阶段的独特优势在于，此时矩阵乘法的结果仍驻留在GPU的片上高速寄存器中，尚未被冲刷到全局显存。这是一个转瞬即逝的优化窗口：若能在此窗口内顺带完成其他计算，就能彻底避免一次“写回再读取”的昂贵数据往返。

CODA的核心贡献正在于此：通过巧妙的代数重参数化，Transformer中许多内存密集型的操作可以被无缝嵌入这个“尾声”窗口执行。

这需要精密的数学设计。以典型的“GEMM-RMSNorm-GEMM”计算模式为例：传统实现需要三个独立算子串行执行，中间结果需两次落地显存。

CODA团队发现，RMSNorm中的行缩放因子r是每行共享的标量，它与后续的矩阵乘法满足交换律。这意味着，可以将r的应用时机从“第二个GEMM之前”推迟到“第二个GEMM的尾声”。如此，第一个GEMM的尾声只需计算局部的“分块均方根”，由一个极轻量的辅助规约内核合并即可，完整的RMSNorm计算过程便被高效“消化”了。

类似的重新参数化技巧同样适用于SwiGLU、RoPE、交叉熵损失等操作，并且其反向传播也能自动继承相同的优化结构。论文中的定理保证了这一点：只要前向传播中的尾声操作是“分块局部”的，反向传播就能保持一致的优化模式。

五类原语与一套组合语言

CODA并非一个固定的融合内核，而是一套可编程的抽象系统。

它固定了经过深度优化的GEMM主循环，并在尾声位置暴露了五类可像乐高积木一样组合的基本原语：

• 逐元素变换（如残差加法、激活函数、RoPE）
• 向量加载与存储（如广播RMSNorm的权重）
• 矩阵分块加载与存储（用于保存中间激活值，供反向传播使用）
• 分块规约（如计算局部均方根、分块的log-sum-exp）
• 有状态变换（如在线归一化所需的max和sum-exp统计量）

利用这五类“积木”，一个标准Transformer层中除注意力和词嵌入之外的几乎所有前向与反向操作，都能被覆盖。

这套抽象对实现者的宽容度颇具启发性。论文评估了两种实现路径：一种由经验丰富的程序员手动编写；另一种则交由Claude Code生成——研究者仅需提供CODA原语说明、示例代码和一份持续更新的实现技巧日志，AI便能完成大部分内核代码，人类进行轻量监督即可。

两种路径的性能均达到了高水平。Tri Dao推文中所说的“LLM以及新手就可以编写光速内核”，正是这一实验结论的生动体现。

性能基准测试

CODA的基准测试选择了颇具挑战的对手：英伟达官方的cuBLAS库配合PyTorch的torch.compile，以及专为LLM优化的Liger Kernel和FlashInfer。

论文对每个内核评估了两种实现：CODA (LLM) 由Claude Code生成；CODA (Human) 由程序员手动编写，但遵循相同的高层重参数化思路。两组结果均与上述优化库对比。

在单算子层面，以“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架构中，除注意力和词嵌入之外的几乎全部计算。包括RMSNorm、残差加法、SwiGLU激活、RoPE位置编码、交叉熵损失及其反向梯度计算。
• 加速效果：在对应1B至70B规模的隐藏维度下，单算子层面相比cuBLAS + torch.compile基线有不同程度提升，反向传播收益最显著（部分内核可达1.6倍以上）；完整Transformer层的前向端到端加速约为5%至20%，模型越大效果越明显。
• 适用对象：CODA基于CuTeDSL实现，支持经验丰富的程序员和AI模型两种内核编写方式，且两种方式均能达到高性能。
• 当前限制：目前仅支持单GPU场景，不涉及分布式训练；其重参数化主要针对标准Transformer架构，其他变体架构的适用性有待验证。

结语

CODA并非一个孤立的技术突破。它代表了一类优化思想的具象化：在GPU上，真正的性能瓶颈往往不在于“计算什么”，而在于“数据如何搬运”。

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

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

从这个角度看，“CODA”这个名字或许别有深意。在古典音乐中，Coda是乐曲末尾收束全篇的段落。在这里，它是GEMM内核的“尾声”——而写好这段尾声，很可能正是开启Transformer训练系统效率提升新篇章的关键所在。