考察点
FlashAttention 是基础而高频的考点，面试官想区分两种候选人：背过"分块 + 重计算"名词的，和真正理解瓶颈在 HBM 与 SRAM 之间的 IO 而非 FLOPs 的。答题主线：标准注意力的显存瓶颈在哪 → tiling 和 online softmax 怎么配合 → 反向传播的重计算 → 效果与局限。追问常落在 online softmax 的数学细节、FlashAttention-2/3 的改进、对上下文长度的影响。
参考答案
先看瓶颈在哪：不是算力，是 IO
标准注意力的实现是三段式：算 S = QK^T（N×N 矩阵）、对 S 做 softmax、再乘 V 得输出。问题是中间矩阵 S 和 softmax 结果 P 都是 N×N，要完整地写进 HBM（显存）再读回来。N=4096 时一个 head 的 S 就是 64MB（fp16），几十个 head 几十层叠起来，注意力的大部分耗时不是乘法运算，而是在 HBM 和计算单元之间来回搬这个 N×N 矩阵。用术语说：标准注意力是 memory-bound 的，FLOPs 是 O(N²d) 但 IO 也是 O(N²)，N 一长 IO 先爆。
GPU 内存层级是理解 FlashAttention 的钥匙：HBM 大（几十 GB）但慢，片上 SRAM 快得多但小（A100 每 SM 约 192KB）。思路自然就是——别把 N×N 矩阵落到 HBM，把它切成能塞进 SRAM 的小块，在片上算完。
做法一：Tiling，分块计算
把 Q、K、V 按行切成块（比如每次取 Q 的一块 Br 行、K/V 的一块 Bc 列，尺寸按 SRAM 容量选）。外层循环遍历 Q 块，内层循环逐个把 K、V 块从 HBM 搬进 SRAM，在片上算这个 Q 块对当前 K 块的局部注意力，累进更新输出，最后只把 O(Nd) 的结果写回 HBM。全程不产生完整的 N×N 矩阵，HBM IO 从 O(N²) 降到 O(N²d²/M)（M 是 SRAM 大小），实测墙钟时间数倍提升，且显存占用从 O(N²) 降到 O(N)——长上下文训练从不可能变可行，靠的就是这个。
做法二：Online Softmax，让分块在数学上成立
分块的麻烦在 softmax：它对整行做归一化，需要全局最大值和全局求和，可 K 是一块一块来的，算第一块时根本不知道后面的值。FlashAttention 用 online softmax（流式 softmax）解决：维护两个运行量——当前见过的最大值 m 和归一化分母 l。每处理一个新块：
- 算该块的局部最大值，更新 m' = max(m_old, m_new)；
- 把之前累积的结果按 e^(m_old - m') 重新缩放（数学上等价于换底），再并入新块的贡献；
- 分母 l 同样缩放累加。
所有块处理完后除以 l 就是精确的 softmax 结果。数学上和原始 softmax 完全等价，不是近似，这一点面试官很可能确认。
反向传播：重计算代替存储
反向需要 softmax 矩阵 P，存它就是 O(N²) 显存。FlashAttention 的做法是不存，反向时把 Q、K、V 从 HBM 重新读进来再算一遍 P。听起来浪费算力，但省下的 IO 时间远超重算的 FLOPs 开销——因为整个算法本来就是 IO 瓶颈，算力闲着也是闲着。前向只需保存 O(N) 的输出和 softmax 归一化统计量（m、l）。
版本演进与使用
FlashAttention-2 主要优化并行粒度和 warp 分工（在序列维度并行、减少非 GEMM 计算占比），再提速约 2 倍；FlashAttention-3 针对 H100 用异步指令和 FP8 进一步压榨硬件。工程上基本不用手写：PyTorch 的 F.scaled_dot_product_attention 会自动dispatch 到 flash kernel，xFormers、各训练/推理框架都已集成。对应用工程师的要求是理解原理 + 知道它的前提（head_dim 有限制、需要 Ampere 以上架构效果最佳）和边界（它优化单机单卡的注意力 IO，不管通信和调度）。
可能的追问
- FlashAttention 改变注意力复杂度了吗？ 计算复杂度仍是 O(N²d)，没变成线性注意力；它降的是 HBM IO 次数和显存占用（O(N²)→O(N)），墙钟时间因此数倍改善。
- 为什么反向重计算反而更快？ 算法瓶颈在 IO 不在 FLOPs：重算 P 的 FLOPs 是在空闲算力上做的，而省掉的 O(N²) 读写是真金白银的时间，净赚。
- 对长上下文的意义是什么？ 显存从 O(N²) 降到 O(N)，训练/推理能容纳的序列长度直接上一个台阶，配合 RoPE 外推等技术才有今天的百 k 级上下文；推理侧 vLLM 的 PagedAttention kernel 也融合了 flash 思想。