How to Install Flash-Attention

Install a prebuilt FlashAttention wheel from the Astral GPU index to avoid compiling CUDA kernels. Route PyTorch to the matching CUDA index too: the PyTorch version alone does not identify its CUDA build.

PyPI publishes a source distribution. Its installer attempts to download a GitHub release wheel before falling back to compilation, but preparing that source distribution also has build prerequisites. The Astral index lets the package resolver select a wheel directly.

Prerequisites

• Linux on an architecture and Python version covered by the chosen index. Python 3.12 is a practical starting point; current PyTorch builds require Python 3.10 or later.
• An NVIDIA Ampere, Ada, or Hopper GPU for the FlashAttention-2 CUDA path. Turing GPUs such as T4 do not support these kernels. For Blackwell, check the specific wheel’s architecture support and use a suitable CUDA build, starting at CUDA 12.8.
• A working NVIDIA driver compatible with the packaged CUDA runtime. The driver stays on the host; wheels do not install it.
• An existing uv project.

A prebuilt wheel needs no local CUDA toolkit or nvcc. The CUDA version displayed by nvidia-smi describes driver capability, while torch.version.cuda identifies PyTorch’s packaged runtime. CUDA minor-version compatibility has limits for newer GPU features and PTX, so the driver number alone does not establish kernel compatibility. See the driver, runtime, and toolkit explanation.

Official CUDA wheels do not target macOS, and upstream Windows support is experimental. AMD ROCm uses separate upstream build instructions, a matching ROCm PyTorch build, and supported AMD hardware; the CUDA indexes in this guide do not cover it.

Route both packages to matching indexes

Choose a CUDA line covered by both PyTorch and the Astral manifest. For CUDA 12.8:

Add these source mappings and explicit indexes to pyproject.toml before adding the dependencies:

[tool.uv.sources]
torch = { index = "pytorch-cu128" }
flash-attn = { index = "astral-cu128" }

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true

[[tool.uv.index]]
name = "astral-cu128"
url = "https://wheels.astral.sh/simple/cu128/"
explicit = true

Then declare both packages directly:

uv add torch flash-attn

Expect a CUDA 12.8 torch build and a flash-attn version containing +cu.12.8.torch. followed by the matching PyTorch minor version. The command updates dependencies, the lockfile, and the environment.

Configure indexes before the first resolve. Changing explicit later does not rewrite sources already recorded in uv.lock.

Both packages must also be declared as direct project dependencies. explicit = true keeps unrelated packages on PyPI; a source mapping makes the named package use that index. An explicit index without a source mapping is not used for ordinary resolution.

Astral wheels constrain the PyTorch minor version in their metadata, such as torch==2.11.*. That excludes an incompatible PyTorch minor version but still accepts either a CUDA 12 or CUDA 13 build of the same release. Routing torch is therefore necessary even when dependency resolution succeeds.

Check the live manifest for the complete Python, CPU architecture, PyTorch, and CUDA combination. Coverage differs between indexes, and the newest upstream source release may have no corresponding Astral wheel. A lockfile reproduces the selected artifacts; it does not make their kernels work on every GPU.

Install outside a project

In an existing virtual environment, install PyTorch first and leave it installed while resolving FlashAttention:

uv pip install torch --index https://download.pytorch.org/whl/cu128
uv pip install flash-attn --index https://wheels.astral.sh/simple/cu128/

Expect prebuilt wheels for both packages. If an existing torch pin has no matching FlashAttention wheel, choose a covered pairing from the manifest rather than switching one package to another CUDA line.

Build images without detecting a GPU

An image build can install these wheels without GPU access. Keep the explicit CUDA indexes and run the kernel check after starting a container with a GPU. Automatic GPU detection during the image build sees the build host, which may have no accelerator.

For a source build, the absence of a GPU is a different problem: the build still needs a CUDA toolkit, compiler, and explicit target architectures for the deployment GPU. Follow the upstream package’s build configuration; a successful CPU-only image build or import is not a FlashAttention execution test.

Verify an attention operation

Save this script on the GPU host. It checks the NVIDIA CUDA backend and executes FlashAttention with float16 tensors:

import torch
from flash_attn import flash_attn_func

assert torch.cuda.is_available(), 'CUDA is unavailable to PyTorch'
assert torch.version.hip is None, 'This check targets NVIDIA CUDA'
assert torch.cuda.get_device_capability()[0] >= 8
print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name())
q = torch.randn(2, 256, 8, 64, dtype=torch.float16, device='cuda')
k, v = torch.randn_like(q), torch.randn_like(q)
out = flash_attn_func(q, k, v, causal=True)
torch.cuda.synchronize()
assert out.is_cuda and out.shape == q.shape
assert torch.isfinite(out).all()
print('FlashAttention CUDA operation OK')

uv run python check_flash_attn.py

uv run executes the command inside the project’s virtual environment, ensuring all dependencies are installed first. Expect the GPU name and FlashAttention CUDA operation OK. An import alone does not establish that a wheel contains a working kernel for the device.

Use a GitHub wheel or source build

When the Astral index has no suitable wheel, inspect the upstream release assets. A GitHub wheel must match Python, platform, CUDA, PyTorch, and the CXX11 ABI. Copy its exact filename from the release rather than guessing a URL. Unlike the Astral builds, GitHub wheels may declare an unpinned torch dependency; install the matching PyTorch first.

For the PyPI source path, install PyTorch plus packaging, ninja, and psutil before using --no-build-isolation. The CUDA build needs a supported C++ compiler and a full CUDA toolkit with nvcc; CUDA_HOME must point to the toolkit. A runtime-only image is insufficient, even if it runs ordinary PyTorch operations. Limit build parallelism with MAX_JOBS when memory is constrained. Follow upstream build requirements for the chosen release.

The source installer constructs CUDA extensions before its GitHub-wheel download attempt. Consequently, taking that path can require a toolkit even when a cached release wheel ultimately exists. A direct wheel URL or Astral wheel avoids that metadata/build step.

Fix a failed operation

• Missing libcudart.so or another CUDA library: check both packages’ CUDA builds and the full loader error. A CUDA 12 extension paired with a CUDA 13 PyTorch build is one cause; library-path conflicts are another.
• An undefined symbol in flash_attn_2_cuda: check the PyTorch minor version and, for upstream release wheels, the CXX11 ABI. Reinstall a matching wheel.
• No matching wheel: check Python, operating system, and CPU architecture against the live manifest. A supported CUDA line does not imply every Python/PyTorch pairing is covered.
• CUDA_HOME or nvcc errors during installation: the installer is using a source path. Select a matching prebuilt wheel, or provide the full build toolchain.
• Unsupported GPU architecture: use a kernel implementation that supports the device. For Turing, upstream links to the separate Turing project; installing a CUDA wheel alone does not add T4 support.

FlashAttention-3 and FlashAttention-4 are separate implementations with different hardware and packaging requirements. Check their upstream instructions rather than applying this FlashAttention-2 workflow to them.