pip install flash-attn: No module named 'torch' (fix)

pip install flash-attn fails with ModuleNotFoundError: No module named 'torch' even though torch is installed. We reproduced that and the three errors that come after it (including ComfyUI's --use-flash-attention message), and checked which prebuilt wheels exist.

1. ComfyUI: --use-flash-attention without flash-attn

ComfyUI v0.37.0 started with python main.py --use-flash-attention in a venv that does not have flash-attn:

[ERROR] 

To use the `--use-flash-attention` feature, the `flash-attn` package must be installed first.
command:
	/path/to/venv/bin/python -m pip install flash-attn

ComfyUI exited with code 255. flash-attn is not in ComfyUI's requirements.txt, and the command it suggests is exactly the one that fails next.

2. pip install flash-attn: No module named 'torch'

flash-attn 2.8.3.post2 is published on PyPI as a source archive only (no wheels). With pip 26.2.2 and torch 2.9.0 already installed in the venv:

python -m pip install flash-attn

        File "<string>", line 22, in <module>
      ModuleNotFoundError: No module named 'torch'
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
ERROR: Failed to build 'flash-attn' when getting requirements to build wheel

Torch is installed, but pip builds the package in an isolated build environment that only contains the build tools, and flash-attn's setup.py imports torch at line 22. That is why its README says to use --no-build-isolation.

3. --no-build-isolation: Cannot import 'setuptools.build_meta'

Without isolation, pip uses whatever is in your venv. Venvs made by uv venv (and some others) contain no setuptools:

pip._vendor.pyproject_hooks._impl.BackendUnavailable: Cannot import 'setuptools.build_meta'

(pip exit code 2.) Install the build requirements first: python -m pip install setuptools wheel packaging psutil ninja.

4. CUDA_HOME environment variable is not set

With those installed and torch 2.9.0+cu128 from PyPI, on a machine without the CUDA toolkit:

python -m pip install flash-attn --no-build-isolation

          raise OSError('CUDA_HOME environment variable is not set. '
      OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> flash-attn

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

The PyPI torch wheel ships CUDA runtime libraries, not the compiler (nvcc). flash-attn's setup.py defines its CUDA extension while pip is still generating metadata, so it needs a CUDA toolkit before it even reaches its "download a prebuilt wheel" step.

5. The trap: FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE

Setting that variable makes the install "succeed":

FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE python -m pip install flash-attn --no-build-isolation

  Building wheel for flash-attn (pyproject.toml): finished with status 'done'
  Created wheel for flash-attn: filename=flash_attn-2.8.3.post2-py3-none-any.whl size=373167 sha256=b8379362a4a69ab969a3ea8ff13f9bc50e51b5fcb7005f9263be51d3af2a81a7
  Stored in directory: ~/.cache/pip/wheels/03/65/4b/0234d6d6db26f61784989e2e0f1782aeee16df3f25dff78c93
Successfully built flash-attn
Installing collected packages: einops, flash-attn

Successfully installed einops-0.8.2 flash-attn-2.8.3.post2

Note the 373 kB py3-none-any wheel: there is no compiled kernel in it. Importing it:

python -c "from flash_attn import flash_attn_func"

  File "<string>", line 1, in <module>
  File "/path/to/venv/lib/python3.11/site-packages/flash_attn/__init__.py", line 3, in <module>
    from flash_attn.flash_attn_interface import (
  File "/path/to/venv/lib/python3.11/site-packages/flash_attn/flash_attn_interface.py", line 15, in <module>
    import flash_attn_2_cuda as flash_attn_gpu
ModuleNotFoundError: No module named 'flash_attn_2_cuda'

Uninstall it (python -m pip uninstall flash-attn) and use one of the fixes below.

6. Fix: install the matching prebuilt wheel

The project's GitHub release v2.8.3.post2 (2026-10-13) has 48 prebuilt wheels. All are Linux x86_64, CUDA 12, torch 2.4 to 2.8, CPython 3.9 to 3.13. There are no Windows wheels and none for torch 2.9. Check which one you need:

python -c "import torch,sys;print(torch.__version__, torch.version.cuda, torch._C._GLIBCXX_USE_CXX11_ABI, 'cp%d%d'%sys.version_info[:2])"

2.9.0+cu128 12.8 True cp311

The wheel name is flash_attn-2.8.3.post2+cu12torch<MAJOR.MINOR>cxx11abi<TRUE|FALSE>-<cpXY>-<cpXY>-linux_x86_64.whl. For example, with torch 2.8, ABI True and Python 3.11:

python -m pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3.post2/flash_attn-2.8.3.post2+cu12torch2.8cxx11abiTRUE-cp311-cp311-linux_x86_64.whl

• torch 2.9 or newer (as in our output above): no prebuilt wheel matches. Either install torch 2.8 in that venv, or build from source: install the CUDA toolkit, set CUDA_HOME, and run step 4's command (setting MAX_JOBS=4 keeps RAM use down; the build takes a long time).
• Windows: the project publishes no Windows wheels; you have to build from source with Visual Studio and the CUDA toolkit, or use third-party wheels at your own risk.
• ComfyUI without flash-attn: just drop --use-flash-attention. ComfyUI falls back to PyTorch attention.

Limits: this box has no NVIDIA GPU, so we measured every failure above but could not run the flash-attn kernel itself.

Related: ComfyUI: sageattention No module named 'triton' · ComfyUI: No module named 'torchsde' · ComfyUI: No module named 'torch_directml'

Written with AI assistance. Measured on 2026-10-24 on Linux x86_64 (no GPU), CPython 3.11.17, pip 26.2.2: flash-attn 2.8.3.post2 from PyPI with torch 2.9.0+cpu and 2.9.0+cu128; ComfyUI v0.37.0 --use-flash-attention. Release assets read from GitHub the same day. The CUDA kernel itself not run. Output pasted unedited except shortened paths. PlainTables is our extension; flash-attn and ComfyUI are free.