KaiSpace
Login
tech

Flash Attention 3/4 笔记

Flash Attention笔记

在 FlashAttention（以及几乎所有现代 GPU 算子）中，计算绝不是一个 token 一个 token 算的，而是用masked attention一个tile一个tile算

  • 硬件瓶颈：GPU 的 Tensor Core（MMA 单元）非常“饥渴”。如果你只给它算一个 token 的注意力，它的计算密度太低，根本跑不满，大部分时间都在等内存传输 。
  • 分块（Tiling）：FlashAttention 的核心就是将 N 个 token 划分成若干个 Tile 。在 Blackwell 架构上，由于第 5 代 Tensor Core 的硬约束，这个 Tile 的大小通常被固定为 128×128 。
  • 并行维度：GPU 会同时启动成百上千个线程块（CTA），每个 CTA 负责处理一个 128×128 的小矩阵块

Flash attention 3

1. 数学定义与性能目标

FlashAttention-3 是对注意力机制的高性能异步实现，旨在 H100 等硬件上达到接近 GEMM 的利用率 。

  • 核心公式：

    O=softmax(QKTd)VO = \text{softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) V

    其中 Q,K,V∈RN×dQ, K, V \in \mathbb{R}^{N \times d}，NN 是序列长度，dd 是 Head 维度 。

  • 计算拆解：

    1. GEMM 0：计算 S=QKTS = Q K^T 。
    2. Softmax：对 SS 进行在线更新，计算归一化权重 PP 。
    3. GEMM 1：计算 O=PVO = P V 。

  • 性能飞跃：在 H100 上，FP16 性能达 740 TFLOPS（理论峰值 75%），FP8 性能接近 1.2 PFLOPS 。

2. 生产者-消费者模式与 Warp 专门化 (Warp-Specialization)

在 Hopper 上，如果让同一个 Warp 交替执行全局内存访问和高强度的张量核心（Tensor Core）计算，极易造成指令发射窗口的停顿和寄存器资源的碎片化。FA-3 通过显式的角色划分解决了这一问题 ：

  • Producer Warpgroup（只负责搬运）： 只需要极少的线程发起异步 TMA load 指令。首先通过 setmaxnreg 释放不必要的寄存器，然后连续将 Q,Kj,VjQ, K_j, V_j 从全局内存搬运至多级 Circular SMEM Buffer，并通过 mbarrier 提交完成信号 。
  • Consumer Warpgroup（只负责计算）： 承接计算任务，使用 setmaxnreg 重新分配大量寄存器池，用于存放计算所需的巨大中间状态。它通过 wait 监听 TMA 的完成信号，然后直接调用 SS-WGMMA（Shared-to-Shared）和 RS-WGMMA（Register-to-Shared）执行矩阵乘 。

这种解耦使得数据搬运和计算完全重叠，隐藏了访存延迟 。

3. Pingpong 调度与 2-Stage 指令级流水线

这是 FA-3 压榨硬件峰值性能的核心所在。在 H100 上，FP16 GEMM 算力高达 989 TFLOPS，但执行 Softmax 中 exp 运算的特殊功能单元（Multifunction Unit）算力仅有 3.9 TFLOPS 。朴素的串行执行会导致 Tensor Core 长时间空转等待 Softmax 。

  • Pingpong 调度（多 Warpgroup 间协同）： FA-3 利用 bar.sync 强制同步同一 CTA 内的两个 Consumer Warpgroup 的执行相位。当 Warpgroup-1 处于 Softmax 计算阶段（占用特殊功能单元）时，Warpgroup-2 刚好在执行 GEMM 阶段（占用 Tensor Core），反之亦然。这种完美的错峰执行，将 FP16 forward 的性能从 570 TFLOPS 直接拉升至 640 TFLOPS 。总之，Pingpong只是用来保证，一个warpgroup在做softmax的时候另一个warpgroup一定不在做softmax。 虽然其实MMA才是真正时间瓶颈，但是由于MMA要在softmax后做完，我们希望将MMA作为唯一瓶颈，于是让softmax交替掩盖所有softmax的时间。

    Figure 1: Pingpong scheduling for 2 warpgroups to overlap softmax and GEMMs.

  • 2-Stage 流水线（单 Warpgroup 内重叠）： 为了打破 GEMM0→Softmax→GEMM1GEMM_0 \rightarrow Softmax \rightarrow GEMM_1 的串行（RAW 依赖）阻塞，FA-3 在单个 Warpgroup 内部署了 2-Stage 流水线 。在第 jj 次迭代时，允许 WGMMA1(j−1)WGMMA_1^{(j-1)}、Softmax(j)Softmax^{(j)} 和 WGMMA0(j+1)WGMMA_0^{(j+1)} 同时活跃 。

    • 寄存器生存期分析代价： 这种指令级重叠的代价是必须在线程块中额外分配一份 SnextS_{next} 的寄存器副本（增加 Br×Bc×sizeof(float)B_r \times B_c \times \text{sizeof(float)} 的寄存器压力）。这也是为什么 Tile 尺寸不能无限增大，必须在指令吞吐和寄存器溢出之间寻找平衡 。由于 NVCC 的重排策略可能会破坏这种理想的指令重叠，通常需要在 SASS 级别进行反汇编验证其流水线编排是否生效 。

    • 时间窗口Warpgroup 1 的行为 (2-stage)Warpgroup 2 的行为 (2-stage)硬件状态
      窗口 T正在执行 SoftmaxjSoftmax_j正在执行 WGMMA1k−1WGMMA1_{k-1} 与 WGMMA0k+1WGMMA0_{k+1}并行 (Multifunction + Tensor Core)
      窗口 T+1正在执行 WGMMA1jWGMMA1_{j} 与 WGMMA0j+2WGMMA0_{j+2}正在执行 SoftmaxkSoftmax_k并行 (交换角色)

4. FP8 的 Layout 转换与精度修复机制

  • 在 Hopper 架构上，FP8 的引入虽然带来了近 2 倍的理论峰值提升，但也给底层的寄存器管理和数值稳定性带来了巨大的挑战。

    由于你对这些机制比较感兴趣，我们可以从 “硬件对齐” 和 “数值摊平” 两个技术维度拆解：

    1) Layout 转换：解决硬件的“强迫症”

    Hopper 的 FP8 WGMMA 指令对操作数（Operand）的排布要求极其严苛。

    VV 矩阵的转置（In-kernel Transpose）

    • 硬约束： FP8 WGMMA 的 BB 操作数必须是 K-major（内层维度连续）。
    • 冲突点： 在 Attention 的第二个矩阵乘法 O=PVO = P V 中，VV 在显存中通常是 Head-major（维度 dd 连续）存储的 (对于 V 矩阵（维度为 N×d），“Head-major”意味着头维度 d 是内存连续的）。直接读取的话，内存布局不符合 WGMMA 指令的要求。
    • FA-3 的解法： 它没有在计算前做一次全局转置，而是利用 Producer Warpgroup 在数据搬运的“影子时间”里顺便把活干了 。
      • 它利用了 Hopper 特有的 ldmatrix.trans 指令，在从 SMEM 读取到寄存器（RMEM）的过程中直接完成转置，或者在写入 SMEM 时进行重排 。

    累加器布局不一致与 Byte-permute

    • 硬约束： GEMM0GEMM_0（计算 QKTQ K^T）的结果 SS 是以 FP32 存储在寄存器里的。然而，FP32 累加器在寄存器中的逻辑序列 {d0,d1,d2,d3...}\{d_0, d_1, d_2, d_3...\} 与紧接着的 GEMM1GEMM_1（计算 SVS V）所要求的 FP8 操作数 AA 的 Layout 完全对不上 。
    • FA-3 的解法： 如果用传统的跨线程 Shuffle 指令来重排，性能开销会爆炸。FA-3 使用了 byte-permute 指令，在单个线程内部直接对寄存器内的字节进行原地重排（例如将 {d0,d1,d4,d5}\{d_0, d_1, d_4, d_5\} 调换位置）。
      • 为了补偿这种逻辑上的列置换，它还在前面提到的 VV 矩阵转置过程中，配套调整了 VV 块的行序，从而实现了“负负得正”的逻辑一致性 。

    2) 精度修复：抑制“离群值”的负面影响

    FP8（特别是 e4m3 格式）的动态范围非常有限，很容易遇到 Outlier（离群值）。这些极少数但数值极大的元素会导致量化精度大幅下降。

    分块量化（Block Quantization）

    • 原理： 传统的做法是对整个 Tensor 取一个 Scaling Factor（缩放因子）。但在长序列中，不同位置的分布差异巨大。
    • 做法： FA-3 针对每个 Br×dB_r \times d 的小分块独立计算量化参数 。
    • 优势： 这种 Tiling 结构天然契合 FlashAttention 的分块计算模型，可以将量化操作融合进 RoPE 或数据搬运过程中，实现零额外开销 。

    非相干处理（Incoherent Processing）

    • 原理： Outlier 往往集中在某些特定的维度上。如果能把这些“尖峰”能量均匀地摊平到所有维度，量化误差就会显著降低 。
    • 数学手段： 它引入了一个随机正交矩阵 MM（由 Hadamard 变换和随机 ±1\pm 1 对角阵组合而成）。
    • 操作： 实际上是计算 (QM)×(KM)T(QM) \times (KM)^T。因为 MM 是正交矩阵，MMT=IM M^T = I，所以数学结果 QKTQ K^T 保持不变 。
    • 结果： 这一步被称为“摊平”。实验证明，这种处理让 FP8 的 RMSE（均方根误差）比普通做法降低了 2.6 倍，使得 FP8 训练在大模型上变得真正可用 。

Flash Attention 4

FlashAttention-4 (FA-4) 的核心设计动机是系统性地解决 NVIDIA Blackwell 架构（如 B200/GB200）上的“非对称硬件缩放”瓶颈 。在 Hopper 架构上，MMA（矩阵乘累加）是主要瓶颈；但在 Blackwell 上，Tensor Core 的吞吐量翻倍，而 SMEM（共享内存）带宽和 MUFU（指数计算单元）的吞吐量却几乎保持不变 。此外，Blackwell 的张量核心强制使用 128×N128\times N 的 tile，并将累加器从寄存器移至了 TMEM (Tensor Memory) 。

为了应对计算瓶颈向 SMEM 流量和 Softmax 指数运算的转移，FA-4 重新协同设计了算法与底层内核流水线 。以下是其核心设计维度的详细拆解：

1. forward pipeline：1-CTA Ping-Pong 与指数软件仿真

Figure 1: FlashAttention-4 forward pipeline.

在 M=N=d=128M=N=d=128 的配置下，MMA 与指数计算（exp）均需要 1024 cycles，形成了并列瓶颈 。FA-4 通过以下机制打破这一限制：

  • Ping-Pong 调度： 设计了交错流水线，将“高 Q tile 的 MMA”与“低 Q tile 的 softmax”重叠执行 。这样可以让 Tensor Core 和 MUFU 始终保持在各自的 cadence 上同时忙碌 。（不分离softmax其实是因为pipeline启动时可以立刻并行算两个softmax）
  • FMA 指数逼近： 单纯依赖 16 ops/cycle 的 MUFU 已经无法满足需求，因此 FA-4 利用大量空闲的 FMA 单元，通过多项式逼近（Cody-Waite 范围归约 + Sollya 拟合）来仿真 2x2^x 。为了避免产生寄存器溢出（Softmax 对寄存器极其敏感），仅对 10-25% 的元素使用 FMA 仿真路径，其余仍由 MUFU 处理 。
  • 条件 Softmax 重缩放： 引入了阈值 τ=8.0\tau = 8.0。只有当最大值的跳变大于 8 bits 时，才会执行重缩放操作，从而跳过了大部分昂贵的校正乘法 。
  • TMEM 传递： 与 Hopper 利用寄存器传递不同，中间结果 PP 通过 TMEM 传递，这使得输出重缩放操作可以被剥离到一个独立的 correction warpgroup 中，移出关键路径 。

2. backward propogation：2-CTA UMMA 与 DSMEM 共享

Figure 2: FlashAttention-4 backward computation graph.

反向传播包含 5 个 MMA 操作（S,dP,dV,dQ,dKS, dP, dV, dQ, dK），此时 SMEM 的流量比 MMA 高出 30%，成为绝对的性能瓶颈 。

  • 突破 1-CTA 限制： 1-CTA 模式受限于 M≤128M \leq 128，无法有效均摊操作数 B 的 staging 成本 。FA-4 转而利用 Blackwell 特有的 2-CTA MMA 模式 。

  • DSMEM 数据交换： 两个 CTA 作为一个 Cluster，将 MM 扩展至 256 。它们通过分布式共享内存 (DSMEM) 交换各自的一半 dSdS tile 。这意味着两个 CTA 可以共同分担操作数 B 在 SMEM 中的 staging 流量，使其近乎减半 。

    Figure 3: 2-CTA backward dQ step using DSMEM exchange.

  • 归约维度重排： 这种设计将 dQdQ 的归约从“K 序列外循环”挪到了“N 维度内”，使得 dQdQ 在写回全局内存时的 atomicAdd 冲突与写入次数同样减半 。

    • 两个 CTA（CTA 0 和 CTA 1）不再各干各的，而是组成一对。它们利用 Blackwell 硬件支持的 DSMEM（分布式共享内存） 交换彼此手里的数据 。

      • 之前：两个 CTA 各自处理 M 行，归约维度（列宽）只有 N。

      • 之后：通过交换数据，每个 CTA 现在只负责处理 M/2 行（行数减半），但它拿到了两个 CTA 共同拥有的数据，使其归约维度（列宽）变成了 2N 。

    • B. 为什么叫“从外循环挪到内维度”？

      • “N 维度内”归约：在矩阵乘法指令执行时，算子内部会自动对这 2N 宽度的数据进行累加。因为这种累加发生在寄存器或 TMEM（Tensor Memory）中，速度极快且不消耗全局内存带宽。

      • “K 序列外循环”减少：因为单个 CTA 一次性处理了原来两倍宽的 K 序列信息（2N 而不是 N），它对全局内存的写入频率自然就降低了 。

      • 在反向传播中，我们要算的是 Query 的梯度 dQ。其公式是：

        dQ=dS⋅KdQ=dS⋅K

        • dS 的形状：(N×N) —— 代表序列中每个 Token 对其他所有 Token 的影响梯度。

        • K 的形状：(N×d) —— 代表 Key 矩阵。

        • dQ 的形状：(N×d) —— 最终我们要得到的 Query 梯度 。

          关键点来了： 如果你要计算 dQ 的某一行（也就是某一个特定的 Query Token 的梯度），根据矩阵乘法规则，你需要用 dS 的整整一行去乘以 K 的整整一列。

        • dS 的这一行有 N 个元素（N 是总序列长度，比如 32K）。

        • 但是你的 CTA 每次处理的分块 M×N 里的 N 只是一个 Tile（方块）的宽度（比如 128） 。

          结论： 当你处理完一个 128×128 的块时，你只计算了这一个 Token 对应的那 32,768 个关联项中的 128 个。剩下的 32,640 个项还在其他的 K 块里等着你。

        2. 为什么要“累加”？（工程层面的操作）

        因为 GPU 没法一次性吞掉整个 32K 长度的 K 矩阵，FlashAttention 必须在“外循环”里不断地换下一批 K 块（Tiles）进来算 。

        1. 第一次循环：CTA 拿 dS 的前 128 列和 K 的前 128 行算，得到一个 M×d 的临时值。
        2. 第二次循环：CTA 拿 dS 的第 129-256 列和 K 的第 129-256 行算，得到另一个临时值。
        3. 累加：你必须把这两个（以及后面所有的）临时值加在一起，才能凑出这个 Token 真正的、完整的梯度 dQ。

        3. FA-4 优化了“哪里的”累加？

        你之前看到的 2-CTA 模式，其实是把这种累加做了一个“合并同类项”的动作 ：

        • 传统做法：每个 CTA 算完自己的 128×128 就急着往显存里写一次 atomicAdd（原子累加）。
        • FA-4 做法：让两个 CTA 组队。它们先在芯片内部（通过 DSMEM 交换数据）把各自的局部结果合在一起，变成一个更宽的 128×256 的局部结果 。

3. Causal 调度策略优化

在因果注意力中，由于序列中的每个 Token 只能看到它之前的 Token，计算量呈“三角形”分布：序列开头的 Token 计算量极小，而末尾的 Token 需要与前面所有的 Key 计算，任务极重 。朴素的按序调度会导致 SM 在处理 Causal 掩码时效率低下。FA-4 采用了精细的调度控制：

  • 混合 LPT 调度： 采用“batch 最外层 + heads 切成不溢出 L2 的 section + mblock 反向遍历”的策略（最长处理时间优先），在解决 SM 负载不均的同时，确保了 KV 块加载的 L2 缓存命中率 。

    • 外层：Batch 遍历。保持 Batch 作为最外层循环，确保基础的局部性 。

      中层：Heads 切片（Sectioning）。将 Heads 切成大小足以塞进 L2 缓存的“Section”。这样在处理这一组 Heads 时，KV 数据能稳稳地留在 L2 中被复用 。

      内层：反向 Mblock 遍历。在 Section 内部，不再按顺序执行，而是采用反向遍历（即 LPT，长任务优先）。先处理计算量最大的块，让长任务尽早开始，从而让所有 SM 尽量在同一时间点完成工作 。

  • 确定性反向传播： 在反向传播（Backward）中，情况更复杂。为了支持强化学习（RL）训练，结果必须是“确定性”的（即多次运行结果完全一致），但这通常需要加锁，会严重拖慢速度 。

    FA-4 引入了 SPT（最短处理时间优先） 调度策略来对消这个开销 ：

    • 原理：在 Causal 模式下，对于 dQ 的写入，FA-4 让“短任务”优先执行 。

    • 逻辑：因为短任务需要的归约次数少，让它们先算完并写入。配合信号量（Semaphore）控制，这种顺序能保证“首个写入永远不会被前驱任务阻塞” 。

    • 效果：即便加了确定性锁，其性能依然能达到非确定性版本的 75% 左右（在 32K 序列下达 950 TFLOPS），为训练提供了极高的可复现性代价比 。

4. CuTe-DSL Python JIT 编译

在工程实现上，FA-4 没有使用 FA-3 的 CUTLASS C++ 模板，而是完全由嵌入 Python 的 CuTe-DSL 编写 。

  • 极致的迭代速度： 零 C++ 代码使得 JIT 编译时间比 C++ 模板缩短了 20-30 倍 。

  • PTX 逃生口： 当 DSL 无法覆盖极底层的硬件行为时，框架保留了直接内联 PTX 汇编的“逃生口”，从而维持了对指令和内存分配的绝对表达力 。

Comments

No comments yet.