手撕Flash Attention！原理解析及代码实现_flashattention 的torch实现

手撕Flash Attention！原理解析及代码实现

大模型智能 2024-05-21 00:04 吉林

大模型智能｜分享

来源 | https://zhuanlan.zhihu.com/p/696850636

作者丨晚安汤姆布利多

01

前言

自 2022 年 11 月 OpenAI 发布 ChatGPT 以来，这一年多来大语言模型 (Large Language Model, LLM) 的发展十分迅速，国内外众多厂商纷纷加入“百模大战”。但是，由于大语言模型的参数量非常巨大（通常为十亿、百亿甚至千亿量级），加之训练语料很庞大，模型的训练成本十分高昂。

当前，Transformer 已经成为了大语言模型的默认网络结构，为了降低大语言模型的训练成本，一些工作尝试对 Transformer 的计算成本进行优化，比如降低注意力运算的时间成本或者显存占用等。

本文介绍 Flash Attention，一种优化的注意力算法。Flash Attention 论文链接如下：

https://arxiv.org/pdf/2205.14135

本文从注意力机制出发，分析