CODA：将 Transformer 块重写为 GEMM-Epilogue 程序

Transformer 训练系统是围绕密集线性代数构建的，但端到端时间的很大一部分花费在周围的内存限制运算符上。归一化、激活、残差更新、约简和相关计算会在全局内存中重复移动大型中间张量，同时执行少量算术，这使得数据移动成为高度优化的训练堆栈中越来越重要的瓶颈。我们引入 CODA，一种 GPU 内核抽象，它将这些计算表示为 GEMM-plus-epilogue 程序。 CODA 基于以下观察：许多作为单独框架内核公开的 Transformer 运算符可以通过代数重新参数化来执行，同时 GEMM 输出块保留在芯片上，然后再写入内存。该抽象修复了 GEMM 主循环，并公开了一小组可组合的尾声原语，用于缩放、缩减、成对变换和累积。这种约束接口保留了专家编写的 GEMM 的性能结构，同时保持足够的表达能力，以涵盖标准 Transformer 块的前向和后向传递中的几乎所有非注意力计算。在代表性的 Transformer 工作负载中，人类编写的 CODA 内核和 LLM 编写的 CODA 内核都实现了高性能，这表明 GEMM 加尾声编程提供了一条将框架级生产力与硬件级效率相结合的实用途径。