CODA革新编程体验无需手写CUDA也能让Transformer极速运行

5月22日，Tri Dao在社交媒体上转发了一篇研究，并附上了引人深思的评论：“经过数学上的巧妙重写，我们发现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风格的10亿参数模型时，许多人会认为时间主要花在矩阵乘法和注意力计算这类“重活”上。这个直觉基本正确：矩阵乘法和注意力机制确实占据了主要的算力开销。

然而，当你使用性能分析工具深入观察，会发现一批“小算子”也在默默消耗着可观的时间：归一化（RMSNorm）、激活函数（SwiGLU）、旋转位置编码（RoPE）、残差加法、跨层规约……它们单个计算量不大，却频繁地迫使庞大的中间张量在显存和GPU计算单元之间来回搬运。

这就是典型的“内存带宽瓶颈”。好比一位顶级大厨，每处理一道食材，都需要从远处的冷库取来，用完再放回去，而不是直接从手边的备料台获取。厨师技艺再精湛，等待搬运的时间也是无法避免的浪费。

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

论文中的数据直观地说明了这一点：在H100上使用TorchTitan训练10亿参数模型时，非矩阵乘法操作占据了相当比例的端到端运行时间。并且，随着FP8精度的引入，这一比例会变得更加突出。

现有的主流深度学习框架对此几乎无能为力。像PyTorch这样的框架，将Transformer的计算表达为一连串独立的算子序列，算子之间有清晰的边界。这种设计虽然便于自动微分，却恰恰阻碍了跨算子的融合优化：每一个算子边界，通常就意味着一次不必要的显存写回操作，带来了额外的延迟。

CODA：在“尾声”中挖掘性能宝藏

CODA的核心理念其实非常直观。

在GPU上，一个高性能的矩阵乘法内核通常由两部分构成：主循环负责核心的分块乘加计算，而“尾声”则负责在结果写回全局显存之前，进行一些收尾处理，例如添加偏置、转换数据类型或执行简单的缩放。

“尾声”阶段的关键在于，此时矩阵乘法的结果仍然驻留在GPU的片上高速寄存器中，尚未落盘到全局显存。这是一个转瞬即逝的黄金窗口：如果能在这个时刻顺手完成一些额外的计算，就能完全避免一次“写入显存再读出来”的昂贵数据往返。

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

这需要一些精妙的数学技巧。以最常见的“GEMM-RMSNorm-GEMM”计算模式为例：一个矩阵乘法的结果，经过残差加法、RMS归一化，然后再进行另一个矩阵乘法。传统做法是三个独立算子串行执行，中间结果需要两次写入显存。

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

类似的重新参数化技巧，同样适用于SwiGLU、RoPE、交叉熵损失等操作，甚至对反向传播也成立。论文中有一个定理证明：只要前向传播中的尾声操作是“分块局部”的，那么反向传播就能自动继承相同的优化结构。具体数学细节可参考原论文。

五类“积木”与一套“乐高语言”

CODA并非一个固定的、写死的融合内核，而是一套高度灵活的编程抽象。

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

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

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

更有趣的是这套抽象对开发者的宽容度。论文在实验中评估了两种实现模式：一种由经验丰富的程序员手动编写；另一种则交由Claude Code这类AI代码助手来生成——研究者只需提供CODA的原语说明、一些示例和一份持续更新的实现技巧日志，AI就能完成大部分内核代码，人类开发者进行轻度监督和验证即可。

两种模式的性能表现都达到了业界领先水平。Tri Dao在推文中说“LLM以及新手就可以编写光速内核”，正是这一实验结果在现实层面的生动体现，极大地降低了高性能GPU内核开发的门槛。

实验结果与性能基准

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架构，其他变体架构（如Mamba、RetNet等）的适用性有待进一步验证。

结语与展望

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

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

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

从这个角度看，“CODA”这个名字或许别有深意。在古典音乐中，Coda是乐曲末尾收束全篇、升华主题的段落。在这里，它是GEMM内核的“尾声”——而写好这段尾声，很可能正是开启Transformer训练系统效率提升新篇章的关键所在，为未来更大规模、更高效的AI模型训练铺平道路。