FlashAttention1. 使用背景Transformer 很强但 attention 一直有一个很现实的问题又慢又吃显存。如果序列长度是NNN标准自注意力里最核心的中间量是SQK⊤ S QK^\topSQK⊤它的形状是N×NN\times NN×N。这意味着随着序列变长attention 的中间结果会迅速膨胀。早期很多人一看到 attention 是二次复杂度就会自然想到是不是只能去做近似、稀疏化、低秩化。但 FlashAttention 提出了一个非常关键的观点很多时候 attention 慢不只是算得多更是因为显存读写太重。也就是说问题不只是 FLOPs不只是数学复杂度还包括一个很重要的系统因素IO输入输出 / 内存访问成本。FlashAttention 就是在这个背景下提出来的。它最核心的思想可以概括成一句话不改变标准 attention 的数学结果而是通过分块、融合和在线 softmax把大规模中间矩阵的显存读写降下来。2. 理论基础1标准 attention 到底在算什么给定输入隐藏状态先通过线性变换得到QXWQ,KXWK,VXWV QXW_Q,\qquad KXW_K,\qquad VXW_VQXWQ,KXWK,VXWV标准 scaled dot-product attention 写成Attention(Q,K,V)softmax(QK⊤d)V \text{Attention}(Q,K,V)\text{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right)VAttention(Q,K,V)softmax(dQK⊤)V如果把它拆开看其实主要有三步先算分数矩阵SQK⊤d S\frac{QK^\top}{\sqrt{d}}SdQK⊤再做 softmax得到注意力权重Psoftmax(S) P\text{softmax}(S)Psoftmax(S)最后和VVV相乘OPV OPVOPV这个流程数学上很清楚但在实现上会有一个大问题中间矩阵SSS和PPP都是N×NN\times NN×N。2为什么这会成为瓶颈当序列长度NNN很大时SSS和PPP的存储、读写都非常重。即使 GPU 算力很强如果每一步都要频繁把这些大矩阵在 HBM高带宽显存和片上 SRAM / cache 之间来回搬运速度也会被拖慢。所以 attention 的瓶颈并不只是“乘法很多”而是中间结果太大内存访问代价太高。3为什么普通实现会浪费在常规实现里通常会先把SQK⊤SQK^\topSQK⊤整个算出来并写到显存再从显存里读出来做 softmax得到PPP再把PPP写回显存最后再读出PPP和VVV去算输出。这个过程最大的问题在于很多中间结果只是临时用一下却被完整写回了大显存。这就是 FlashAttention 最想解决的地方。3. FlashAttention 的核心思想1一句话理解FlashAttention 最核心的想法就是不要把整个 attention matrix 显式落到显存里而是在片上做小块计算边算边归约直接得到输出。也就是说它不是先完整构造SQK⊤,Psoftmax(S) SQK^\top,\qquad P\text{softmax}(S)SQK⊤,Psoftmax(S)再去算输出而是把这几个步骤融合在一起以分块方式完成。2为什么叫 IO-awareFlashAttention 论文里一个非常关键的词就是IO-awareness。传统算法通常更关注算术复杂度比如时间复杂度是不是O(N2)O(N^2)O(N2)FlashAttention 则进一步强调在现代 GPU 上显存和片上缓存之间的数据搬运本身就是主要开销。所以 FlashAttention 并不是去近似 attention而是去优化 attention 的数据流。3最重要的一点它是 exact attention这一点特别重要。FlashAttention 不是 Performer、Linformer 那种近似 attention 方法。它并没有改 attention 的数学定义而只是改了实现方式。所以它的核心价值不是“牺牲精度换速度”而是在保持精确等价的前提下把显存访问和运行时间都降下来。4. FlashAttention 是怎么做的1分块tiling假设把Q,K,VQ,K,VQ,K,V都按块切开。比如把QQQ按行分成若干 query block把K,VK,VK,V按列分成若干 key/value block。那么 attention 就不再是一次性拿整块大矩阵去算而是按 block 逐块处理。直观上理解就是整张大表一次性装不下那就分小块搬进片上缓存里一块一块算。2问题softmax 不是一个纯局部操作如果只是矩阵乘法分块很好理解。但 attention 里还有 softmax而 softmax 涉及整行归一化softmax(si)jesij∑kesik \text{softmax}(s_i)_j\frac{e^{s_{ij}}}{\sum_k e^{s_{ik}}}softmax(si)j∑kesikesij这里的分母要看整行所有元素所以它不像普通矩阵乘法那样天然可以局部分开。这也是 FlashAttention 真正巧的地方它用在线online方式维护 softmax 所需的统计量。3在线 softmax对于某一行 attention score常规 softmax 会先拿到整行所有值再求最大值、再求指数和。FlashAttention 则是在分块扫描时逐步维护当前见过的最大值当前对应的归一化系数当前输出累积值。这样一来即使一整行分散在多个 block 中也不需要把整行完整写回显存后再统一 softmax。所以 FlashAttention 的关键点不只是“分块”而是分块 在线 softmax 输出累积。5. 在线 softmax 的直观理解1为什么 softmax 通常要先减最大值普通 softmax 为了数值稳定常写成softmax(si)jesij−mi∑kesik−mi \text{softmax}(s_i)_j\frac{e^{s_{ij}-m_i}}{\sum_k e^{s_{ik}-m_i}}softmax(si)j∑kesik−miesij−mi其中mimaxjsij m_i\max_j s_{ij}mijmaxsij这样做是为了防止指数爆炸。所以如果想分块做就必须一边扫描 block一边正确维护这个最大值和归一化项。2FlashAttention 的核心统计量对于某一行FlashAttention 会在扫描 block 的过程中不断更新行最大值mmm行归一化和lll输出向量累计值。每看到一个新的 key block就先算当前 block 对这一行的局部分数再和旧的m,lm,lm,l合并更新。这样到最后虽然没有显式存整个SSS和PPP但最终得到的输出和标准 attention 完全一致。3本质上在做什么如果用一句话概括这个过程我觉得最合适的是把原来“先得到完整注意力矩阵再做 softmax”这件事改写成“扫描中间块时同步完成归一化和加权求和”。所以它不是把 attention 简化了而是把计算顺序重新排了一遍。6. FlashAttention 到底省了什么1不是把二次复杂度变成线性这是最容易误解的一点。FlashAttention 并没有把 attention 的理论计算复杂度从O(N2)O(N^2)O(N2)变成O(N)O(N)O(N)。从数学上讲所有 query-key 对之间的交互还是要算。它真正优化的是中间矩阵的显式物化和反复显存访问。2它主要省的是 HBM 访问在 GPU 上HBM 访问代价远高于片上 SRAM / register 的访问。FlashAttention 通过块级计算让很多中间值留在片上而不是反复写回大显存。所以它真正省的不是“注意力不用算了”而是很多原本没必要写出的大矩阵不再落到 HBM。3为什么这会带来显著提速因为在长序列下attention 往往不是单纯 compute-bound而是明显受内存访问限制。一旦把显存读写压下去GPU 的有效利用率就能明显提高。所以 FlashAttention 的加速来源本质上是更合理的数据流而不是近似。7. FlashAttention 和普通 fused kernel 有什么区别1不是简单 kernel fusion看到这里很多人会觉得“这不就是把几个 kernel 融合一下吗”其实比普通 fusion 更进一步。普通 kernel fusion 更多是把几个连续操作合并以减少 launch overhead 和部分中间读写FlashAttention 则是从 attention 的算法结构本身出发重新设计了 block-wise 的 IO 流程。所以它不只是“把 kernel 拼起来”而是把 attention 的执行顺序按 GPU 内存层级重新组织了一遍。2为什么这很重要因为 attention 里的问题不只是 kernel 多而是中间矩阵太大。单纯 fusion 并不能自动解决“整块 attention matrix 落显存”的问题。FlashAttention 真正解决的是这个更深层的问题。8. FlashAttention 的训练和推理意义1训练里最直接的收益在训练阶段FlashAttention 可以明显降低 attention 中间激活带来的显存占用。这意味着能撑更长序列能用更大 batch或者在同样显存下训练更大的模型。所以 FlashAttention 一开始最重要的影响其实是在训练侧。2推理里也很重要在推理阶段尤其是 prefill 阶段长上下文 attention 本身依然很重。FlashAttention 仍然能带来明显收益。不过在 decode 阶段很多时候瓶颈又会更多落在 KV Cache 访问上。所以在推理里FlashAttention 和 KV Cache 往往是互补关系而不是替代关系。3和 KV Cache 怎么配合理解我觉得最简单的理解方式是FlashAttention 主要优化“单次 attention 怎么算”KV Cache 主要优化“历史 K/V 怎么复用”。二者关注点不同但在大模型推理里通常会一起出现。9. FlashAttention-2 在改什么1为什么还要有 FlashAttention-2初代 FlashAttention 已经把 IO 问题解决得很好了但它并不意味着 GPU 利用率就已经到头。FlashAttention-2 继续往前推进关注的重点更偏工程实现和并行划分怎样让 GPU 上的线程块、warp 和 matmul 工作分配得更合理。2核心改进方向FlashAttention-2 的主线大致可以概括成三件事减少不必要的非 matmul FLOPs提高不同线程块之间的并行度优化单个线程块内部不同 warp 的工作划分减少共享内存通信。所以它相对初代的重点已经有点变化初代更像“把 attention 变成 IO-aware”二代更像“在这个框架下把 GPU 并行度和吞吐继续榨出来”。3为什么它还是很重要因为很多时候真正影响实际速度的不是“有没有一个好算法”而是“这个算法在硬件上吃得够不够满”。FlashAttention-2 做的就是把这件事继续推进。10. FlashAttention-3 又在改什么1为什么还会有 3到 FlashAttention-3 时重点已经更偏向新硬件特性尤其是 Hopper 架构 GPU。它关注的是怎样利用新一代 GPU 的异步执行能力、Tensor Memory AcceleratorTMA、以及低精度能力把 attention 再往前推。2主线变化如果很粗地说FlashAttention 1解决 IO-aware exact attentionFlashAttention 2优化并行划分和吞吐FlashAttention 3进一步拥抱新硬件特性与低精度执行。所以这一条线越往后越明显能看出来它已经不只是一个“attention 小技巧”而是在逐渐变成 attention 的高性能实现主线。11. FlashAttention 为什么这么重要1它改变了大家看 attention 的方式早期一提 attention很多人第一反应就是二次复杂度然后想到各种近似。FlashAttention 提醒大家在真实硬件上算法快不快不只取决于 FLOPs也取决于 IO。这一点其实很重要因为它把问题从“数学复杂度”推进到了“系统复杂度”。2它证明了 exact 也可以很快以前很多人会下意识觉得想让 attention 快就必须近似。FlashAttention 的价值就在于它说明了一件事即使不近似只要实现方式对exact attention 也能快很多。这个结论对后面的很多工作都很有启发。3它已经变成基础设施级组件现在很多训练框架、推理引擎、长上下文模型一提高性能 attention基本都会碰到 FlashAttention。这说明它已经不只是某篇论文里的方法而是逐渐变成了 Transformer 生态中的基础设施级组件。12. FlashAttention 和近似注意力怎么理解1它们解决的是不同问题近似注意力方法的核心通常是从数学形式上减少 attention 的计算或存储开销。比如稀疏化、低秩化、核方法等都在试图减少理论复杂度。FlashAttention 则不同它基本保留标准 attention 数学形式重点放在如何把标准 attention 在真实 GPU 上算得更高效。2所以它不是“替代一切近似方法”FlashAttention 很强但它并不意味着以后所有长上下文问题都只靠 FlashAttention 就够了。当序列真的极长时二次计算本身仍然是问题。所以在更极端的长度下稀疏化、压缩、窗口化等方法仍然有意义。更准确地说FlashAttention解决的是“标准 attention 该怎么高效实现”近似方法解决的是“当标准 attention 本身太贵时要不要改数学形式”。13. 从更高一层看FlashAttention 到底是什么1它不是新的注意力定义我觉得理解 FlashAttention 最关键的一点是它不是新的 attention 公式而是新的 attention 执行方式。这个区别很重要。因为它并没有改模型定义而是改了模型在硬件上的落地方式。2它本质上是“attention 的 IO 重排”如果只用一句话概括 FlashAttention我觉得最准确的是FlashAttention 是一种基于分块和在线 softmax 的 IO-aware exact attention 实现。这个定义比单纯说“更快 attention”更准确。3它代表了一种非常重要的研究思路我觉得 FlashAttention 这条线真正厉害的地方不只是把 attention 做快了而是它代表了一种思路算法设计不能只盯数学公式还得盯硬件的数据流。这个思路在大模型时代会越来越重要。14. 一点理解1FlashAttention 最漂亮的地方我觉得 FlashAttention 最漂亮的地方在于它不是靠“偷懒少算一点”来变快而是靠“别把没必要的东西搬来搬去”来变快。这个想法其实非常朴素但也非常有效。2它为什么看起来像工程优化却又很有方法味因为它既有工程味也有算法味。如果只看表面你会觉得它像 kernel 优化但再往里看它其实重新设计了 attention 的执行顺序和 softmax 归约方式。所以它不是纯工程 patch而是一个很完整的方法。3怎么记 FlashAttention如果只是为了学习我觉得可以把 FlashAttention 记成四句话标准 attention 慢不只是因为二次复杂度还因为大矩阵显存读写很重FlashAttention 不显式存整个 attention matrix而是按块计算并在线完成 softmax它保持和标准 attention 数学等价因此是 exact attention所以它本质上是在做 attention 的 IO 优化而不是近似替代。15. 参考鸣谢FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesshttps://arxiv.org/abs/2205.14135FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioninghttps://arxiv.org/abs/2307.08691FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisionhttps://arxiv.org/abs/2407.0860816. 注这篇主要是个人学习整理重点放在主线理解文中没有展开很多实现细节比如 backward pass 的重计算策略、causal mask 的具体 block 处理方式、以及不同 GPU 架构下 kernel 设计的差异才疏学浅欢迎批评、指导和交流有错误望大家及时指正
FlashAttention
发布时间：2026/7/8 4:11:52
FlashAttention1. 使用背景Transformer 很强但 attention 一直有一个很现实的问题又慢又吃显存。如果序列长度是NNN标准自注意力里最核心的中间量是SQK⊤ S QK^\topSQK⊤它的形状是N×NN\times NN×N。这意味着随着序列变长attention 的中间结果会迅速膨胀。早期很多人一看到 attention 是二次复杂度就会自然想到是不是只能去做近似、稀疏化、低秩化。但 FlashAttention 提出了一个非常关键的观点很多时候 attention 慢不只是算得多更是因为显存读写太重。也就是说问题不只是 FLOPs不只是数学复杂度还包括一个很重要的系统因素IO输入输出 / 内存访问成本。FlashAttention 就是在这个背景下提出来的。它最核心的思想可以概括成一句话不改变标准 attention 的数学结果而是通过分块、融合和在线 softmax把大规模中间矩阵的显存读写降下来。2. 理论基础1标准 attention 到底在算什么给定输入隐藏状态先通过线性变换得到QXWQ,KXWK,VXWV QXW_Q,\qquad KXW_K,\qquad VXW_VQXWQ,KXWK,VXWV标准 scaled dot-product attention 写成Attention(Q,K,V)softmax(QK⊤d)V \text{Attention}(Q,K,V)\text{softmax}\left(\frac{QK^\top}{\sqrt{d}}\right)VAttention(Q,K,V)softmax(dQK⊤)V如果把它拆开看其实主要有三步先算分数矩阵SQK⊤d S\frac{QK^\top}{\sqrt{d}}SdQK⊤再做 softmax得到注意力权重Psoftmax(S) P\text{softmax}(S)Psoftmax(S)最后和VVV相乘OPV OPVOPV这个流程数学上很清楚但在实现上会有一个大问题中间矩阵SSS和PPP都是N×NN\times NN×N。2为什么这会成为瓶颈当序列长度NNN很大时SSS和PPP的存储、读写都非常重。即使 GPU 算力很强如果每一步都要频繁把这些大矩阵在 HBM高带宽显存和片上 SRAM / cache 之间来回搬运速度也会被拖慢。所以 attention 的瓶颈并不只是“乘法很多”而是中间结果太大内存访问代价太高。3为什么普通实现会浪费在常规实现里通常会先把SQK⊤SQK^\topSQK⊤整个算出来并写到显存再从显存里读出来做 softmax得到PPP再把PPP写回显存最后再读出PPP和VVV去算输出。这个过程最大的问题在于很多中间结果只是临时用一下却被完整写回了大显存。这就是 FlashAttention 最想解决的地方。3. FlashAttention 的核心思想1一句话理解FlashAttention 最核心的想法就是不要把整个 attention matrix 显式落到显存里而是在片上做小块计算边算边归约直接得到输出。也就是说它不是先完整构造SQK⊤,Psoftmax(S) SQK^\top,\qquad P\text{softmax}(S)SQK⊤,Psoftmax(S)再去算输出而是把这几个步骤融合在一起以分块方式完成。2为什么叫 IO-awareFlashAttention 论文里一个非常关键的词就是IO-awareness。传统算法通常更关注算术复杂度比如时间复杂度是不是O(N2)O(N^2)O(N2)FlashAttention 则进一步强调在现代 GPU 上显存和片上缓存之间的数据搬运本身就是主要开销。所以 FlashAttention 并不是去近似 attention而是去优化 attention 的数据流。3最重要的一点它是 exact attention这一点特别重要。FlashAttention 不是 Performer、Linformer 那种近似 attention 方法。它并没有改 attention 的数学定义而只是改了实现方式。所以它的核心价值不是“牺牲精度换速度”而是在保持精确等价的前提下把显存访问和运行时间都降下来。4. FlashAttention 是怎么做的1分块tiling假设把Q,K,VQ,K,VQ,K,V都按块切开。比如把QQQ按行分成若干 query block把K,VK,VK,V按列分成若干 key/value block。那么 attention 就不再是一次性拿整块大矩阵去算而是按 block 逐块处理。直观上理解就是整张大表一次性装不下那就分小块搬进片上缓存里一块一块算。2问题softmax 不是一个纯局部操作如果只是矩阵乘法分块很好理解。但 attention 里还有 softmax而 softmax 涉及整行归一化softmax(si)jesij∑kesik \text{softmax}(s_i)_j\frac{e^{s_{ij}}}{\sum_k e^{s_{ik}}}softmax(si)j∑kesikesij这里的分母要看整行所有元素所以它不像普通矩阵乘法那样天然可以局部分开。这也是 FlashAttention 真正巧的地方它用在线online方式维护 softmax 所需的统计量。3在线 softmax对于某一行 attention score常规 softmax 会先拿到整行所有值再求最大值、再求指数和。FlashAttention 则是在分块扫描时逐步维护当前见过的最大值当前对应的归一化系数当前输出累积值。这样一来即使一整行分散在多个 block 中也不需要把整行完整写回显存后再统一 softmax。所以 FlashAttention 的关键点不只是“分块”而是分块 在线 softmax 输出累积。5. 在线 softmax 的直观理解1为什么 softmax 通常要先减最大值普通 softmax 为了数值稳定常写成softmax(si)jesij−mi∑kesik−mi \text{softmax}(s_i)_j\frac{e^{s_{ij}-m_i}}{\sum_k e^{s_{ik}-m_i}}softmax(si)j∑kesik−miesij−mi其中mimaxjsij m_i\max_j s_{ij}mijmaxsij这样做是为了防止指数爆炸。所以如果想分块做就必须一边扫描 block一边正确维护这个最大值和归一化项。2FlashAttention 的核心统计量对于某一行FlashAttention 会在扫描 block 的过程中不断更新行最大值mmm行归一化和lll输出向量累计值。每看到一个新的 key block就先算当前 block 对这一行的局部分数再和旧的m,lm,lm,l合并更新。这样到最后虽然没有显式存整个SSS和PPP但最终得到的输出和标准 attention 完全一致。3本质上在做什么如果用一句话概括这个过程我觉得最合适的是把原来“先得到完整注意力矩阵再做 softmax”这件事改写成“扫描中间块时同步完成归一化和加权求和”。所以它不是把 attention 简化了而是把计算顺序重新排了一遍。6. FlashAttention 到底省了什么1不是把二次复杂度变成线性这是最容易误解的一点。FlashAttention 并没有把 attention 的理论计算复杂度从O(N2)O(N^2)O(N2)变成O(N)O(N)O(N)。从数学上讲所有 query-key 对之间的交互还是要算。它真正优化的是中间矩阵的显式物化和反复显存访问。2它主要省的是 HBM 访问在 GPU 上HBM 访问代价远高于片上 SRAM / register 的访问。FlashAttention 通过块级计算让很多中间值留在片上而不是反复写回大显存。所以它真正省的不是“注意力不用算了”而是很多原本没必要写出的大矩阵不再落到 HBM。3为什么这会带来显著提速因为在长序列下attention 往往不是单纯 compute-bound而是明显受内存访问限制。一旦把显存读写压下去GPU 的有效利用率就能明显提高。所以 FlashAttention 的加速来源本质上是更合理的数据流而不是近似。7. FlashAttention 和普通 fused kernel 有什么区别1不是简单 kernel fusion看到这里很多人会觉得“这不就是把几个 kernel 融合一下吗”其实比普通 fusion 更进一步。普通 kernel fusion 更多是把几个连续操作合并以减少 launch overhead 和部分中间读写FlashAttention 则是从 attention 的算法结构本身出发重新设计了 block-wise 的 IO 流程。所以它不只是“把 kernel 拼起来”而是把 attention 的执行顺序按 GPU 内存层级重新组织了一遍。2为什么这很重要因为 attention 里的问题不只是 kernel 多而是中间矩阵太大。单纯 fusion 并不能自动解决“整块 attention matrix 落显存”的问题。FlashAttention 真正解决的是这个更深层的问题。8. FlashAttention 的训练和推理意义1训练里最直接的收益在训练阶段FlashAttention 可以明显降低 attention 中间激活带来的显存占用。这意味着能撑更长序列能用更大 batch或者在同样显存下训练更大的模型。所以 FlashAttention 一开始最重要的影响其实是在训练侧。2推理里也很重要在推理阶段尤其是 prefill 阶段长上下文 attention 本身依然很重。FlashAttention 仍然能带来明显收益。不过在 decode 阶段很多时候瓶颈又会更多落在 KV Cache 访问上。所以在推理里FlashAttention 和 KV Cache 往往是互补关系而不是替代关系。3和 KV Cache 怎么配合理解我觉得最简单的理解方式是FlashAttention 主要优化“单次 attention 怎么算”KV Cache 主要优化“历史 K/V 怎么复用”。二者关注点不同但在大模型推理里通常会一起出现。9. FlashAttention-2 在改什么1为什么还要有 FlashAttention-2初代 FlashAttention 已经把 IO 问题解决得很好了但它并不意味着 GPU 利用率就已经到头。FlashAttention-2 继续往前推进关注的重点更偏工程实现和并行划分怎样让 GPU 上的线程块、warp 和 matmul 工作分配得更合理。2核心改进方向FlashAttention-2 的主线大致可以概括成三件事减少不必要的非 matmul FLOPs提高不同线程块之间的并行度优化单个线程块内部不同 warp 的工作划分减少共享内存通信。所以它相对初代的重点已经有点变化初代更像“把 attention 变成 IO-aware”二代更像“在这个框架下把 GPU 并行度和吞吐继续榨出来”。3为什么它还是很重要因为很多时候真正影响实际速度的不是“有没有一个好算法”而是“这个算法在硬件上吃得够不够满”。FlashAttention-2 做的就是把这件事继续推进。10. FlashAttention-3 又在改什么1为什么还会有 3到 FlashAttention-3 时重点已经更偏向新硬件特性尤其是 Hopper 架构 GPU。它关注的是怎样利用新一代 GPU 的异步执行能力、Tensor Memory AcceleratorTMA、以及低精度能力把 attention 再往前推。2主线变化如果很粗地说FlashAttention 1解决 IO-aware exact attentionFlashAttention 2优化并行划分和吞吐FlashAttention 3进一步拥抱新硬件特性与低精度执行。所以这一条线越往后越明显能看出来它已经不只是一个“attention 小技巧”而是在逐渐变成 attention 的高性能实现主线。11. FlashAttention 为什么这么重要1它改变了大家看 attention 的方式早期一提 attention很多人第一反应就是二次复杂度然后想到各种近似。FlashAttention 提醒大家在真实硬件上算法快不快不只取决于 FLOPs也取决于 IO。这一点其实很重要因为它把问题从“数学复杂度”推进到了“系统复杂度”。2它证明了 exact 也可以很快以前很多人会下意识觉得想让 attention 快就必须近似。FlashAttention 的价值就在于它说明了一件事即使不近似只要实现方式对exact attention 也能快很多。这个结论对后面的很多工作都很有启发。3它已经变成基础设施级组件现在很多训练框架、推理引擎、长上下文模型一提高性能 attention基本都会碰到 FlashAttention。这说明它已经不只是某篇论文里的方法而是逐渐变成了 Transformer 生态中的基础设施级组件。12. FlashAttention 和近似注意力怎么理解1它们解决的是不同问题近似注意力方法的核心通常是从数学形式上减少 attention 的计算或存储开销。比如稀疏化、低秩化、核方法等都在试图减少理论复杂度。FlashAttention 则不同它基本保留标准 attention 数学形式重点放在如何把标准 attention 在真实 GPU 上算得更高效。2所以它不是“替代一切近似方法”FlashAttention 很强但它并不意味着以后所有长上下文问题都只靠 FlashAttention 就够了。当序列真的极长时二次计算本身仍然是问题。所以在更极端的长度下稀疏化、压缩、窗口化等方法仍然有意义。更准确地说FlashAttention解决的是“标准 attention 该怎么高效实现”近似方法解决的是“当标准 attention 本身太贵时要不要改数学形式”。13. 从更高一层看FlashAttention 到底是什么1它不是新的注意力定义我觉得理解 FlashAttention 最关键的一点是它不是新的 attention 公式而是新的 attention 执行方式。这个区别很重要。因为它并没有改模型定义而是改了模型在硬件上的落地方式。2它本质上是“attention 的 IO 重排”如果只用一句话概括 FlashAttention我觉得最准确的是FlashAttention 是一种基于分块和在线 softmax 的 IO-aware exact attention 实现。这个定义比单纯说“更快 attention”更准确。3它代表了一种非常重要的研究思路我觉得 FlashAttention 这条线真正厉害的地方不只是把 attention 做快了而是它代表了一种思路算法设计不能只盯数学公式还得盯硬件的数据流。这个思路在大模型时代会越来越重要。14. 一点理解1FlashAttention 最漂亮的地方我觉得 FlashAttention 最漂亮的地方在于它不是靠“偷懒少算一点”来变快而是靠“别把没必要的东西搬来搬去”来变快。这个想法其实非常朴素但也非常有效。2它为什么看起来像工程优化却又很有方法味因为它既有工程味也有算法味。如果只看表面你会觉得它像 kernel 优化但再往里看它其实重新设计了 attention 的执行顺序和 softmax 归约方式。所以它不是纯工程 patch而是一个很完整的方法。3怎么记 FlashAttention如果只是为了学习我觉得可以把 FlashAttention 记成四句话标准 attention 慢不只是因为二次复杂度还因为大矩阵显存读写很重FlashAttention 不显式存整个 attention matrix而是按块计算并在线完成 softmax它保持和标准 attention 数学等价因此是 exact attention所以它本质上是在做 attention 的 IO 优化而不是近似替代。15. 参考鸣谢FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesshttps://arxiv.org/abs/2205.14135FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioninghttps://arxiv.org/abs/2307.08691FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisionhttps://arxiv.org/abs/2407.0860816. 注这篇主要是个人学习整理重点放在主线理解文中没有展开很多实现细节比如 backward pass 的重计算策略、causal mask 的具体 block 处理方式、以及不同 GPU 架构下 kernel 设计的差异才疏学浅欢迎批评、指导和交流有错误望大家及时指正
