novitalabs/flash-attention

novitalabs/flash-attention

Description: Fast and memory-efficient exact attention

License: BSD-3-Clause

Stars: 2

Forks: 0

Open issues: 0

Created: 2025-10-27T08:40:34Z

Pushed: 2025-10-26T20:49:25Z

Default branch: main

Fork: yes

Parent repository: Dao-AILab/flash-attention

Archived: no

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-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:

import flash_attn_interface
flash_attn_interface.flash_attn_func()

Installation and features

Requirements:

• CUDA toolkit or ROCm toolkit
• PyTorch 2.2 and above.
• packaging Python package (pip install packaging)
• 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). Support for Turing GPUs (T4, RTX 2080) is coming soon, please use FlashAttention 1.x for Turing GPUs for now. 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. MI200 or MI300 GPUs. 2. Datatype fp16 and bf16 3. Both forward's and backward's head dimensions up to 256.

Triton Backend

The Triton implementation of the Flash Attention v2 is currently a work in progress.

It supports AMD's CDNA (MI200, MI300) and RDNA GPU's using fp16, bf16 and fp32 datatypes.

These features are supported in Fwd and Bwd 1) Fwd and Bwd with causal masking 2) Variable sequence lengths 3) Arbitrary Q and KV sequence lengths 4) Arbitrary head sizes 5) Multi and grouped query attention 6) Dropout 7) Rotary embeddings 8) ALiBi

We are working on the following things 1) Paged Attention 2) Sliding Window 3) FP8 4) Performance Improvements

##### Getting Started To get started with the triton backend for AMD, follow the steps below.

First install the recommended Triton version

pip install triton==3.2.0

Then install Flash Attention with the flag FLASH_ATTENTION_TRITON_AMD_ENABLE set to "TRUE".

cd flash-attention
git checkout main_perf
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" python setup.py install

To test that things are working, you can run our tests. These tests take hours so you don't need to run the full thing.

FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pytest tests/test_flash_attn_triton_amd.py

You can use autotune for better performance by using this flag FLASH_ATTENTION_TRITON_AMD_AUTOTUNE="TRUE"