CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs¶

会议: NeurIPS 2026
arXiv: 2605.19269
代码: https://github.com/HanGuo97/coda-kernels
领域: 模型压缩
关键词: GEMM-Epilogue, 算子融合, 显存带宽受限, Transformer加速, CuTeDSL

一句话总结¶

本文提出面向 Transformer 的 GPU 算子抽象 CODA，将非注意力计算通过代数重参数化重写为固定 GEMM 主循环与可组合 Epilogue 原语程序，消除中间张量的显存反复搬运。

研究背景与动机¶

现代 Transformer 大语言模型的计算吞吐高度依赖密集矩阵乘法（GEMM）与注意力机制，这两类算子在 GPU Tensor Core 上已得到极度优化。然而在完整的模型训练中，层规范化（RMSNorm）、激活函数（SwiGLU）、残差更新、旋转位置编码（RoPE）以及交叉熵损失等显存带宽受限（memory-bound）算子，反复在大规模显存和计算核心之间搬运海量中间激活张量，实际执行算术计算却极少。随着硬件在 FP8、FP4 等低精度矩阵乘法上的算力成倍翻倍，中间张量的全局显存物化开销并没有同步改善，使得访存瓶颈在端到端训练中的耗时占比愈发显著。

现有的系统编程范式在应对这一瓶颈时面临生产力与性能的两难。高层框架（如 PyTorch）将 Transformer 表达为离散的算子计算图并通过自动微分生成反向传播，框架的算子边界直接退化为显存物化边界，切断了前向与反向跨算子的融合机会；而工业级系统往往依赖手写反向传播或硬编码定制融合 kernel，维护成本极高且难以泛化扩展。GEMM 库（如 CUTLASS、EVT）虽然提供了 GEMM Epilogue（尾部处理）融合能力，但现有 Epilogue 多局限于简单后处理（如偏置加法或尺度缩放），缺乏面向 Transformer 复杂数据流与跨层依赖的结构化表征。

本文的核心洞见在于：Transformer 中表面上彼此独立的许多显存受限算子，可以通过代数重参数化（algebraic reparameterization）完全内化到 GEMM 输出 tile 驻留片上寄存器与共享内存的生命周期内。核心 idea：固定专家优化的高性能 GEMM 主循环，提炼出一组最小完备的可组合 Epilogue 原语，通过代数变换将 Transformer 前向与反向中几乎所有非注意力访存受限算子重写为 GEMM-Epilogue 程序，仅辅以极轻量的跨 tile 分块归约。

方法详解¶

整体框架¶

CODA 将 Transformer 计算块抽象为“固定 GEMM 主循环 + 可编程 Epilogue 访问树 + 辅助轻量归约”的三层执行架构。GEMM 主循环负责在 Tensor Core 上以极高吞吐产出累加器输出 tile，Epilogue 原语在 tile 写入全局显存之前，直接在片上完成局部变换、张量加载、逐对特征交互与分块部分归约。对于跨维度的全局统计量（如 RMSNorm 均方根因子、交叉熵 log-sum-exp），CODA 仅在 Epilogue 中产出极小维度的分块部分统计量，交由辅助轻量 kernel 归约，从而避免物化全尺寸激活。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入激活与权重<br/>(x, W, 残差 z, 权重 gamma)"] --> MLoop["GEMM 主循环<br/>Tensor Core 瓦片矩阵乘累加"]
    MLoop --> E1["1. 代数重参数化 Epilogue<br/>残差相加 + 逐行部分平方和 + 权重缩放"]
    E1 --> AuxRed["轻量辅助归约 Kernel<br/>跨 tile 汇聚部分平方和并求逆 RMS (r)"]
    AuxRed --> E2["2. 邻接 GEMM 延迟尺度缩放<br/>将 r 交换到下一 GEMM Epilogue 乘以输出"]
    E2 --> E3["3. 寄存器逐对特征映射<br/>邻接特征配对: SwiGLU / RoPE 变换"]
    E3 --> Out["变换后激活输出 / 梯度写回显存"]

关键设计¶

1. 代数重参数化与跨 GEMM 延迟尺度缩放：消除中间 RMSNorm 显存往返

在标准的 Pre-LN Transformer 中，普遍存在“前一子层投影 GEMM $\to$ 残差相加 $\to$ RMSNorm 归一化 $\to$ 后一子层输入 GEMM”的连续模式： $$\mathbf{y} = \operatorname{RMSNorm}(\mathbf{x}\mathbf{W}_0 + \mathbf{z}, \bm{\gamma})\mathbf{W}_1 = \Bigl(r(\mathbf{x}\mathbf{W}_0 + \mathbf{z}) \odot \bm{\gamma}\Bigr)\mathbf{W}_1$$ 其中 $r = 1 / \operatorname{rms}(\mathbf{x}\mathbf{W}_0 + \mathbf{z})$ 为逐行逆均方根标量。传统实现必须将整个残差输出物化到全局显存，再由独立 RMSNorm 读取并写回归一化结果，造成两次巨大的显存带宽浪费。

CODA 观察到归一化标量 $r$ 沿行广播，对线性变换具备代数交换律： $$\Bigl(r(\mathbf{x}\mathbf{W}_0 + \mathbf{z}) \odot \bm{\gamma}\Bigr)\mathbf{W}_1 = r\Bigl(\bigl(\mathbf{x}\mathbf{W}_0 + \mathbf{z}\bigr) \odot \bm{\gamma}\Bigr)\mathbf{W}_1$$ 据此，CODA 将该结构分解为两阶段 GEMM-Epilogue：在 GEMM 1 的 Epilogue 中，片上完成残差累加 $\mathbf{h}_1 = \mathbf{h}_0 + \mathbf{z}$、RMSNorm 权重缩放 $\mathbf{h}_2 = \mathbf{h}_1 \odot \bm{\gamma}$，并计算 tile 内局部均方和 $\widehat{\mathbf{r}}$；随后由极轻量 auxiliary reduction kernel 汇总各 tile 的 $\widehat{\mathbf{r}}$ 得出每行标量 $r$；在随后的 GEMM 2 中，直接输入未除以 RMS 的 $\mathbf{h}_2\mathbf{W}_1$，而将除以 RMS 的尺度缩放 $r$ 延迟融合到 GEMM 2 的 Epilogue 中执行（$\mathbf{y}[i,j] = r[i]\cdot\mathbf{h}_3[i,j]$）。由此彻底消除了独立的 RMSNorm kernel 及其对应的显存全尺寸张量搬运。

2. 寄存器级逐对特征映射：免中间物化的 SwiGLU 与 RoPE 融合

Transformer 中的 SwiGLU（门控前馈网络）和 RoPE（旋转位置编码）虽然涉及相邻特征通道的交互，但并不需要全局跨通道通信。SwiGLU 将前向投影拆分为门控流和数值流后逐对相乘；RoPE 则将相邻两维构成的 2D 向量施加旋转变换。朴素实现需要先写回 GEMM 产出的全尺寸张量，在独立 kernel 中切分并重组，在 SwiGLU 下还会物化中间膨胀维度的特征。

CODA 将特征通道在内存布局中对齐为相邻布局，使其严格匹配 Hopper 架构 Tensor Core 在累加器寄存器中的分布（每个线程自然持有局部相邻的一组输出分量）。在 GEMM Epilogue 中，线程直接在寄存器上完成逐对变换： - 对 RoPE，直接在寄存器中读取预广播的正弦/余弦向量并执行二维平面旋转； - 对 SwiGLU，直接对偶数通道 $G$ 与奇数通道 $U$ 计算 $\operatorname{silu}(G) \odot U$，将维度减半并在写入共享内存前完成局部布局压缩紧凑； - 在反向传播时，单路输入梯度在 Epilogue 寄存器中直接计算出对应成对输入的梯度并打包。完全避免了中间张量的全局显存落盘。

3. 对偶性反向传播 Epilogue 与跨算子行统计量推导：反向结构自洽保真

对于任意前向逐 tile Epilogue $\mathbf{x}_{\ell}[i,j] = \mathbf{f}_{\ell}[i,j](\mathbf{h}_{\ell}[i,j])$，CODA 证明其反向激活梯度计算天然继承相同的 GEMM-Epilogue 对偶结构：反向传播仅需将前向“生成输入的 GEMM 尾部融合”转变为“生成输出梯度的 GEMM 尾部融合”，且局部 tile 规则 $\mathbf{g}_{\ell}$ 保持严格的 tile 内独立性，不引入额外的跨 tile 通信。

针对 RMSNorm 反向传播中需要计算行内积统计量 $\mathbf{s} = \frac{1}{d}\operatorname{sum}_{\mathrm{cols}}(\nabla_{\mathbf{h}_2}\mathcal{L} \odot \mathbf{h}_2)$ 的难题，CODA 通过线性代数恒等变换，利用 $\nabla_{\mathbf{h}_2}\mathcal{L} = \nabla_{\mathbf{y}}\mathcal{L}\mathbf{W}_1^\top$ 与 $\mathbf{y} = \mathbf{h}_2\mathbf{W}_1$ 证明： $$\mathbf{s} = \frac{1}{d}\operatorname{sum}_{\mathrm{cols}}\left((\nabla_{\mathbf{y}}\mathcal{L}\mathbf{W}_1^\top) \odot \mathbf{h}_2\right) = \frac{1}{d}\operatorname{diag}\left(\nabla_{\mathbf{y}}\mathcal{L}(\mathbf{h}_2\mathbf{W}_1)^\top\right) = \frac{1}{d}\operatorname{sum}_{\mathrm{output}}(\nabla_{\mathbf{y}}\mathcal{L} \odot \mathbf{y})$$ 这意味着 RMSNorm 反向所必需的行内积统计量，既可以在 GEMM 之前计算，也可以等价地在生成 $\nabla_{\mathbf{y}}\mathcal{L}$ 的后向 GEMM Epilogue 中计算。因此，反向传播执行时，每个 GEMM Epilogue 可以顺带为其前序模块累加并输出分块偏置统计量 $\widehat{\mathbf{s}}_{\mathrm{prev}} = \operatorname{reduceTile}_{\mathrm{cols}}(\mathbf{x} \odot \nabla_{\mathbf{x}}\mathcal{L})$，消除了独立的 RMSNorm 反向 kernel。

4. 结构化 Epilogue 原语抽象与 LLM 代码生成闭环

CODA 将 Epilogue 的编写收敛至五大正交原语集合：(1) 逐元素与逐对映射（Maps）；(2) 向量级加载与存储（Rank-1 Loads/Stores，用于偏置与规范化权重）；(3) Tile 级加载与存储（Rank-2 Loads/Stores，基于 Hopper TMA 异步搬运）；(4) Tile 级局部归约（Partial Reductions）；(5) 状态累加变换（Stateful Transforms，如在线 log-sum-exp 维护）。

底层实现基于 CuTeDSL 并适配 NVIDIA Hopper 架构的 TMA（Tensor Memory Accelerator）与 Warp 归约机制。由于搜索空间被约束在“固定高性能 GEMM 主循环 + 结构化 Epilogue 原语组合”，CODA 大幅降低了高性能算子生成的难度。开发者或大语言模型（如 Claude Code）无需探索庞大的调度策略与底层硬件指令，只需通过少量原语示范组合即可生成高性能 GPU 代码，兼顾了极致性能与工程生产力。

实验关键数据¶

主实验¶

在单张 NVIDIA H100 GPU 上，针对典型 Transformer 维度设置（$M=N=K \in \{4096, 8192\}$，RoPE 针对 QKV 投影维度设为 $3N$，词表大小设为 $32768$），CODA 与 PyTorch cuBLAS + torch.compile、Liger Kernel、FlashInfer 进行了算子级性能对比（加速比基准为 cuBLAS + torch.compile = 1.00×）：

算子类型	测试规格 ($M \times N \times K$)	cuBLAS + compile	Liger / FlashInfer	CODA (Human)	CODA (LLM)
GEMM + RoPE	$4096 \times 4096 \times 4096$ ($3N$)	1.00×	1.10×	1.22×	1.21×
GEMM + RoPE	$8192 \times 8192 \times 8192$ ($3N$)	1.00×	1.12×	1.28×	1.27×
GEMM + SwiGLU	$4096 \times 4096 \times 4096$	1.00×	1.14×	1.26×	1.25×
GEMM + SwiGLU	$8192 \times 8192 \times 8192$	1.00×	1.16×	1.32×	1.31×
GEMM + Cross-Entropy	$4096 \times 32768 \times 4096$	1.00×	1.25×	1.45×	1.43×
GEMM + Cross-Entropy	$8192 \times 32768 \times 8192$	1.00×	1.28×	1.51×	1.49×

消融实验¶

论文在完整 Transformer Block（包含两个连续 GEMM-Residual-RMSNorm-GEMM 结构，分别覆盖 RoPE 与 SwiGLU 路径，并纳入轻量级辅助归约与胶水开销）上，针对 1B ($d=2048$)、7B ($d=4096$)、70B ($d=8192$) 模型规模及批大小 16K tokens 评估端到端 Block 加速比：

模型规约 / 评估阶段	隐藏维度 $d$	cuBLAS + compile 基线	独立融合算子 (Liger等)	CODA Block 级加速比
1B 规模 - 前向块 (Forward Block)	2048	1.00×	1.18×	1.38×
7B 规模 - 前向块 (Forward Block)	4096	1.00×	1.15×	1.31×
70B 规模 - 前向块 (Forward Block)	8192	1.00×	1.11×	1.24×
1B 规模 - 完整层 (Full Layer Fwd+Bwd)	2048	1.00×	1.21×	1.42×
7B 规模 - 完整层 (Full Layer Fwd+Bwd)	4096	1.00×	1.18×	1.35×
70B 规模 - 完整层 (Full Layer Fwd+Bwd)	8192	1.00×	1.14×	1.27×

关键发现¶

小模型受显存瓶颈制约更严重：在隐藏维度较小的 1B 模型上，非 GEMM 访存算子的耗时占比更高，CODA 在 1B 前向/后向中取得了最高达 1.42× 的加速提升；即使在计算极其密集的 70B 规模下，CODA 仍能保持约 1.24×~1.27× 的端到端层加速。
LLM 代码生成逼近人类专家性能：利用 Claude Code 在 CODA 约束原语空间下生成的 Kernel 代码，性能普遍达到手写专家代码（CODA Human）的 98%~99.5%，验证了结构化抽象对自动化代码生成的赋能价值。
数值精度稳步提升：在 BF16 精度下，与 FP32 黄金标准对比，CODA 将 RMSNorm 缩放因子延迟到后续 GEMM Epilogue 计算，由于减少了中间量化截断与显存存取带来的精度损失，其实测相对数值误差甚至低于 PyTorch 标准实现路径。

亮点与洞察¶

代数交换消除依赖：巧妙利用行归一化因子与线性矩阵乘法的数学可交换性，打破了“必须先归一化再投影”的框架序列依赖，将中间张量物化彻底降解为两级局部归约。
反向统计量前置转移：揭示了 RMSNorm 反向传播所需的行内积内积统计量可以从前序激活与输出梯度的交互中计算，从而使反向传播完整维持了纯粹的 GEMM-Epilogue 拓扑。
DSL 与 LLM 协同的代码生成新范式：不让 LLM 凭空生成复杂的 GPU 汇编或硬件流水线，而是将不变的专家 GEMM 框架固化，暴露极小且语义正交的 Epilogue 接口，为 AI 辅助生成高性能系统内核提供了教科书般的切分思路。

局限与展望¶

适用架构局限：目前的代数重参数化主要针对标准 Pre-LN Transformer 结构（RMSNorm + SwiGLU + RoPE），对新兴的混合专家架构（MoE）、状态空间模型（Mamba）或带门控机制的复杂注意力结构的覆盖有待进一步验证。
单卡与分布式通信正交：当前原型实现局限于单 GPU 内的算子融合，未将跨卡张量并行（Tensor Parallelism All-Reduce）或流水线并行通信的 Epilogue 融合纳入统一抽象。
打破框架抽象的调试成本：代数重写打破了传统深度学习框架中模块（Layer/Module）的封装边界，将前一层的尾部与后一层的头部混洗在一起，给计算图自动微分与代码调试带来了更高的认知负荷。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙的 Transformer 代数重参数化与完备的反向 Epilogue 映射理论。
实验充分度: ⭐⭐⭐⭐☆ 覆盖了单算子、全层及反向，在 H100 上评测详实，若能补全跨卡分布式训练表现则更佳。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严密，从动机、数学证明到工程实现层次分明。
价值: ⭐⭐⭐⭐⭐ 为低精度大模型时代的显存带宽墙瓶颈提供了极具通用性与工程可落地性的系统级解法。

算子类型	测试规格 (\(M \times N \times K\))	cuBLAS + compile	Liger / FlashInfer	CODA (Human)	CODA (LLM)
GEMM + RoPE	\(4096 \times 4096 \times 4096\) (\(3N\))	1.00×	1.10×	1.22×	1.21×
GEMM + RoPE	\(8192 \times 8192 \times 8192\) (\(3N\))	1.00×	1.12×	1.28×	1.27×
GEMM + SwiGLU	\(4096 \times 4096 \times 4096\)	1.00×	1.14×	1.26×	1.25×
GEMM + SwiGLU	\(8192 \times 8192 \times 8192\)	1.00×	1.16×	1.32×	1.31×
GEMM + Cross-Entropy	\(4096 \times 32768 \times 4096\)	1.00×	1.25×	1.45×	1.43×
GEMM + Cross-Entropy	\(8192 \times 32768 \times 8192\)	1.00×	1.28×	1.51×	1.49×