FlashAttention算法详解

这篇文章的目的是详细的解释Flash Attention，为什么要解释FlashAttention呢？因为FlashAttention 是一种重新排序注意力计算的算法，它无需任何近似即可加速注意力计算并减少内存占用。所以作为目前LLM的模型加速它是一个非常好的解决方案，本文介绍经典的V1版本，最新的V2做了其他优化我们这里暂时不介绍。因为V1版的FlashAttention号称可以提速5-10倍，所以我们来研究一下它到底是怎么实现的。

介绍

论文的标题是：

“FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness”

内存的效率与普通注意力相比（序列长度是二次的，O(N²)），FlashAttention是次二次的/线性的N (O(N))。并且它不是注意力机制的近似值(例如，稀疏或低秩矩阵近似值方法)-它的输出与“传统”注意力机制相同。与普通的注意力相比，FlashAttention的注意力是”有感知“的。

它利用底层硬件的内存层次知识(例如gpu，但其他AI加速器也应该工作，我这里使用gpu作为示例)。一些[近似]方法在序列长度上将计算要求降低到线性或近线性，但其中许多方法专注于减少FLOP，而忽略内存访问(IO)的开销。

经过多年的发展gpu的FLOPS的增长速度一直在以比内存吞吐量(TB/s)更快。内存的瓶颈应该引起重视。FLOPS和内存吞吐量需要紧密结合，由于硬件上的差距，我们就需要软件层面上的工作进行平衡。

根据计算和内存访问之间的比率，操作可以分为以下两种:

计算约束：矩阵乘法
内存约束:元素操作(激活，dropout，masking)，归并操作(softmax， layer norm，sum等)

在当前的AI加速器（GPU）上是受内存大小限制的。因为它“主要由元素操作组成”，或者更准确地说，注意力的算术密度不是很高。

我们看看这个图:

可以看到，masking，softmax和dropout是占用大量时间的操作，而不是矩阵乘法(即使大部分FLOPS是在matmul中)。内存不是一个单一的工件，它在本质上是分层的，一般的规则是:内存越快，越昂贵，容量越小。

我们在上面说的，FlashAttention的注意力是”有感知“的可以归结为利用SRAM比HBM(高带宽内存)快得多来确保减少两者之间的通信。

以A100为例：

A100 GPU有40-80GB的高带宽内存(HBM)，带宽为1.5-2.0 TB/s，而每108个流处理器有192KB的SRAM，带宽估计在19TB/s左右。

可以看到大小小了很多，但是速度却提升了10倍，所以如何高效的利用SRAM是提速的关键，让我们看看标准注意力实现背后的计算:

标准实现如何显示对HW操作方式不大尊重。它基本上将HBM加载/存储操作视为0成本(它不是“io感知”)。

我们首先考虑如何使这个实现更有效(时间和内存方面)。最简单的方法是删除冗余的HBM读/写。

如何把S写回HBM只是为了(重新)加载它来计算softmax，那么我们可以将其保存在SRAM中，执行所有中间步骤，然后将最终结果写回HBM。

内核基本上是“GPU操作”的一种奇特的说法（参考我们以前发布的CUDA入门，往简单了说就是一个函数）。融合则可以将多个操作融合在一起。所以只从HBM加载一次，执行融合的op，然后将结果写回来。这样做可以减少通信开销。

这里还有一个专业名词术语是“materialization”（物化/实体化）。它指的是，在上面的标准注意力实现中，已经分配了完整的NxN矩阵(S, P)。下面我们将看到如何直接将内存复杂度从O(N²)降低到O(N)。

Flash attention基本上可以归结为两个主要观点:

Tiling (在向前和向后传递时使用)-基本上将NxN softmax/scores矩阵分块成块。

Recomputation (仅在向后传递中使用)

算法如下:

上面我们提到了很多名词，你可能还不了解。没关系下面我们开始逐行解释算法。

FlashAttention算法

让Tiling方法的主要障碍是softmax。因为softmax需要将所有的分数列耦合在一起。

看到分母了吗?这就是问题所在。

要计算输入序列中的特定第i个标记对序列中其他标记的关注程度，需要在SRAM中随时可用所有这些分数(这里用z_j表示)。

但是SRAM的容量是有限的。N(序列长度)可以是1000甚至100000个令牌。所以N²爆炸得很快。所以论文使用了一个技巧：把softmax的计算分成更小的块，最终仍然得到完全相同的结果。

我们可以只获取前一个B分数(x_1到x_B)并为它们计算softmax。然后通过迭代，“收敛”到正确的结果。以一种聪明的方式组合这些每块部分softmax的数字，这样最终的结果实际上是正确的。方法如下:

基本上，为了计算属于前2个块(大小为B)的分数的softmax，必须要跟踪每个块的2个统计数据:m(x)(最大分数)和l(x) (exp分数总和)。然后就可以用归一化系数将它们无缝地融合在一起。

这里主要是一些基本的代数运算，通过展开f(x)和l(x)项并与e^x相乘一些项会相互抵消，这里就不写了。这个逻辑递归地一直持续到最后一个(N/B)块，这样就得到了N维正确的softmax输出!

为了详细的介绍这个算法，假设有一个大小为1的批处理(即单个序列)和单个注意力头，稍后会扩展它(通过简单地跨GPU的并行化-稍后会详细介绍)。我们暂时忽略了dropout和masking，因为稍后再添加。

我们开始计算：

初始化:HBM的容量以GB为单位测量(例如RTX 3090有24 GB的VRAM/HBM, A100有40-80 GB等)，因此分配Q, K和V不是问题。

第1步

计算行/列块大小。为什么ceil(M / 4 d) ?因为查询、键和值向量是d维的，所以我们还需要将它们组合成输出的d维向量。所以这个大小基本上允许我们用q k v和0个向量最大化SRAM的容量。

比如说，假设M = 1000, d = 5。那么块大小为(1000/4*5)= 50。所以一次加载50个q, k, v, o个向量的块，这样可以减少HBM/SRAM之间的读/写次数。

对于B_r，我也不太确定他们为什么要用d执行最小运算?如果有人知道，请评论指教!

第2步：

用全0初始化输出矩阵O。它将作为一个累加器，l也类似它的目的是保存softmax的累积分母——exp分数的总和)。M(保存逐行最大分数)初始化为-inf，因为我们将对其进行Max运算符，因此无论第一个块的Max是什么-它肯定大于-inf 。

第3步：

步骤1中的块大小将Q, K和V分成块。

第4步：

将O, l, m分割成块(与Q的块大小相同)。

第5步：

开始跨列循环，即跨键/值向量(上图中的外部循环)。

第6步：

将K_j和V_j块从HBM加载到SRAM。在这个时间点上我们仍然有50%的SRAM未被占用(专用于Q和O)。所以SRAM是这样的：

第7步：

开始跨行内部循环，即跨查询向量。

第8步：

将Q_i (B_r x d)和O_i (B_r x d)块以及l_i (B_r)和m_i (B_r)加载到SRAM中。

这里需要保证l_i和m_i能够载入SRAM(包括所有中间变量)，这块可能是CUDA的知识，我不太确定如何计算，所以如果你有相关的信息，请留言

第9步：

计算Q_i (B_r x d)和K_j转置(d x B_c)之间的点积，得到分数(B_r x B_c)。并没有将整个nxns(分数)矩阵“物化”。

假设外部循环索引为j (j=3)，内部循环索引为i (i=2)， N为25，块大小为5，下面就是刚刚计算的结果(假设以1为基础的索引):

也就是输入序列中标记11-15的标记6-10的注意力得分。这里的一个要点是，这些都是精确的分数，它们永远不会改变。

第10步：

使用上一步计算的分数计算m_i_j、l*i_j和P~*i_j。M ~_i_j是按行计算的，找到上面每一行的最大元素。

然后通过应用元素运算得到P~_i_j:

归一化-取行最大值并从行分数中减去它，然后EXP

l~_i_j是矩阵P的逐行和。

第11步：

计算m_new_i和l_new_i。同样非常简单，可以重复使用上面的图表:

M_i包含之前所有块的逐行最大值(j=1 & j=2，用绿色表示)。M _i_j包含当前块的逐行最大值(用黄色表示)。为了得到m_new_i我们只需要在m_i_j和m_i之间取一个最大值，l_new_i也类似。

第12步(最重要):

这是算法中最难的部分。

它允许我们用矩阵的形式做逐行标量乘法。如果你有一列标量s (N)和一个矩阵a (NxN)如果你做diag(s)* a你基本上是在用这些标量做a行的元素乘法。

公式1(为了方便再次粘贴在这里):

第12步的第一项所做的(用绿色下划线)是:更新了在同一行块中当前块之前的块的当前softmax估计。如果j=1(这是这一行的第一个块。

第一项乘以diag(l_i)是为了抵消之前迭代中除以的相同常数(这个常数隐藏在O_i中)。

表达式的第二项(黄色下划线)是不需要消去的，因为可以看到我们直接将P~_i_j矩阵与V向量块(V_j)相乘。

e^x项是用来修改矩阵P~_i_j & O_i的，方法是消去前一次迭代中的m，用最新的估计(m_new_i)来更新它，该估计包含到目前为止逐行最大值。

以下是我的逐步分析(实际上只需要5分钟，希望能有所帮助!)

重点是这些外面的e项和P/O矩阵里面的e项消掉了，所以总是得到最新的m_new_1估计!

第三次迭代也是类似的，得到了正确的最终结果!

回想一下:这只是对最终O_i的当前估计。只有在我们遍历上图中的所有红色块之后，我们才能最终得到确切的结果。

第13步

将最新的累加到统计数据(l_i & m_i)写回HBM。注意它们的维数是B_r。

第13、14、15、1步

嵌套的for循环结束，O (Nxd)将包含最终结果:每个输入令牌的注意力加权值向量!

简单汇总

算法可以很容易地扩展到“block-sparse FlashAttention”，这是一种比FlashAttention快2-4的稀疏注意力算法，扩展到64k的序列长度!通过使用一个块形式的掩码矩阵，可以跳过上面嵌套的for循环中的某些加载/存储，这样我们可以按比例节省稀疏系数，比如下图

现在让我们简单地讨论一下复杂性。

复杂度分析

空间:在HBM中分配了Q, K, V, O (Nxd)， l和m (N)。等于4Nd + 2*N。去掉常量，并且知道d也是一个常量并且通常比N小得多(例如d={32,64,128}， N={1024，…，100k})，可以得到O(N)的空间，这有助于扩展到64k序列长度(再加上一些其他“技巧”，比如ALiBi)。

时间:这里不会严格地进行时间复杂度分析，但是我们将使用一个好的指标:HBM访问的数量。

论文的解释如下：

他们是怎么得到这个数字的?让我们来分析嵌套的for循环:

我们的块大小是M/4d。这意味着向量被分割成N/(M/4d)块。取它的2次方(因为要遍历行/列块)得到O(N²d²/ M²)

我们不能一次获取整个块，如果做一个大O分析，可能会让我们认为这并不比标准注意力好多少，但对于典型的数字，这导致访问次数减少了9倍(根据上面的论文截图)。

我们的伪算法集中在一个单头注意力，假设批处理大小为1。下面我们就开始进行扩展了

多头注意力

要扩展到batch_size > 1和num_heads > 1实际上并不难。

算法基本上是由单个线程块(CUDA编程术语)处理的。这个线程块在单个流多处理器(SM)上执行(例如，A100上有108个这样的处理器)。为了并行化计算，只需要在不同的SMs上并行运行batch_size * num_heads线程块。该数字与系统上可用的SMs数量越接近，利用率就越高(理想情况下是多个，因为每个SM可以运行多个线程块)。

反向传播

对于GPU内存的占用，另外一个大头就是反向传播，通过存储输出O (Nxd)和softmax归一化统计数据(N)，我们可以直接从SRAM中的Q, K和V (Nxd)块中反向计算注意力矩阵S (NxN)和P (NxN) !从而使内存保持在O (N)。这个比较专业了，我们了解以下就可以了，所以需要详细的内容请看原论文。

代码实现

最后，让我们看看在使用flash attention时可能出现的一些问题。因为涉及到显存的操作，所以我们只能深入CUDA，但是CUDA又比较复杂。

这就是OpenAI的Triton等项目的优势(参见他们的FlashAttention实现)。Triton基本上是一种DSL(领域特定语言)，介于CUDA和其他领域特定语言(例如TVM)之间的抽象级别。可以编写超级优化的Python代码(一旦编译)，而不必直接处理CUDA。这样Python代码可以部署在任意的加速器上(这是Triton任务)。

另外一个好消息是Triton最近已经与PyTorch 2.0集成了。

另外对于某些用例，比如对于超过1K的序列长度，一些近似注意方法(如Linformer)开始变得更快。但是flash attention的块稀疏实现优于所有其他方法。