Interview: Flash Attention没有改变数学计算结果，为什么能加速2-4倍？瓶颈到底在哪？
题目解析
Flash Attention是近年来最有影响力的系统优化工作之一。它不改变任何数学运算，输出与标准Attention完全一致，却实现了显著加速。这道题考察候选人对GPU计算架构和内存层次结构的理解——现代深度学习的瓶颈往往不在计算而在访存。
解答思路
标准Attention的计算流程：先算Q·K^T得到N×N的注意力矩阵，存到HBM（显存），再做softmax，再存回HBM，最后与V相乘。三次读写HBM的N×N矩阵是巨大开销。Flash Attention的核心思想是分块计算（tiling）：将Q、K、V分成小块加载到SRAM（片上内存），在SRAM中完成所有中间计算，只将最终结果写回HBM，避免了N×N注意力矩阵的显存物化。
关键要点
瓶颈在于内存带宽而非算力：A100的计算能力是312 TFLOPS（BF16），但HBM带宽仅2TB/s。标准Attention的算术强度（FLOPs/byte）很低，是典型的memory-bound操作。Flash Attention通过增加额外的重计算（online softmax的数值稳定处理）来换取减少HBM访问。关键技术：online softmax算法允许分块计算softmax而无需全局归一化常数。
加分回答
深入讨论GPU内存层次：SRAM约20MB（A100），带宽约19TB/s；HBM约40-80GB，带宽约2TB/s。Flash Attention将IO复杂度从O(N²)降到O(N²d/M)，其中M是SRAM大小。还可以提到Flash Attention 2的进一步优化：减少non-matmul FLOPs、更好的warp级并行。Flash Attention 3利用Hopper架构的异步特性实现更大加速。
常见踩坑
最常见的错误是说Flash Attention减少了计算量——实际上它的FLOPs与标准实现相同甚至略多（因为online softmax需要额外计算）。另一个误区是认为Flash Attention改变了注意力的数学定义或引入了近似——它是精确计算。也有人混淆Flash Attention和稀疏Attention，后者才是通过减少计算量来加速。