【GPU 算子工程】FlashAttention:在线 softmax 与 IO-aware 注意力
FlashAttention 把注意力重写成分块的在线 softmax,不落地 N×N 分数矩阵,用重算换访存。本文推导算法、给出实测正确的简化实现(误差 4e-7、避免 16.8MB 分数矩阵),并引用原论文的加速与显存数据。
Linux 内核、存储与网络、可观测性、系统架构与大模型基础设施的工程笔记:机制拆解、踩坑复盘与可核对证据,少空谈。
共 2 篇文章 · 返回首页
FlashAttention 把注意力重写成分块的在线 softmax,不落地 N×N 分数矩阵,用重算换访存。本文推导算法、给出实测正确的简化实现(误差 4e-7、避免 16.8MB 分数矩阵),并引用原论文的加速与显存数据。
归约类算子是 memory-bound 的典型。讲 softmax 的数值稳定写法(减最大值、在线 softmax)、LayerNorm 的 Welford 单遍方差,以及逐元素融合:实测把 scale+bias+GELU 三个 kernel 融成一个,提速 2.94 倍。