arXiv:2605.19269cs.LG2026-05被引 4

将Transformer中的冗余计算重写为高效矩阵乘法后处理,显著减少内存瓶颈。

CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

  • 把归一化等操作转化为矩阵乘法的后处理步骤,避免频繁内存访问。
  • 在多个Transformer任务中,性能接近手工优化代码,且支持自动生成。
  • 适合追求框架易用性与硬件效率平衡的研究者和工程师。

Transformer训练系统依赖密集线性代数运算,但大量时间耗费在内存密集型的外围操作上。归一化、激活函数、残差更新、规约等计算反复在全局内存中搬运大型中间张量,却执行少量算术运算,导致数据移动成为高性能训练堆栈中的主要瓶颈。本文提出CODA,一种将这些计算表达为GEMM-后处理程序的GPU内核抽象。其核心思想是:许多以独立框架内核形式暴露的Transformer操作,可代数重构为在GEMM输出块仍驻留在芯片上时执行,从而避开内存写入开销。该抽象固定了GEMM主循环,仅暴露一组可组合的后处理原语,用于缩放、规约、成对变换和累加。这一受限接口在保持专家级GEMM性能结构的同时,仍具备足够表达力覆盖标准Transformer块前向与反向传播中的几乎所有非注意力计算。在典型Transformer工作负载上,无论是人工编写还是大模型生成的CODA内核均表现优异,表明GEMM-后处理编程为兼顾框架生产力与硬件效率提供了一条可行路径。

原文摘要 · Abstract (English)

Transformer training systems are built around dense linear algebra, yet a nontrivial fraction of end-to-end time is spent on surrounding memory-bound operators. Normalization, activations, residual updates, reductions, and related computations repeatedly move large intermediate tensors through global memory while performing little arithmetic, making data movement an increasingly important bottleneck in otherwise highly optimized training stacks. We introduce CODA, a GPU kernel abstraction that expresses these computations as GEMM-plus-epilogue programs. CODA is based on the observation that many Transformer operators exposed as separate framework kernels can be algebraically reparameterized to execute while a GEMM output tile remains on chip, before it is written to memory. The abstraction fixes the GEMM mainloop and exposes a small set of composable epilogue primitives for scaling, reductions, pairwise transformations, and accumulation. This constrained interface preserves the performance structure of expert-written GEMMs while remaining expressive enough to cover nearly all non-attention computation in the forward and backward pass of a standard Transformer block. Across representative Transformer workloads, both human- and LLM-authored CODA kernels achieve high performance, suggesting that GEMM-plus-epilogue programming offers a practical path toward combining framework-level productivity with hardware-level efficiency.

TransformerGEMMGPU优化内存瓶颈

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。