前言

最近在学习 FlashAttention，看到一份不错的手稿分享下🤗

manuscript：From Online Softmax to FlashAttention

0. Abstract

FlashAttention 的关键创新是使用类似于 Online Softmax 的思想来 tile 分块 self-attention 的计算，这样可以融合整个多头注意力层，而无需多次重复访问 GPU 的全局内存来获取临时变量和注意力分数矩阵 $A$。本文将简要解释如何 tiling 分块 self-attention 的计算，以及如何从 online softmax 中推导出 FlashAttention 计算

1. The Self-Attention

Self-Attention 的计算可以描述为：

$$\begin{array}{c}O=softmax\left(Q{K}^T\right)V\end{array}$$

其中 $Q,K,V,O\in {\mathbb{R}}^{L\times D}$，$L$ 是序列长度，$D$ 是每个头的维度，softmax 将按列应用于 $Q{K}^T$

Note：这里我们忽略了多头、多 batch，因为在这些维度上的计算是完全并行的，也就是说我们只关注单头、单 batch 的计算就行。另外为了简单起见，我们也忽略了注意力掩码以及缩放因子 $\frac{1}{\sqrt{D}}$

计算 self-attention 的标准方法是将其分解为以下几个阶段：

$$\begin{array}{rl}X& =Q{K}^T\\ A& =softmax(X)\\ O& =AV\end{array}$$

我们称 $X$ 矩阵为 pre-softmax logits，$A$ 矩阵为注意力分数（attention score），$O$ 矩阵为输出

FlashAttention 的一个惊人之处在于，我们不需要在全局内存（global memory）上实现 $X$ 和 $A$ 矩阵，而是将公式 (1) 中的整个计算融合到单个 CUDA kernel 中。这要求我们设计一种能够精心管理片上内存（on-chip memory）的算法，因为 NVIDIA GPU 的共享内存（shared memory）非常小

对于矩阵乘法等经典算法，我们通常会使用 tiling 技术来确保片上内存不超过硬件限制。下图提供了一个例子，在 kernel 执行期间，无论矩阵形状如何，只有 $3T^2$ 个元素存储在片上。这种 tiling 方法之所以有效，是因为加法是关联的，允许将整个矩阵乘法分解为许多 tile-wise 矩阵乘法的总和

然而，Self-Attention 包含一个不直接关联的 softmax 运算符，因此很难像下图那样简单地对 Self-Attention 进行 tile，那有没有办法让 softmax 具有关联性呢？🤔

上图简要说明了如何对矩阵乘法 $C=A\times B$ 的输入和输出矩阵进行 tile（分块），矩阵被划分为 $T\times T$ 个小块。对于每个输出块，我们从左到右遍历 $A$ 中的相关块，从上到下遍历 $B$ 中的相关块，并将它们的值从全局内存（global memory）加载到片上内存（on-chip memory）（以蓝色表示，整体片上内存占用为 $O(T^2)$）

接着逐块进行矩阵乘法，对于位置 $(i,j)$，我们先从片上内存中加载块内所有 $i$ 行的元素即 $A[i,k]$ 以及所有 $j$ 列的元素即 $B[k,j]$（以红色表示），然后利用 $A[i,k]$ 和 $B[k,j]$ 计算得到片上内存中的 $C[i,j]$。以此循环，直到完成一个块的计算后，我们再将片上 $C$ 块的结果写回主内存并继续处理下一个块

实际 tiling 的应用要复杂得多，大家可以参考：cutlass 在 A100 上实现矩阵乘法

2. (Safe) Softmax

我们先回顾一下 softmax 算子，下面是 softmax 计算的通用公式：

s o f t m a x ( { x 1 , … , x N } ) = { e x i ∑ j = 1 N e x j } i = 1 N \begin{equation} \mathrm{softmax}(\{x_1,\ldots,x_N\})=\left\{\frac{e^{x_i}}{\sum_{j=1}^Ne^{x_j}}\right\}_{i=1}^N \end{equation} softmax({x1​,…,xN​})={∑j=1N​exj​exi​​}i=1N​​​

值得注意的是，当 $x_i$ 比较大时 $e^{x_i}$ 很容易溢出，例如 float16 可以支持的最大值是 65536，这意味着当 $x⩾11$ 时，$e^x$ 将超过 float16 的有效范围

为了解决这个问题，像 pytorch、tensorflow 等框架通常会使用一种被称为 “safe” 的 softmax 计算方法：

$$\begin{array}{c}\frac{e^{x_i}}{{\sum }_{j=1}^N{e}^{x_j}}=\frac{e^{x_i-m}}{{\sum }_{j=1}^N{e}^{x_j-m}}\end{array}$$

其中 $m={\max }_{j=1}^N(x_j)$

我们将 $x_i$ 减去它们中的最大值，这样可以确保每个元素 $x_i-m⩽0$，从而再做指数运算时可以确保其值不会溢出

我们可以将 safe softmax 的计算总结为下面的 3-pass 算法：

Algorithm 3-pass safe softmax

NOTATIONS

• { m i } \{m_i\} {mi​}： max ⁡ j = 1 i { x j } \max_{j=1}^i\{x_j\} maxj=1i​{xj​}，初始值 $m_0=-\infty$
• { d i } \{d_i\} {di​}：${\sum }_{j=1}^i{e}^{x_j-m_N}$，初始值 $d_0=0$，$d_N$ 是 safe softmax 的分母
• { a i } \{a_i\} {ai​}：最终的 softmax 值

BODY

for $i\leftarrow 1,N$ do

$$\begin{array}{c}{m}_i\leftarrow \max (m_{i-1},x_i)\end{array}$$

end

for $i\leftarrow 1,N$ do

$$\begin{array}{c}{d}_i\leftarrow d_{i-1}+e^{x_i-m_N}\end{array}$$

end

for $i\leftarrow 1,N$ do

$$\begin{array}{c}{a}_i\leftarrow \frac{e^{x_i-m_N}}{d_N}\end{array}$$

end

该算法要求我们对 $[1,N]$ 进行 3 次迭代，在 Transformer 的 self-attention 中，每个 $x_i$ 是通过 $Q{K}^T$ 计算得到的。如果我们无法将所有的 $x_i$ 都缓存到片上内存（SRAM）中（实际上我们也没有足够大的片上内存 SRAM 来容纳所有的 $x_i$），那么在每次迭代时不得不访问 $Q,K$ 以动态重新计算 $x_i$，这种频繁的访问会导致大量的 I/O 操作，从而降低整体效率

3. Online Softmax

如果我们能够在一个循环中融合公式 (7)、(8)、(9)，那么就可以将 global memory 的访问时间降低 3 倍，不幸的是，我们不能在同一个循环中融合公式 (7)、(8)，因为公式 (8) 的计算需要依赖 $m_N$，而 $m_N$ 只有在第一个循环完成之后才能够确定

既然数据之间存在着依赖关系，那我们可以创建另外一个序列 $d_i^{\prime }:={\sum }_{j=1}^i{e}^{x_j-m_i}$ 作为原始序列 $d_i^{\prime }:={\sum }_{j=1}^i{e}^{x_j-m_N}$ 的替代，以消除对 $N$ 的依赖，并且这两个序列的第 $N$ 项是相同的即 $d_N=d_N^{\prime }$，因此我们可以安全地用 $d_N^{\prime }$ 替换公式 (9) 中的 $d_N$

Note：在数学和计算机科学中，符号 $:=$ 用于表示 “定义为” 或 “被定义为”，例如当我们写下 $x:=y$ 时，我们是在说明 “我们将 $x$ 定义为 $y$” 或 “$x$ 等于 $y$（这是它的定义）”

此外，我们还可以找到 $d_i^{\prime }$ 和 $d_{i-1}^{\prime }$ 之间的递归关系：

$$\begin{array}{c}\begin{array}{rl}{d}_i^{\prime }& =\sum _{j=1}^i{e}^{x_j-m_i}\\ & =\left(\sum _{j=1}^{i-1}{e}^{x_j-m_i}\right)+e^{x_i-m_i}\\ & =\left(\sum _{j=1}^{i-1}{e}^{x_j-m_{i-1}}\right)e^{m_{i-1}-m_i}+e^{x_i-m_i}\\ & =d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}+e^{x_i-m_i}\end{array}\end{array}$$

这个递归公式只依赖于 $m_i$ 和 $m_{i-1}$，此外我们还可以在一个循环中同时计算 $m_j$ 和 $d_j^{\prime }$

Algorithm 2-pass online softmax

for $i\leftarrow 1,N$ do

$$\begin{array}{rl}{m}_i& \leftarrow \max (m_{i-1},x_i)\\ d_i^{\prime }& \leftarrow d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}+e^{x_i-m_i}\end{array}$$

end

for $i\leftarrow 1,N$ do

$$a_i\leftarrow \frac{e^{x_i-m_N}}{d_N^{\prime }}$$

end

这是 Online Softmax 论文中提出的算法，但它仍然需要两次循环才能完成 softmax 的计算，我们能否将循环次数减少到 1 次来最小化全局的 I/O 呢？🤔

4. FlashAttention

不幸的是，对于 softmax 来说，答案是否定的，但在 Self-Attention 中，我们最终的目标并不是注意力分数矩阵 $A=softmax(Q{K}^T)$，而是输出矩阵 $O=A\times V$，既然无法找到 softmax 的一次递归形式，那我们换个思路思考下能否找到矩阵 $O$ 的一次递归形式呢？

下面我们尝试下先将 Self-Attention 计算的第 $k$ 行公式转化为递归算法：

Note：由于所有行的计算都是独立的，为了简单起见，这里我们只解释一行的计算

Algorithm Multi-pass Self-Attention

NOTATIONS

• $Q[k,:]$：查询矩阵 $Q$ 的第 $k$ 行向量
• $K^T[:,i]$：键矩阵 $K^T$ 的第 $i$ 列向量
• $O[k,:]$：输出矩阵 $O$ 的第 $k$ 行
• $V[i,:]$：值矩阵 $V$ 的第 $i$ 行
• ${\mathbf{o}}_i$：${\sum }_{j=1}^i{a}_{j}V[j,:]$，存储 $A[k,:i]\times V[:i,:]$ 结果的行向量

BODY

for $i\leftarrow 1,N$ do

$$\begin{array}{rl}{x}_i& \leftarrow Q[k,:]K^T[:,i]\\ m_i& \leftarrow \max (m_{i-1},x_i)\\ d_i^{\prime }& \leftarrow d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}+e^{x_i-m_i}\end{array}$$

end

for $i\leftarrow 1,N$ do

$$\begin{array}{rl}{a}_i& \leftarrow \frac{e^{x_i-m_N}}{d_N^{\prime }}\\ {\mathbf{o}}_i& \leftarrow {\mathbf{o}}_{i-1}+a_{i}V[i,:]\end{array}$$

end

$$O[k,:]\leftarrow {\mathbf{o}}_N$$

我们将公式 (12) 中的 $a_i$ 替换公式 (11) 中的定义，有：

$$\begin{array}{r}{\mathbf{o}}_i:=\sum _{j=1}^i\left(\frac{e^{x_j-m_N}}{d_N^{\prime }}V[j,:]\right)\end{array}$$

该公式仍然依赖 $m_N$ 和 $d_N$ 变量，这两个值在前一个循环完成之前无法确定。但我们可以再次使用第 3 节中的替代技巧，即创建替代序列 ${\mathbf{o}}^{\prime }$：

$${\mathbf{o}}_i^{\prime }:=\left(\sum _{j=1}^i\frac{e^{x_j-m_i}}{d_i^{\prime }}V[j,:]\right)$$

$\mathbf{o}$ 和 ${\mathbf{o}}^{\prime }$ 的第 $N$ 个元素相同即 ${\mathbf{o}}_N^{\prime }={\mathbf{o}}_N$ ，并且我们可以发现 ${\mathbf{o}}_i^{\prime }$ 和 ${\mathbf{o}}_{i-1}^{\prime }$ 之间存在如下的递归关系：

$$\begin{array}{c}\begin{array}{rl}{\mathbf{o}}_i^{\prime }& =\sum _{j=1}^i\frac{e^{x_j-m_i}}{d_i^{\prime }}V[j,:]\\ & =\left(\sum _{j=1}^{i-1}\frac{e^{x_j-m_i}}{d_i^{\prime }}V[j,:]\right)+\frac{e^{x_i-m_i}}{d_i^{\prime }}V[i,:]\\ & =\left(\sum _{j=1}^{i-1}\frac{e^{x_j-m_{i-1}}}{d_{i-1}^{\prime }}\frac{e^{x_j-m_i}}{e^{x_j-m_{i-1}}}\frac{d_{i-1}^{\prime }}{d_i^{\prime }}V[j,:]\right)+\frac{e^{x_i-m_i}}{d_i^{\prime }}V[i,:]\\ & =\left(\sum _{j=1}^{i-1}\frac{e^{x_j-m_{i-1}}}{d_{i-1}^{\prime }}V[j,:]\right)\frac{d_{i-1}^{\prime }}{d_i^{\prime }}{e}^{m_{i-1}-m_i}+\frac{e^{x_i-m_i}}{d_i^{\prime }}V[i,:]\\ & ={\mathbf{o}}_{i-1}^{\prime }\frac{d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}}{d_i^{\prime }}+\frac{e^{x_i-m_i}}{d_i^{\prime }}V[i,:]\end{array}\end{array}$$

它仅依赖于 $d_i^{\prime },d_{i-1}^{\prime },m_i,m_{i-1},x_i$ ，因此我们可以在单个循环中融合 Self-Attention 中的所有计算

Algorithm FlashAttention

for $i\leftarrow 1,N$ do

$$\begin{array}{rl}{x}_i& \leftarrow Q[k,:]K^T[:,i]\\ m_i& \leftarrow \max \left(m_{i-1},x_i\right)\\ d_i^{\prime }& \leftarrow d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}+e^{x_i-m_i}\\ o_i^{\prime }& \leftarrow o_{i-1}^{\prime }\frac{d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}}{d_i^{\prime }}+\frac{e^{x_i-m_i}}{d_i^{\prime }}V[i,:]\end{array}$$

end

$$O[k,:]\leftarrow {\mathbf{o}}_N^{\prime }$$

状态量 $x_i,m_i,d_i^{\prime },{\mathbf{o}}_i^{\prime }$ 占用的内存都很小，可以非常轻松的放入 GPU 的 shared memory 中。另外由于此算法中的所有操作都是关联的，因此它与 tiling 兼容，如果我们逐个 tiling 计算，则该算法可以表示为：

Algorithm FlashAttention（Tiling）

NEW NOTATIONS

• $b$：tile 的 block size 大小
• $\text{\#tiles}$：每行 tile 的数量，$N=b\times \text{\#tiles}$
• ${\mathbf{x}}_i$：存储第 $i$ 个 tile 的 $Q[k]K^T$ 值的向量 $[(i-1)b:ib]$
• $m_i^{(\mathrm{local})}$： ${\mathbf{x}}_i$ 内部的局部最大值

BODY

for $i\leftarrow 1,\text{\#tile}$ do

$$\begin{array}{rl}{\mathbf{x}}_i& \leftarrow Q[k,:]K^T[:,(i-1)b:ib]\\ m_i^{\text{(local)}}& =\underset{j=1}{\overset{b}{\max }}({\mathbf{x}}_i[j])\\ m_i& \leftarrow \max (m_{i-1},m_i^{\text{(local)}})\\ d_i^{\prime }& \leftarrow d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}+\sum _{j=1}^b{e}^{{\mathbf{x}}_i[j]-m_i}\\ {\mathbf{o}}_i^{\prime }& \leftarrow {\mathbf{o}}_{i-1}^{\prime }\frac{d_{i-1}^{\prime }{e}^{m_{i-1}-m_i}}{d_i^{\prime }}+\sum _{j=1}^b\frac{e^{{\mathbf{x}}_i[j]-m_i}}{d_i^{\prime }}V[j+(i-1)b,:]\end{array}$$

end

$$O[k,:]\leftarrow {\mathbf{o}}_{N/b}^{\prime }$$

下图说明了如何将该算法应用到硬件上：

上图说明了 FlashAttention 在硬件上的计算方式，其中蓝色块表示在 SRAM 中的 tile，而红色块对应于 tile 中的第 $i$ 行，$L$ 表示序列长度，$D$ 表示每个头的维度，$B$ 表示 block size 的大小

值得注意的是，整体 SRAM 的内存占用仅仅取决于 $B$ 和 $D$，与 $L$ 无关，因此，该算法可以扩展到长上下文而不会遇到内存问题（GPU 的共享内存很小，例如 H100 架构中为 228kb/SM）。在计算过程中，我们从左到右遍历 $K^T$ 和 $A$，从上到下遍历 $V$，并相应地更新 $m,d,O$ 的状态

结语

这篇文章主要分享了 FlashAttention 的思想，FlashAttention 想要做的就是在单个循环中完成 self-attention 的整个计算

我们先从 safe softmax 出发分析了 safe softmax 算法需要循环迭代 3 次，第一次循环求取序列中的最大值 $m_N$，第二次循环求取 safe softmax 的分母 $d_N$，第三次循环求取最终的 softmax 值 $a_i$。由于 SRAM 片上内存空间有限，我们无法将所有的 $x_i$ 缓存，导致我们在每次迭代时需要重新访问 $Q,K$ 以动态重新计算 $x_i$，这造成了大量的 I/O 操作，我们并不希望看到

在 online softmax 中提出了一个有意思的想法，那就是创建一个新的序列 $d_i^{\prime }$ 来替换原始的 $d_i$，主要区别在于将原来 $d_i$ 定义中的 $m_N$ 替换为了 $m_i$，以消除对 $N$ 的依赖，并且我们发现 $d_N=d_N^{\prime }$，那么整个 safe softmax 的计算就只需要迭代 2 次即可，第一次循环迭代可以同时计算 $m_i$ 和 $d_i^{\prime }$

FlashAttention 想要把 self-attention 中的循环次数降低到一次，但是发现对于其中的 softmax 而言无法实现，因此它另辟蹊径试图寻找输出矩阵 $O$ 的一次迭代形式。经过推导我们发现输出行向量 ${\mathbf{o}}_i$ 仍然依赖于 $m_N$ 和 $d_N$ 变量，这两个值在前一个循环完成之前无法确定，因此仍需要两次循环

FlashAttention 借助 online softmax 的替代思想，将序列 $\mathbf{o}$ 替换为 ${\mathbf{o}}^{\prime }$，替换之后发现 ${\mathbf{o}}_i^{\prime }$ 不再依赖于 $m_N$ 和 $d_N$ 等变量，而仅依赖于 $d_i^{\prime },d_{i-1}^{\prime },m_i,m_{i-1},x_i$，从而可以在单个循环中融合 self-attention 中的所有计算

最后，由于 flashattention 算法中的所有操作都是关联的，因此可以逐个 tiling 分块计算，并且整体 SRAM 的内存占用与序列长度无关，因此可以扩展到长上下文而不用担心遇到内存问题

OK，以上就是关于这篇手稿的全部内容了，大家感兴趣的可以看看，作为 FlashAttention 的入门还是没有问题的😄

参考

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Gtc 2020: developing cuda kernels to push tensor cores to the absolute limit on nvidia a100
• Online normalizer calculation for softmax