← 返回 AI 开源导航

Flash Attention

训练与微调

GitHub 仓库

一、项目简介

FlashAttention 是高效的注意力计算实现，通过 IO 感知的分块与重计算策略，将 Transformer 注意力的显存占用从平方级降到线性级，并显著加速训练与推理，已成为大模型训练的标配组件。

二、核心特性

  • IO 感知分块：减少 HBM 与 SRAM 间数据搬运；
  • 显存线性化：从 O(n²) 降到 O(n)，支持更长上下文；
  • 训练推理加速：显著提升吞吐；
  • 主流框架集成：Transformers、vLLM 等；
  • 持续迭代（FlashAttention-2/3）；
  • 数值稳定性经过验证。

三、适用场景

  • 大模型训练加速与显存优化；
  • 长上下文推理（RAG、长文档）；
  • 显存受限的本地部署。

四、快速上手

pip install flash-attn --no-build-isolation
# 或在加载模型时：attn_implementation="flash_attention_2"

五、注意事项

  • 需 NVIDIA GPU（Ampere 及以上更佳）；
  • 安装需编译，注意 CUDA 与 PyTorch 版本；
  • 老旧显卡可能不支持。