'pip install flash-attn' fails or hangs for hours (Building wheel for flash-attn, 'OSError: CUDA_HOME environment variable is not set', out of memory while compiling, or ImportError / undefined symbol in flash_attn_2_cuda after install). How do I install FlashAttention, or fall back safely?

What it is: flash-attn (FlashAttention) is a compiled CUDA/ROCm extension. Its setup first tries to download a prebuilt wheel from the project's GitHub releases that exactly matches your Python, PyTorch, CUDA version and C++ ABI; if none matches, it compiles from source, which needs the CUDA toolkit (nvcc), a working ninja and lots of RAM. Common failures: building in pip's isolated environment where torch isn't visible, no CUDA toolkit or CUDA_HOME (PyTorch's extension builder raises 'CUDA_HOME environment variable is not set'), a broken ninja making the build take hours, too many parallel compile jobs running out of RAM, or an ImportError / undefined symbol later because flash-attn was built against a different PyTorch than the one now installed. Am I affected: - Check requirements from the project README: Linux (Windows is not officially tested), PyTorch 2.2+, CUDA 12.0+ for NVIDIA, and for FlashAttention-2 an Ampere, Ada or Hopper GPU (A100, RTX 3090/4090, H100...). Older GPUs such as T4 or V100 aren't supported by the main package. - python -c "import torch; print(torch.__version__, torch.version.cuda)" and nvcc --version show your PyTorch and CUDA toolkit versions. Fix: 1. Install PyTorch first, then: pip install packaging psutil ninja and pip install flash-attn --no-build-isolation (the flag lets the build see your installed torch). This often picks up a matching prebuilt wheel and finishes quickly. 2. If it compiles: make sure ninja works (ninja --version should exit with 0; if not, pip uninstall -y ninja && pip install ninja), install a CUDA toolkit matching your PyTorch's CUDA major version and set CUDA_HOME to it, and limit parallel jobs on machines with less RAM: MAX_JOBS=4 pip install flash-attn --no-build-isolation. 3. Easiest environment: the NVIDIA PyTorch container, which the project recommends because it has all build tools. 4. After upgrading PyTorch, reinstall flash-attn (pip install --force-reinstall --no-build-isolation flash-attn) to fix undefined symbol errors caused by an ABI mismatch. 5. If you can't install it, you often don't need it: in transformers, load the model with attn_implementation="sdpa" (PyTorch's built-in scaled dot product attention) instead of "flash_attention_2"; it is fast on modern GPUs and needs no extra package. Verify: python -c "import flash_attn; print(flash_attn.__version__)" works, and the model loads with attn_implementation="flash_attention_2" (or runs with "sdpa"). Sources: https://github.com/Dao-AILab/flash-attention https://huggingface.co/docs/transformers/attention_interface https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html