Skip to content

flash_attn / diffusers Startup Error Fix — Technical Note

Choose a tag to compare

View all tags
@ussoewwin ussoewwin released this 03 Apr 05:57
· 319 commits to main since this release
1.4.6
c4bdaaa

This document describes Flash Attention 2 (flash_attn) related startup failures that appeared after upgrading PyTorch to 2.11.0+cu130, and the launcher-side changes made to address them.

1. Errors Observed (Representative Messages)

1.1 Native extension load failure (ties directly to the root cause)

ImportError: DLL load failed while importing flash_attn_2_cuda: The specified procedure could not be found.

(On non-English Windows locales the same failure may appear with a localized message instead of the English text above.)

What this means

  • The flash_attn package loads not only Python code but also CUDA native modules (on Windows, typically something like flash_attn_2_cuda.cp313-win_amd64.pyd).
  • If that .pyd does not match the ABI and symbols of the currently installed PyTorch (and CUDA runtime), loading the DLL fails as above.
  • “The specified procedure could not be found” typically means the expected C++ exports (symbols) are missing from torch_cuda or related DLLs (e.g. loading an extension built for an older torch on top of a newer torch).

1.2 Cascading error through diffusers

If importing flash_attn fails, some diffusers modules (e.g. attention_dispatch) run from flash_attn import ... at import time, so the failure may surface wrapped in a RuntimeError like:

RuntimeError: Failed to import diffusers.models.transformers.transformer_flux because of the following error (look up to see its traceback):
DLL load failed while importing flash_attn_2_cuda: ...

The stack trace may show nunchaku → diffusers → flash_attn, but the first failure is loading flash_attn_2_cuda.

1.3 pip errors on reinstall (Windows-specific)

When trying to overwrite a broken or outdated build, you may see:

ERROR: Could not install packages due to an OSError: [WinError 5] Access is denied.
'...\site-packages\flash_attn_2_cuda.cp313-win_amd64.pyd'
Check the permissions.

What this means

  • On Windows, while another process has the .pyd mapped, overwriting or deleting the same file can fail.
  • If the Web UI or another Python process is running, pip install -U or pip uninstall may stop with WinError 5.

1.4 Launcher message

If pip fails during dependency installation, the launcher still prints:

Installing flash_attn
Failed to install flash_attn; Please manually install it

That means the pip command exited with a non-zero status. Causes may include access denied, network issues, bad URL, etc.

2. Root Cause Analysis (Why It Happened)

2.1 Direct cause: torch upgrade vs. flash_attn build mismatch

  • Windows wheels for flash_attn are separate binaries built for a specific PyTorch version, CUDA version, and sometimes C++ ABI.
  • If PyTorch was upgraded to 2.11.0+cu130 but the venv still contains a flash_attn .pyd built for an older torch (e.g. 2.9.x + different CUDA), metadata may still say flash_attn 2.8.3 while runtime DLL loading breaks.

2.2 Detection gap: “package exists” ≠ “can be used”

Older launcher logic often treated the install as complete when:

  • importlib.metadata finds the distribution name, or
  • importlib.util.find_spec("flash_attn") returns a path.

Neither guarantees that the native extension loads successfully.
They cannot detect “Python package present but .pyd built for another torch”.

2.3 diffusers detection gap

In diffusers utils/import_utils.py, optional availability uses _is_package_available("flash_attn"), which is effectively find_spec + version string — it does not guarantee that import flash_attn succeeds.

The failure chain:

  1. _flash_attn_available is True
  2. models/attention_dispatch.py at import time runs from flash_attn import ...
  3. DLL load fails → whole app crashes on startup

2.4 Why pip -U -I alone is not always enough on Windows

  • If the target .pyd is held by another process, pip may fail with WinError 5.
  • You may need uninstall then reinstall, or stop other processes first.

3. Files Changed and Their Roles

File Role
modules/launch_utils.py First-run dependency install. Verifies with a real import; on failure, force uninstall → reinstall.
modules_forge/transformers_cache_compat.py Right after diffusers.utils.import_utils loads, re-validates flash_attn with a real import; on failure, sets _flash_attn_available to False (avoids module-level import crash).

4. modules/launch_utils.py — Change Details

4.1 _verify_native_import(package)

Goal: Decide success by whether import actually works in a subprocess, not merely “is the package registered”.

def _verify_native_import(package: str) -> bool:
    """Verify a package can actually be imported (catches DLL / native extension failures)."""
    if not is_installed(package):
        return False
    return check_run_python(f"import {package}")

Meaning

  • check_run_python runs python -c "..." from the venv in a separate process.
  • The main process is not polluted by import side effects, and you learn whether loading the C extension succeeds.
  • Intended for packages that depend on native code, e.g. flash_attn / sageattention.

4.2 _force_uninstall(package)

Goal: Even if pip uninstall fails (e.g. locked .pyd), remove leftovers under site-packages as far as possible so a later pip install can succeed.

def _force_uninstall(package: str) -> None:
    """Uninstall a package, force-removing residual files if pip fails (e.g. locked .pyd on Windows)."""
    import shutil
    site_packages = os.path.join(venv_dir, "Lib", "site-packages")
    try:
        run(f'"{python}" -m pip uninstall -y {package}', f"Uninstalling broken {package}")
    except RuntimeError:
        print(f"pip uninstall failed for {package}; attempting manual removal")
    for entry in os.listdir(site_packages):
        if entry.startswith(package) or entry.startswith(package.replace("-", "_")):
            path = os.path.join(site_packages, entry)
            try:
                if os.path.isdir(path):
                    shutil.rmtree(path)
                else:
                    os.remove(path)
            except OSError as e:
                print(f"Cannot remove {path}: {e} — close other Python processes and retry")

Meaning

  • First try normal pip uninstall -y.
  • On failure, do not abort the whole startup with RuntimeError; scan site-packages and remove files/dirs whose names start with flash_attn / flash_attn-* (or the sage package name).
  • If removal still fails, another process likely holds the .pyd — the message asks the user to close other Python processes.

Caveat: Prefix matching is simple; in theory a similarly named package could collide. For current names (flash_attn, sageattention) this is unlikely in practice.

4.3 Install conditions for --flash / --sage

Before (conceptually)

  • Run pip install only when flash_attn is “not installed”.

After (conceptually)

  • When _verify_native_import("flash_attn") is False (not installed or import fails):

    1. _force_uninstall("flash_attn")
    2. run_pip(f"install --no-deps {flash_package}", "flash_attn")

The same pattern applies to sageattention.

Meaning

  • After a torch upgrade, an old build can remain; the next launch can reinstall from the correct FLASH_PACKAGE / SAGE_PACKAGE.
  • Instead of relying only on -U -I, uninstall (and remove leftovers) first, then install — this reduces overwrite failures on Windows.

Overriding the wheel URL with FLASH_PACKAGE / SAGE_PACKAGE still works as before.

5. modules_forge/transformers_cache_compat.py — Change Details

5.1 Background

This module already provides a builtins.__import__ hook for transformers 5.x compatibility (HybridCache, no_init_weights, etc.). The diffusers-related logic is an add-on on top of that.

5.2 _patch_diffusers_flash_attn_flag(mod)

When: Immediately after the module diffusers.utils.import_utils is loaded — once only.

Flow

  1. Only if _flash_attn_available exists and is True, run validation.
  2. Try import flash_attn using the hook’s saved original __import__.
  3. On any exception: set
    mod._flash_attn_available = False and mod._flash_attn_version = "N/A"
    so is_flash_attn_available() returns False afterward.

Meaning

  • attention_dispatch.py evaluates _CAN_USE_FLASH_ATTN = is_flash_attn_available() and ... at import time.
  • If a real import would fail, fixing the flag to False first avoids entering from flash_attn import ... and prevents a full startup crash.
  • This is not “permanently disabling the feature”; it is a guardrail until the next launch with --flash lets launch_utils reinstall the correct wheel.

5.3 Hook wiring

if name == "diffusers.utils.import_utils":
    _patch_diffusers_flash_attn_flag(mod)

Meaning

  • transformers_cache_compat.apply() must run before importing diffusers (this repo calls it from launch_utils.start() and initialize.shush_nunchaku()).

6. End-to-End Flow

flowchart TD
    A[Startup: prepare_environment] --> B{_verify_native_import flash_attn?}
    B -->|OK| C[Continue]
    B -->|NG| D[_force_uninstall flash_attn]
    D --> E[pip install correct wheel]
    E --> C
    C --> F[start: apply import hook]
    F --> G[import diffusers...]
    G --> H[import_utils loads]
    H --> I{_patch_diffusers: import flash_attn}
    I -->|OK| J[FA available as usual]
    I -->|NG| K[_flash_attn_available = False]
    K --> L[attention_dispatch skips FA branch]

7. Checklist If It Still Fails on Your Machine

  1. No other Python / Web UI running (check Task Manager for python.exe).
  2. FLASH_PACKAGE points to a wheel that matches your Python (e.g. cp313), torch, and CUDA.
  3. With the venv activated, does
    python -c "import flash_attn; print('ok')"
    succeed?

8. Related Environment Variables (Recap)

Variable Meaning
FLASH_PACKAGE URL of the Windows flash_attn wheel (default: e.g. Hugging Face cu130 + torch 2.11 build)
SAGE_PACKAGE URL of the Windows sageattention wheel
TORCH_COMMAND / PYTORCH_VERSION PyTorch install / version parsing (inside launch_utils)

9. Summary

  • Error: flash_attn_2_cuda DLL load failure — usually torch/CUDA vs. extension build mismatch.
  • Detection gaps: Package presence alone is not enough; diffusers can also mis-detect using find_spec-style logic.
  • Mitigations:
    • At startup: _verify_native_import + _force_uninstall + reinstall
    • After import_utils loads: re-validate with a real import and fix the flag

Together, this makes it easier to detect and replace a broken flash_attn after a PyTorch upgrade, and if import still fails, reduces hard crashes through diffusers on startup.

Assets 2
Loading