CODA:将 Transformer 模块重写为 GEMM 后处理程序
针对Transformer训练系统中因频繁内存访问导致的性能瓶颈,研究团队提出了名为CODA的新型GPU内核抽象方案。传统Transformer训练依赖密集线性代数,但大量时间消耗在归一化、激活函数、残差更新等内存受限算子上。这些操作在传输大型中间张量时计算量极低,使得数据移动成为优化系统的最大障碍。 CODA的核心创新在于将此类非注意力机制的计算重写为"GEMM(通用矩阵乘法)加后处理(Epilogue)”程序。该方案基于关键观察:许多原本作为独立框架内核运行的算子,经过代数重参数化后,可在GMM输出块驻留芯片期间直接执行,无需先写入全局内存再读取。CODA固定了GEMM的主循环,并暴露了一组可组合的后处理原语,涵盖缩放、归约、成对变换及累加等操作。这种受控接口既保留了专家编写的高性能GEMM结构,又具备足够表达力,能覆盖标准Transformer模块前向和反向传播中几乎所有的非注意力计算。 在代表性Transformer工作负载测试中,无论是人工编写还是由大语言模型自动生成的CODA内核,均实现了优异的性能表现。这一成果表明,采用GEMM加后处理的编程范式,是一条兼具框架级开发效率与硬件级运行效率的实用路径。该研究通过重新定义算子执行方式,有效缓解了数据搬运带来的延迟,为未来高效能AI训练系统提供了新的技术方向。相关成果已发布于arXiv,编号为2605.19269,属于机器学习领域的重要探索。
