Dao-AILab /flash-attention

Fast and memory-efficient exact attention

★ 25,264⑂ 3,214
PythonCreated 2022-05-19
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 

Snapshot: 2026-11-10T09:35:25Z · Data license Apache-2.0 · JSON