Flash Attention 3 指南：GPU IO 优化的终极武器
为什么需要 Flash Attention？ 标准 Attention 的瓶颈不在计算，而在显存 IO。考虑 $n = 8192$, $d = 128$ 的注意力计算： 操作 HBM 读写量 SRAM 计算 $S = QK^T$ 读 Q(8MB) + 读 K(8MB) + 写 S(512MB) ~16 GFLOPs $P = \text{softmax}(S)$ 读 S(512MB) + 写 P(512MB) ~8 MFLOPs $O = PV$ 读 P(512MB) + 读 V(8MB) + 写 O(8MB) ~16 GFLOPs 总计 ~1.6 GB ~32 GFLOPs A100 的 HBM 带宽 2 TB/s，SRAM 带宽 19 TB/s。计算只需 ~0.02ms，但 IO 需要 ~0.8ms——IO 是瓶颈的 40 倍。 ...