Source Code Compilation and Build

Obtain the FlashAttention source code, build a wheel from the source code, and install the local build product.

Obtain the FlashAttention 2.8.3 source code.

git clone --branch v2.8.3 --depth 1 https://github.com/Dao-AILab/flash-attention.git

Ensure that PyTorch is a CUDA version.

cd flash-attention
python3 - <<'PY'
import torch  

print("torch_version=" + torch.__version__) 
print("torch_cuda=" + str(torch.version.cuda)) 
print("cuda_built=" + str(torch.backends.cuda.is_built())) 
print("cuda_available=" + str(torch.cuda.is_available())) 
assert torch.backends.cuda.is_built() 
assert torch.cuda.is_available() 
PY

Build the wheel from the source code.

export CUDA_HOME=/usr/local/cuda-13.0 
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}
export FLASH_ATTENTION_FORCE_BUILD=TRUE 
export FLASH_ATTN_CUDA_ARCHS=80 
export MAX_JOBS=8 
export NVCC_THREADS=1  

python3 -m pip wheel --no-build-isolation --no-deps --wheel-dir dist .

Install the local build product.

python3 -m pip install --force-reinstall --no-deps dist/flash_attn-2.8.3-*.whl

The following is an example of the build product.

flash_attn-2.8.3-cp311-cp311-linux_aarch64.whl

Parent topic: Developer Guide