FlashAttention 四代演进：从 IO 感知到 CuTeDSL，H100 与 B200 上的注意力加速方案

它解决什么问题：精确注意力在长序列下的显存与速度瓶颈

标准 attention 需要将 Q 和 K 的乘积矩阵完整写入显存，序列长度一长，显存占用随平方增长，训练和推理都会卡在带宽上。FlashAttention 的核心是 IO 感知，它不改变注意力计算结果，而是把计算分块，让中间矩阵在片上 SRAM 中完成聚合，避免反复读写 HBM。这个仓库是 Tri Dao 等人两篇论文的官方实现，覆盖 FlashAttention 和 FlashAttention-2，并附带 FlashAttention-3 的 beta 版本和 FlashAttention-4 的 CuTeDSL 版本。它的目标用户是训练大模型、处理超长上下文的工程师，以及在 MLPerf 这类基准上追求吞吐的团队。

从 FlashAttention-2 到 FlashAttention-4：四代实现的架构差异

FlashAttention-2 在 Ampere、Ada 和 Hopper 上支持 fp16 与 bf16，头维度最高 256，backward 在无 dropout 时也支持 256。FlashAttention-3 转向 Hopper 专用，只开放 FP16/BF16 的 forward 和 backward，以及 FP8 的 forward，要求 H100 或 H800，CUDA 12.3 以上。FlashAttention-4 用 CuTeDSL 编写，目标覆盖 Hopper 和 Blackwell，例如 H100 和 B200。三者的接口不同：FlashAttention-2 的接口在 src/flash_attention_interface.py，FlashAttention-3 需要从 hopper 子目录导入 flash_attn_3，FlashAttention-4 则从 flash_attn.cute 导入。这种分裂意味着你选定的版本决定了你的代码写法，升级并不只是换一个 pip 包名。

安装路径：pip 预编译、源码编译与 uv 配置

最直接的安装方式是 pip install flash-attn --no-build-isolation，但 README 强调编译依赖 ninja，否则编译时间可能长达 2 小时。ninja 正常工作时，64 核机器编译只需 3 到 5 分钟。若机器内存低于 96GB，需设置 MAX_JOBS=4 来限制并行编译任务。FlashAttention-3 需要进入 hopper 子目录执行 python setup.py install，然后设置 PYTHONPATH 才能运行测试。FlashAttention-4 可以直接 pip install flash-attn-4，若在 CUDA 13 环境，建议安装 flash-attn-4[cu13]。用 uv 时，需要在 pyproject.toml 中声明 no-build-isolation = true，并把 flash-attn-3 指向 git 仓库的 hopper 子目录。

硬件边界与数据类型限制：哪些卡能用，哪些不能用

FlashAttention-2 的 CUDA 后端只支持 Ampere、Ada 和 Hopper，例如 A100、RTX 3090、RTX 4090、H100。Turing 卡（T4、RTX 2080）不在支持列表里，需要另找 flash-attention-turing 这个独立仓库，它只覆盖核心子集。bf16 需要 Ampere 及以上，fp16 则 Turing 也可能可用，但官方没有承诺。ROCm 后端有两个分支：composable_kernel 默认支持 MI200 到 MI355x 以及 RDNA 3/4，但只支持 fp16 和 bf16；Triton 后端支持 CDNA 和 RDNA，数据类型扩展到 fp32，还支持 MQA/GQA、dropout、rotary embeddings 和 ALiBi。滑动窗口注意力在 Triton 后端仍在开发中，这算一个明确的功能缺口。

接口对比：从 flash_attn_func 到 flash_attn_3 再到 flash_attn.cute

FlashAttention-2 的经典调用是 flash_attn_func(q, k, v, causal=True)，但 FlashAttention-3 的接口是 flash_attn_interface.flash_attn_func()，FlashAttention-4 则从 flash_attn.cute 导入同名函数。这三个函数虽然名字相似，但参数和返回类型未必一致。README 没有给出 FlashAttention-2 的具体函数签名，也没有说明 FlashAttention-3 的完整参数列表。实际使用时，你需要查阅各自子目录的源码或测试文件，不能假设接口兼容。这种接口碎片化是该项目的一个真实痛点，尤其当你在同一份代码里尝试在不同 GPU 上切换版本时。

测试与验证：pytest 跑通不代表生产可用

FlashAttention-3 的测试命令是 pytest -q -s test_flash_attn.py，FlashAttention-4 的测试套件没有在 README 中给出具体命令。Triton 后端的测试命令是 FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pytest tests/test_flash_attn_triton_amd.py，但 README 警告完整测试套件需要数小时。这个时间成本意味着你不能在每次代码改动后都跑全量测试，需要按功能模块筛选。另外，README 建议使用 NVIDIA 的 PyTorch 容器来安装，因为它自带所需工具，这暗示了在裸系统上编译可能遇到环境问题。

维护与升级成本：beta 版本与版本碎片化

FlashAttention-3 明确标注为 beta 版本，用于测试和基准，尚未与主仓库集成。FlashAttention-4 的版本号是 fa4-v4.0.0.beta28，说明它仍在迭代，且发布频率很高，例如 beta26 到 beta28 相隔一周左右。这种快速迭代意味着你依赖的接口可能在下个版本变化，需要跟随 release note 调整代码。许可证是 BSD-3-Clause，允许自由使用和修改，但 README 要求引用论文和注明来源。维护成本方面，源码编译依赖 ninja 和 packaging 等包，且 Windows 支持不完整，从 v2.3.2 开始才有零星成功报告，官方明确说 Windows 编译仍需更多测试。

替代方案与适用判断：什么时候不该选它

如果你的 GPU 是 Turing 或更老，FlashAttention 官方不支持，替代方案是 flash-attention-turing 仓库，但它只覆盖核心特性，且不是官方维护。另一个替代是使用 PyTorch 自带的 scaled_dot_product_attention，它内置了内存高效的实现，但性能优化程度取决于你的硬件和 PyTorch 版本。FlashAttention 的优势在于精确性和对特定 GPU 的深度调优，而 PyTorch 原生实现更通用，但可能牺牲峰值性能。若你的场景是短序列、显存充足，或者你不想处理编译依赖，那么标准 attention 或 SDPA 可能足够，不必引入 FlashAttention 的安装复杂度。

编辑结论

适合拥有 Ampere 及以上 NVIDIA GPU 或 AMD MI200 系列，且需要在不牺牲精度的前提下加速长序列训练的团队。不适合需要 Turing 卡完整功能支持、Windows 环境或非 Linux 平台的用户，也不适合希望开箱即用、不愿处理编译依赖的场合。采用前应核实：你的 CUDA 版本是否不低于 12.0，PyTorch 是否 2.2 以上，以及是否愿意接受 3 到 5 分钟（64 核）的源码编译。若你的 GPU 是 H100 或 B200，且追求最新性能，可直接安装 flash-attn-4 并配置 uv 的 no-build-isolation；否则建议先从 flash-attn 稳定版入手。

官方来源

1. Official README
2. Project repository
3. Release notes