flash-attention

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-2: Faster Attention with Better Parallelism and Work Partitioning
Tri Dao

Paper: https://tridao.me/publications/flash2/flash2.pdf

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

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:

cd hopper
python setup.py install

To run the test:

export PYTHONPATH=$PWD
pytest -q -s test_flash_attn.py

Once the package is installed, you can import it as follows:

from flash_attn_3 import flash_attn_interface
flash_attn_interface.flash_attn_func()

To install using uv, in your pyproject.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:

pip install flash-attn-4

If you're on CUDA 13, we recommend installing with the cu13 extra for best performance:

pip install "flash-attn-4[cu13]"

Once installed, you can use it as follows:

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 $? should return exit code 0). If not (sometimes ninja --version then echo $? returns a nonzero exit code), uninstall then reinstall ninja (pip uninstall -y ninja && pip install ninja). Without ninja, compiling can take a very long time (2h) since it does not use multiple CPU cores. With ninja compiling takes 3-5 minutes on a 64-core machine using CUDA toolkit.

To install:

pip install flash-attn --no-build-isolation

Alternatively you can compile from source:

python setup.py install

If your machine has less than 96GB of RAM and lots of CPU cores, ninja might run too many parallel compilation jobs that could exhaust the amount of RAM. To limit the number of parallel compilation jobs, you can set the environment variable MAX_JOBS:

MAX_JOBS=4 pip install flash-attn --no-build-isolation

Interface: src/flash_attention_interface.py

NVIDIA CUDA Support

Requirements:

• CUDA 12.0 and above.

We recommend the Pytorch container from Nvidia, which has all the required tools to install FlashAttention.

FlashAttention-2 with CUDA currently supports:

1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). For Turing GPUs (T4, RTX 2080), see the separate flash-attention-turing repo, which supports a core subset of FlashAttention features on Turing.
2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).
3. All head dimensions up to 256. Head dim > 192 backward requires A100/A800 or H100/H800. Head dim 256 backward now works on consumer GPUs (if there's no dropout) as of flash-attn 2.5.5.

AMD ROCm Support

ROCm version has two backends. There is composable_kernel (ck) which is the default backend and a Triton backend. They provide an implementation of FlashAttention-2.

Requirements:

• ROCm 6.0 and above.

We recommend the Pytorch container from ROCm, which has all the required tools to install FlashAttention.

Composable Kernel Backend

FlashAttention-2 ROCm CK backend currently supports:

1. MI200x, MI250x, MI300x, MI355x, and RDNA 3/4 GPUs.
2. Datatype fp16 and bf16
3. Both forward's and backward's head dimensions up to 256.

Triton Backend

The Triton implementation of Flash Attention supports AMD's CDNA (MI200, MI300) and RDNA GPUs using fp16, bf16, and fp32 datatypes. It provides forward and backward passes with causal masking, variable sequence lengths, arbitrary Q/KV sequence lengths and head sizes, MQA/GQA, dropout, rotary embeddings, ALiBi, paged attention, and FP8 (via the Flash Attention v3 interface). Sliding window attention is currently a work in progress.

The Triton backend kernels are provided by the aiter package, included as a git submodule at third_party/aiter and automatically installed during setup.

To install, first get PyTorch for ROCm from https://pytorch.org/get-started/locally/, then install Flash Attention:

cd flash-attention
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .

To use a specific aiter commit (e.g., for testing or development):

cd flash-attention
cd third_party/aiter && git fetch origin && git checkout <commit-sha> && cd ../..
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .

To run the tests (note: full suite takes hours):

FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pytest tests/test_flash_attn_triton_amd.py

The Triton backend uses a default kernel configuration optimized for determinism and reasonable performance across workloads. For peak throughput, enable FLASH_ATTENTION_TRITON_AMD_AUTOTUNE="TRUE" to search for optimal settings, which incurs a one-time warmup cost.

Alternativly, if not autotuning, FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON may be used to set a single triton config overriding the hardcoded defaults for attn_fwd. E.g.

FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON='{"BLOCK_M":128,"BLOCK_N":64,"waves_per_eu":1,"PRE_LOAD_V":false,"num_stages":1,"num_warps":8}'

For a quick start with Docker:

FROM rocm/pytorch:latest

WORKDIR /workspace

# build flash attention with triton backend
RUN git clone https://github.com/Dao-AILab/flash-attention &&\ 
    cd flash-attention &&\
    FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .

# set working dir
WORKDIR /workspace/flash-attention

# set env variable to use triton backend
ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"

Build and run:

docker build -t flash-attn-triton .
docker run -it --network=host --user root --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --shm-size 16G --device=/dev/kfd --device=/dev/dri flash-attn-triton

How to use FlashAttention

The main functions implement scaled dot product attention (softmax(Q @ K^T * softmax_scale) @ V):

from flash_attn import flash_attn_qkvpacked_func, flash_attn_func

…

…

…

To see how these functions are used in a multi-head attention layer (which includes QKV projection, output projection), see the MHA implementation.

Using with 🤗 Kernels

If your hardware environment belongs to any of the above-mentioned, you can also use the kernels library to use Flash Attention 2 and 3 right away.

# pip install kernels

from kernels import get_kernel

# FA2
fa_module = get_kernel("kernels-community/flash-attn2", version=1)
flash_attn_func = fa_module.flash_attn_func

# FA3
fa3_module = get_kernel("kernels-community/flash-attn3", version=1)
flash_attn_func = fa3_module.flash_attn_func

Changelog

2.0: Complete rewrite, 2x faster

Upgrading from FlashAttention (1.x) to FlashAttention-2

These functions have been renamed:

• flash_attn_unpadded_func -> flash_attn_varlen_func
• flash_attn_unpadded_qkvpacked_func -> flash_attn_varlen_qkvpacked_func
• flash_attn_unpadded_kvpacked_func -> flash_attn_varlen_kvpacked_func

If the inputs have the same sequence lengths in the same batch, it is simpler and faster to use these functions:

flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)

flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)

2.1: Change behavior of causal flag

If seqlen_q != seqlen_k and causal=True, the causal mask is aligned to the bottom right corner of the attention matrix, instead of the top-left corner.

For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
v2.0:
1 0 0 0 0
1 1 0 0 0
v2.1:
1 1 1 1 0
1 1 1 1 1

If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
v2.0:
1 0
1 1
1 1
1 1
1 1
v2.1:
0 0
0 0
0 0
1 0
1 1
If the row of the mask is all zero, the output will be zero.

2.2: Optimize for inference

Optimize for inference (iterative decoding) when query has very small sequence length (e.g., query sequence length = 1). The bottleneck here is to load KV cache as fast as possible, and we split the loading across different thread blocks, with a separate kernel to combine results.

See the function flash_attn_with_kvcache with more features for inference (perform rotary embedding, updating KV cache inplace).

Thanks to the xformers team, and in particular Daniel Haziza, for this collaboration.

2.3: Local (i.e., sliding window) attention

Implement sliding window attention (i.e., local attention). Thanks to Mistral AI and in particular Timothée Lacroix for this contribution. Sliding window was used in the Mistral 7B model.

2.4: ALiBi (attention with linear bias), deterministic backward pass.

Implement ALiBi (Press et al., 2021). Thanks to Sanghun Cho from Kakao Brain for this contribution.

Implement deterministic backward pass. Thanks to engineers from Meituan for this contribution.

2.5: Paged KV cache.

Support paged KV cache (i.e., PagedAttention). Thanks to @beginlner for this contribution.

2.6: Softcapping.

Support attention with softcapping, as used in Gemma-2 and Grok models. Thanks to @Narsil and @lucidrains for this contribution.

2.7: Compatibility with torch compile

Thanks to @ani300 for this contribution.

Performance

We p