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」频道最新
- AI 圈梗图：「偷来神的火种就为刷性能」调侃算力狂热 — bronzeagepapi · 2026-09-20
- SemiAnalysis：Neocloud 应加入 Nvidia 禁运制安全披露计划 — typewriters · 2026-09-20
- 7900XTX 本地跑 27B 模型 40tok/s,搞定稳定 Agentic 编码 — W61k3r · 2026-09-20
- 「推理负载远超你优化所想」：一句 AI 圈流行语道破长尾需求 — charles_irl · 2026-09-20
- 断供约两月后,Kimi 订阅服务恢复开放 — ChrisGPT · 2026-09-20
- HN 热议：OpenAI 用自家 LLM 辅助设计自研 Jalapeño 芯片 — petrusenko_max · 2026-09-20