Flash Attention Setup

Flash Attention is a memory-efficient attention algorithm for transformer models that reduces memory usage for long sequences. Installing it requires PyTorch with CUDA support and careful dependency management. This guide covers two installation paths using uv: from a pre-built wheel (faster, no compilation), or built from source when no compatible wheel is available.

Prerequisites

Make sure to read the following sections of the documentation before using this example:

• Quick Start
• Running your code
• uv

Other resources:

• Flash Attention GitHub repository

Installation

Flash Attention provides pre-built wheels for common combinations of CUDA, PyTorch, and Python on the Flash Attention releases page.

Some dashboards are available online to help find the right wheel for your configuration, like this one, where you can copy an URL or download the wheel to your local machine.

If a compatible wheel is listed, use the From a pre-built wheel section. Otherwise, use Building from source to compile Flash Attention on a compute node.

From a pre-built wheel

Pre-built wheels are available for common CUDA, PyTorch, and Python combinations. Wheel filenames encode the target configuration, like :

flash_attn-2.8.3.post1+cu126torch2.7cxx11abiFALSE-cp312-cp312-linux_x86_64.whl

This wheel targets:

• flash_attn-2.8.3.post1 — flash attention 2.8.3.post1
• cu126 — CUDA 12.6
• torch2.7 — PyTorch 2.7
• cp312 — Python 3.12
• linux_x86_64 — Linux x86_64

Tip

Pre-built wheels require no compilation and can be installed on login nodes.

Add the wheel URL as a dependency source in pyproject.toml, replacing the URL with the one matching the target configuration:

[project]
name = "flash-attn-setup"
version = "0.1.0"
requires-python = ">=3.12, <3.14"
dependencies = [
    "numpy",
    "torch<=2.8.0",
    "flash-attn>=2.7",
]

[tool.uv.sources]
# Look at pre-builds wheels urls from https://github.com/Dao-AILab/flash-attention/releases
flash-attn = { url = "<URL_to_wheel>" }

Install the dependencies:

# Create the virtual environment and install all dependencies
uv sync

Reusing a locally built wheel

A wheel built from source (see theBuilding from source) can also be used. In the sources section of the pyproject.toml, replace url by path, like:

[tool.uv.sources]
flash-attn = { path = "<path_to_wheel>" }

Building from source

Build Flash Attention from source when no pre-built wheel matches the target CUDA, PyTorch, and Python combination.

Warning

Building from source requires significant memory and time. Do not run the build on login nodes, submit a dedicated job to a compute node instead.

Use the following configuration and job script:

ConfigurationJob script

[project]
name = "flash-attn-setup"
version = "0.1.0"
requires-python = ">=3.12, <3.14"
dependencies = [
    "numpy",
    "ninja",
    "torch<=2.8.0",
    "flash-attn>=2.7",
]

[tool.uv.extra-build-dependencies]
flash-attn = [
    { requirement = "torch", match-runtime = true },
    { requirement = "ninja", match-runtime = true },
]

[tool.uv.extra-build-variables]
flash-attn = { MAX_JOBS = "4", FLASH_ATTENTION_SKIP_CUDA_BUILD = "0", TORCH_CUDA_ARCH_LIST = "9.0" }

Warning

MAX_JOBS = "4" limits parallel compilation to prevent out-of-memory errors during the build. Increase this value only when the compute node has sufficient memory for more parallel jobs.

FLASH_ATTENTION_SKIP_CUDA_BUILD = "0" ensures that Flash Attention is compiled with CUDA support.

Tip

Adapt TORCH_CUDA_ARCH_LIST to the compute capability of the target GPU. Find compute capabilities on the NVIDIA website. "9.0" targets the H100. To support multiple architectures, separate values with semicolons: TORCH_CUDA_ARCH_LIST="9.0;8.0;...".

#!/bin/bash
#SBATCH --ntasks=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=128G
#SBATCH --time=03:00:00

# Exit on error
set -e

# Echo time and hostname into log
echo "Date:     $(date)"
echo "Hostname: $(hostname)"

# Get access to the CUDA libraries
module load cuda/13.2

# Create the virtual environment and install all dependencies
uv sync

Submit the build using the job script:

sbatch build_from_source/job.sh

Once the build completes, uv caches the compiled wheel, which can be reused in other projects without recompiling. To reuse the wheel directly, (for example, on an other cluster), locate the cached wheel with:

find ~/.cache/uv -name "flash_attn*.whl" 2>/dev/null

Example

The full source code for this example is available on the mila-docs GitHub repository.

Job scriptPython code

#!/bin/bash
#SBATCH --ntasks=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=1
#SBATCH --gpus-per-task=l40s:1
## Or --gpus-per-task=h100:1    to request a different GPU model
#SBATCH --mem-per-gpu=16G
#SBATCH --time=00:15:00

# Exit on error
set -e

# Echo time and hostname into log
echo "Date:     $(date)"
echo "Hostname: $(hostname)"

# Execute Python script
# Use `uv run --offline` on clusters without internet access on compute nodes.
uv run python main.py

import torch
from flash_attn import flash_attn_interface


def main():
    assert torch.cuda.is_available(), "Flash Attention requires a CUDA-capable GPU."
    device = torch.device("cuda")

    # Flash Attention requires float16 or bfloat16 inputs.
    dtype = torch.bfloat16

    batch_size = 2
    seq_len = 4096
    num_heads = 8
    head_dim = 64

    q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
    k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
    v = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)

    out, _ = flash_attn_interface.flash_attn_func(
        q,
        k,
        v,
        dropout_p=0.0,
        softmax_scale=None,
        causal=True,
    )

    print(f"GPU:          {torch.cuda.get_device_name(device)}")
    print(f"Input  shape: {q.shape}  dtype: {q.dtype}")
    print(f"Output shape: {out.shape}  dtype: {out.dtype}")


if __name__ == "__main__":
    main()

Running this example

Submit the job with sbatch:

sbatch job.sh

Comments