25,264
stars
+19+25+156
3,214
forks
–
watchers
–
prs
–
issues
–
releases
–
discussions
README
FlashAttention
This repository provides the official implementation of FlashAttention and
FlashAttention-2 from the
following papers.
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
Paper: https://arxiv.org/abs/2205.14135
IEEE Spectrum article about our submission to the MLPerf 2.0 benchmark using FlashAttention.
!FlashAttention
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Tri Dao
Paper: https://tridao.me/publications/flash2/flash2.pdf
!FlashAttention-2
Usage
We've been very happy to see FlashAttention being widely adopted in such a short
time after its release. This page
contains a partial list of places where FlashAttention is being used.
FlashAttention and FlashAttention-2 are free to use and modify (see LICENSE).
Please cite and credit FlashAttention if you use it.
FlashAttention-3 beta release
FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
Blogpost: https://tridao.me/blog/2024/flash3/
Paper: https://tridao.me/publications/flash3/flash3.pdf
!FlashAttention-3 speedup on H100 80GB SXM5 with FP16
This is a beta release for testing / benchmarking before we integrate that with
the rest of the repo.
Currently released:
- FP16 / BF16 forward and backward, FP8 forward
Requirements: H100 / H800 GPU, CUDA = 12.3.
We highly recommend CUDA 12.8 for best performance.
To install:
sh
cd hopper
python setup.py install
To run the test:
sh
export PYTHONPATH=$PWD
pytest -q -s test_flash_attn.py
Once the package is installed, you can import it as follows:
python
from flash_attn_3 import flash_attn_interface
flash_attn_interface.flash_attn_func()
To install using uv, in your pyproject.toml:
toml
[project]
dependencies = [
"flash-attn-3"
]
[tool.uv]
no-build-isolation = true
[tool.uv.sources]
flash-attn-3 = { git = "https://github.com/Dao-AILab/flash-attention", subdirectory = "hopper" }
FlashAttention-4 (CuTeDSL)
FlashAttention-4 is written in CuTeDSL and optimized for Hopper and Blackwell GPUs (e.g. H100, B200).
To install:
sh
pip install flash-attn-4
If you're on CUDA 13, we recommend installing with the cu13 extra for best performance:
sh
pip install "flash-attn-4[cu13]"
Once installed, you can use it as follows:
python
from flash_attn.cute import flash_attn_func
out = flash_attn_func(q, k, v, causal=True)
Installation and features
Requirements:
- CUDA toolkit or ROCm toolkit
- PyTorch 2.2 and above.
- packaging Python package (pip install packaging)
- psutil Python package (pip install psutil)
- ninja Python package (pip install ninja)
- Linux. Might work for Windows starting v2.3.2 (we've seen a few positive reports) but Windows compilation still requires more testing. If you have ideas on how to set up prebuilt CUDA wheels for Windows, please reach out via Github issue.
\ Make sure that ninja is installed and that it works correctly (e.g. ninja
--version then echo