FlashAttention原理：从原始Attention到FlashAttention

最新推荐文章于 2025-04-27 18:34:33 发布
最新推荐文章于 2025-04-27 18:34:33 发布
阅读量3.2k 收藏 28
点赞数 16
分类专栏： llm 文章标签： flash-attention llm
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/qq_43527718/article/details/142302113
版权

一、背景概述

以前的attention加速方法旨在减少attention的计算和内存需要，如sparse-attention、low-rank approximation等，但由于它们主要关注FLOP reduction，且倾向于忽略内存访问的开销，所以都没有达到wall-clock speedup。

FlashAttention比普通attention的HBM（GPU high bandwidth memory）访问量更少，并适用于一系列SRAM大小。
FlashAttention旨在让attention算法变得IO-aware，即仔细计算对于不同级别的memory的读写（如fast GPU on-chip SRAM和relatively slow GPU high bandwidth memory即HBM），各memory层级如下图。这是在普通的深度学习框架（如PyTorch、Tensorflow）所不天然支持的
请添加图片描述

二、具体原理

  1. 宏观目标：compute exact attention with far fewer memory accesses
    具体来说就是避免从HBM中读/写attention矩阵，这需要从两方面入手：
    • computing the softmax reduction without access to the whole input
    • backward时不存储大型的中间矩阵
  2. 做法：提出两个策略
    • 重塑attention计算：将input分成若干blocks，然后incrementally进行softmax reduction（即tilinng）
    • 存储forward中的softmax normalization factor，以便在backward中快速recompute attention on-chip（原来是从HBM中读取中间矩阵）

一、attention是什么

attention的本质：带权特征筛选。注意力机制描述了（序列）元素的加权平均值，其权重是根据输入的query和元素的键值进行动态计算的。

先有seq2seq模型，seq2seq其实就是输入序列通过一个RNN，得到最终编码后，再通过一个RNN，每次输出翻译的单词。如下图
在这里插入图片描述
2015年提出attention，此时attention就被绑定在上面这个seq2seq模型上，2016年提出self-attention，2017年提出transformer。因此attention并不能和transformer划等号，它只是transformer的最核心的零件。

输入的时候是一起全部输入的，输出的时候还是一个一个输出的

经典传统文本翻译模型seq2seq是基于循环神经网络RNN的，而transformer丢掉了循环神经网络，只用到了注意力

transformer

  1. 最基础的注意力：q、k、v
    如下图，给定一张图，什么问题都不问，人对它的本来印象就是k（如第一行的数字，表示的是本来印象），然后提出问题，问图里的动物有哪些。此时q（query）就是“动物”。q和k相乘得出来的就是每个元素的权重（第二行）。本例中右上角其实有个海鸥，但由于一开始根本没有注意到，因此即使乘了q也不会有任何权重。v就是客观上的各个元素，没有任何权重。
    在这里插入图片描述
    k和v都是由x（原本要翻译的这句话中的每一个词向量）分别经过不同矩阵相乘而得到的结果
    在这里插入图片描述
    在这里插入图片描述
    也就是输入向量x，分别乘不同的权重矩阵得到向量V和向量K，然后向量Q乘向量K转置得到矩阵 α \alpha α，最终矩阵 α \alpha α右乘向量V得到向量b
  2. 自注意力：自注意力就是上面的Q、K、V全部都来源于输入x（前面是只有k和v）。如下图，自注意力输入N个，也输出N个。BERT就是自注意力。并行计算的
    在这里插入图片描述
  3. 多头注意力：多头就是有多组Q、K、V。如上图是一组（transformer是8头），然后再来一组，就会生成第二组 b 1 , b 2 , . . . , b N b_1,b_2,...,b_N b1​,b2​,...,bN​。最终生成的多组b，连接在一起就是结果。说白了多头就是多组参数，就是让模型复杂一些、让参数更多一些。
  4. 位置编码：如上述多头注意力计算过程可知，调换任意两个x的位置，各x对应的结果是不变的。因此我们若要考虑位置信息，需要引入位置编码
  5. transformer：
    在这里插入图片描述
    • 左边encoder，右边decoder。（假设任务是中翻英）encoder的作用就是把要翻译的中文进行充分理解，充分获取里面的信息。decoder就是把理解的内容进行翻译版本的输出

二、FlashAttention原理解读

  1. 贡献要点：Fast and Memory-Efficient Exact Attention with IO-Awareness
    • fast：加快模型训练速度
    • memory-efficient：减少显存占用
    • exact attention：和标准attention计算得到的结果完全一致，精度不变
    • IO-awareness：算法是改进IO的效率
  2. 以一个attention的标准计算过程为例进行存储方面的说明：从下图中可看到随着序列长度N的增长，中间矩阵的缓存是以 N 2 N^2 N2进行增长的。而且这些中间矩阵是必须要存的，因为反向传播时需要它们来计算梯度
    请添加图片描述
    pytorch在实际显卡上的attention实现：
    矩阵Q、K、V（N*d，N是序列长度，d是特征维度）存储在HBM中
    • 从HBM加载Q、K到SRAM
    • 计算出 S = Q K T S=QK^T S=QKT
    • 将S写到HBM
    • 将S加载到SRAM
    • 计算 P = s o f t m a x ( S ) P=softmax(S) P=softmax(S)
    • 将P写到HBM
    • 从HBM加载P和V到SRAM
    • 计算 O = P V O=PV O=PV
    • 把O写到HBM
    • 返回O
  3. 模型训练时的两个制约因素
    • Compute-bound：速度瓶颈在于运算。比如对于大的矩阵乘法需要多channel的卷积运算，这些运算需要的数据不多，但是计算很复杂
    • Memory-bound：速度瓶颈在于HBM的读取速度，即算力等待数据。这样的操作包含两类，这些操作都是需要数据很多，但是计算相对简单：
      • 按位操作：Relu、Dropout
      • 规约操作：sum、softmax
    • 对于Memory-bound的优化一般是进行fusion融合操作，通过不对中间结果缓存来减少GBM的访问。但中间结果在方向传播时是有用的，这样不保存的话只能在反向传播时重新计算
  4. memory层级：需要尽可能增加芯片内（SRAM）的访存、减少芯片外的HBM的访存
    请添加图片描述
  5. Flash Attention原理：避免attention matrix从HBM的读写
    • 通过分块计算，融合attention内的操作，不缓存中间结果到HBM，从而加快速度
    • 反向传播时，重新计算中间结果，以此来解决不缓存后梯度无法计算的问题
    • 具体实现：
      请添加图片描述
      请添加图片描述
      请添加图片描述
      如上图， K T K^T KT、V不变，分别读取Q的前、中间、后两行，计算出O的前、中间、后两行（非最终结果）；然后取 K T K^T KT和V的后三列、后三行，并分别和Q的前、中、后两行进行计算，附加到O上，得到O的最终结果。

      上述去掉了softmax。Softmax是按行进行的，如下图所示，要让矩阵分块对attention多步进行融合计算得以进行的前提就是要解决softmax的分块问题：
      请添加图片描述
  6. softmax的分块计算：如下图2所示，softmax是可以分块计算的，但是就是不如上面的简单，softmax的分块计算需要额外保存几个变量（即 m ( x 1 ) , m ( x 2 ) , p ( x 1 ) , p ( x 2 ) , l ( x 1 ) , l ( x 2 ) m(x^1), m(x^2), p(x^1), p(x^2), l(x^1), l(x^2) m(x1),m(x2),p(x1),p(x2),l(x1),l(x2)）、以及在分块的合并时需要增加一些计算量。但这些计算量相对于减少的IO时间都是非常划算的
    原softmax：
    请添加图片描述
    分块版softmax：
    请添加图片描述
    伪代码：
    请添加图片描述
  7. FlashAttention2的大致思想与FlashAttention1类似，只不过增加了以下一些工程上的优化：
    • 减少了非矩阵乘法计算，可利用TensorCore加速
    • 调整了内外循环，Q为外层循环、KV为内层循环。进一步减少了HBM读写
    • 如果一个block处于矩阵上三角部分（也就是被mask掉的部分），就不进行attention计算，进一步减少计算量

三、FlashAttention更详细的细节剖析

可见这篇文章，写得很好，一篇就足够了。
笔记很详细，见实验室-FlashAttention修改-flashAttention原理

flash attention原理
强化学习曾小健
06-06 1612
在深度学习中，注意力机制（Attention Mechanism）广泛用于自然语言处理（NLP）、图像处理等领域。注意力机制的核心思想是根据输入序列中的每个元素的重要性来动态调整它们的权重，从而更好地捕捉序列中的依赖关系。自注意力（Self-Attention）是其中的一种常见形式，它用于计算输入序列中每个元素对其他元素的影响。Flash Attention 是通过引入分块计算、内存访问优化和并行计算等技术来提高注意力机制效率的一种方法。
Attention优化｜2w字原理&图解: 从Online-Softmax到FlashAttention V1/V2/V3
强化学习曾小健
05-25 4166
本文通过原理分析和图解的方式，通俗易懂地FlashAttention系列算法。FlashAttention V1/V2在LLM领域的应用已经非常广泛，相关的论文也反复读了几遍。FA1和FA2论文非常经典，都推荐读一下（不过FA2论文中公式错误不少本文大约1.8w字，包括以下内容：0x07 分布式训推使用FlashAttention0x09 FlashAttention中MQA/GQA以及Causal Mask处理。
FlashAttention的核心优化原理解析（快收藏）
Jason160918的博客
01-29 1055
这个代码实现了一个基本的注意力机制，其中包括将输入分割成多个头、计算查询和键之间的得分、应用 softmax 函数进行归一化、结合值，以及最后的输出线性层。请注意，这不是一个完整的、优化的实现，而是为了演示其核心思想而设计的。在传统的注意力机制中，计算过程需要频繁地读写内存，这会导致显著的延迟。它对矩阵乘法（Transformer 中注意力机制的关键操作）进行了优化，以减少不必要的计算和提高并行处理能力。对于长序列，FlashAttention 可能采用分块处理的方法，将长序列分成较小的部分进行并行计算。
Flash Attention介绍
weixin_42035282的博客
06-13 2287
Flash Attention是一种优化Transformer模型中注意力机制的技术，旨在提高计算效率并减少内存使用。
【LLMs篇】10：让大模型更快更省显存的魔法FlashAttention
u013010473的博客
04-22 1028
FlashAttention旨在加速注意力计算并减少内存占用。FlashAttention利用底层硬件的内存层次知识，例如GPU的内存层次结构，来提高计算速度和减少内存访问开销。
LLM底层架构---手撕flashattention1
最新发布
看llm论文和code看得相似
04-27 1077
手撕flashattention底层原理
Flash attention入门
weixin_40777649的博客
05-15 2246
flashAtention其加速的原理是非常简单的，也是最基础和常见的系统性能优化的手段，即通过利用更高速的上层存储计算单元，减少对低速更下层存储器的访问次数，来提升模型的训练性能。首先，从HBM中读取完整的Q和K矩阵（每个大小为N x d），计算点积得到相似度得分S（大小为N x N），需要进行O(Nd + N^2)次HBM访问。将原始的注意力矩阵分解成更小的子矩阵，然后分别对这些子矩阵进行计算，只要这个子矩阵的大小可以在SRAM内存放，就可以在计算过程中只访问SRAM。
FlashAttention/ PagedAttention原理,大模型加速
stephen147的博客
07-10 2502
1.1 GPU 硬件特点 由于 FlashAttention 计算 self-attention 的主要关键是有效的硬件使用，所以了解GPU内存和各种操作的性能特征是很有必要的。 以 A100 (40GB HBM) 为例，下面显示其内存层次结构的粗略图。SRAM内存分布在108个流式多处理器(SMs)上，每个处理器192KB。片上SRAM比HBM快得多，但比HBM小得多，在计算方面，使用Tensor Core的BFLOAT16 的理论峰值吞吐量为 312 TFLOPS。GPU 的典型操作方式是使用大量的线程
FlashAttention的原理及其优势
weixin_63866037的博客
01-13 1335
FlashAttention 是一种高效且内存优化的注意力机制实现方法，通过分块计算和内存访问优化，显著降低了注意力机制的计算开销。它不仅适用于现有的Transformer模型，还为未来更大规模的模型提供了可能性。随着深度学习模型的不断扩展，FlashAttention 将成为解决计算和内存瓶颈的重要工具。
通透理解FlashAttention(含其2代和第3代)：全面降低显存读写、加快计算速度
热门推荐
结构之法 算法之道
10-06 7万+
因此，可以确认：在 MQA 中，除了 query 向量还保存着 8 个头，key 和 value 向量都只剩 1 个「公共头」了，这也正好印证了论文中所说的「所有 head 之间共享一份 key 和 value 的参数」所以，上面讲到计算注意力的主要瓶颈是显存访问，因此减少对HBM的读写次数，有效利用更高速的SRAM来进行计算是非常重要的，而GPU有大量的线程来执行某个操作，称为。(需要注意的是，模型训练通常会影响到算子融合的效果，因为为了后向传递计算梯度，通常需要将某些中间结果写入到HBM中)
Flash Attention原理讲解
周同学的博客
03-16 1436
Flash Attention原理讲解
图解Attention优化: 从Online-Softmax到FlashAttention V1/V2/V3
lqfarmer的博客
07-28 998
本文首先从Online-Softmax的角度切入，由浅入深地讲解了3-pass Safe-Softmax、2-pass Online-Softmax以及1-pass FlashAttention的原理；然后，进一步详细讲解了FlashAttention-1和FlashAttention-2算法中各自的优化点、FlashAttention IO复杂度分析以及适用场景、FlashAttention在分布式训推中的应用；
【深度学习】FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness【二】
q742971636的博客
06-18 903
提出 FlashAttention 算法FlashAttention 是一种 IO 感知的精确注意力算法，通过分块技术减少了 GPU 高带宽内存（HBM）与 GPU 片上 SRAM 之间的内存读写次数。该算法通过避免对大型中间注意力矩阵的读写，显著减少了 HBM 访问次数。IO 复杂度分析论文详细分析了 FlashAttention 的 IO 复杂度，证明其在减少 HBM 访问次数方面比标准注意力算法更高效。
大模型训练加速之FlashAttention系列：爆款工作背后的产品观
强化学习曾小健
04-15 1409
FA把优化目标是单个Head的Attention计算内，N是seqence length长度、d是hidden dimension大小。输入是。
NLP高频面试题（二十）——flash attention原理
WeLearnNLP
03-27 581
FlashAttention是一种针对Transformer模型中自注意力机制的优化算法，旨在提高计算效率并降低内存占用，特别适用于处理长序列任务。在Transformer架构中，自注意力机制的计算复杂度和内存需求随着序列长度的平方增长。这意味着当处理较长序列时，计算和内存负担会显著增加，导致模型训练和推理的效率降低。
FlashAttention
maxzcl的博客
05-03 1534
paper: https://arxiv.org/abs/2205.14135an informal talk by the author Tri Dao: https://www.youtube.com/watch?v=FThvfkXWqtEcode repo: GitHub - HazyResearch/flash-attention: Fast and memory-efficient exact attention introduction to transformer: Transformer a
FlashAttention：Fast and Memory-Efficient Exact Attention with IO-Awareness
qq128252的博客
06-29 3466
给 Transformer 配备更长的上下文仍然很困难，因为 Transformer 核心自注意力模块的时间复杂度以及内存复杂度在序列长度上是二次方的。FlashAttention是一种新的注意力机制，旨在解决Transformer在处理长序列时速度慢且内存需求大的问题。FlashAttention的创新之处在于引入了IO（输入/输出）感知的设计原则，专注于减少GPU内存（高带宽内存）和GPU片上内存（SRAM）之间的读写次数。
【大模型训练】Flash Attention详解
weixin_43799388的博客
09-13 3461
FlashAttention系列工作，是一种加速注意力计算方法，目前已经应用在：GPT-3、Falcon2（阿联酋大模型）、Llama2、Megatron-LM、GPT-4等流行LLM上。并且FlashAttention2已经集成到了pytorch2.0中，可以很便捷的调用。1. FlashAttention动机，可以看出由于Transformer中self-attention 的时间和内存复杂度是序列长度的二次方，所以序列过长时，算法速度会变慢，需要消耗很高的内存，导致低效的。
加速attention计算的工业标准：flash attention 1和2算法的原理及实现
bug404
10-07 7413
transformer模型大火，flash attention技术同时也被提出用于加速attention的计算，目前已经被pytorch、huggingface、paddlepaddle等集成到其框架中。
写一个flashattention算子
03-23
<think>嗯，用户让我介绍如何写一个FlashAttention算子。首先，我需要理解FlashAttention是什么。我记得FlashAttention是一种高效计算注意力机制的算法，由斯坦福团队提出，主要优化了GPU内存的使用，特别是减少对高带宽内存（HBM）的访问，从而加快训练速度。好的，现在需要把这个过程讲清楚，但用户可能想了解具体实现的步骤，而不仅仅是概念。 用户可能有一定的深度学习背景，尤其是对Transformer模型和注意力机制有所了解。他们可能正在实现自己的模型，遇到了性能瓶颈，或者对优化感兴趣。所以需要详细说明实现的关键点，比如分块计算、循环处理、重计算等。 接下来，我需要分解实现步骤。首先，分块处理输入矩阵，将大的Q、K、V矩阵分成小块，这样可以利用GPU的SRAM，减少HBM的访问次数。然后，循环加载这些块到SRAM，计算局部注意力分数，接着处理softmax的分母部分，可能需要在线计算最大值和求和来保持数值稳定。还要处理反向传播时的重计算，避免存储中间结果，从而节省内存。 需要注意的是，用户可能对具体代码实现感兴趣，但直接写代码可能太复杂，而且不同框架（如PyTorch或CUDA）实现方式不同。所以应该重点讲解算法逻辑和关键步骤，而不是具体代码。 另外，内存管理的细节很重要，比如如何分块，块的大小如何确定，这会影响性能。同时，数值稳定性是softmax计算中的关键点，需要解释在线计算max和sum的方法，以及如何合并不同块的结果。 可能还需要比较传统注意力机制和FlashAttention的区别，突出FlashAttention的优势，比如内存效率和计算速度。但用户的问题主要是如何编写这个算子，所以重点放在实现步骤上。 最后，要确保回答结构清晰，分步骤说明，每个步骤的关键点，可能遇到的挑战，比如如何处理块间的依赖，如何高效地管理内存，以及如何实现反向传播的重计算。需要提醒用户在实际实现中需要处理多线程、内存对齐等底层优化，可能建议参考原始论文或现有开源实现，如Tri Dao的代码。</think>好的，我将分步骤为您介绍如何实现一个简化版的FlashAttention算子： --- ### 一、FlashAttention核心思想 1. **目标**：高效计算注意力矩阵 $QK^TV$ 2. **核心优化**： - 通过分块计算(tiling)减少GPU HBM访问次数 - 在线计算softmax避免存储中间矩阵 - 反向传播时重计算(Recompute)代替存储 --- ### 二、关键实现步骤 #### 步骤1：矩阵分块 ```python # 将Q, K, V分块 (示例维度) Q = [B, H, N, d] # Batch, Head, SeqLen, Dim K = [B, H, N, d] V = [B, H, N, d] # 分块大小根据SRAM容量确定 block_size = 64 # 典型值 ``` #### 步骤2：外循环处理K/V块 ```python for j in range(0, N, block_size): Kj = K[:, :, j:j+block_size, :] # 加载K的块到SRAM Vj = V[:, :, j:j+block_size, :] ``` #### 步骤3：内循环处理Q块 ```python for i in range(0, N, block_size): Qi = Q[:, :, i:i+block_size, :] # 加载Q的块到SRAM # 计算局部注意力分数 Sij = einsum('bhid,bhjd->bhij', Qi, Kj) / sqrt(d) ``` #### 步骤4：在线Softmax计算 ```python # 计算每块的最大值和求和项 m_ij = max(Sij, dim=-1) # 当前块的最大值 p_ij = exp(Sij - m_ij) # 稳定化指数 l_ij = sum(p_ij, dim=-1) # 分母部分 # 合并到全局统计量 m_new = max(m_prev, m_ij) l_new = exp(m_prev - m_new) * l_prev + exp(m_ij - m_new) * l_ij ``` #### 步骤5：增量更新输出 ```python # 更新输出矩阵 O_i = (exp(m_prev - m_new) * l_prev * O_prev + exp(m_ij - m_new) * p_ij @ Vj) / l_new ``` #### 步骤6：反向传播重计算 ```python # 前向时不存储S矩阵，反向时重新计算 def backward(ctx, grad_output): recompute_Sij() # 重新分块计算注意力分数 # 计算梯度... ``` --- ### 三、关键优化技术 1. **内存管理** - 将计算分解为多个块，确保每个块能在SRAM中完成计算 - 典型内存消耗：$O(N)$ → $O(\sqrt{N})$ 2. **数值稳定性** - 在线计算softmax分母（使用对数域合并） - 增量更新输出矩阵 3. **硬件适配** - 使用CUDA kernel实现 - 利用Tensor Cores加速矩阵乘 --- ### 四、与传统注意力对比 | 特性 | 传统Attention | FlashAttention | |---------------|------------------|-----------------| | 内存复杂度 | $O(N^2)$ | $O(N)$ | | HBM访问次数 | $Ω(N^2)$ | $O(N^2/M)$ | | 最大序列长度 | ~4k tokens | ~64k tokens | --- ### 五、实现建议 1. 参考官方CUDA实现：https://github.com/HazyResearch/flash-attention 2. 使用Triton语言简化开发： ```python @triton.jit def flash_attention_kernel(...) ``` 3. 注意内存对齐和bank conflict避免 需要更具体的代码实现细节，可以进一步讨论某个模块的实现！
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值