PyTorch 新增 flex_gemm API:普通函数写 GEMM 融合 epilogue

tri_dao · x · 2026-10-23

Tri Dao 转发了 drisspg 在 PTC 的演讲:flexgemm 是 PyTorch 中基于 Quack 构建的新 GEMM epilogue API,允许用一个普通的 PyTorch 函数把 epilogue 融合进 GEMM,单 kernel 内在寄存器中的累加器上直接计算 epilogue(如 GELU+bias)再写回。

作者开源了一组示例( transformernuggets 仓库的 flexgemm 目录),涵盖 pointwise、多输出、SwiGLU、reduction、MXFP 量化、backward、随机舍入、函数变换与 autotune 等 9 个模式;每个脚本会用 torch.compile 对照验证正确性并打印 GPU 耗时。作者表示会持续补充更多示例。