FlashAttention 的第一原则

我相信你已经听过很多次了：“注意力机制在序列长度为 N 时会变得很差，确切地说是 O(N²)”。但究竟是什么导致了序列长度的变差呢？为什么会变差？问题真的出在序列长度上吗？在本系列文章中，我们将深入探讨Tri Dao 等人 (2022) 撰写的 FlashAttention 论文^[1]，探讨为什么注意力机制在现代 GPU 上如此缓慢且耗费资源，以及 FlashAttention 是如何解决这个问题的。

这原本只是一篇简短的论文评论，但随着我逐渐深入挖掘，我发现市面上大多数信息都非常肤浅，而且大多数资源都只是要求你了解注意力机制、GPU 架构和 CUDA 编程模型的细节，并且还要掌握大量的数学知识。因此，我决定将文章分成两部分。

在第一部分中，我们将介绍注意力机制、GPU 架构和 CUDA 编程模型，为什么注意力在现代 GPU 上如此缓慢且耗费资源，以及是什么阻碍了我们应用 CUDA 编程中常见的优化。

在第二部分中，我们将介绍 FlashAttention，以第一部分获得的知识为基础，深入研究前向和后向传递的算法细节，一步一步地讲解。

1. 我们为什么要关心？

在深入探讨细节之前，我们先来思考一下为什么我们应该关注这个话题。注意力机制是 Transformer 架构的关键组成部分，而 Transformer 架构又是当今几乎所有深度学习领域的支柱。随着模型规模越来越大，训练和推理模型的成本也越来越高，无论是在时间上还是金钱上。因此，如果我们能让注意力机制更快、更高效，就能在更短的时间内以更少的成本训练和推理更大的模型，甚至可以训练以前无法训练的模型。

我认为下图很好地说明了这一点。它展示了 FlashAttention 与之前方法相比的性能，以及它如何随着序列长度N的变化而变化。

如你所见，FlashAttention 比以前的方法快得多，并且它可以处理以前的方法内存不足的序列长度。

考虑到这一点，我们首先简单解释一下注意力机制。

2. 简要解释注意力

注意力机制是 Transformer 架构的关键组成部分，由 Vaswani 等人在 2017 年发表的《注意力就是一切》（Attention is All You Need）一文中提出。从宏观层面来看，注意力机制使模型能够动态地关注输入序列的不同部分，并在处理特定 token 时忽略其他部分。token 是输入序列中的单个元素，可以是单词、像素块，甚至是图中的一个节点。

对于输入序列中的每个标记，都会根据序列中的所有其他标记计算一个注意力分数。

最初，注意力机制是为了处理文本序列而引入的，但后来它被应用于各种任务和数据模式，例如图像、视频和图形，成为一种用于捕获序列中元素之间依赖关系的通用机制。

看一下下面的例子：

非常简单，我喜欢把注意力想象成一种向数据提问的机制，而注意力分数则是该问题的答案。注意力分数越高，该信息对当前标记就越重要。多头注意力（MHA）本质上就是同时针对同一个标记提出多个问题，然后将答案组合起来。稍后会详细介绍。

下面更详细地了解一下 Transformer 编码器架构和注意力机制。

Transformer 的编码器由多个相同的阶段堆叠而成。每个阶段包含两个子阶段：多头自注意力 (MHA) 和多层感知器 (MLP)。

MHA 子阶段实现多个并行注意力头，然后将它们连接起来并线性转换为预期的隐藏维度。

深入研究多头注意力层的单个头，我们可以看到注意力机制的架构，可以用以下等式来描述：

请注意，Dropout 和 masking 是可选的，但由于 FlashAttention 提到了它们，所以我也会将它们包括在这里。

由于本文并非讨论注意力机制本身，因此我不会在此进行更详细的介绍。如果你有兴趣了解更多关于注意力机制的知识，我建议你阅读Vaswani 等人（2017）的原始论文^[2]或Jay Alammar 的著名博客文章^[3]。

我们感兴趣的是，为什么注意力机制在现代 GPU 上如此缓慢且耗费资源。为了理解这一点，我们先来研究一下注意力机制中涉及的矩阵的形状。

这里的关键点在于，注意力机制在中间结果中具体化了NxN矩阵，而我们感兴趣的输出仅仅是Nxd形状，其中N是序列长度，d 是隐藏层维度。正因如此，Transformer 的二次运算和内存复杂度才被宣称为O(N²)，相对于序列长度N而言。虽然这个例子提到的是N=4,096的典型值（当时），但论文将其放大到N=65,536，这将导致中间结果值达到 42.9 亿！

请注意，图 5 展示了单层中单个头的注意力机制。实际上，Transformer 会使用多个头和多层，这进一步大幅增加了计算复杂度。

为了进一步理解为什么注意力机制在现代 GPU 上如此缓慢且耗费资源，我们首先需要简单介绍一下 GPU 的硬件架构以及一些 CUDA 编程的基础知识。我保证，我们会尽量简短地讲解。😉

3. 现代 GPU 和 CUDA 设备架构

3.1. GPU架构

我们先来比较一下 CPU 和 GPU 的架构。CPU 由几个针对顺序处理进行优化的核心组成。相比之下，GPU 则拥有大规模并行架构，由数千个更小、更高效的核心组成，旨在同时处理多个任务。CPU 和 GPU 都拥有大小和速度各异的内存层次结构。这一点对于后面讨论 Flash Attention 如何实现大幅加速至关重要。

这种通用 GPU 架构以及 CUDA 编程语言是由 Nvidia 的 Tesla 架构于 2006 年推出的，其许多核心概念至今仍存在于现代 GPU 中。我强烈建议你阅读Nvidia关于 Tesla 架构^[4]的白皮书。

GPU 架构围绕流多处理器 (SM) 构建，SM 由多个 CUDA 和 Tensor 核心、共享内存和寄存器组成。每个 SM 可以并行执行多个线程块。同一块内的线程可以使用共享内存相互通信，而不同块中的线程则不能通信。

稍微透露一下这能做什么：我们现在有如此多的线程，实际上是数千个线程，我们可以通过将每个线程分配给一小块数据的一小部分计算来并行化矩阵乘法和卷积等繁重的计算。

3.2. CUDA 编程模型

Nvidia GPU 使用 CUDA 编程模型进行编程，这是 C++ 编程语言的扩展。它允许开发人员编写 CUDA 内核，这些内核是由每个线程执行的小函数。这些线程被组织成块，然后这些块被组织成网格。

这种线程层次结构是 CUDA 的一个巧妙特性，因为它模拟了我们正在处理的数据的形状，即向量、矩阵和张量。此外，每个线程都知道其在网格和块中的位置，这使我们能够计算出其在正在处理的数据结构中的位置。

我们还可以将 CUDA 编程模型映射到 GPU 架构，以查看数据存储在层次结构的哪个级别以及线程如何访问它。

为了在 FlashAttention 中引入内核融合的概念，我们需要了解 CPU 如何与 GPU 交互，以及数据如何流动。从异构编程模型开始，简单来说，这意味着我们使用不同的设备（即 CPU 和 GPU）来执行代码。

CPU 执行串行代码，负责在主机和设备之间传输数据，并启动内核。另一方面，GPU 通过在设备上的线程块网格内启动多个线程来并行执行这些内核。结果被读回 CPU，然后 CPU 继续进行串行处理，并最终启动另一个内核。这里的关键点在于，在简单的实现中，每个内核都需要与设备之间传输数据。将其置于我们的注意力机制中，这些顺序执行的内核可以是矩阵乘法、缩放、掩码、SoftMax 和 Dropout。

放大单个内核执行的数据流，我们可以看到数据如何在 CPU 和 GPU 之间流动，以及内核（线程）如何访问 GPU 内存（全局内存）中的数据。

毋庸置疑，上图极其简化，但应该能让你大致了解。关键在于，数据需要多次传递，而且核心经常处于等待状态。

呼，信息量好大啊。我们来看一个例子来更清楚些。最突出的例子是矩阵乘法，它是深度学习中的一个关键运算。

3.3. CUDA 中的矩阵乘法

大家肯定知道，要将两个矩阵A和B相乘，我们需要对矩阵A的每一行与矩阵B的每一列进行点积。在一个简单（且极其低效）的实现中，我们可以在 CUDA 中复制此算法，为每个线程分配计算结果矩阵C中一个元素的任务。

请注意，计算结果矩阵C的整行需要矩阵A的一行和矩阵B的所有列。

再次，我们必须稍微绕个弯，了解一下内核的执行方式以及它在不同的内存带宽上的表现。每次启动内核时，都会构建线程块网格，将内核加载到各个核心中，并将数据从全局内存加载到处理器的寄存器中。然后执行内核，并将结果存储回全局内存中。一般来说，内存越大，距离处理器越远，速度也越慢。因此，全局内存是层次结构中最大但速度最慢的内存，而各个核心的寄存器是层次结构中最小但速度最快的内存。

为了给你一个概念（毕竟 A100 的 1.5 TB/s 听起来还不错），看看需要加载的绝对最小值是多少，才能将两个大小为NxN的矩阵相乘，其中N=4096。我们假设值是 float32 格式，这意味着每个元素为 4 个字节。

• 我们需要从矩阵A中读取4096x4096 个元素，从矩阵B中读取4096x4096个元素：2x4096x4096x4 字节 = 134MB。
• 我们需要在矩阵C中存储4096x4096 个元素：4096x4096x4 字节 = 64MB。

如果我们现在按照矩阵乘法的简单实现，我们实际上是为矩阵A的每一行加载整个矩阵B。这意味着我们要加载4096x4096个元素4096次，也就是134MB x 4096 = 549GB，这可是个不小的流量😅。而且这还只是针对我们N=4096 个token 的小上下文窗口而言的……

❓那么，了解了这些之后，有什么技巧可以让 GPU 上的计算更快、更高效呢？

3.4. 内核执行优化

你可能已经猜到了，我们主要关注的是减少对全局内存的读写次数，因为它是层次结构中最慢的内存，并且我们的内核经常等待数据加载。

有几种技术可以实现这一点，包括：

• 共享内存——将常用数据保存在靠近核心的位置。
• 块平铺——重新使用已加载的数据来计算多个输出。
• 核融合——避免写入和读取中间结果。

请注意，为了简化理解 FlashAttention 的贡献，我们将省略一些细节。如果你对 CUDA 中矩阵乘法的更详细解释以及可以应用哪些优化感兴趣，我推荐你阅读这篇精彩且图文并茂的Simon Boehm 博客文章^[5]。

但现在，讨论前面提到的优化，以及如何将它们应用于矩阵乘法。

共享内存

将经常使用的数据保存在靠近核心的位置。

我们要做的第一件事是利用层次结构中的不同内存级别。共享内存是一种小型、快速的内存，由同一块中的线程（跨单个流多处理器内的所有核心）共享。正如 FlashAttention 论文中针对 A100 GPU 所述，全局内存（40GB）的带宽为 1.5TB/s，而共享内存（20MB）的带宽为 19TB/s。

我们知道，块中的线程可以使用共享内存相互通信，并且块可以在不同的流多处理器上并行执行。

因此，对于我们处理的每个块，我们将数据从全局内存加载到共享内存中，然后块中的线程可以从共享内存中访问数据。在处理下一个块时，我们重复此过程。

块状平铺

重新使用已加载的数据来计算多个输出。

接下来，我们在一个线程中计算结果矩阵C的多个元素。这被称为“平铺”，这是一种减少全局内存读写次数的常用技术。其基本思想是将已经从共享内存加载到寄存器的数据在线程中重复使用，用于多次计算。

需要明确的是，我们执行的计算数量保持不变，只是从内存中读取的数据更少，重复使用了已加载的数据。这意味着我们等待数据的时间更少，但计算时间更多。我们改进了所谓的 “算术强度”，即操作次数与内存访问次数的比率。

内核融合

避免写入和读取中间结果。

最后，我们可以将多个内核融合成一个内核，以避免将中间结果写入全局内存，只需在下一个内核中再次读取它们。

重要提示：只有当核操作于同一块数据时，它们才能融合。当我们稍后讨论注意力机制时，记住这一点很重要。

有了关于注意力机制、CUDA 和 GPU 设备架构的所有知识，我们现在可以理解为什么注意力机制在现代 GPU 上如此缓慢且耗费资源！

4. 现代 GPU 上的注意力

4.1. 为什么注意力机制如此缓慢且耗费资源

你现在可能已经猜到了，我们在全局内存中写入和读取了太多数据。

再看一下注意力机制和所涉及的矩阵的形状，但这次是通过 GPU 架构和 CUDA 编程模型的视角。

我们可以在此做出几点观察：

1. 我们有 3 个输入矩阵（Q、K和V），每个矩阵的大小为Nxd，以及一个大小为Nxd的输出矩阵O。
2. 注意力机制在中间结果中实现了巨大的NxN矩阵，这些矩阵都需要存储起来，然后在执行下一个内核时再次读取。
3. 一些运算（例如矩阵乘法和 SoftMax）会将其各自输入的相同数据重复用于输出中的单个元素。
4. 其他操作（例如 Dropout 和 masking）都是逐元素应用的。

正如你所看到的，全局内存中存在大量昂贵的读写操作，最糟糕的是，我们甚至对它们不感兴趣，因为这些都是中间结果。

那么，看看如何将前面提到的优化应用到注意力机制中。或者更好的是，看看为什么我们不能应用它们（在深入研究 FlashAttention 是如何实现的之前）。

4.2. 如何改进注意力机制？

注意力机制的最大问题是，我们必须在中间结果中具体化NxN矩阵，然后必须将其写入和读取到全局内存中。那么，为什么不简单地将所涉及的核融合在一起，以避免写入和读取中间结果呢？

这是一个好主意，但是存在两个问题：

1. 并非所有这些内核都对同一个数据块进行操作，这是前面提到的内核融合所必需的。
2. 我们需要反向传播中的中间结果来计算梯度。

稍微放慢速度，我们来看看第一个问题。SoftMax 函数在实际应用中是如何实现的？我们先来看看NxN矩阵。

首先，请注意 SoftMax 是按行应用的，而不是整个矩阵！其次，更重要的是，要计算 SoftMax 矩阵P的单个元素，我们需要读取输入矩阵S相应行的所有元素。

现在，回想一下，矩阵乘法核应用了分块平铺和滑动窗口，将输入加载到共享内存中以计算部分和。这与 SoftMax 函数不兼容，因为我们需要该行的所有元素来计算输出中的单个元素。

每行应用 SoftMax 函数背后的直觉是什么？这些 NxN 行的每一行都与输入序列中的一个 token 相关。计算每行的 SoftMax 函数会得到一个分数，该分数表示序列中其他 token 对当前 token 的“关注”程度。换句话说，其他 token 对于当前 token 解决当前任务的重要性。

假设我们可以将 SoftMax 函数应用于与矩阵乘法核相同的数据块。这样我们就可以融合这些核，避免写入和读取中间结果。但我们仍然面临第二个问题：我们需要在反向传播中使用中间结果来计算梯度。

为了训练模型，我们通常会输入一批数据，计算针对某个目标的损失，然后应用反向传播来更新模型参数，以便下次输入数据时，模型能够表现得更好（希望如此）。应用链式法则，我们可以计算模型中每个单独操作的梯度。从整个模型中提取一个简化的注意力机制：

请注意，需要前向传播的中间结果来计算梯度。

所以，我们需要将中间结果保存在内存中，这意味着我们必须将它们写入全局内存，并在反向传播中再次读取它们。这正是我们一开始就想要避免的。所以，我们又回到了原点……

幸运的是，FlashAttention 解决了这些问题。不过现在，我们先把第一部分总结一下，让自己消化一下这些信息。

这只是深入探究 FlashAttention 及其背后工程的开始。我很想听听其他使用注意力机制或 GPU 优化的朋友们的意见——你们遇到了哪些挑战？你们有什么见解或探索过的替代方法吗？欢迎在下方留言。

如果你觉得这篇分析有用，点个赞就能帮助其他人也发现它。

5. 总结

我们想要的是：

• 使得注意力机制更加快速、高效。

❓为什么它这么慢并且占用大量资源？

• 我们必须在中间结果中实现NxN矩阵，然后将其写入全局内存并从中读取。

❓为什么我们不能简单地将内核融合在一起以避免写入和读取中间结果？

• 并非所有这些内核都对同一个数据块进行操作，这是内核融合所必需的。
• 我们需要反向传播中的中间结果来计算梯度。

在第 2 部分中，我们将介绍 FlashAttention，以我们在本部分获得的知识为基础，深入研究前向和后向传递的算法细节，一步一步地讲解。