flash attention: Fast and Memory-Efficient Exact Attention with IO-Awareness

mex_wayne

已于 2024-08-12 10:22:15 修改

阅读量645

点赞数 15

分类专栏： transformer 文章标签：人工智能 transformer 自然语言处理 machine translation

于 2024-07-18 23:08:35 首次发布

本文链接：https://blog.csdn.net/mikhailbran/article/details/140446863

版权

transformer 专栏收录该内容

5 篇文章 0 订阅

订阅专栏

FlashAttention 学习笔记，不会按照文章结构走，理解为王

前言

目前transformer 相关应用非常广泛，因此分享一篇关于flash attention的文章。这里为什么先分享flash attention？
首先，之前的 attention优化都是基于计算和稀疏性，这篇文章直接从硬件构架角度思考，减少硬件的开销。
其次，本文的效果好，直接从 $N^2$ 的复杂度降为N，在不改变attention 结构的情下加速显然，甚至因为能增加输入的长度，使得效果有提升，所以在这里做一篇分享，欢迎交流。
在这里插入图片描述

1、transformer 基础

网上关于transformer的讲述非常多也讲的很好，我这里不进行赘述，祭出一个公式，后面用的比较多
$attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt(d_k)})V$ , 其中 $Q, K, V \in N \times d$
N：sequence length
d：head dimension

2、GPU memory 结构

FlashAttention 有个最大的优点就是基于硬件讨论。作者在文中指出，文章的出发点就是减少IO的访问(access，论文中作者后面的复杂度都指 access) 。下图是论文作者分享的一个比较重要的memory 构架图其中标注了硬件名称和bandwith 我会详细阐明。
在这里插入图片描述
这里推荐链接：https://developer.nvidia.com/zh-cn/blog/nvidia-ampere-architecture-in-depth/
讲明了英伟达A100的构架，都是干货

2.1、Memory hierarchy

（1）GPU SRAM
sram 一直是static random access memory，它不是寄存器也不能当做寄存器（部分博主视频讲解会把它认为是寄存器）。准确的理解是片上高速缓存，不会直接参与计算，用于存放数据。由于访问sram的delay非常小，因此将其存放反复访问的数据。寄存器的耗时delay是1cycle，sram的访问delay是2~5cycle级别。
如果在DSP构架中，SRAM是TCM（Tightly-Coupled Memory）可以用于存放指令或数据。因为几乎与内核同频，因此叫做tightly-coupled。在GPU构架中SRAM也可以做L1/L2缓存/share memory。要注意的是L2cache 和 share memory 访问速度会慢一些，访问delay的数量级大概十几个cycle或者share memory会更慢一些。

（2）GPU HBM
HBM 是 high bandwidth memory，它是由dram做的。dram全名是Dynamic Random Access Memory。访问速
会更慢，访问delay大概在几十个cycle，我们常说的 DDR就是由DRAM构成，论文这里HBM也是由DRAM构成，只是NVIDIA 加入了自己的技术TSV（Through-Silicon Vias, TSV），TSV 是贯穿每个 DRAM 芯片的垂直导电通道，使得数据可以在不同的 DRAM 层之间传输。TSV 技术使得多个 DRAM 芯片可以堆叠在一起，形成一个紧凑的内存模块。
TSV 提供了高速数据传输路径，显著提高了内存带宽。同样造价更高，因此整体尺寸小于DDR。

（3）Main Memory
Main Memory就是指的我们常说的DDR或者广义一些就是内存，也是由dram构成。相对于HBM少了 TSV因造价便宜一些使得我们可以将其的尺寸定的更大一些。

综观上述三个结构，可见sram上的数据要和计算单元紧密合作所以要求访问速度快，而HBM主打带宽高方便计算大量吞吐，而cpu内存就是主打量大。所以我们不难看出来这个内存尺寸一定是类似金字塔型的。

2.2 bandwidth

论文中有个出处：
[43] Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study on optimizing transformers. Proceedings of Machine Learning and Systems, 3: 711–732, 2021.
这个出处详细分析了roofline对transformer的影响。roofline主要讲解的是什么时候计算被memorysize限制，什么时候被计算带宽限制，维基百科和知乎的一篇文讲的很好：
https://en.wikipedia.org/wiki/Roofline_model
https://zhuanlan.zhihu.com/p/34204282
我这里仅贴图，不赘述。
在这里插入图片描述

引用flash attention作者的结论就是 transformer 的瓶颈是：memory access.
原文：GPUs, compute speed has out-paced memory speed [61, 62, 63], and most operations in Transformers are
bottlenecked by memory accesses [43].

3、online softmax

online soft max其实是nvidia很早之前发的一篇论文：Online normalizer calculation for softmax。没有华丽的辞藻也没有高深的数学模型，就是实打实的将如何优化softmax，步步可推敲且严谨。

3.1 、safe softmax

这里有softmax原始公式：[y] = softmax([x])定义为：
$y_i=\frac{e^{x_i}}{\sum^{V}_{j=1}{e^{x_j}}}$ ，其中 $x,y∈\R^V$ , 即都是长度为V的实数域向量。
上述公式按照online softmax 作者描述，是2输入1输出因此，写成伪代码则是：

Algorithm 1 Naive softmax
1: d0 = 0
2: for j in range(1, V):
3:     d[j] = d[j−1] + e^x[j]
5: for i in range(1, V):
6:     y[i] ← (e^x[i])/dV

但实际电脑上的代码则为了保护溢出值防止计算误差变大，
因为 $e^x$ 中档x>64，当我们遇到token数量为1000甚至100000时（尤其fp16时经常遇到因为计算误差导致梯度消失不收敛），那么这个过程就非常不精确。因此这里将其减去最大，将 $e^x$ 拉到0~1之间，公式为：
$y_i=\frac{e^{x_i-\max^V_{k=1}x_k}}{\sum^{V}_{j=1}{e^{x_j-\max^V_{k=1}x_k}}}$ ，其中 $x,y∈\R^V$ , 即都是长度为V的实数域向量。
且该文的作者指出，tensoflow、ptyroch、caffe2等主流构架都是用该数学公式。
因此我们有伪代码：

Algorithm 2 Safe softmax
1: m0 = −∞
2: for k in range(1, V):
3:     m[k] = max(m[k−1], x[k])
4: d[0] = 0
5: for j in range(1, V):
6:     d[j] = d[j−1] * e^m[j-1] + e^(x[j]−m[V])
7: for i ← 1, V do
8:     y[i] = e^(x[i]−m[V])/d[V]

对于这个algorithm 2，计算向量中的一个元素，就要访存4次 memory。因此需要提升这个softmax。

3.2、online softmax

为什么要看online softmax，因为FlashAttention作者提到了这个。
online softmax作者灵感来自于：
[18] B. P. Welford. Note on a method for calculating corrected sums of squares and products. Technometrics, 4(3):419–420, 1962. URL https://amstat.tandfonline.com/doi/abs/10.1080/00401706.1962.10490022.
将一个元素的4次memory访问降为3次，伪代码如下：

Algorithm 3 Safe softmax with online normalizer calculation
1: m0 = −∞
2: d[0] = 0
3: #for k in range(1, V):#被优化掉了
4: #    m[k] = max(m[k−1], x[k])#被优化掉了
5: for j in range(1, V):
6:     m[j] = max(m[j-1], x[j-1])#计算优保留，但是不用直接访存
7:     d[j] = d[j−1] * e^m[j-1] + e^(x[j]−m[V])
7: for i ← 1, V do
8:     y[i] = e^(x[i]−m[V])/d[V]

同时，作者也证明了这个过程是safe的。写到这里 online softmax其实就是通过一个递推公式(代码第7行)： $d[j] = d[j−1] * e^{m[j-1]} + e^{(x[j]−m[V])}, d[0]=0$ , 通过该公式将第N个结果算出而不是通过访问memory读出。

在此之上，该文作者也提出了如何并行化上述代码,其实就是SIMD向量（说的简单些就是一次将多个同类型数据送入vector register进行并行计算）并行：
在这里插入图片描述
以及针对inference过程中的topK其实也用的是softmax，因此可以有新的优化算法从上面提到的3-pass（访问3次）降为2-pass（访问2次），于是又有如下伪代码：

Algorithm 4 Online softmax and top-K
1: m0 = −∞
2: d0 = 0
3: u = {−∞, −∞, . . . , −∞}T# u ∈ R^(K+1) The 1st K elems will hold running TopK values
4: p = {−1, −1, . . ., −1}T # p ∈ Z^(k+1) The 1st K elems will hold running TopK values' indices 
5: for j in range( 1, V):
6:      m[j] = max (m[j−1], x[j] )
7:      d[j] = d[j−1] * (e^(m[j−1]−m[j]) + e^(x[j]−m[j]))
8:      u[K+1] = x[j] # Initialize K + 1 elem with new value from input vector
9:      p[K+1] = j# Initialize K + 1 elem with new value from input vector's index
10:
11:     # Sort u in descending order, permuting p accordingly. The first K elements are 
12:     # already sorted, so we need just a single loop, inserting the last element in 
13:     # the correct position.
14:     _k = K 
15:     while k ≥ 1 and u[_k] < u[_k+1]:
16:         swap(u[_k], u[_k+1])
17:         swap(p[_k], p[_k+1])
18:         _k = _k − 1
19: for i in range( 1, K): # The algorithm stores only K values and their indices
20:     v[i] = e^(u[i]−m[V])/d[V]
21:     z[i] = p[i]

基本上追上原始softmax的效率
在这里插入图片描述

4、selective gradient checkpointing

本小节来自于论文作者的引用[10]
[10] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016.
~~[34] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.~~[34]不公开，且技术点在[10]中也有说明，我这里用[10]举例说明。这里不是attention 的一个主要特点，因此提一下。

还有，为什么要单独开这么一节，因为这些数据非常亮眼：
1 内存消耗度直接干成O $\sqrt{n}$ ,部分case变成O(logn)
2 残差网络48g干成7g

4.1 memory sharing

在这里插入图片描述
这篇引用论文的作者思路就是，在反向通路中，我们分析不同数据流的依赖关系，进而得到哪些节点的数据输入输出是一样的，那么这个节点依赖的memory就是share memory。

4.2 gradient checkpointing

这篇引用论文作者也点出了技术点：gradient checkpointing。该技术其实是在反向传播中挑选用于计算梯度的关键点。这些关键点就是 check point，如果全储存所有的中间结果代价较大，选取基哥点那么可以大大减少内存负担。
那么问题来了，如何选取？在[10]的引用文章《Implementation of Checkpointing for the Reverse or Adjoint Mode of Computational Differentiation》中的 revolve过程说明了该算法。因为我们重点是将flash attention详细展开会非常多，我这里仅做说明。

5、FlashAttention

基于上面的基础，作者首先分析了transformer 的瓶颈在于访问IO的次数，也就是频繁IO增加了耗时。因此角度都从这个方面切入。

5.1、flash attention IO complexity：

按照第1小节的公式，我们有qkv在硬件上的实现方法基本就是先算qk 接着softmax，最后乘以V，访存过程都放入HBM，因此有如下复杂度分析：
（1） $Q,K,V ∈ \R^{N×d}$ 因此 $Q×K^T$ 的计算复杂度为O(N×N×d)，
（2）因为SRAM的宽度为M，因此访存次数为 $\frac{N^2d}{M}$
（2）softmax 如果用英伟达的 online softmax 则是倍数乘以3，因为这里是线性倍数不是指数倍数，因此可以暂时忽略。
（3）softmax之后的结果再乘以V，其计算复杂度再次增加N×d，但我们会这样写：

for(d) #d for O
    # calc the S  ################ start ####
	for(N) #N for Q
		for(N) #N for trans(K)
			for(d) # dot(Q,K), d is for both Q and trans(K)
	# calc the S ################# end  #####
	P = softmax(S)
	for(N) # N for P
		for(d) #  d for P*V

在这里插入图片描述

(4) 不考虑是否能恰好对齐sram长度，以及其他情况，假设每次都可以度满sram，因此访问sram的次数需要除以M，因为从HBM一次最多读M大小的数据。
因此得到文中的flash attention 的 IO复杂度（IO complexity）为O( $\frac{N^2d^2}{M}$ )。

5.2、standard attention IO complexity：

$S= Q*K^T, P = softmax(S) ，O = PV ∈ \R^{N × d}$ 在这里插入图片描述
对于standard IO 复杂度，按照作者提供的这个过程，我们有：
（1）load Q,K ，write S： $Nd + Nd + N^2$ 。
（2）read S, softmax(S), write P: read S 不增加，因为本身就在 HBM中，接着用了 softmax（s）就是 *3，即 $N^2*3$ ，注：线性系数不用出现在复杂度，因此这里是 $N^2$ 。
（3）load P and V compute PV write O： $N^2 + Nd$
（4）return O：略
所以作者说standard attention IO复杂度是 $Ω(Nd + N^2)$ ，即下限。

5.3、IO complexity compare：

flash attention 的 IO 复杂度为：O( $\frac{N^2d^2}{M}$ )
而standard 的 IO复杂度为： $Ω(Nd + N^2)$
看起来 flash attention的更多一些，但这里需要注意：
在这里插入图片描述
文中举出来了一种典型数据尺度：Nsequence lenght ：1024， d head dimension：64~128，M为100KB。
`这种数量级起步，因此这中IO减少的思路还是非常可观的。作者也贴出了数据，证明效果。论文作者特别喜欢用`wall-clock时间`，wall-clock时间就是实际运行时间不关心中间过程如何也不用芯片频率做换算，就是`实打实，我们看墙上的表开始然后结束的时间`。下图为原文截图，作者以此说明flash attention的炸裂效果。
在这里插入图片描述

5.4、flash attention的优化手段

在描述self attention 算法前，我这里需要明确一下过程：

5.4.1 tiling

tiling 就是大家经常做的分tile。这里按照行（raw）分。
论文中：
sram_size: 原文是M，很容易和向量混淆，我这里 sram_size 替代更清晰。
如果 $x∈\R^B$ ,那么有如下定义：

$m(x) = max(x) = max(x_1, x_2, x_3, ....)$ 是标量。
$f(x) = [e^{x_1-m(x)}, ...., e^{x_B - m(x)}]$ 是向量。
$\ell(x) = \sum_i{f(x)_i} \\=e^{x_1-m(x)} + e^{x_2-m(x)} + ......$ 符号i有迷惑性，根据上下文就是求和。
$softmax=\frac{f(x)}{\ell(x) }$

如果有 $x=[x^{(1)}\quad x^{(2)}]$ , 其中 $x^{(1)}, x^{(2)} ∈\R^B$ , 那么上面的4条又有如下定义：

$m(x) = max([x^{(1)}, x^{(2)}]) \\= max(max(x^{(1)}), max(x^{(2)})) \\=max(max(x_1^{(1)}, x_2^{(1)}, x_3^{(1)}, ....), max(x_1^{(2)}, x_2^{(2)}, x_3^{(2)},....)$ 是标量。
$[e^{m( x^{(1)} ) - m(x) } f( x^{(1)} )\quad e^{m( x^{(2)} ) - m(x) } f( x^{(2)} ) ]$ 这里不加逗号而是空格，我理解作者其实希望两个 $x^{(1)}$ 和 $x^{(2)}$ 写在一起不增加维度所以采用了个这么写法。
其中： $e^{m( x^{(1)} ) - m(x) } f( x^{(1)} ) \\=e^{m( x^{(1)} ) - m(x) } [e^{x_1-m( x^{(1)} )}, ...., e^{x_B - m( x^{(1)} )}] \\=e^{m( x^{(1)} ) - m(x) }e^{x_1-m( x^{(1)} )}, .....,e^{m( x^{(1)} ) - m(x) }e^{x_B - m( x^{(1)} )} \\=e^{ \cancel{ m( x^{(1)} ) } - m(x) } e^{ x_1 - \cancel{m( x^{(1)} )} }, .....,e^{ \cancel{ m( x^{(1)} ) } - m(x) } e^{ x_B - \cancel{ m( x^{(1)} ) } } \\=e^{ x_1-m(x) }, .....,e^{x_B - m(x) }$ 也是和定义相同。
$\ell(x) = \ell( x^{(1)} \quad x^{(2)} ) \\= [e^{m( x^{(1)} ) - m(x) } \ell( x^{(1)} )\quad e^{m( x^{(2)} ) - m(x) } \ell( x^{(2)} ) ]$ 这里推导和2.中过程一样，不赘述。
$softmax=\frac{f(x)}{\ell(x) }$ 同理，不赘述。

通常我们分tile可以减少横跨地址带来的cache miss以及提高memory的利用率。但是对于flash attention最主要的目的是将输入减小。减小后可以一次放在 sram上做完一个tile的 $softmax(K_iQ^T_iV_i)$ ，然后再把每个tile的结果拼在一起。至于作者为什么分为4d，个人认为就是一个经验值，也可以8d或16d，就像transform里的 $\sqrt{d_k}$ 。但是这里的4是fp32类型，因为sram的单位是byte，github上工程使用fp32实现，是4个byte后面再求得尺寸时可以得到验证。
在这里插入图片描述

5.4.2 recomputation：

这里作者首先说这里是为在反向传播时，不存储 $O(N^2)$ 尺寸的数据。如果没有优化需要存储S和P， $S,P ∈ \R^{N × N}$ 才能计算出有关KQV的梯度。但是如果存储输出O和矩阵归一化统计信息 $\ell)$ 可以重新计算出有关KQV的 S 和 P。这种方法可以看成是selective gradient checkpointing，在第4章节说过这个方法。

5.4.3 implementation：

在这里插入图片描述
Require:
$Q,K,V ∈ \R^{N × d}，S= Q*K^T$ ， $P = so f t ma x (S)$ ， $O = PV ∈ \R^{N × d}$
line 1：
block size有两种:
(1) $B_c = celling(sram_size / 4d)$
(2) $B_r = min(celling(sram_size / 4d), d)$

line 2:
初始化 $O=(0)_{N×d}∈\R^{N × d}$ ， $\ell=(0)_{N×d}∈\R^{N}$ ， $m=(-∞)_{N}∈\R^{N}$
这里注意的是m，m维度是N。且上述内存都在HBM中。

line3：
(1) Q分为 $T_r个=celling(N/B_r)$ 个 blocks $Q_1,Q_2,....,Q_{Tr}$ 共计 $B_r×d$ 个。
(2) K和V分 $T_c个=celling(N/B_c)$ 个blocks $K_i$ 和 $V_i$ 都是 $B_c×d$ 个。
(3) 注意这里有个tricky的地方： $B_c$ 是sram_size/4划分，但是原文中却这样描述：
在这里插入图片描述
这里是因为sram_size是byte大小，因此按照M/4d划分。但是运行数据是fp32(4个btye)所以是 $B_r×d$ 。

line4：
同理划分O, $\ell$ 各为 $B_r×d$ ，m 分为 $B_r$

line5：
for Tc 循环：

line6：
将 $K_j和$ V_i$从HBM上加载到sram中。

line7：
for Tr 循环：

line8：
从HBM中读取 $Q_i,O_i,\ell_i,m_i$
line9：
计算 $S_{ij}=Q_iK^T_j ∈ \R^{B_r × B_c}$
line10:
(1) 计算 $\widetilde m_ij=rowmax(S_{ij})∈\R^{B_r}$ 到这里，我用一张图说明过程：
在这里插入图片描述
(2) 计算 $\widetilde P_{ij}=exp(S_{ij} - \widetilde m_{ij}) ∈ \R^{B_r × B_c}$ ,注意这里算出来的 $\widetilde m_{ij}$ 是需要在行方向复制扩充才能匹配 $S_{ij}$ 的尺寸，做的工作就是
第3章节中说的online softmax的safe过程。同样依照online softmax，这里需要求和得到softmax的分母。
在这里插入图片描述
line11:

因为循环是不断迭代的，对于第一次，计算出来的最大值和 sum求和都是局部的最大值和求和，因此我们需要5.4.1 中tiling部分的2个公式：
$\ell(x) = \ell( x^{(1)} \quad x^{(2)} ) = [e^{m( x^{(1)} ) - m(x) } \ell( x^{(1)} )\quad e^{m( x^{(2)} ) - m(x) } \ell( x^{(2)} ) ]$
$[e^{m( x^{(1)} ) - m(x) } f( x^{(1)} )\quad e^{m( x^{(2)} ) - m(x) } f( x^{(2)} ) ]$
随着循环不断更新，更新的数值是靠前面乘的系数。因此我们每次循环都会从新得到 $m_i^{new}$ 和 $\ell_i^{new}$ 。
同时也得到 $P = so f t ma x (S)$

line12:
计算出 $O = PV ∈ \R^{N × d}$
** line13：**
将算出的 $\ell_i$ 和 $m_i$ 存到HBM中，因为反向传播要用。

5.6、flash attention IO complexity ：

作者在写完上述算法后，又进行了理论分析。并且对反向传播也写了详细的优化算法。
我将在下一篇文章进行分享

mex_wayne

关注

15
点赞
踩
10

收藏

觉得还不错? 一键收藏
1
评论
flash attention: Fast and Memory-Efficient Exact Attention with IO-Awareness

目前transformer 相关应用非常广泛，因此分享一篇关于flash attention的文章。这里为什么先分享flash attention？首先，之前的 attention优化都是基于计算和稀疏性，这篇文章直接从硬件构架角度思考，减少硬件的开销。其次，本文的效果好，直接从N2N^2N2的复杂度降为N，在不改变attention 结构的情下加速显然，甚至因为能增加输入的长度，使得效果有提升，所以在这里做一篇分享，欢迎交流。
复制链接

扫一扫