deepinfra/flash-attention

deepinfra/flash-attention

Description: Fast and memory-efficient exact attention

Language: Python

License: BSD-3-Clause

Stars: 0

Forks: 0

Open issues: 0

Created: 2025-02-19T22:18:06Z

Pushed: 2025-02-20T00:29:54Z

Default branch: main

Fork: yes

Parent repository: vllm-project/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

These features are supported in Fwd for now. We will add them to backward soon. 1) Multi and grouped query attention 2) ALiBi and matrix bias

These features are in development 1) Paged Attention 2) Sliding Window 3) Rotary embeddings 4) Dropout 5) Performance Improvements

Getting Started

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

First install the recommended Triton commit.

git clone https://github.com/triton-lang/triton
cd triton
git checkout 3ca2f498e98ed7249b82722587c511a5610e00c4
pip install --verbose -e python

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

export FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
cd flash-attention
python setup.py install
pytest tests/test_flash_attn.py

How to use…