3. FlashAttention原理概览：分块计算、重计算、IO感知算法，理解其如何减少HBM访问

在Transformer推理优化中，注意力机制的计算往往成为吞吐与延迟的瓶颈。传统注意力实现将完整的 S×S 注意力矩阵（S 为序列长度）写入高带宽内存（HBM），再从中读取以完成 Softmax 归一化与加权求和。当序列长度达到 4K、8K 甚至 32K 时，注意力矩阵的显存占用以平方级增长，HBM 带宽成为主要约束。FlashAttention 的核心思想并非改变注意力数学公式，而是通过 分块计算（Tiling）、重计算（Recomputation） 与 IO感知调度（IO-Awareness） 三种机制，将注意力计算重新组织为对 HBM 访问次数极少的融合内核。本章从原理层面拆解这三个关键机制，并给出量化分析与工程判断方法。

上图展示了 FlashAttention 与传统注意力在 HBM 访问模式上的本质差异。传统实现中，每一步中间结果（S、P）都需要写回 HBM，再读入 SRAM 进行下一步计算，导致 HBM 访问量随序列长度平方增长。FlashAttention 则将 Q、K、V 切分为块，在 SRAM 中完成块级注意力计算，并通过在线 Softmax 算法维护全局归一化统计量，从而避免中间矩阵落盘。下面逐一深入三个核心机制。

3.1 分块计算：将平方级矩阵拆解为 SRAM 可容纳的块

分块计算是 FlashAttention 的基石。GPU 的存储层次中，SRAM（共享内存/寄存器）带宽远高于 HBM，但容量极小（例如 A100 为 192KB，H100 为 228KB）。FlashAttention 将 Q 按行切分为大小为 Bᵣ 的块，将 K、V 按列切分为大小为 B꜀ 的块，使得每个块的计算所需的数据（Qᵢ块、Kⱼ块、Vⱼ块）能够完全放入 SRAM。对于每个 Qᵢ 块，算法遍历所有 K/V 块，在 SRAM 内完成 Sᵢⱼ = QᵢKⱼᵀ、Softmax 归一化以及 PᵢⱼVⱼ 的累加。关键点在于：块内计算产生的中间矩阵 Sᵢⱼ 和 Pᵢⱼ 只存在于 SRAM 中，从不写回 HBM。只有最终的输出块 Oᵢ 在完成所有 K/V 块遍历后写回 HBM。

块大小的选择直接影响 HBM 访问量。设 SRAM 容量为 M，头维度为 d，则 Bᵣ 和 B꜀ 需满足 Bᵣ·d + B꜀·d + Bᵣ·B꜀ ≤ M（分别对应 Q 块、K/V 块和 S 块）。在 A100 上，典型配置为 Bᵣ=64，B꜀=64，d=64 时，所需 SRAM 约为 64×64×4 + 64×64×4 + 64×64×4 = 48KB，留有余量。更大的块能减少 K/V 块的遍历次数，但受限于 SRAM 容量；更小的块则增加遍历开销。工程上需要根据具体 GPU 的 SRAM 大小和头维度进行调优。

3.2 重计算：用 FLOPs 换取 HBM 带宽

在反向传播中，传统注意力需要保存前向的 S 和 P 矩阵来计算梯度。FlashAttention 选择不保存这些中间结果，而是在反向传播时重新计算。具体而言，反向传播需要 ∂S 和 ∂P 时，从 HBM 重新读取 Q、K、V 块，在 SRAM 中重新执行前向的块级计算，得到 S 和 P 后再计算梯度。这看似浪费了计算量，但实际上 GPU 的 FLOPs 远多于 HBM 带宽——以 A100 为例，FP16 算力 312 TFLOPS，HBM 带宽 2TB/s，算术强度（FLOPs/Byte）约为 156。重计算增加的 FLOPs 远小于节省的 HBM 带宽带来的收益。

重计算策略需要权衡：保存哪些中间结果？FlashAttention 的典型做法是保存每个块的 Softmax 统计量（行最大值 m 和归一化因子 l），因为这些量很小（每个 Q 块仅 2×Bᵣ 个标量），且反向传播时需要它们来恢复正确的归一化。而 S 和 P 矩阵则完全重计算。这种“部分保存 + 部分重计算”的策略在显存占用和计算开销之间取得了平衡。

在推理场景中，重计算同样有意义。虽然推理不需要反向传播，但 KV Cache 的显存占用与序列长度成正比。FlashAttention 的分块计算使得推理时无需为注意力矩阵分配额外显存，从而为 KV Cache 留出更多空间，支持更长的上下文窗口。

3.3 IO感知算法：以 HBM 访问量为第一优化目标

IO感知算法的核心思想是：在 GPU 上，HBM 带宽往往比 FLOPs 更稀缺。传统内核优化以计算吞吐为目标，而 FlashAttention 将 HBM 访问量作为首要优化指标。这体现在三个层面：

1. 内核融合（Kernel Fusion）：将 QKᵀ、Softmax、PV 三个操作融合为单个内核，避免多次内核启动和中间结果写回 HBM。传统实现需要 3 次内核调用，FlashAttention 只需 1 次。
2. 块调度（Block Scheduling）：按照 SRAM 容量计算最优块大小，并安排遍历顺序，使得每个数据块从 HBM 加载后尽可能被多次复用。Q 块被加载一次后，遍历所有 K/V 块；K/V 块被加载后，被多个 Q 块复用（取决于遍历顺序）。
3. 计算与访存重叠（Overlap）：通过 CUDA 流或异步拷贝，在计算当前块的同时预取下一块数据，隐藏 HBM 延迟。

IO感知算法的量化收益可以通过“算术强度”来评估。传统注意力的算术强度约为 O(S)（每字节数据对应 S 次操作），而 FlashAttention 的算术强度提升至 O(S·M/d)。当 S=4096，d=64，M=192KB 时，算术强度提升约 48 倍，意味着 HBM 带宽需求降低 48 倍。

3.4 案例拆解：从 8K 序列到 32K 序列的显存与速度优化

我们以一个实际推理优化案例来量化 FlashAttention 的收益。假设模型配置：头维度 d=128，8 个注意力头，序列长度 S=8192，批量大小 B=1，使用 FP16 精度（2 字节/元素）。

初始现象与目标：使用 PyTorch 原生注意力实现，推理时显存占用为 14.2GB，其中注意力矩阵占用约 8×8192×8192×2 = 1.07GB，加上 KV Cache 和激活值，总显存接近 A100 80GB 的 18%。目标是将显存占用降低 30% 以上，并将单层注意力延迟从 3.2ms 降至 1.5ms 以下。

关键参数与日志：通过 torch.cuda.max_memory_allocated() 观察到峰值显存出现在注意力层。使用 nsys profile 发现 QKᵀ 和 Softmax 两个内核的 HBM 读取量分别为 1.07GB 和 1.07GB，写入量分别为 1.07GB 和 1.07GB，总计约 4.3GB 的 HBM 流量。

分析判断过程：传统实现中，S 矩阵（8×8192×8192×2=1.07GB）被写入 HBM 一次，Softmax 读一次写一次，PV 再读一次。仅 S 矩阵就产生了 4 次 HBM 访问（写 1 次 + 读 3 次），共 4.3GB。而 FlashAttention 将 S 矩阵完全限制在 SRAM 内，HBM 访问量降为 Q/K/V 读取（3×8192×128×2=6.3MB）加上输出写入（8192×128×2=2.1MB），总计约 8.4MB，降低 500 倍以上。

实际操作步骤：将注意力替换为 FlashAttention 内核（使用 flash_attn_func），保持模型其他部分不变。调整块大小：在 A100 上设置 Bᵣ=128，B꜀=128，确保 SRAM 使用不超过 192KB 的 80%。启用 CUDA Graph 捕获内核启动开销。

结果对比：

风险与代价：FlashAttention 在反向传播中增加了约 1 倍的前向 FLOPs，但推理场景无反向传播，因此不受影响。对于训练场景，需要评估增加的 FLOPs 是否被 HBM 带宽节省所抵消——在长序列（S>2048）下通常收益显著，短序列下收益不明显。此外，FlashAttention 要求头维度 d 为 16 的倍数，且不支持某些自定义注意力变体（如相对位置编码需要额外处理）。

复盘经验：IO感知优化的核心是“先算 HBM 流量，再调内核”。在优化任何注意力实现前，先用 profiler 测量 HBM 读写量，如果中间矩阵的读写占比超过 50%，则分块融合是必然选择。块大小的调优应遵循“SRAM 利用率 70%-90%”的经验法则，过高会导致 bank conflict，过低则无法充分复用数据。

3.5 总结：三种机制如何协同减少 HBM 访问