---
author: Hung-yi Lee
date: '2026-03-23'
guest: ''
layout: post.njk
source: https://www.youtube.com/watch?v=vXb2QYOUzl4
speaker: Hung-yi Lee
tags:
  - flash-attention
  - gpu-optimization
  - language-model-inference
  - transformer-acceleration
  - memory-hierarchy
title: 加速大型语言模型：Flash Attention深度解析与GPU底层优化
summary: 本课程深入讲解了如何通过Flash Attention技术加速大型语言模型（LLM）的推理过程。核心在于优化GPU内存访问层级，减少HBM与SRAM之间的数据搬运次数。Flash Attention作为一种不改变计算结果、模型无关的即插即用技术，能显著提升Transformer架构的运行效率，尤其在处理长序列时效果更佳，避免了CUDA内存溢出问题。课程还通过实例代码对比了Flash Attention与传统方法的性能差异。
insight: ''
draft: true
series: ''
category: ai-ml
area: tech-engineering
project: []
people: []
companies_orgs:
  - Hugging Face
products_models:
  - Yi-34B
media_books: []
status: evergreen
---
大家好，我们现在开始上课。今天的课程旨在讲解如何加快**语言模型**（Language Model: 用于理解和生成人类语言的AI模型）的生成速度。本课预设大家已非常清楚**语言模型**内部的运作原理，例如**Transformer**（Transformer: 一种基于自注意力机制的神经网络架构）的工作方式。在已知**Transformer**运作原理的前提下，我们将探讨如何加速其生成过程，而不是训练过程。换言之，本课假设一个**语言模型**已被训练好，我们将聚焦于如何加速其推理阶段，也就是我们常说的**推论**（Inference: 语言模型训练完成后，用于生成新内容或预测的过程）。

首先，我们用两页投影片快速回顾一下**语言模型**的内部运作。众所周知，**语言模型**本质上在做“文字接龙”，你给它一个未完成的句子，它预测接下来应该产生哪个**token**（Token: 文本在模型中处理的最小单元，可以是词、字或字符）。当前主流**语言模型**的神经网络架构是**Transformer**。**Transformer**由多层组成，每层都包含一个名为**自注意力机制**（Self-Attention: Transformer 中用于处理序列中不同位置信息的核心机制）的机制。**自注意力机制**使得**Transformer**能够考虑整个输入**Sequence**（序列: 输入到模型中的一系列数据，如文本中的词语顺序）中的所有信息。

**自注意力机制**的工作原理是：输入一排向量（例如 X1 到 X5），它将输出另一排向量（O1 到 O5）。具体运作方式是，首先将 X1 到 X5 各自乘上三个不同的**Transformation Matrix**，将 X1 转换为 V1、K1、Q1，X2 转换为 V2、K2、Q2，以此类推。在大多数课程中，我们通常会省略这个从 X 转换为 QKV 的过程，直接假设你已了解输入的 X 会变成 QKV 三个向量。

接下来，以生成第四个位置的输出为例进行计算（其他位置的计算方式相同）。我们将第四个位置的**查询向量**（Query: 在自注意力机制中，用于查询其他所有键的向量）q4 与前面所有的**键向量**（Key: 在自注意力机制中，用于与查询向量进行匹配的向量）k1 到 k4 进行**点积**（Dot Product: 向量之间的一种乘法运算），计算它们的内积。这些内积计算结果我们记作 a1 到 a4。由于这些内积值可能从负无穷大到正无穷大，我们会通过**Softmax函数**（Softmax Function: 将任意实数向量压缩为概率分布的函数）对其进行**归一化**（Normalization: 将数据缩放到特定范围，使其符合某种分布），使其数值介于 0 到 1 之间且总和为 1。经过**Softmax**后的**注意力权重**（Attention Weight: 表示输入序列中各部分对当前输出的贡献程度）我们记作 A1 hat 到 A4 hat。我们将 A 视为**点积**后得到的**注意力权重**，而 A hat 表示经过**归一化**处理的**注意力权重**。有了 A1 Hat 到 A4 Hat 后，每个 A Hat 都会乘上其对应的**值向量**（Value: 在自注意力机制中，承载实际信息的向量），然后进行**加权和**（Weighted Sum: 各项乘以其权重再求和），这就是**注意力层**在第四个位置的最终输出。以上这些都是我们预设你已经了解的**Transformer**基础知识。

### 加速Transformer：代价与策略

今天，我们将探讨一系列加速**Transformer**运算的方法。在评估任何加速方法时，我们需要问：“其代价是什么？”任何号称能加速的方法，背后往往都付出了某种代价，即用你可能不那么关心的资源，来换取**Transformer**的加速。常见的代价包括：
*   **改变自注意力计算**：某些方法可能并非真正计算**自注意力**，而是一种**近似**（Approximation），导致计算结果与原始**自注意力**有所不同。
*   **模型绑定**：有些方法是**模型绑定**的，要求你必须训练特定的模型或对模型进行**定制化**（Customization）才能使用，并非**即插即用**（Plug-and-Play）的方法。
*   **其他代价**：如果以上两点都不是代价，那么它可能付出了其他形式的代价来实现加速。

我们将通过一个表格来总结这些方法。今天课程的重点在于详细解释**Flash Attention**这项技术。事实上，现在大多数人在使用**语言模型**时，即便不自知，也已经在利用**Flash Attention**。

此外，第二部分会介绍一系列与**KV Cache**（Key-Value Cache: 存储Transformer层中Key和Value向量以加速推理的技术）相关的方法，虽然方法众多，但每个方法都会快速带过。有一个今天不会讲的有效加速推理方法是**推测解码**（Speculative Decoding: 一种通过草稿模型预生成序列，再用大模型验证以加速推理的方法），因为它已在之前的课程（2024年生成式AI导论第16讲）中讲解过，并包含在作业中。

### Flash Attention核心机制：GPU存储层级与计算瓶颈

接下来，我们将深入探讨**Flash Attention**。这项技术最早来源于2022年的一篇论文，可谓是“上古时代”的技艺，但其威力巨大。**Flash Attention**不会改变**注意力机制**的计算结果，其输出与传统方法完全一致，并非**近似**。更重要的是，它是一个**即插即用**的方法，可以直接应用于任何使用**自注意力机制**的**Transformer**，不与特定模型绑定。它的代价非常小。

**Flash Attention**之所以能实现加速，其核心思想在于它考虑了**GPU**（Graphics Processing Unit: 图形处理器，擅长并行计算）运算的底层逻辑。我们需要先了解**GPU**运算时发生了什么，以及需要关注的关键部分。需要强调的是，以下描述并非极其精确，而是为了便于理解而进行的简化，且并非只针对**GPU**，许多高速运算都有类似概念。

在**GPU**中，存在大量负责运算的**执行单元**（Execution Unit: GPU内部负责实际计算的部件）。这些**执行单元**并非单一存在，而是数量众多，可以想象成拥有多个分身，能同时执行任务。它们强大之处在于数量多，运算速度极快。然而，它们的弱点是**工作台**（即**SRAM**，Static Random-Access Memory: 静态随机存取存储器，GPU的工作台，速度快容量小）太小，能存放的数据有限。如果你有一个非常大的矩阵或向量，无法直接放在**工作台**上，每次只能放置少量数值进行运算。

大量的数据则存放在一个**仓库**（即**HBM**，High Bandwidth Memory: 高带宽内存，GPU仓库，容量大速度相对慢）中。虽然**仓库**并非无限大，但相较于**工作台**，**仓库**的容量要大得多。因此，在实际运算时，**执行单元**需要从**仓库**中将数据搬到**工作台**上，每次只能搬运正好能放入**工作台**的数据。一旦数据进入**工作台**，运算速度极快，可在瞬间完成各种运算，然后将结果搬回**仓库**。所以，只要数据上了**工作台**，就能进行高速运算，**工作台**上的一切运算几乎瞬间完成。然而，**搬运数据**是耗时的，是拖慢运算的瓶颈。

**Flash Attention**的目标正是改变**自注意力机制**的计算算法，在不改变计算结果的前提下，通过调整数值计算顺序，减少**搬运数据**的次数，从而缓解这个较慢的瓶颈，使整体运算更快。

### 朴素Attention的计算流程与瓶颈

在讲解**Flash Attention**之前，我们先来看看传统**注意力机制**（朴素Attention）是如何运算的。我们知道**查询向量**（Query）、**键向量**（Key）等大量数值都存放在**仓库**中，而**工作台**（用青色区域表示）非常小。现在我们要计算**注意力**。

假设我们有一个**查询向量**需要与一长串**键向量**进行**点积**计算。实际上，**GPU**可以同时处理多个**查询向量**，但为简化理解，我们只考虑一个**查询向量**的情况（多**查询向量**的情况可以类推）。我们无法将所有**键向量**一次性放入**工作台**，因为**键向量**数量太多了。你需要想象，**键向量**的数量就是当前要处理的**序列长度**（Sequence Length）。一个**语言模型**作为AI代理使用时，其输入**序列**往往非常长，可能是上万、十万甚至百万个**token**。因此，我们必须将整排**键向量**（长度记作大L）切分成一个个**块**（Chunk），每个**块**包含N个**键向量**。每个**块**的大小取决于**GPU****工作台**的大小，**工作台**大则**块**可大，**工作台**小则**块**亦小。

当包含 N 个**键向量**的**块**被放入**工作台**，**查询向量**也被放入**工作台**后，**查询向量**与所有**键向量**的**点积**计算瞬间完成。这些**点积**结果（灰色框表示 Q 与 K 的**点积**）计算完成后，会被放回**仓库**，无法长期停留在**工作台**上，因为需要清出空间进行其他操作。第一个**块**读取进来，计算出**查询向量**与每个**键向量**的**点积**；接着读取第二个**块**，再读取第三个**块**，以此类推。每次可以计算出一个**块**中所有**键向量**与**查询向量**的**点积**。

### Attention权重的归一化：Amax的作用

我们将每个灰色的框框记作 A_i，它代表**注意力权重**。但我们真正需要用于后续运算的是经过**归一化**（Normalized）的**注意力权重**，我们记作 A Hat。**归一化**的公式是将每个 A_i 取**指数**（Exponential），作为分子，分母则是所有 A_i 取**指数**后的数值总和。这样就能得到**归一化**后的**注意力权重** A Hat。

在实际实现中，通常会额外加入一个 Amax 项。在取**指数**之前，会将 A_i 的数值减去 Amax。Amax 是在**点积**计算过程中，所有计算出的**点积**数值（A_i）中的最大值。将 A_i 减去 Amax 的目的是确保**指数**项的最大值是 0。因为 A_i 的数值可能非常大，如果在**指数**项中放入过大的数值，计算时可能会导致**溢出**（Overflow）或其他问题。所以，通过减去 Amax，可以确保**指数**项的最大值是 0。接下来，我们假设使用这个带有 Amax 的公式将 A_i 转换为 A hat。

那么，如何将 A_i 转换为 A hat 呢？你可能会想，我们能否将整个**序列**中所有的 A_i 一次性放入**工作台**，然后**工作台**非常强大，瞬间就能算出 A hat 并将其读出。然而，这种操作是不允许的。虽然在这个例子中只有 9 个数值（看起来像**标量**），但你需要注意，**工作台**不能存放任何与**序列长度** L 相关的东西。你之所以觉得这些数值可以放入**工作台**，是因为在这个例子中只展示了 9 个数值。但在实际**语言模型**运算中，**序列长度**可能达到 10 万甚至 100 万。即使只是 100 万个**标量**，也无法放入**工作台**。

既然不能存放与**序列长度** L 相关的东西，那么从 A_i 转换为 A hat 的过程就变得非常繁琐。首先，我们必须找到 Amax。如果我们只加载了一个**块**中的数值，我们根本不知道哪个是 Amax，因为我们没有遍历所有数字。所以，寻找 Amax 的方法是：一次先加载一个**块**中的 A_i，找出这个**块**中最大的 A_i，假设为 D1。然后将 D1 这个数值存放在**工作台**的一个区域，我们称之为 D。D 中存储着 D1，它是当前**块**中最大的 A_i，而非整个**序列**中最大的 A_i。

第一个**块**处理完后，保留 D1 的值，接着从第二个**块**中读取**注意力权重**。然后比较这些新读取的**注意力权重**与 D 中当前存储的值（D1）哪个更大。也就是说，比较上一个**块**中最大的值 D1 与当前读取的**注意力权重** a_{n+1} 到 a_{2n}（假设每个**块**有 n 个 a_i）哪个更大。将最大的那个值保存为 D2。现在 D 中存储的是 D2。这个过程反复进行，总共执行大 B 次，其中大 B 等于**序列长度** L 除以**块**的大小 N。

直到最后第 B 次操作时，比较 D 中存储的数值 D_{b-1} 与最后一个**块**中读取的**注意力权重**，将最大的存入 D 中，这个最大的数值就是 D_B。此时，D 中存储的值就是 Amax。因为我们每次都将当前看到的最大 A_i 放入 D 中，所以当所有**块**都遍历一遍后，D 中存储的便是 Amax。

### Flash Attention简化版：一次性找出Amax与分母

找到 Amax 后，我们将 Amax 的值仍保留在**工作台**上。然后从第一个**块**开始读取，将读取到的 A_i 减去 D（即 Amax），再取**指数**。这样我们就得到了分子项。对每个**块**都执行相同的操作。我们将 Exponential(A_i - D)（其中 D 是 Amax）记作 A_i prime。

现在我们已经计算出了分子项，接下来要计算分母项。分母项需要将所有分子项，即所有 A_i prime 加起来。寻找总和的方法与之前寻找最大值的方法非常相似：每次读取一个**块**，将**块**中的所有数值加起来存入 S 中。第一次存入的数值是 S_1，即第一个**块**中所有 A_i prime 的总和。然后保留 S 的值，再读取第二个**块**中的内容。将 S 中原有的 S_1 加上当前读取的第 N+1 到 2N 个 A_i prime 的值，得到 S_2。将 S_2 存放在**工作台**上。这个过程反复进行，直到读取到最后一个**块**时，S 中存储的数值是 S_{B-1}。再加上最后一个**块**中所有 A_i prime 的总和，得到 S_B。这个 S_B 就是所有 A_i prime 的总和，即整个**序列**中所有 A_i prime 的总和。这个 A_i prime 就是 Exponential(A_i - Amax)，也就是我们所需 A hat 的分母。

现在我们已经求出了分母。有了分母项和 A_i prime（分子项），接下来计算 A hat 似乎是轻而易举。我们将分子除以分母，然后读出结果，就得到了 A hat。经过这一番折腾，我们终于从 A_i 转换到了 A hat，中间已经多次读写**仓库**，将**工作台**上的数据存入**仓库**，这是一个非常耗时的过程。最终算出 A hat 后，任务仍未结束。此时，你需要从**仓库**中搬运**值向量** V。将 A hat 与 V 进行**加权和**。每次同样只能读取一个**块**的数据到**工作台**上。将这批 V 读取到**工作台**，将 A hat 读取到**工作台**，然后计算 V 与 A hat 的**加权和**，得到 O。O 是我们最重要的输出，也就是 V 对 A hat 的**加权和**。

一开始我们只能计算第一个**块**的 O，所以它并非最终答案。有了第一个**块**的 O 后，接着将 O_1 加上第二个**块**中 A hat 乘以 V 的**加权和**。这里总共有 N 个数值，N 个**加权和**的结果，加上 O_1 得到 O_2。这个过程反复进行，直到最后 O_{B-1} 加上当前**块**中 A hat 乘以 V_i 的**加权和**，得到 O_B。这个 O_B 就是**注意力层**在某个位置的最终输出，即 Q、K、V 进行**注意力**计算后最终得到的结果。

在整个过程中，我们对**仓库**进行了多次读写。**Flash Attention**提出的问题是：真的需要这么多次读写吗？有没有减少读写次数的方法？我们先介绍一个**Flash Attention**的简化版，稍后再讲真正的**Flash Attention**。

简化版的想法是：从 A_i 到 A hat 的多次读取，能否减少读取次数？这个 Amax 和分母的**总和项**（Summation）可以一次性找出。也就是说，将 A_i 读取到内存一次，就能同时算出 Amax 和这个**总和项**。之前我们必须将这两项分开计算，因为**总和项**依赖于 Amax 的值，没有 Amax 就无法计算**总和项**。但神奇的是，即使**总和项**依赖于 Amax，也有办法同时找出 Amax 和**总和项**的结果。

具体做法是：首先读取一个**块**中的数值，找出**块**中最大的数值，记作 D1。我们暂时假设 D1 就是真正的 Amax。然后计算 Exponential(A_i - D1) 并将其总和记作 S1。这里我们假设 D1 就是 Amax。如果 D1 确实是 Amax，那么计算结果就是正确的。但实际上并非如此，我们来看看**Flash Attention**如何解决这个问题。

我们现在计算出了 D，它是这堆数值中最大的。我们计算出了 S，它是假设 Amax 被替换为 D1 时的总和。接着进入下一个**块**，其中包含上一个**块**在**工作台**上留下的 D 和 S 的值。我们读取第二个**块**的数值，并与 D 中原有的值 D1 进行比较。你可能会找到更大的数值，记作 D2。即使没有找到更大的数值，即 D1 等于 D2，也不会影响后续计算结果。这里我们假设找到了一个不同的值 D2，即当前**块**中有比第一个**块**更大的数值。

显然，我们应该将 D2 作为 Amax，而之前将 D1 作为 Amax 是错误的，因为 D2 至少比 D1 大，更接近 Amax。所以我们将 D2 作为 Amax，并计算 Exponential(A_i - D2)。然后，我们将 Exponential(A_i - D2) 在第 N+1 到 2N 个**块**中的每个数值的总和计算出来。接下来如何处理呢？直观的方法可能是将其与之前计算出的总和 S1 直接相加，得到 S2。但问题是，S1 是假设 D1 是 Amax 时计算出的，而这里我们假设 D2 是 Amax。我们现在已经知道 D1 是 Amax 是错误的，因为 D2 比 D1 大。所以 S1 显然是错误的。前一个**块**中计算出的数值是错误的。

但怎么办呢？你不能将前一个**块**的这些数值重新读取进来重新计算，这会浪费读写时间。所以需要想个办法，直接根据现有的 S1 数值进行调整，让错误变为正确，仿佛错误从未发生过。这里的调整是一个公式：S1 乘以 Exponential(D1 - D2)。这样就好像什么都没发生过，直接抹去了之前将 D1 作为 Amax 的痕迹，转变为将 D2 作为 Amax。

为什么将 S1 乘以 Exponential(D1 - D2) 就能移除 D1 的影响，并将其视为以 D2 为 Amax 的计算呢？如果你将 S1 的表达式代入，很容易就能看出：S1 是 summation over i 等于 1 到 N 的 Exponential(A_i - D1)。再乘以 Exponential(D1 - D2)，其中 Exponential(-D1) 和 Exponential(D1) 可以抵消，结果变为 summation over i 等于 1 到 N 的 Exponential(A_i - D2)。这与我们在第 N+1 到 2N 的**总和**计算中（Exponential(A_i - D2)）一致。所以最终，S2 变成了 summation over i 等于 1 到 2N 的 Exponential(A_i - D2)。这就好像我们以 D2 为 Amax 计算了 i 等于 1 到 2N 的总和。

因此，我们可以在不回溯过去数值的情况下，直接对 S1 进行调整，就好像我们以 D2 为 Amax 进行计算一样。这就是**Flash Attention**的核心技巧，稍后会反复利用此技巧减少读写次数。

### Flash Attention完整版：直接计算输出O

现在我们已经了解在一个**块**中如何运作。我们以第一个和第二个**块**为例，后续**块**的算法是完全相同的。我们只是将刚才讲述的内容重新叙述一遍。在第 k 个**块**时，我们做的事情是：读取第 k 个**块**中的 A_i，与 D 中存储的值进行比较，找出最大的那个，记作 D_k。然后计算 Exponential(A_i - D_k)。接着，我们将 Exponential(A_i - D_k) 在第 k 个**块**中所有数值的总和计算出来。

之前存储在 S 中的值 S_{k-1} 需要进行一个变化，才能与新的 D_k 配合。所以 S_{k-1} 需要乘以 Exponential(D_{k-1} - D_k)，才能匹配新的 Exponential(A_i - D_k) 这些项，然后得到 S_k。这个操作反复进行，直到你到达最后一个**块**（大 B 个**块**）。接下来的操作都与之前讲过的类似：读取一串 A_i 进来，比较大小。最大的那个就是 D_B，而这个 D_B 实际上就是 Amax。在我们遍历完所有**块**之后，最终会放在这个区域的数值。我们每次都会用这个区域的数值去与当前所有**块**中读取到的数值进行比较。当所有**块**都比较过后，这个蓝色区域中存储的便是 Amax。

我们现在已经有了 Amax，接下来计算 Exponential(A_i - D_B)。由于 D_B 已经等于 Amax，所以这样计算出的数值是正确的。然后，我们将 L - N 加 1 到 L（即最后一个**块**的数值）的所有 Exponential(A_i - D_B) 取出，并对前面留下的 S_{B-1} 乘以 Exponential(D_{B-1} - D_B) 进行变化和调整，再加上后面这一项，得到 S_B。这个 S_B 就是我们所需的分母项，因为现在 D_B 已经是 Amax，所以这个计算会得到 summation over i 等于 1 到 N 的 Exponential(A_i - Amax)。

经过这一系列操作，S 中存储的数值已经就是分母项了。所以我们现在有了 Amax，有了分母项。接下来计算 A hat 似乎是举手之劳。现在**工作台**上有 D 和 S，D 就是 Amax，S 就是分母项。你要做的事情就是将 A_i 的值读取进来，然后取 Exponential(A_i - D)（D 就是 Amax），再除以分母项 S（即 summation over i 等于 1 到大 L）。你对每个**块**的数值都做反复同样的操作，就能得到 A hat。所以，经过上述操作后，我们已经将从 A_i 到 A hat 这一连串本来需要多次读写的任务，变为只需两次读写。刚才只需从**仓库**中搬运两次 A_i 的数值进来，就能将 A_i 转换为 A hat。因此，我们可以将多次读写减少到两次读写。

你可能会觉得这样已经很不错了。前面先将 Q 和 K 进行一次**点积**得到 A_i，然后两次读写得到 A hat，再进行**加权和**。这比原始未经优化的方法效率更高。但这并非**Flash Attention**的全部。**Flash Attention**最神奇的地方在于，它提出了一个关键的灵魂拷问：一定要算出**注意力权重**才能计算最终的**加权和** O 吗？它跳过了真正计算出 A hat 的步骤，直接得到了 O。

因此，如果你使用**Flash Attention**，一个神奇之处在于你无法直接读取到**注意力权重**。许多人在分析**Transformer**时，会想绘制**注意力矩阵**来查看**注意力权重**的数值。但如果你选择使用**Flash Attention**计算**注意力**，例如使用**Hugging Face**的模块，它会报错，告诉你没有**注意力权重**可以读取。这是因为**Flash Attention**这个方法可以跳过计算真正的**注意力权重** A hat，直接给出 O。而且它得到的 O 的结果，与先计算 A hat 再得到 O 的结果，在理论上是完全一致的。

### Flash Attention的实现与性能优势

**Flash Attention**是如何做到这一点的呢？它的做法是：它希望**一步到位**。将 Q、K、V 全部放入**工作台**后，进行一次运算就得到最终结果。我们来看看这一连串的运算最终是如何实现的。

前面的操作与刚才大致相同：Q 需要与一个**块**中的 K 都进行**点积**计算，然后找出**点积**中最大的数值，记作 D。我们这里将 D1 存入 D 中，它是这些数值中最大的那个。然后计算 Exponential(A_i - D1) 作为 S。S1 就是 summation over i 等于 1 到 N 的 Exponential(A_i - D1)。这部分操作我们刚才已经见到过。

最后，它要做的事情是，在这个时间点直接将 V 读取进来，计算**加权和**。我们直接将 V_i 乘以 Exponential(A_i - D1) 除以 S1（也就是 S1 中存储的数值），然后 summation over i 等于 1 到 N，得到 O1。你可能会想，这个**加权和**对吗？这个**加权和**的权重根本不是我们想要的**注意力**，它根本不是 A hat。没关系，稍后会想办法修复这个问题。所以这里是“将错就错”，继续计算下去。要记住，这里这项**注意力权重**是错误的，它不是真正的 A hat。例如，它的分子是错误的，因为 D1 不是真正的 Amax。它的分母也是错误的，因为 S1 只对 i 等于 1 到 N 进行了**总和**，它的 Amax 也是错误的。它只是对这个区间进行了**总和**，还没有对整个**序列**进行**总和**，但它已经被用作分母了。

接下来我们来看看**Flash Attention**如何修复之前的错误。它的做法是这样的：读取第二个**块**，同样计算出 D2，它是到目前为止看到的**点积**中数值最大的那个。然后，我们需要计算 S2。这与刚才在讲解**Softmax**简化版时所说的一样：原有的 S1 需要进行处理和修复，然后才能加上这个 Exponential(A_i - D2) 得到 S2。刚才我们也提到，如果你一直进行这种修复，最终 S 中存储的就会是正确的分母项，即正确 A hat 的分母项。但实际上，对于真正的**Flash Attention**来说，这件事已经不重要了。它真正关心的是最终能否正确计算出 O。

那么，如何正确计算 O 呢？刚才我们已经计算出了 O1，它是前一个**块**中 V 对一个错误的**注意力权重**的**加权和**。现在我们计算出了 S2。这个 S2 即使它可能仍然不是一个正确的分母，但它至少比 S1 更正确一点。所以，我们现在新的**注意力权重**可以写成 Exponential(A_i - D2) 除以 S2。尽管它不是很正确，但也许比这一项更正确一点。所以我们用第二个**块**中的 V 来对 Exponential(A_i - D2) 除以 S2 进行**加权和**。

那前面的 O1 怎么办呢？前面的 O1 已经算错了，所以我们需要抹去 S1 和 D1 存在的痕迹。我们要将 S1 想办法换成 S2，将 D1 想办法换成 D2。我们要做的事情就是将 S1 分之 Exponential(A_i - D1) 乘以 V_i，换成 S2 分之 Exponential(A_i - D2) 乘以 V_i。

那么，这个置换如何实现呢？你看它们的分母不同，一个是 S1，一个是 S2。我就将 O1 乘以 S1 再除以 S2。乘以 S1 可以抵消这一项，除以 S2 可以将这一项加回来。所以我们就弥补了 S1 和 S2 之间的差异。然后**指数项**（Exponential Term）不同，如何弥补这个差异呢？其实和前面讲的一样，就是乘以 Exponential(D1 - D2)。乘以 Exponential(D1 - D2) 就可以弥补**指数项**的差异。所以，将 O1 乘以 S1 除以 S2，再乘以 Exponential(D1 - D2) 后，O1 与这一项计算的差异，即 D1、D2、S1、S2 造成的差异，就被抹平了。你得到的结果是 summation over i 等于 1 到 2N 的 S2 分之 Exponential(A_i - D2) 乘以 V_i。这就是**Flash Attention**的精神。我们真正关心的是最终得到的 O。

这个过程反复进行。在第 k 个**块**时，操作与刚才一样：计算**注意力权重**。这里我们没有画出**查询向量**，你可以想象有一个**查询向量**进来，与当前**块**中的 K 进行**点积**计算，得到一堆**点积**结果后，找出最大的就是 D_k。然后我们需要计算这个**总和**。需要对原有的 S_{k-1} 进行修正，然后再加上新的 Exponential(A_i - D_k) 得到 S_k。接着我们计算 O。O_k 是原有的 O_{k-1} 经过修正项（这个是修正项）再加上我们新得到的 V_i 的**加权和**。所以我们的过程就是不断修正前面的错误，再加入新的东西。前面有错没关系，因为接下来总能修正回来。

所以在最后第 B 个**块**的操作是：读取一个**块**的数据进来，计算**加权和**，找出最大的就是 D_B。然后计算 S_B。S_{B-1} 需要进行修正，再加上新的数值。O 也一样，最终计算出的 O_B 就是前面的 O_{B-1} 经过一系列修正后再加入新的数值。最终 D 中存储的是 Amax，S 中存储的是 A hat 的分母项。但你根本不在乎这两项，你想要的是 O。这个 O 是 summation over i 等于 1 到大 L（即整个**序列**）的 V_i 乘以正确的**注意力权重**的**加权和**。这一连串的修正会让我们得到正确的**注意力**与 V_i 的**加权和**结果，尽管在整个计算过程中**注意力权重**从未被真正计算出来。这就是**Flash Attention**。

### Flash Attention实战与性能评估

以上就是**Flash Attention**算法的简单讲解，稍后在助教的讲解和作业中会有更详尽的说明。这里提供了一个示例程序，让大家感受一下**Flash Attention**实际使用时的效果。我们将比较使用**Flash Attention**和不使用**Flash Attention**时的速度差异。

首先，这是Colab的典型使用方式：先导入一些必要的库，然后确认**GPU**环境。这里我们使用的是A100，其内存为80GB。请注意，这80GB的内存指的是**仓库**容量，而**工作台**（SRAM）通常只有十几MB。

接着，我们实现两种不同的**注意力机制**：一种使用**Flash Attention**，另一种是**朴素Attention**（Naive Attention），也就是你直觉上的算法。这两种实现都调用了PyTorch的**scaled_dot_product_attention**模块。**scaled_dot_product_attention**允许你设置使用哪种**注意力**算法，默认情况下，如果没有特殊设置，通常会直接使用**Flash Attention**。所以，现在你在运行**Transformer**时，如果没有做特殊处理，很可能已经默认使用了**Flash Attention**。

这两个函数都接收QKV作为输入，并且`causal`参数表示我们正在处理的是一个**语言模型**，它只关注前面的信息，而不是双向**注意力**。QKV读取进来后，唯一的不同在于第二行。第三行完全一样，都是执行PyTorch的函数来计算**注意力**。`run_Naive_Attention`中设置`attn_mask_config={'is_causal': True, 'attention_implementation': 'eager'}`，表示使用朴素**注意力**；而`run_Flash_Attention`中设置`attn_mask_config={'is_causal': True, 'attention_implementation': 'flash_attention'}`，表示使用**Flash Attention**。

接下来是数值验证，目的是证明**Flash Attention**和**朴素Attention**计算出的数值几乎是一样的。我们初始化QKV，设置几个参数：B是**Batch Size**（批处理大小: 一次处理的样本数量），H是**注意力头**（Attention Head: 自注意力机制中的并行计算单元）的数量，L是**序列长度**，D是QKV向量的**维度**。这里分别设置为4、8、256和64。256的**序列长度**相对较小，实际使用中通常会更长。

然后，随机初始化QKV这三个对象。这里没有绑定任何**语言模型**，QKV的数值是随机设置的，因为我们现在只关注计算速度，而非真正运行一个**语言模型**。在衡量计算速度之前，我们分别执行**朴素Attention**和**Flash Attention**，看看它们计算出的数值是否一致。一个结果存储在`out_naive`中，另一个存储在`out_flash`中。我们检查这两个输出的最大差异。执行后会发现最大差异非常小，约为10的负7次方。在进行这种不同运算时，由于底层运算机制的差异，总会存在一点点差异，但这个差异非常小，始终在10的负7次方左右。

### 实际案例分析：长序列加速效果

接下来，我们真正比较使用**Flash Attention**与不使用**Flash Attention**的时间差异。我们编写了一个名为`benchmark_attention`的函数，用于计算不同**注意力机制**所花费的时间。它接收QKV作为输入，还接收一个`mode`参数，表示当前要运行**朴素Attention**还是**Flash Attention**。另外两个参数`warm_up`和`repeat`，你可以自行调整。`warm_up`表示在真正计算之前先试运行一次。因为你可能不清楚Colab背后**GPU**的调度和使用情况，所以先运行一次，让它加载所需的库。至于需要预热多少次，你可以自行尝试。为了加速，这里只运行一次。

然后，我们开始计时，从这一行代码开始计时。我们将同一种**注意力机制**运行三次（你也可以运行更多次，计算结果可能更准确）。然后结束计时，计算每次**注意力**需要多长时间，以毫秒为单位。我们使用不同的**序列长度**：64、128，直到4096，来观察在不同长度下，**Flash Attention**能加速多少。

运行后，你会发现它运行得很快。对于不同的**序列长度**，例如64、128、256、512等，都会显示**朴素Attention**和**Flash Attention**分别花费的毫秒数，并将两者相除，告诉你**Flash Attention**相对于没有使用它加速了多少。你会发现，当**序列长度**设置为4096时，它最多可以加速到9倍左右。**Flash Attention**仅仅通过减少数据搬运次数，就可以实现8到9倍的加速。所以这是一项非常有用的技术。它的代价很小，只是算法复杂一些，并且算法中为了修正进行了一些额外的运算。但瑕不掩瑜，减少数据搬运次数确实非常划算，可以大幅提升速度。

到目前为止，我们只是做了一个玩具示例，使用假的QKV进行计算。接下来，我们将使用一个真正的**语言模型**来比较安装**Flash Attention**与否的差异。这里我们使用Yi-34B模型（其他模型的结果也类似）。首先，我们从**Hugging Face**下载这个模型，这需要一些时间。

在模型下载过程中，我们来看看接下来要做什么。我们会先存储一个非常长的字符串，这将在稍后作为**语言模型**的**Prompt**（提示: 输入给语言模型，引导其生成内容的文本）。我们称之为`long_text`，因为它确实非常长。我们只是将一个无聊的句子反复重复很多次。我们将这个句子重复100次。稍后，我们将使用**Hugging Face**的`pipeline`来调用**语言模型**。这个**序列**有7300个**token**。

接下来，我们开始进行测试。首先测试没有使用**Flash Attention**的情况。我们使用**Hugging Face**的`pipeline`函数来调用**语言模型**。这里有一个`pipeline eager`，其设置告诉我们**注意力机制**的实现要使用`eager`方法。`eager`方法就是不使用**Flash Attention**。`eager`虽然有“急切”的意思，但在这里它代表没有使用**Flash Attention**。如果你没有特殊设置，默认通常是使用**Flash Attention**。所以，如果你不想使用**Flash Attention**，你需要特别进行设置。

现在，我们使用没有**Flash Attention**的`pipeline`进行**推论**。我们只让它输出一个**token**。输入一个非常长的**序列**，只预测下一个**token**，看看花了多少时间。结果花了0.15秒。然后，我们换成使用**Flash Attention**，这里是`pipe_sdpa`，表示使用了**Flash Attention**。如果使用了**Flash Attention**，同样只输出一个**token**，你会发现速度并没有快多少。为什么呢？因为我们的**序列**太短了。

在一个**语言模型**中，除了**注意力机制**外，还有非常多不同的组成部分，例如大量的**前馈网络**（Feed Forward Network）。甚至在计算时，仅仅将文字转换为**嵌入向量**（Embedding: 将离散的符号（如单词）映射到连续的向量空间中）也需要时间，**嵌入向量**转回文字也需要时间。许多部分都需要时间。所以，当**序列**太短时，**Flash Attention**的效果并不明显。

我们将**序列长度**改为1000，给它7万个**token**看看。结果，没有使用**Flash Attention**时花了2秒，而使用**Flash Attention**时花了1.3秒。所以你看，在**序列**较长时，使用**Flash Attention**可以显著加快速度。

你可能会想，如果再长一点呢？我们刚才使用了1000，如果使用1万个**token**呢？那将有73万个**token**。运行一下，看看会发生什么。

哇，无法运行！出现了`CUDA out of memory`。这个`CUDA out of memory`并非**工作台**的内存不足，而是**仓库**的内存不足。所以我刚才说，**仓库**虽然很大，但它也有极限。当**序列**太长时，最终会撑爆**仓库**。至于**序列**太长为何会撑爆**仓库**，我们将在下一堂课讲解**KV Cache**时详细说明。