图解Attention优化: 从Online-Softmax到FlashAttention V1/V2/V3

文章预览

作者：DefTruth，AI Infra，暨南大学(JNU) 主页：github.com/DefTruth 声明：本文只做分享，版权归原作者，青稞AI整理 原文：https://zhuanlan.zhihu.com/p/668888063 前言 本文通过原理分析和图解的方式，通俗易懂地FlashAttention系列算法。FlashAttention V1/V2在LLM领域的应用已经非常广泛，相关的论文也反复读了几遍。FA1和FA2论文非常经典，都推荐读一下（不过FA2论文中公式错误不少）。本文大约1.8w字，包括以下内容： • Standard Self-Attention • (Safe) Softmax: 3-pass • Online Softmax: 2-pass • FlashAttention V1 • FlashAttention V2 • Analysis: IO Complexity of FlashAttention • 分布式训推使用FlashAttention • Memory-Efficient Attention • FlashAttention中MQA/GQA以及Causal Mask处理 • FlashAttention V3: 比V2更快、支持Hopper FP8 (TODO) • 总结 本篇文章，主要是记录些FlashAttention论文阅读笔记，温故知新，不追求独 ………………………………