FlashNorm 两行代数换 33-35% 提速,CUDA 竞态曾让模型倒放输出

AI Engineer · youtube · 2026-09-20

AI Engineer 视频访谈 Filip Makraduli(与 Nils Graef 合写论文)讲解 FlashNorm。要点:

问题:RMS norm 层几乎不做算术,但一次 decode 可能启动它约 33 次;GPU 擅长计算、不擅长启动/搬运/等待,这个间隙就是优化空间。
方法:①把 norm 的 gain 离线折叠进投影权重,一个矩阵吸收两者;②推迟标量除法,让矩阵单元和向量单元并行而非互相空等;③对连续两次归一化的新架构,利用 scale invariance 直接删掉第二次。
效果:norm+投影操作提速 33-35%,折叠后的 checkpoint 兼容 torch compile 和量化模型。
踩坑:除法推迟需 CUDA kernel 实现,Python 层做不到。最初把 matmul 放张量核心、RMS reduction 放 CUDA 核心并行跑,单元测试和困惑度都正常,但长生成时模型开始以一步之差复读过去的输出——原因是两条流之间的 join 是隐式的,一条未完成,post scale 读到了未完成乘法的过期 buffer。修复是显式标记每条流的结束,让 post scale 等待两条流。
部署:修改 checkpoint 的 kernel 手术无法在租用的推理端点上做,后半段需要 Superlinked 的开源推理引擎。

视频含完整时间戳,论文与仓库见作者主页。

「Infra」频道最新