FlashAttention算法之美：极简推导版

文章预览

知乎：方佳瑞（已授权） 地址：https://zhuanlan.zhihu.com/p/4264163756 FlashAttention（FA）是大模型训练和推理性能优化最重要的组件。从并行计算角度，FA算法设计是可以写进教科书的。通过利用简单数学知识，等价变化任务的计算流程，从而将算法并行执行起来，实现最佳的内存效率，这无疑是并行计算Phd心中最完美的idea。 FA的算法流程也以复杂著称，原始论文中公式包含纷繁变量，复杂计算流程图，让普通人很难理解。对于FA也有很多非官方解读版本，比如Zihao Ye的《From Online Softmax to FlashAttention》（2023年5月作为UW研究生作业发布）。Ye从FA的演化历史入手，囊括了完毕的前置知识，让读者读此一文即可搞懂FA。但是Ye版本力求完备，公式保持完整严谨，让普通人很那一遍理解，需要反复揣摩。Zhihao Ye也是FlashInfer的作者。 去年此时，FA v2刚刚更新，我写过 ………………………………