---

一、先给一个总览

FlashAttention 是大模型时代最重要的 CUDA 优化之一，由 Tri Dao 等人在 2022 年提出。

它的核心思想是：

通过 tiling 和 online softmax，把 Attention 的计算分块进行，避免显式存储完整的 $N \times N$ 注意力矩阵，从而同时降低显存占用和提升计算效率。

标准 Attention 的问题在于：

• 需要显式存储 $N \times N$ 的注意力矩阵；
• 显存占用 $O(N^2)$；
• 大量 HBM 读写，成为瓶颈。

FlashAttention 通过 IO 感知的设计，把 Attention 从“显存受限”变成“计算受限”，在现代 GPU 上大幅加速。

一句话总结：FlashAttention 用分块计算和在线 softmax，把 Attention 的 HBM 读写从 $O(N^2)$ 降到 $O(N)$，实现了显存和速度的双赢。

---

二、先看标准 Attention 的问题

1. 标准 Attention 的计算

对于输入 $Q, K, V \in \mathbb{R}^{N \times d}$，标准 Attention 为：

$$ S = QK^T $$

$$ P = \text{softmax}(S) $$

$$ O = PV $$

其中：

• $S \in \mathbb{R}^{N \times N}$ 是注意力分数；
• $P \in \mathbb{R}^{N \times N}$ 是注意力概率；
• $O \in \mathbb{R}^{N \times d}$ 是输出。

2. 显存问题

标准实现需要显式存储：

• $S$：$N \times N$；
• $P$：$N \times N$。

显存占用为：

$$ O(N^2) $$

当 $N = 8192$ 时，$N^2 = 6.7 \times 10^7$，FP16 下约 128MB；当 $N = 32768$ 时，约 2GB。这还只是一个 attention head，多头、多层、多 batch 后更夸张。

3. 带宽问题

更严重的是 HBM 读写：

• 写 $S$：$O(N^2)$；
• 读 $S$ 做 softmax：$O(N^2)$；
• 写 $P$：$O(N^2)$；
• 读 $P$ 做 $PV$：$O(N^2)$。

大量时间花在 HBM 读写上，而不是计算上。

GPU 的算力增长远快于 HBM 带宽增长，所以 Attention 成为显存带宽受限的典型。

---

三、FlashAttention 的核心优化

1. Tiling（分块计算）

核心思想：

把 $Q, K, V$ 分成小块，每次只加载一小块到 SRAM（shared memory），在 SRAM 内完成计算。

具体做法：

• 把 $Q$ 分成 $T_r$ 个块，每块大小 $B_r \times d$；
• 把 $K, V$ 分成 $T_c$ 个块，每块大小 $B_c \times d$；
• 外层循环遍历 $K, V$ 块，内层循环遍历 $Q$ 块；
• 每次只把当前块加载到 shared memory。

这样：

• 不需要一次性加载完整的 $Q, K, V$；
• 不需要存储完整的 $S, P$；
• HBM 读写大幅减少。

2. Online Softmax

Tiling 带来的问题是：softmax 需要全局的归一化因子，但分块后无法一次看到所有分数。

Online softmax 解决这个问题：

在遍历 $K, V$ 块的过程中，动态维护当前最大值 $m$ 和归一化因子 $\ell$，逐步更新输出。

标准 softmax：

$$ \text{softmax}(s_i) = \frac{e^{s_i - m}}{\sum_j e^{s_j - m}} $$

其中 $m = \max_j s_j$。

Online softmax 的更新规则：

假设当前已处理到第 $j$ 块，维护：

• $m_j$：到目前为止的最大值；
• $\ell_j$：到目前为止的指数和；
• $O_j$：到目前为止的加权输出。

处理新块时：

$$ m_{j+1} = \max(m_j, \tilde{m}) $$

$$ \ell_{j+1} = e^{m_j - m_{j+1}} \ell_j + e^{\tilde{m} - m_{j+1}} \tilde{\ell} $$

$$ O_{j+1} = e^{m_j - m_{j+1}} O_j + e^{\tilde{m} - m_{j+1}} \tilde{O} $$

其中 $\tilde{m}, \tilde{\ell}, \tilde{O}$ 是新块的局部统计量。

这样，不需要存储完整的 $S$，也不需要两次遍历，就能得到正确的 softmax 结果。

3. 重计算（Recomputation）

反向传播时，需要 $P$ 来计算梯度。

FlashAttention 不保存 $P$，而是在反向时：

• 用保存的 $m$ 和 $\ell$；
• 重新计算 $S$ 和 $P$。

这就是 activation checkpointing 的思想，用计算换显存。

因为重计算在 SRAM 内进行，速度很快，代价可接受。

4. 减少 HBM 读写

综合以上：

• 标准 Attention 的 HBM 读写为 $O(N^2)$；
• FlashAttention 的 HBM 读写为 $O(N^2 / M)$，其中 $M$ 是 SRAM 大小；
• 当 $M$ 较大时，HBM 读写大幅减少。

实际上，FlashAttention 的 HBM 访问量约为：

$$ O\left(\frac{N^2 d^2}{M}\right) $$

相比标准的 $O(N^2 d)$，有显著下降。

5. Kernel Fusion

FlashAttention 把多个操作融合进一个 kernel：

• $QK^T$；
• softmax；
• $PV$。

避免了中间结果的 HBM 读写，减少了 kernel launch 开销。

---

四、为什么能同时降低显存和提升速度

1. 降低显存

• 不存储 $S$ 和 $P$，显存从 $O(N^2)$ 降到 $O(N)$；
• 只存储输出 $O$ 和统计量 $m, \ell$；
• 反向时重计算，不保存中间激活。

显存节省非常显著，尤其长序列下。

2. 提升速度

• HBM 读写大幅减少，缓解带宽瓶颈；
• 计算在 SRAM 内进行，速度快；
• kernel fusion 减少 launch 开销；
• 更好的 GPU 利用率。

在 A100 上，FlashAttention 相比标准实现可以提速 2–4 倍，长序列下更明显。

3. 从带宽受限到计算受限

标准 Attention 是带宽受限：

• 算力用不满；
• 时间花在等 HBM。

FlashAttention 变成计算受限：

• HBM 读写减少；
• 算力充分利用；
• 更接近 GPU 峰值。

---

五、FlashAttention 的版本演进

1. FlashAttention-1

• 提出 tiling + online softmax；
• 支持 FP16/BF16；
• 主要针对训练。

2. FlashAttention-2

• 优化了工作划分，减少非矩阵乘操作；
• 更好的并行化；
• 提升约 2 倍；
• 支持更多 head 维度。

3. FlashAttention-3

• 针对 Hopper 架构优化；
• 利用 TMA、WGMMA 等新特性；
• 支持 FP8；
• 进一步提升。

---

六、和标准 Attention 的对比

维度	标准 Attention	FlashAttention
显存	$O(N^2)$	$O(N)$
HBM 读写	$O(N^2)$	大幅减少
中间矩阵	存储 $S, P$	不存储
反向	用保存的 $P$	重计算
瓶颈	带宽受限	计算受限
速度	较慢	快 2–4 倍

---

七、局限和注意点

1. 实现复杂

• 需要精细的 shared memory 管理；
• 需要处理边界；
• 需要针对不同 GPU 架构优化。

2. 对 head 维度有要求

• 早期版本对 head 维度有限制；
• 后续版本逐步放宽。

3. 不改变数学结果

• FlashAttention 是精确 Attention，不是近似；
• 输出和标准 Attention 数值上一致，只是计算方式不同。

4. 适合长序列

• 序列越长，收益越大；
• 短序列下收益有限。

---

八、简单评价这道面试题

1. 这道题的价值

• 这是大模型 CUDA 优化中最经典的一道题；
• 一道题能考察 tiling、online softmax、重计算、HBM 读写、kernel fusion 等多个核心概念；
• 区分度好，初级可能只知道“FlashAttention 快”，中级能说清 tiling 和 online softmax，高级能讲清 IO 复杂度和版本演进；
• 贴近实战，几乎所有大模型训练和推理都会用到。

2. 这道题的不足

• 题目偏难，如果没有 CUDA 和 Attention 基础，容易答不上来；
• 没有涉及具体的 shared memory 大小和分块策略；
• 对 FlashAttention-2/3 的细节考察不够；
• 没有涉及和其他 Attention 变体的对比。

3. 建议的追问方向

• Online softmax 的更新公式怎么推导？
• FlashAttention 的 HBM 访问量怎么估算？
• 为什么 FlashAttention 是精确的而不是近似的？
• FlashAttention-2 相比 1 做了哪些改进？
• FlashAttention 和 PagedAttention 有什么关系？
• 在什么情况下 FlashAttention 收益不明显？

4. 总体评价

这是一道很好的大模型 CUDA 优化面试题，适合考察候选人对 Attention 底层实现和 IO 优化的理解。它既能考察 tiling、online softmax 等核心算法，也能延伸到 HBM 读写、kernel fusion 和版本演进。如果候选人能把 FlashAttention 的 tiling、online softmax、重计算、HBM 读写减少以及为什么能同时省显存和提速都讲清楚，基本可以判断其具备大模型算子优化的扎实基础。

---

---