【LLM系列】深入理解FlashAttention V4：从计算链路演进看非对称硬件时代的注意力优化

自 2022 年 FlashAttention 横空出世以来这一系列注意力算子几乎成为了大模型训练与推理的算力基石。每一代 GPU 架构升级FlashAttention 都会同步迭代精准击穿当下的性能瓶颈。2026 年发布的 FlashAttention V4更是将优化思路从 “访存优化”“搬运优化” 推向了算法与内核流水线协同设计的新阶段在 NVIDIA Blackwell 架构上将注意力计算的张量核心利用率推至 71%逼近纯矩阵乘法的资源效率。要真正理解 FlashAttention V4 的突破不能只看零散的技术点而要顺着 “计算链路→硬件瓶颈→代际优化” 的脉络看清楚每一次升级到底解决了哪一环的问题。一、先搞懂两个基础计算链路与硬件层级所有性能优化都建立在两个底层逻辑之上自注意力本身的计算步骤以及 GPU 的内存 - 计算资源分布。1.1 自注意力的完整计算链路标准缩放点积注意力的公式如下其中Q∈Rn×dQ \in \mathbb{R}^{n \times d}Q∈Rn×d、K∈Rn×dK \in \mathbb{R}^{n \times d}K∈Rn×d、V∈Rn×dV \in \mathbb{R}^{n \times d}V∈Rn×dnnn为序列长度ddd为头维度。把公式拆成可执行的计算链路一共包含 4 个核心步骤且天然串行计算注意力分数执行矩阵乘法SQK⊤S QK^\topSQK⊤得到尺寸为n×nn \times nn×n的分数矩阵数值稳定处理对每行减去最大值mmax⁡(Srow)m \max(S_{\text{row}})mmax(Srow​)防止指数运算溢出Softmax 归一化逐元素计算指数exp⁡(S)\exp(S)exp(S)按行求和得到分母最终归一化得到权重矩阵PPP加权求和输出执行矩阵乘法OPVO PVOPV得到最终输出。原生实现会将SSS和PPP完整存储在显存中带来两个致命问题一是显存占用随序列长度平方增长长序列直接爆显存二是所有中间结果往返读写高带宽显存HBM计算单元大部分时间在等待数据算力利用率通常仅 10%~20%。1.2 GPU 的内存 - 计算金字塔GPU 的资源遵循 “容量越小、速度越快” 的金字塔规律这是所有算子优化的底层依据。内存层级从慢到快HBM 显存全局大仓库容量几十 GB访问延迟高、带宽有限是最容易成为瓶颈的资源共享内存每个计算单元SM自带的高速缓存容量仅几十到上百 KB速度是 HBM 的数十至上百倍寄存器线程私有存储空间容量 KB 级速度最快直接喂给计算单元。计算单元每个 SM 内的三类执行单元张量核心专用矩阵乘法加速器算力极强是大模型计算的核心动力新一代 GPU 的算力提升基本来源于此FMA 乘加单元CUDA 核心通用计算单元可执行各类加减乘除单精度算力远低于张量核心但数量多、通用性强SFU 特殊功能单元专门负责指数、对数、三角函数等特殊运算数量极少很容易成为隐性瓶颈。性能优化的本质就是不断识别当前链路中最慢的 “短板”然后通过算法重构让它不再拖后腿。FlashAttention 每一代的演进都精准对应了当时 GPU 上的核心瓶颈。二、三代演进瓶颈是如何一步步转移的从原生注意力到 FlashAttention V3优化主线始终围绕 “减少访存→重叠搬运” 展开瓶颈也从 HBM 带宽逐步转移到了非计算单元。2.1 原生实现被 HBM 拖垮的算力原生实现的完整链路是从 HBM 读 Q/K/V → 算 S → 写 S 回 HBM → 读 S 算 softmax → 写 P 回 HBM → 读 P 和 V 算 O → 写 O 回 HBM。全程所有中间结果都要进出 HBM访存量为O(n2)O(n^2)O(n2)张量核心绝大多数时间都在等数据属于典型的访存瓶颈。2.2 FlashAttention V1/V2分块计算把计算搬进车间V1 和 V2 的核心思路是分块计算 online softmax用少量重计算换取大幅访存减少将 Q 按行切分为多个小块K、V 按列切分为对应小块每次只加载一个 Q 块和一个 KV 块到共享内存中完成计算循环遍历所有 KV 块计算过程中在寄存器内维护每行的最大值、指数和动态重缩放保证数值稳定全程不保存完整的 S 和 P 矩阵。这一优化将 HBM 访存量从O(n2)O(n^2)O(n2)降至O(n)O(n)O(n)直接击穿了 HBM 带宽瓶颈张量核心利用率提升至 40%~50%。但此时数据加载与计算仍是串行的加载下一块数据时计算单元处于闲置状态。2.3 FlashAttention V3异步流水线边计算边搬运针对 Hopper 架构新增的 TMA张量内存加速器和 WGMMA 指令V3 实现了计算与数据搬运的完全重叠计算当前分块的同时TMA 在后台异步加载下一分块的数据形成 “计算当前块→预加载下一块” 的软件流水线将数据搬运的等待时间完全隐藏在计算中。V3 解决了搬运调度的瓶颈利用率进一步提升至 50%~60%。但此时也出现了新的矛盾矩阵乘法已经足够快softmax 指数运算、数值重缩放、共享内存读写这些 “周边操作”开始拖慢整体节奏。这个问题在 Hopper 上尚不突出直到 Blackwell 架构出现被彻底放大。三、FlashAttention V4面向非对称硬件的全链路重构Blackwell 架构B200/G200的硬件升级呈现显著的非对称缩放特征张量核心 MMA 分块从 64×128 升级为 128×128单周期计算量翻倍理论峰值算力大幅提升每个 SM 新增 256KB 张量内存TMEM专门用于存储张量核心的中间结果但 SFU 特殊功能单元、共享内存带宽、HBM 带宽并未随张量核心同比例增长。简单来说主力生产线张量核心效率翻倍了但辅助工序指数计算、数据同步、共享内存读写的产能没跟上导致生产线经常停工待料。FlashAttention V4 的所有优化都瞄准了这些非计算瓶颈。3.1 乒乓调度让张量核心永不闲置V4 最核心的创新是双 Q 块乒乓调度Ping-pong Schedule从根本上解决了 “主计算等待杂活” 的问题。在 V3 及之前单个 Q 块遍历 KV 分块的流程是串行的算完 QK 矩阵乘法 → 算 softmax → 乘 V 累加 → 加载下一个 KV 块。其中 softmax、重缩放等操作会占用主链路时间让张量核心出现短暂闲置。V4 让每个计算单元同时维护 2 个 Q 分块QA、QB交替占用张量核心时间片 1张量核心执行QA×Ki⊤Q_A \times K_i^\topQA​×Ki⊤​矩阵乘法时间片 2张量核心切换执行QB×Ki⊤Q_B \times K_i^\topQB​×Ki⊤​与此同时FMA 单元在后台完成 QA 上一轮结果的 softmax 计算、重缩放与结果累加时间片 3张量核心切回 QA 的下一轮计算同时处理 QB 的后处理。两个 Q 块像打乒乓球一样交替占用主力计算资源所有 softmax、重缩放、数据搬运的杂活全部隐藏在另一块的计算时间中张量核心几乎全程满载。同时 V4 还增设了独立的校正线程组将 softmax 重缩放、数值校验等操作从主计算线程中剥离异步执行进一步不占用主计算路径。3.2 软件近似指数破解 SFU 算力瓶颈Softmax 中的指数运算exp⁡(x)\exp(x)exp(x)原本完全依赖 SFU 特殊功能单元。在 Blackwell 上张量核心算力翻倍但 SFU 数量没有同步增长指数计算直接成为前向传播的限速环节 —— 矩阵乘法算完了还要等 SFU 把指数算完才能继续。V4 的解法非常巧妙用 FMA 单元通过多项式拟合近似计算指数。既然 FMA 单元数量多、且有大量空闲就用多项式逼近的方式模拟 exp 运算把特种技工的活拆分给大量普通工人并行完成整体指数计算吞吐量直接翻倍。配合条件式重缩放优化传统实现每个分块都会执行一次 “减最大值、重缩放” 的数值稳定操作V4 会动态判断当前分块的数值范围仅在可能溢出时才执行重缩放进一步减少无效计算。3.3 TMEM 复用与 CTA 配对破解反向共享内存瓶颈反向传播需要存储注意力权重、中间 softmax 结果等大量中间数据此前全部放在共享内存中共享内存的容量与带宽很容易成为反向瓶颈。V4 结合 Blackwell 的硬件特性做了两层优化TMEM 复用中间结果将注意力权重矩阵 P 等大尺寸中间结果直接存储在 256KB 的 TMEM 中不占用共享内存空间也减少了共享内存的读写流量2-CTA 配对计算两个协作线程阵列CTA组成配对各自加载一半中间数据通过共享内存交换互补部分。单个 CTA 的共享内存需求直接减半同时原子规约操作次数减少 50%大幅缓解共享内存拥堵。此外 V4 还在保证性能的前提下支持确定性执行模式满足训练可复现的工业需求。3.4 全链路工程优化全异步数据通路不仅 Q/K/V 加载注意力掩码生成、输出结果回写、中间数据搬运全部改为异步执行全程无需 CUDA 核心干预动态分块调度针对因果掩码、变长序列场景动态分配分块任务平衡各 SM 负载避免边缘分块的资源浪费CuTe DSL 开发基于 NVIDIA CUTLASS 的 CuTe 领域特定语言开发内核开发效率相比手写 C 提升数十倍同时保持顶尖性能。四、性能表现与代际对比根据官方在 NVIDIA B200 GPU、BF16 精度下的测试数据峰值算力达到 1605 TFLOPs/s对应张量核心理论峰值的71% 利用率接近纯矩阵乘法的资源效率相比 cuDNN 9.13 的官方注意力实现提速 1.3 倍相比 Triton 实现的注意力算子提速 2.7 倍在长序列、因果掩码等复杂场景下性能优势进一步扩大。五、技术意义与生态展望FlashAttention V4 标志着注意力算子优化正式进入非对称硬件协同设计的新阶段。在此之前注意力优化的主线是 “减少访存”“优化搬运”本质都是让数据更快送到计算单元面前而 V4 第一次将优化重心放在了 “匹配非均匀增长的硬件资源” 上 —— 当算力增长不再均匀单纯移植旧算法已经无法发挥新硬件性能必须从算法层面重构计算流程让每一类硬件单元都保持高负载。目前 FlashAttention V4 已完全开源PyTorch 官方已将其作为 FlexAttention 的高性能后端同时深度融入主流大模型训练框架与推理引擎成为 Blackwell 时代大模型基建的核心组件。对于长上下文大模型、大尺度预训练等场景V4 带来的算力利用率提升将直接转化为训练时长缩短与推理成本下降。从 V1 到 V4FlashAttention 系列的演进也印证了一个朴素的道理顶尖的算子优化永远是算法逻辑与硬件特性的双向奔赴。