Flash Attention原理与实现
一句话总结
Flash Attention通过分块计算和IO感知的内存管理策略，在不牺牲精度的前提下将注意力计算加速2-4倍并大幅降低显存占用。
核心概念
标准注意力需要将完整的N*N注意力矩阵存储在GPU HBM中，内存复杂度O(N^2)。Flash Attention的核心思想是tiling(分块)：将Q、K、V按块加载到GPU SRAM中计算，利用online softmax技巧在不完整信息下逐块更新softmax的分母，最终结果与标准注意力数学上等价。这样避免了完整注意力矩阵的物化存储，将内存从O(N^2)降到O(N)。Flash Attention-2进一步优化了并行策略和warp级别的任务分配。
为什么重要
GPU计算速度远快于内存带宽(计算强度高)，标准注意力的瓶颈在于HBM读写而非计算本身。Flash Attention通过减少HBM访问将注意力从IO密集型变为计算密集型，充分利用GPU算力。这使得训练长序列(如16K、32K)在现有硬件上变得可行。
实践要点
Flash Attention已集成在PyTorch 2.0+(通过scaled_dot_product_attention)和HuggingFace Transformers中。使用时需确保数据类型为fp16或bf16。Flash Attention-2支持不等长序列打包(variable-length)，非常适合训练场景。Flash Attention-3在H100上进一步利用了FP8和异步执行。
常见误区
误区一：认为Flash Attention是注意力的近似算法。它是精确计算，结果与标准实现完全一致。误区二：认为Flash Attention只节省内存不提速。实际上通过减少IO操作也显著提升了计算速度。误区三：在短序列场景下强行使用Flash Attention，此时标准实现可能同样高效。