FlashAttention

按 GPU 显存层级重排计算的注意力实现，结果不变但更快更省显存

FlashAttention 由 Tri Dao 等人在 2022 年提出，是注意力机制的一种 GPU 实现，算出的结果和标准注意力完全一样（精确注意力，不是近似）。标准实现要把 N×N 的注意力矩阵写进显存再读回，序列一长，显存读写就成了瓶颈。FlashAttention 把 Q、K、V 切成小块放进 GPU 片上高速缓存里算完再写回，不再存完整注意力矩阵，显存占用随序列长度线性增长，速度也明显提升。之后又有 FlashAttention-2、针对 Hopper 架构的 FlashAttention-3。VLA、视频世界模型这类长序列模型训练和推理时普遍依赖它，PyTorch 的 scaled_dot_product_attention 也内置了类似后端。

例子训练 VLA 时在 Hugging Face Transformers 里设置 attn_implementation=「flash_attention_2」，同样显存下可以用更大的批大小或更多图像 token。

也叫

FlashAttention-2、FlashAttention-3

相关

注意力机制、Transformer、算子、CUDA、推理延迟、混合精度训练

来源

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Dao-AILab/flash-attention (GitHub)