find the right prebuilt flash-attention wheel from mjun0812/flash-attention-prebuild-wheels

find_flash_attn_wheel.py

#!/usr/bin/env python3
"""
find_flash_attn_wheel.py — find the right prebuilt flash-attention wheel from
https://github.com/mjun0812/flash-attention-prebuild-wheels

Detects (or takes via flags) your Python / PyTorch / CUDA / platform, fetches the
repo's full wheel index (doc/packages.md), and prints every matching wheel with
a recommendation and a ready-to-run `pip install <url>` command.

Usage:
  python find_flash_attn_wheel.py                     # auto-detect everything
  python find_flash_attn_wheel.py --python 3.12 --torch 2.8 --cuda 12.8
  python find_flash_attn_wheel.py --fa 2.8.3          # pin flash-attn version
  python find_flash_attn_wheel.py --fa3               # only Flash-Attention 3 wheels
  python find_flash_attn_wheel.py --loose-cuda        # allow same-major CUDA mismatch

Stdlib only. No GitHub API calls (no rate limits) — one GET of raw markdown.
"""

import argparse
import platform
import re
import sys
import urllib.request

INDEX_URL = (
    "https://raw.githubusercontent.com/mjun0812/"
    "flash-attention-prebuild-wheels/main/doc/packages.md"
)

# flash_attn-2.8.3+cu128torch2.8-cp312-cp312-linux_x86_64.whl
# flash_attn_3-3.0.0+cu126torch2.13gite2743ab-cp39-abi3-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
WHEEL_RE = re.compile(
    r"(?P<pkg>flash_attn(?:_3)?)-"
    r"(?P<fa>[\d.]+(?:\.post\d+)?)"
    r"\+cu(?P<cuda>\d+)torch(?P<torch>\d+\.\d+)(?:git[0-9a-f]+)?"
    r"-cp(?P<py>\d+)t?-(?P<abi>cp\d+t?|abi3)-"
    r"(?P<plat>[\w.]+?)\.whl"
)


def detect_environment():
    """Best-effort detection of python / torch / cuda / platform tags."""
    env = {}
    vi = sys.version_info
    freethreaded = bool(__import__("sysconfig").get_config_var("Py_GIL_DISABLED"))
    env["python"] = f"{vi.major}.{vi.minor}" + ("t" if freethreaded else "")

    try:
        import torch  # noqa: PLC0415

        env["torch"] = ".".join(torch.__version__.split("+")[0].split(".")[:2])
        env["cuda"] = torch.version.cuda  # None on CPU-only builds
        try:
            if torch.cuda.is_available():
                cap = torch.cuda.get_device_capability(0)
                env["sm"] = cap[0] * 10 + cap[1]
                env["gpu"] = torch.cuda.get_device_name(0)
        except Exception:
            pass
    except ImportError:
        env["torch"] = None
        env["cuda"] = None

    mach = platform.machine().lower()
    if sys.platform.startswith("win"):
        env["os_arch"] = "win_amd64"
    elif mach in ("arm64", "aarch64"):
        env["os_arch"] = "linux_aarch64"
    else:
        env["os_arch"] = "linux_x86_64"
    return env


def platform_ok(plat_tag: str, os_arch: str) -> bool:
    """Both plain linux_* and manylinux_* wheels run on a normal glibc Linux."""
    if os_arch == "win_amd64":
        return "win_amd64" in plat_tag
    if os_arch == "linux_aarch64":
        return "aarch64" in plat_tag or "arm64" in plat_tag
    return "x86_64" in plat_tag and "win" not in plat_tag


def python_ok(wheel_py: str, wheel_abi: str, user_py: str) -> bool:
    # user_py like "3.12" or "3.14t"
    ft = user_py.endswith("t")
    base = user_py.rstrip("t")
    user_cp = base.replace(".", "")  # "312"
    if wheel_abi == "abi3":
        # abi3 wheels (cp39-abi3) work on any CPython >= 3.9 with GIL
        return not ft and int(user_cp) >= int(wheel_py)
    wheel_ft = wheel_abi.endswith("t")
    return wheel_py == user_cp and wheel_ft == ft


def cuda_ok(wheel_cu: str, user_cuda: str, loose: bool) -> bool:
    # wheel_cu like "128"; user_cuda like "12.8"
    u = user_cuda.replace(".", "")[:4]
    # normalize: "12.8" -> "128", "13.0" -> "130"
    umaj, umin = user_cuda.split(".")[:2]
    u = f"{umaj}{umin}"
    if wheel_cu == u:
        return True
    if loose:
        wmaj = wheel_cu[:-1] if len(wheel_cu) == 3 else wheel_cu[:2]
        # major match, wheel minor <= user minor (CUDA minor-version compat)
        wmin = wheel_cu[len(wmaj) :]
        return wmaj == umaj and int(wmin) <= int(umin)
    return False


def fa_sort_key(v: str):
    parts = []
    for p in v.replace(".post", ".").split("."):
        parts.append(int(p) if p.isdigit() else 0)
    return tuple(parts)


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[1])
    ap.add_argument("--python", help='e.g. 3.12 (use "3.14t" for free-threaded)')
    ap.add_argument("--torch", help="e.g. 2.8")
    ap.add_argument("--cuda", help="e.g. 12.8 (the CUDA your torch build uses)")
    ap.add_argument("--fa", help="pin a flash-attn version, e.g. 2.8.3")
    ap.add_argument(
        "--fa3",
        action="store_true",
        help="only Flash-Attention 3 (flash_attn_3) wheels",
    )
    ap.add_argument("--os-arch", choices=["linux_x86_64", "linux_aarch64", "win_amd64"])
    ap.add_argument(
        "--loose-cuda",
        action="store_true",
        help="accept wheels built for a lower CUDA minor within the same major",
    )
    ap.add_argument(
        "--all",
        action="store_true",
        help="print every match, not just the top pick per FA version",
    )
    args = ap.parse_args()

    env = detect_environment()
    py = args.python or env["python"]
    torch_v = args.torch or env["torch"]
    cuda_v = args.cuda or env["cuda"]
    os_arch = args.os_arch or env["os_arch"]

    print("── Your target ─────────────────────────────")
    print(f"  Python : {py}")
    print(f"  PyTorch: {torch_v or 'NOT DETECTED (pass --torch)'}")
    print(
        f"  CUDA   : {cuda_v or 'NOT DETECTED (pass --cuda; must match your torch build, see torch.version.cuda)'}"
    )
    print(f"  Arch   : {os_arch}")
    if "sm" in env:
        print(f"  GPU    : {env.get('gpu', '?')} (SM{env['sm']})")

    if not (torch_v and cuda_v):
        sys.exit("\nMissing torch/cuda version. Pass --torch and --cuda explicitly.")

    print(f"\nFetching wheel index ({INDEX_URL.split('/')[-1]}) ...")
    with urllib.request.urlopen(INDEX_URL, timeout=30) as r:
        text = r.read().decode()

    seen, matches = set(), []
    for m in WHEEL_RE.finditer(text):
        d = m.groupdict()
        # Reconstruct the download URL from the surrounding markdown link
        start = text.rfind("(https://", max(0, m.start() - 200), m.start())
        url = None
        if start != -1:
            end = text.find(")", start)
            url = text[start + 1 : end]
        if not url or url in seen:
            continue
        seen.add(url)

        if args.fa3 and d["pkg"] != "flash_attn_3":
            continue
        if args.fa and d["fa"] != args.fa:
            continue
        if not platform_ok(d["plat"], os_arch):
            continue
        if not python_ok(d["py"], d["abi"], py):
            continue
        if d["torch"] != torch_v:
            continue
        if not cuda_ok(d["cuda"], cuda_v, args.loose_cuda):
            continue
        matches.append((d, url))

    if not matches:
        print("\nNo matching wheel found.")
        print(
            "Try --loose-cuda, or check the search page: https://mjunya.com/flash-attention-prebuild-wheels/"
        )
        sys.exit(1)

    # Group by (pkg, fa version); prefer manylinux over plain linux within a group
    groups = {}
    for d, url in matches:
        groups.setdefault((d["pkg"], d["fa"]), []).append((d, url))
    for g in groups.values():
        g.sort(
            key=lambda x: ("manylinux" not in x[0]["plat"], x[1]),
        )

    print(f"\n── Matching wheels ({len(matches)}) ────────────────────")
    ordered = sorted(groups, key=lambda k: (k[0], fa_sort_key(k[1])), reverse=True)
    for pkg, fa in ordered:
        rec = ""
        if pkg == "flash_attn" and fa.startswith("2.8.3"):
            rec = "   <-- recommended (stable FA2 line)"
        if pkg == "flash_attn_3":
            rec = "   (FA3: requires Hopper SM90+ GPU, imports as flash_attn_3)"
        print(f"\n{pkg} {fa}{rec}")
        items = groups[(pkg, fa)] if args.all else groups[(pkg, fa)][:1]
        for d, url in items:
            print(f"  [{d['plat']}]")
            print(f"  pip install {url}")

    # Final single recommendation
    pick = None
    for pkg, fa in ordered:
        if pkg == "flash_attn" and fa.startswith("2.8.3"):
            pick = groups[(pkg, fa)][0][1]
            break
    if not pick:
        fa2 = [k for k in ordered if k[0] == "flash_attn"]
        pick = groups[fa2[0]][0][1] if fa2 else groups[ordered[0]][0][1]

    print("\n── Recommended install ─────────────────────")
    print(f"pip install {pick}")
    if "sm" in env and env["sm"] >= 90 and any(k[0] == "flash_attn_3" for k in ordered):
        print(
            "\nYour GPU is Hopper+ — Flash-Attention 3 wheels above are usable too "
            "(separate package, `import flash_attn_3`)."
        )


if __name__ == "__main__":
    main()