FlashAttention - Ming's Blog

Note

高性能算子系列博客

在介绍 FlashAttention 之前，先统一本文使用的符号。

符号	含义
$N$	Sequence Length，序列长度
$D$	Head Dimension，每个 Attention Head 的特征维度
$Q$	Query 矩阵
$K$	Key 矩阵
$V$	Value 矩阵

为什么 Attention 很慢？#

Transformer 中最重要的计算之一就是 Self-Attention, 从公式上看，Attention 主要由两个矩阵乘法和一个 Softmax 组成：

\begin{aligned} S &= Q K^T \\ P &= \operatorname{Softmax}(S) \\ O &= PV \end{aligned}

Q,K,V\in\mathbb{R}^{N\times D}

而：

S,P\in\mathbb{R}^{N\times N}

最终：

O\in\mathbb{R}^{N\times D}

1 / 3

Self-Attention 计算过程

而问题在于这个 $N \times N$ 的中间矩阵，当 Sequence Length $N$ 增大时，Attention Matrix 的大小按照 $N^2$ 增长，当序列长度从 1024 增长到 8192 时，这个中间矩阵的元素数量会增加 64 倍。

FlashAttention 的核心思想并不是减少 Attention 的 FLOPs 计算量，而是：

Important

不要把巨大的 Attention Matrix 写回 GPU 的 HBM，而是把计算拆成小块，让数据尽可能留在 GPU 的片上高速存储中。

在 Standard Attention 实现中，中间矩阵 $S$ ， $P$ 只是为了下一个计算阶段使用，却被反复写入和读取 HBM，后面的 FlashAttention 就是围绕如何减少这些 HBM IO 展开的。

1 / 4

Standard Attention 实现

IO-Aware Algorithm：Tiling#

FlashAttention 所指的 IO-Aware 思想实际上就是 Tiling。在矩阵乘法的优化中，我们介绍了 Tiling，比如矩阵 $A$ 乘以矩阵 $B$ ，我们将 $A$ 矩阵和 $B$ 矩阵进行分块，逐块加载到 shared memory 中，读取 shared memory 完成计算，而不是直接读取 global memory (HBM)。

1 / 5

矩阵乘法 Tiling

对于 Attention 计算 $S = QK^T$ 我们也可以使用 Tiling，一块一块的计算 Score Matrix。

Online Softmax#

先回顾普通的 Safe Softmax，三次遍历分别求最大值、求指数和以及计算结果。

FlashAttention 不希望保存完整的 Score Matrix，因为它很占空间，而标准的 Softmax 实现需要完整的 Score Matrix 的一行，因此我们需要优化 Softmax 的实现让它也可以一块一块地处理。

假设已经处理：

x_1,\cdots,x_k

定义两个变量 running max $m$ 以及 running sum $l$ ：

m_k=\max_{1\le i\le k}x_i

l_k= \sum_{i=1}^{k}e^{x_i-m_k}

Important

$l_N$ 就是我们想要的 Softmax 的分母。

现在来了新的 block：

x_{k+1},\cdots,x_{k+b}

新的最大值：

m_{new} = m_{k + b} = \max \left( m_k, \max(x_{k+1},\cdots,x_{k+b}) \right) = \max \left( m_k, m_{block} \right)

由于旧的 $l_k$ 是按照 $m_k$ 计算的，因此需要引入纠正系数：

\text{correction factor} = e^{m_k-m_{new}}

不难得到：

l_{new} = l_{k + b} = l_k e^{m_k-m_{new}} + \sum_{j = k + 1}^{k + b} e^{x_j-m_{new}} = l_k e^{m_k-m_{new}} + l_{block}

这样，我们只需要维护：

\boxed{m,l}

就可以将 Softmax 的前两次循环，求最大值循环和求和循环，压缩为一次循环。

我们实际上把三次循环缩减到了两次循环，而这就是 Online Softmax。

Flash Attention#

Online Softmax 还不够，我们真正需要的是 $O = PV$ ，我们希望 $O$ 也可以在线更新。

假设我们将 $Q \in \mathbb{R}^{N \times D}$ 切分为一个一个的 $Q_{block} \in \mathbb{R}^{\text{tile\_q} \times D}$ ，同理将 $K \in \mathbb{R}^{N \times D}$ 和 $V \in \mathbb{R}^{N \times D}$ 分别切分为一个一个的 $K_{block} \in \mathbb{R}^{\text{tile\_k} \times D}$ 和 $V_{block} \in \mathbb{R}^{\text{tile\_k} \times D}$ 。

我们先考虑 $\text{tile\_q} = 1, \text{tile\_k} = b$ 的情况，这个时候我们的 $O_{block} \in \mathbb{R}^{D}$ 和 $Q_{block} \in \mathbb{R}^{D}$ 变成了一维向量，记作 $o$ 与 $q$ 。在 CUDA 实现中，每一个线程块处理到一个 $q$ ，我们后续推导专注一个线程块上的计算，所以默认 $q$ 是不变的。

假设我们处理到某个 $S_{block}$ 时已经处理了：

x_1, \cdots, x_k

其中 $x_i = q k^T_i$ ，其中 $k^T_i$ 是 $K^T$ 矩阵中的一个列向量。

我们在 online softmax 中定义了 running max $m$ ，running row sum $l$ ，也就是:

m_k=\max_{1\le i\le k}x_i

l_k= \sum_{i=1}^{k}e^{x_i-m_k}

当处理完 $S_{block}$ 时会引入：

x_{k + 1}, \cdots, x_{k + b}

我们可以得到新的 $m$ 和 $l$ :

m_{new} = m_{k + b} = \max \left( m_k, m_{block} \right)

l_{new} = l_{k + b} = l_k e^{m_k-m_{new}} + l_{block}

现在考虑 $o$ ，它是：

o = \sum_{j = 1}^N \frac{e^{x_j - m_N}}{l_{N}} v_j

其中 $v_j \in \mathbb{R}^D$ 是 $V$ 中的一个行向量。

我们想让 $o$ 也能在线更新，于是我们定义新的变量 $o_i$ ：

o_i = \sum_{j = 1}^i e^{x_j - m_i} v_j

那么自然就有

o = o_N / l_N

现在我们推导 $o_k$ 和 $o_{k + b}$ 的关系：

\begin{aligned} o_{k + b} &= \sum_{j = 1}^{k + b} e^{x_j - m_{k + b}} v_j \\ &= \sum_{j = 1}^{k} e^{x_j - m_{k + b}} v_j + \sum_{j = k + 1}^{k + b} e^{x_j - m_{k + b}} v_j \\ &= \left(\sum_{j = 1}^{k} e^{x_j - m_{k}} v_j \right) e^{m_{k} - m_{k + b}} + \sum_{j = k + 1}^{k + b} e^{x_j - m_{k + b}} v_j \\ &= o_k \times e^{m_{k} - m_{k + b}} + \sum_{j = k + 1}^{k + b} e^{x_j - m_{k + b}} v_j \\ &= o_k \times e^{m_{k} - m_{new}} + \begin{bmatrix} e^{x_{k+1} - m_{\text{new}}} & e^{x_{k+2} - m_{\text{new}}} & \cdots & e^{x_{k+b} - m_{\text{new}}} \end{bmatrix} \begin{bmatrix} v_{k+1} \\ v_{k+2} \\ \vdots \\ v_{k+b} \end{bmatrix} \end{aligned}

太好了，我们将 $o_{old}$ 和 $o_{new}$ 联系到一起了，现在我们可以在线更新 $o$ 了!

而且我们可以很轻松地扩展到 $\text{tile\_q} \neq 1$ 的情况，Flash Attention 的计算过程如下：

完整的 Flash Attention 实现#

见我的实现。

为什么 Attention 很慢？#

IO-Aware Algorithm：Tiling#

Online Softmax#

Flash Attention#

完整的 Flash Attention 实现#

参考#

支持与分享

目录