训过大模型的人都知道一件反直觉的事:你以为时间都花在 matmul 上,profile 一拉,norm、激活、residual、各种 reduction 这些"小算子"能吃掉一大块端到端时间。
不是因为它们算得多,是因为它们一直在搬数据。
中间 tensor 从 SRAM 写回 HBM,下一个 kernel 再读回来,循环往复。算力够,带宽不够,卡就空转。
CODA 这篇的思路挺干净:既然 GEMM 输出 tile 还在片上,那就别急着写回去,顺手把 norm / scaling / pairwise / reduction 这些都作为 epilogue 跟着算完再落盘。
🔍 它的取舍是把 GEMM 主循环锁死,只暴露一小撮可组合的 epilogue 原语。好处是能继承 expert 手写 GEMM 的性能结构,坏处是表达力受限——不过作者论证了 Transformer 里除 attention 之外几乎都能塞进这个壳。
💡 跟前阵子那批 deep fusion kernel 是一个方向,但 CODA 把它抽象成了一个"编程范式",连 LLM 写的 kernel 都能跑出不错的性能,这点对工程落地比较关键。
📌 不过我有点保留:epilogue 寄存器压力一上来,occupancy 容易掉,论文 figure 看着漂亮,真上到 70B 训练 stack 里替换 RMSNorm + residual 那条链,会不会被 register spill 反噬,得自己跑过才知道。
做推理/训练优化的可以读一下,思路比结论值钱。 2605.19269
发布于 中国香港
