自注意力的计算复杂度是多少？FlashAttention 为什么能加速？

简化版

自注意力对序列长度 n 的算术复杂度为 O(n²d)；朴素实现若物化 n×n 分数/概率矩阵，中间激活显存为 O(n²)。FlashAttention 是 IO 感知的精确注意力算法：把 Q、K、V 分块放入片上存储，用在线 softmax 累积输出，不把完整注意力矩阵写回 HBM。它保持标准注意力的数学定义与 O(n²d) 算术量，同时把额外中间存储降到随序列近线性增长，并减少 HBM 访问；实际是否更快取决于形状、硬件、dtype、mask 和 kernel。

详细版

复杂度从哪来

标准注意力：

S = Q·Kᵀ          # (n×d)·(d×n) = n×n，  计算 O(n²·d)
P = softmax(S)     # n×n
O = P·V            # (n×n)·(n×d)，        计算 O(n²·d)

计算复杂度 O(n²·d)：两次 n×n 规模的矩阵乘。
显存复杂度 O(n²)：要存下 n×n 的 S 和 P。序列一长（如 n=32K），n² 达十亿级，显存和访存都爆炸。

FlashAttention 的关键洞察

朴素拆算子实现的一个关键瓶颈是 IO：n×n 中间矩阵要在 HBM 与计算单元之间往返。FlashAttention 针对这种内存层次设计；但具体工作负载可能是访存受限，也可能因序列、头维和硬件转为计算受限，不能把“注意力永远 memory-bound”当定律。

FlashAttention 三招：

分块（tiling）：把 Q、K、V 切成小块，逐块加载进快得多的 SRAM（片上缓存）计算。
在线 softmax：softmax 需要全局最大值和归一化和，FlashAttention 用「边遍历边维护running最大值和running和」的数值稳定算法，分块累加也能得到正确结果，无需先算完整行。
不落地中间矩阵：完整的 n×n 分数矩阵从不写回 HBM，显存降到 O(n)；反向传播时靠重计算（recompute）而非存储中间矩阵。

收益

显存：O(n²) → O(n)，能训更长上下文。
速度：减少 HBM 读写，特定论文配置曾报告显著加速；项目中必须实测，不能承诺固定倍数。
精确算法：计算同一个标准注意力算子，不做稀疏/线性近似；浮点结合顺序变化仍可能产生容差内数值差异，并非逐 bit 相同。

完整版教学

一、把 O(n²) 说清楚：贵在哪一步

注意力要让每个 token 和其他所有 token 算相关度，于是天然产生一个 n×n 的分数矩阵 S=QKᵀ。

算 S：n×n 个分数，每个是 d 维点积，O(n²·d) 计算。
softmax(S)：对 n×n 逐行归一化。
乘 V：又是 O(n²·d)。

所以计算是 O(n²·d)、显存是 O(n²)（存 S 和 P）。d 是固定的头维，真正随序列爆炸的是 n²。当上下文从 2K 涨到 32K，n² 涨 256 倍——这就是长上下文又慢又占显存的根源，也是各种「高效注意力」要攻克的目标。

二、一个反直觉的事实：注意力慢，主要慢在搬数据

很多人只数乘法，却忽略张量在存储层次间搬运的成本。FlashAttention 的核心贡献是把 IO 纳入算法设计，而不是宣称所有注意力形状都只受带宽限制。

原因是 GPU 有两级存储：

HBM（高带宽显存）：容量大（几十 GB）但带宽相对慢。
SRAM（片上缓存）：极快但容量极小（每个 SM 只有几十 KB～几 MB）。

朴素的分离算子会把 S、P 写入并重新读出 HBM，产生大量流量。融合与分块能减少这种流量；当该流量是瓶颈时收益最大，而在很短序列、小头维或其他算子主导时收益可能有限。

三、FlashAttention 三板斧详解

① 分块（Tiling）——把计算搬进 SRAM 把 Q、K、V 沿序列维切成小块。外层循环遍历 Q 的块，内层循环遍历 K/V 的块，每次只把当前小块加载进 SRAM 计算局部注意力。因为块很小，能装进快速的 SRAM，避免把大矩阵放 HBM。

② 在线 softmax（Online Softmax）——分块也能算对归一化 难点在于 softmax 需要「整行的最大值和求和」来归一化，而分块时看不到整行。FlashAttention 借用「在线/流式 softmax」技巧：遍历每个 K/V 块时，维护当前的running 最大值 m 和 running 归一化和 l，每来一个新块就用「重新缩放旧结果 + 累加新块」的方式更新，数学上保证最终等于对整行做 softmax。这样不必先看完整行，边走边算就得到精确结果。（减去最大值是为了数值稳定，防止 exp 溢出。）

③ 不落地 + 反向重计算——省显存 前向从不把完整 n×n 矩阵写回 HBM，显存需求降到 O(n)。反向传播本来需要中间的 P 矩阵，FlashAttention 选择不存它、需要时按块重新算（recomputation）——用一点额外计算换巨大的显存节省，而由于瓶颈本就在访存，这笔交换非常划算。

简要说抓住精髓：FlashAttention 不改变二次算术阶，而是通过“分块进片上存储 + 在线 softmax + 不落地大矩阵”减少 HBM 访问；它计算同一注意力定义，速度收益需按具体硬件实测。

四、FlashAttention vs 其他”高效注意力”

历史上为对付 O(n²) 有很多路线，FlashAttention 属于「精确、IO 优化」这一支，和「近似」路线本质不同：

路线	代表	是否精确	思路
稀疏注意力	Longformer、BigBird	近似	只算局部+少量全局，跳过大部分对
线性注意力	Performer、Linear Transformer	近似	用核技巧把 O(n²) 降到 O(n)
低秩/压缩	Linformer	近似	把 K/V 投影到低维
IO 感知精确	FlashAttention	精确	不改数学，优化访存与分块

近似方法能改变计算阶数，却也改变注意力模式或引入近似误差；FlashAttention 保持标准注意力定义，主要优化执行与存储。工程框架可能根据设备、dtype、mask 和张量布局在多种 kernel 间选择，不能只看到 API 名称就断言实际走了某一版 FlashAttention。

五、复杂度到底降没降

要点清楚：FlashAttention 的计算复杂度仍是 O(n²·d)（浮点运算数没变），它降的是显存复杂度（O(n²)→O(n)）和实际墙钟时间（靠减少 HBM 访问）。想在计算阶数上降到 O(n) 得靠稀疏/线性注意力那类近似方法。这个区别是高频考点，别答错。

六、如何验证 FlashAttention 的真实收益

固定模型与输入 -> 切换普通/FlashAttention kernel
                  -> 记录峰值显存、step time、数值误差
                  -> 按 2K/8K/32K 长度画曲线

比较时保持 dtype、batch、序列长度、mask 与 dropout 一致，并先确认实际 kernel 没有因不支持的形状回退。以单头、n=8192、bf16 为例，若朴素实现单独物化一个 n×n 概率矩阵，其大小约为 8192²×2=128 MiB；32 个头仅该矩阵就约 4 GiB，还未计 logits、反向与 batch。FlashAttention 不物化完整矩阵，因此主要中间存储不再按 n² 增长。

速度则不能从这 4 GiB 直接推出。应预热后统计中位数与 P95 step time，同时检查输出/梯度在设定容差内一致；短序列可能被调度开销抵消，长序列也可能转成算力瓶颈。正确结论应是“在本硬件与形状下峰值下降 X%、吞吐上升 Y%”，而不是背固定加速倍数。

七、常见误区与追问

误区：FlashAttention 降低了计算复杂度。 没有，计算仍 O(n²)；它降的是显存和访存，靠 IO 优化提速，结果精确。
误区：它是一种近似注意力。 它计算同一数学算子；浮点顺序不同可带来容差内差异，不等于 bitwise identical。
追问：为什么朴素实现可能慢？ n×n 中间矩阵的 HBM 往返可能形成 IO 瓶颈，但是否 memory-bound 需由 profile 判断。
追问：在线 softmax 解决什么？ 让分块计算也能得到正确的全局 softmax 归一化，无需先算完整行。
追问：反向为什么要重计算？ 前向没存 n×n 中间矩阵以省显存，反向需要时按块重算，用算力换显存。
追问：它和 GQA/MQA 冲突吗？ 不冲突，正交——FlashAttention 优化注意力的算法/访存，GQA/MQA 优化 KV 存储量，常一起用。
追问：还有哪些降长序列成本的手段？ 稀疏/线性注意力（降阶但近似）、滑动窗口、MoE（与序列无关）、KV 压缩等。

八、加强记忆

复杂度要说完整：标准自注意力算术是 O(n²d)；朴素物化分数矩阵的中间显存是 O(n²)。FlashAttention 用分块、在线 softmax 和反向重算避免完整矩阵落 HBM，保持同一注意力定义与二次算术阶，把主要中间存储降为近线性并减少 IO。它不是稀疏近似，也不保证逐 bit 相同或固定倍数加速；是否受 IO 限制、实际节省多少，必须在目标形状与硬件上 profile。

简化版

详细版