Flash Attention 原理？为何能加速又省显存？

题目

Flash Attention 已成为大模型训练与推理的标准组件。请说明其原理，以及它为何能同时实现加速与省显存。

参考答案

痛点：标准 Self-Attention 需把完整的 $N \times N$ 注意力矩阵写入 HBM（显存），再读回做 softmax。$N$ 大时（长上下文）：

• 显存爆炸：$N \times N \times \text{bytes}$，128K 上下文下数百 GB。
• IO 瓶颈：显存带宽远低于 SRAM，反复读写 HBM 成为速度瓶颈。

Flash Attention 核心思想：IO-aware 分块计算，避免把完整 $N \times N$ 矩阵物化到 HBM。

两大技术：

1. 分块（Tiling）：把 Q、K、V 切成块加载到 SRAM（快速缓存），在 SRAM 内算该块的注意力，直接写回结果到 HBM，不存中间矩阵。
2. 在线 softmax（Online Softmax）：传统 softmax 需先扫一遍求最大值再归一化。Flash Attention 用流式算法，逐块累加归一化因子，一遍扫描完成 softmax，无需物化中间矩阵。

数学关键：分块后用”rescaling”把各块结果合并：

$$\text{softmax}(QK^T)V = \text{merge}(\text{softmax}(QK^T_{\text{block}})V_{\text{block}})$$

每块的局部 softmax 经缩放后可正确合并为全局结果。

为何加速：HBM 读写次数从 $O(N^2)$ 降到 $O(N^2 / M)$（$M$ 为 SRAM 大小），IO 大幅减少 → 速度提升 2–4 倍。

为何省显存：不再物化 $N \times N$ 注意力矩阵，显存从 $O(N^2)$ 降到 $O(N)$。

进化版本：

• Flash Attention v2：优化并行度与分块策略，速度再提升约 2 倍。
• Flash Attention v3：针对 H100 异步拷贝优化，进一步压榨硬件。

实践影响：

• 几乎所有现代大模型训练/推理都默认用 Flash Attention。
• 是长上下文训练可行化的关键——没有它，128K 上下文训练几乎不可行。

面试加分点：

• 强调 Flash Attention 是算法等价（输出与标准 attention 数学相同），不是近似——无损加速。
• 痛点是 IO 而非计算，所以优化目标是减少 HBM 访问，这点反直觉但关键。
• 在线 softmax 的”流式归一化”思想也可用于其他需在线聚合的场景。

出处：CSDN《2026 大模型 LLM 面试通关秘籍：啃透”三位一体”指南》、Flash Attention 论文《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》。