Files
myWiki/concepts/flash-attention.md

1.4 KiB
Raw Permalink Blame History

title, created, updated, type, tags, sources
title created updated type tags sources
FlashAttention 2025-04-15 2026-05-01 concept

FlashAttention

IO感知的精确注意力优化,由 Dao 等 2022 年提出,是注意力计算效率的最大单次突破。

核心洞见

传统注意力实现的瓶颈不是计算FLOPs而是GPU 内存层级之间的数据传输IO。标准实现需要多次读写 HBM高带宽内存而 HBM 带宽远低于 SRAM。

关键创新

  1. Tiling: 将注意力矩阵分块计算,每块保持在快速 SRAM 中
  2. Recomputation: 反向传播时重新计算 softmax 而非存储中间结果,节省内存
  3. IO-Aware: 算法设计以最小化 HBM↔SRAM 数据传输为核心目标

性能收益

  • 速度: 比标准注意力快 2-4x
  • 内存: 内存占用量从 O(n²) 降至 O(n)
  • 精度: 数值精确(非近似),无精度损失

版本演进

  • FlashAttention-1 (2022): Tiling + Recomputation
  • FlashAttention-2 (2023): 更好的并行化和 work partitioning
  • FlashAttention-3 (2024): 异步计算 + 低精度 (FP8)

相关概念