CUDA FlashAttention：从在线 Softmax 到长序列 IO 极限

大模型把上下文从 2K 推到 8K、32K 以后，Attention 的问题不再只是“两个 GEMM 很贵”。标准实现会把 $QK^T$ 写回显存，经过 mask 和 Softmax 后再读回来与 $V$ 相乘。计算量仍是二次的，真正先让实现失去扩展性的往往是这个 $N\times N$ 中间矩阵。

本文只讨论精确的 dense causal attention，不用稀疏近似换速度。目标是从一个可核验的 CUDA 数据流出发，解释 FlashAttention 为什么有效、如何分层到 CTA/warp/MMA、在什么 shape 下应回退，以及生产实现和教学骨架之间还差哪些工作。

1. 问题、shape 与性能指标#

设

Q,K,V\in\mathbb{R}^{B\times H\times N\times D},

大模型中的 scaled dot-product attention 为

O=\operatorname{softmax}\left(\frac{QK^T}{\sqrt D}+M\right)V,

其中 causal mask $M_{ij}=0$ 当 $j\le i$ ，否则为 $-\infty$ 。本文实测使用 $B=1,H=32,D=128$ ，这对应常见大模型的 head geometry； $N$ 从 512 扩展到 8192。

非 causal 前向的两个矩阵乘法约需 $4BHN^2D$ FLOP。causal kernel 若能跳过上三角 tile，实际有效计算约为

F_{\mathrm{causal}}\approx 2BHN^2D.

性能不能只看 TFLOP/s，还要同时记录：端到端 kernel 延迟、输入之外的峰值显存、长序列是否 OOM，以及 FP16 输出相对 FP32 Softmax 参考路径的误差。

2. 精度语义：低精度乘法不等于低精度 Softmax#

本文的 $Q,K,V$ 和输出为 FP16，点积与输出 accumulator 使用 FP32。在线 Softmax 的行最大值 $m$ 、指数和 $l$ 也必须是 FP32：如果 $l$ 用 FP16，长行会更早丢失小概率项；如果不减行最大值，exp 会溢出。

BF16 有更大的指数范围，但尾数更短，仍应保留 FP32 的 $m/l/O$ accumulator。TF32 是 FP32 GEMM 的硬件折中，不是本文 FP16 大模型路径的输入格式。FP8/INT8 可以降低 Q/K/V 流量，但需要 per-tensor 或 per-block scale，且 Softmax 状态仍应提升到 FP32。FP64 吞吐和容量代价都不适合作为大模型默认路径，只用于高精度离线校验。

3. 朴素实现为何先撞上显存#

朴素路径是三个独立阶段：

GEMM 写出 $S=QK^T/\sqrt D$ ；
mask 与 Softmax 读写 $S$ ；
GEMM 读取 $P=\operatorname{softmax}(S)$ 和 $V$ ，写出 $O$ 。

即使只把 $S$ 保存成 FP16，中间矩阵也占

M_S=2BHN^2\ \text{bytes}.

在 $B=1,H=32,N=8192$ 时，仅一份 $S$ 就是 4096 MiB。安全 Softmax 还可能生成 mask、FP32 临时量或额外输出，因此真实峰值高于这个下界。OneFlow 对 Softmax 的拆解给出了同样的工程信号：exp、reduce 和 broadcast 分成多个 kernel 会反复读写整行，优化应先减少全局访存，而不是先抠一条指令。

从 roofline 看，两个大 GEMM 本身有较高算术强度，但 GEMM 之间的 $N^2$ 物化流量没有产生新的模型信息。FlashAttention 不减少 dense attention 的渐近计算量，而是利用片上 SRAM 重排计算，消掉这部分 HBM 往返。

4. 在线 Softmax：不保存整行仍保持精确#

把一行 score 分成 tile。处理第 $t$ 个 tile 前保存运行状态：最大值 $m^{(t-1)}$ 、指数和 $l^{(t-1)}$ 、未归一化输出向量 $o^{(t-1)}$ 。新 tile 的局部最大值和局部指数和为

\tilde m=\max_j s_j,\qquad \tilde l=\sum_j e^{s_j-\tilde m}.

合并时先更新最大值

m^{(t)}=\max\left(m^{(t-1)},\tilde m\right),

再重标定旧状态：

\alpha=e^{m^{(t-1)}-m^{(t)}},\qquad \beta=e^{\tilde m-m^{(t)}},

l^{(t)}=\alpha l^{(t-1)}+\beta\tilde l,

o^{(t)}=\alpha o^{(t-1)}+ \sum_j e^{s_j-m^{(t)}}v_j.

最后输出 $O=o^{(T)}/l^{(T)}$ 。这是代数重排，不是近似；只要每个 tile 都进入合并，结果与整行安全 Softmax 相同，差异来自浮点加法顺序。

图 1：Q tile 常驻寄存器，K/V tile 通过双缓冲 shared memory 流过；每轮只保留 FP32 的 $m/l/O$ 状态，不把 score tile 写回 HBM。

5. 从能跑到高吞吐的 CUDA 数据流#

5.1 CTA tile：让 K/V 被多个 query 复用#

一个 CTA 负责 $B_r$ 行 query。 $Q_i$ 只加载一次并留在寄存器或 shared memory；K/V 按 $B_c$ 个 token 分块。单个 K/V tile 从 HBM 读取一次，却服务 $B_r$ 行 query，避免“一行 query 一个 block”重复读取整个上下文。

shared memory 预算近似为

M_{\mathrm{smem}}\approx 2B_cD\,s+M_Q+M_P,

其中 $s=2$ bytes 对应 FP16。增大 $B_c$ 提高复用，却会增加 shared memory 和寄存器占用，减少每个 SM 的并发 CTA。不存在对所有 $D,N$ 都最优的 tile。

5.2 warp/MMA：把乘加留给 Tensor Core#

$QK^T$ 和 $PV$ 应映射到 mma.sync 支持的 fragment，而不是逐元素 FP32 FMA。warp 先协作装载 FP16 fragment，Tensor Core 以 FP32 累加；Softmax 的 max/sum 则使用 warp shuffle 和少量 shared memory 做跨 warp 合并。

FlashAttention-2 的关键不是换一个 Softmax 公式，而是减少非矩阵乘 FLOP、在单 head 内沿 query 方向增加 CTA 并行，并调整 warp 分工以减少 shared-memory 通信。对 $B\times H$ 很小而 $N$ 很长的情况，这种并行度尤其重要。

5.3 K 维流水与 bank conflict#

Ampere 的 cp.async 可以把 global-to-shared copy 与当前 tile 的 MMA 重叠：stage 0 计算时预取 stage 1，随后交换 buffer。K/V shared layout 要同时满足 128-bit 合并访存和 MMA fragment 的访问方式；简单的二维行主序可能让多个 lane 命中同一 bank。常见处理是 padding 或 swizzle，而不是盲目增大 shared memory。

5.4 向量化、对齐与尾部#

head dim 为 64/128 且首地址 16-byte 对齐时，使用 uint4 或等价的 128-bit load 搬运 8 个 FP16。尾部不整齐、跨页或变长序列必须走 masked load。向量化减少指令数，但不能越界，也不能假设框架给出的任意 stride 都连续。

6. 核心 CUDA 骨架#

下面只展示在线状态更新和 tiled 主循环的边界。生产 kernel 还需要真正的 MMA fragment、swizzle layout、causal 边界和架构特化；benchmark 不使用这段教学骨架冒充工业实现。

1
struct OnlineState {
2
  float row_max;
3
  float row_sum;
4
};
5

6
__device__ __forceinline__ OnlineState merge_softmax_tile(
7
    OnlineState old_state, float tile_max, float tile_sum) {
8
  const float next_max = fmaxf(old_state.row_max, tile_max);
9
  const float old_scale = __expf(old_state.row_max - next_max);
10
  const float tile_scale = __expf(tile_max - next_max);
11
  return {
12
      next_max,
13
      old_state.row_sum * old_scale + tile_sum * tile_scale,
14
  };
15
}
16

17
template <int Br, int Bc, int D>
18
__global__ void flash_attention_forward(
19
    const half* __restrict__ q,
20
    const half* __restrict__ k,
21
    const half* __restrict__ v,
22
    half* __restrict__ out,
23
    int sequence,
24
    float scale) {
25
  extern __shared__ half smem[];
26
  half* k_stage[2] = {smem, smem + Bc * D};
27
  half* v_stage[2] = {smem + 2 * Bc * D, smem + 3 * Bc * D};
28

29
  // Q fragments and FP32 m/l/O accumulators stay live across the K/V loop.
30
  OnlineState state{ -CUDART_INF_F, 0.0F };
31
  float output_fragment[D / 32] = {};
32

33
  for (int col = 0, stage = 0; col < sequence; col += Bc, stage ^= 1) {
34
    // Production code issues aligned cp.async into the next stage here.
35
    load_kv_tile_async<Bc, D>(k, v, col, k_stage[stage], v_stage[stage]);
36
    cp_async_wait_and_sync();
37

38
    // mma_qk returns FP32 score fragments; masked tiles become -infinity.
39
    auto scores = mma_qk<Br, Bc, D>(q, k_stage[stage], scale);
40
    const float tile_max = warp_block_max(scores);
41
    const float next_max = fmaxf(state.row_max, tile_max);
42
    rescale_output(output_fragment, __expf(state.row_max - next_max));
43
    const float tile_sum = exp_scores_and_mma_pv(
44
        scores, next_max, v_stage[stage], output_fragment);
45
    state = {next_max,
46
             state.row_sum * __expf(state.row_max - next_max) + tile_sum};
47
  }
48
  store_normalized(out, output_fragment, state.row_sum);
49
}

代码中最容易写错的是状态顺序：旧输出必须按新最大值重标定后再累加当前 $PV$ ；causal mask 必须在 tile max 之前生效；全 mask 行不能让 $-\infty-(-\infty)$ 产生 NaN。

7. shape、mask 与大模型 dispatch#

工业实现至少应按以下维度分派：

$D=64/128$ ：使用专用 MMA/layout；其他 head dim 走通用 kernel 或库回退。
短 $N$ ：kernel launch 和 online bookkeeping 占比高，普通 fused SDPA 或 GEMM 可能更合适。
长 $N$ 、小 $B\times H$ ：沿 K 维做 split，多个 CTA 产生局部 $(m,l,o)$ ，再用同一合并公式归约。
causal：跳过完全位于上三角的 tile；对角 tile 才做逐元素 mask。
GQA/MQA：query head 映射到共享 KV head，避免物理复制 K/V；但多个 query head 的复用要和 CTA 调度一起设计。
dropout/反向：训练必须保存或重算 log-sum-exp，并保证随机数 offset 可复现，不能直接套推理 kernel。

极长序列还要考虑 32-bit 索引溢出、grid 维度、变长 batch 的有效 token 数，以及 sequence parallel/context parallel 的跨卡归约。单卡 FlashAttention 解决的是片上/显存 IO，不自动解决多卡通信。

8. 独占 GPU benchmark#

测试环境为 RTX 3060 Laptop 6 GiB、SM86、driver 581.57、PyTorch 2.11.0+cu130。输入为 causal FP16， $B=1,H=32,D=128$ 。每个 shape 先 warmup，再做 7 组 CUDA Event 测量，每组内重复 5 到 30 次，报告 7 组均值的中位数。

测试前 nvidia-smi 和 pmon 均无 compute process。运行时临时把允许的 power limit 从 100 W 提到 130 W，监控峰值为 100% SM、7000 MHz 显存、1965 MHz 核心、102 W、76°C；结束后恢复 100 W。监控进程只采样硬件计数，不提交 CUDA 工作。

对比项强制为 PyTorch SDPBackend.FLASH_ATTENTION 与 SDPBackend.MATH。因此下表验证的是生产 FlashAttention 数据流，而不是上一节的简化骨架。

$N$	FlashAttention	Math SDPA	加速	Flash 峰值额外显存	Math 峰值额外显存
512	0.149 ms	2.214 ms	14.82x	8.1 MiB	109.0 MiB
1024	0.447 ms	7.878 ms	17.61x	16.1 MiB	364.0 MiB
2048	1.704 ms	30.276 ms	17.77x	32.3 MiB	1312.1 MiB
4096	6.573 ms	155.493 ms	23.66x	64.5 MiB	4960.1 MiB
8192	25.854 ms	OOM	-	129.0 MiB	OOM

在 $N\le4096$ 的可比区间，最大绝对误差均为 $1.953125\times10^{-3}$ ，相对 $L_2$ 误差从 $2.54\times10^{-4}$ 到 $2.68\times10^{-4}$ 。这是 FP16 输出和不同归约顺序的预期量级，不应要求 bitwise 一致。

可复现入口如下：

1
from torch.nn.attention import SDPBackend, sdpa_kernel
2
import torch.nn.functional as F
3

4
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
5
    output = F.scaled_dot_product_attention(
6
        query, key, value, is_causal=True, dropout_p=0.0
7
    )

这里的 14.82x 到 23.66x 不能外推到 H100，也不能解释为“算法少算了同样倍数的 FLOP”。Math 路径承担了中间张量物化、mask 和多 kernel 流量；硬件、框架版本、shape 或非 causal 模式改变后比例会变化。

9. 局限与回退路径#

本文只验证 forward causal FP16。没有测试 backward、dropout、变长 packed sequence、RoPE 融合、FP8 scale 或多 GPU context parallel。SM86 也没有 Hopper 的 TMA、WGMMA 和更大的片上资源。

生产中优先调用框架 SDPA、FlashAttention 官方实现或 cuDNN，并让后端根据 dtype、head dim 和 mask 选择 kernel。只有自定义 bias、特殊稀疏模式或框架不覆盖的融合边界，才值得维护手写 CUDA；不支持的 shape 应明确回退，而不是静默给出错误结果。

音乐

音乐