【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×d，nnn为序列长度，ddd为头维度。
把公式拆成可执行的计算链路，一共包含 4 个核心步骤，且天然串行：

计算注意力分数：执行矩阵乘法S=QK⊤S = QK^\topS=QK⊤，得到尺寸为n×nn \times nn×n的分数矩阵；
数值稳定处理：对每行减去最大值m=max⁡(Srow)m = \max(S_{\text{row}})m=max(Srow)，防止指数运算溢出；
Softmax 归一化：逐元素计算指数exp⁡(S)\exp(S)exp(S)，按行求和得到分母，最终归一化得到权重矩阵PPP；
加权求和输出：执行矩阵乘法O=PVO = PVO=PV，得到最终输出。

原生实现会将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 到 V4，FlashAttention 系列的演进也印证了一个朴素的道理：顶尖的算子优化，永远是算法逻辑与硬件特性的双向奔赴。