install pre-built flash atten

app.py

@@ -2,12 +2,73 @@ import os
 import subprocess
 import sys
 
-# Install flash-attn before other imports
-subprocess.run(
-    [sys.executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
-    check=True,
-)
+def install_flash_attn():
+    """Auto-detect CUDA, PyTorch, Python versions and install the matching pre-built flash-attn wheel."""
+    import torch
+
+    # Python version (e.g., "cp310")
+    py_major = sys.version_info.major
+    py_minor = sys.version_info.minor
+    cp_tag = f"cp{py_major}{py_minor}"
+
+    # PyTorch version (e.g., "2.4" from "2.4.0")
+    torch_version = torch.__version__.split("+")[0]  # strip +cu121 if present
+    torch_major_minor = ".".join(torch_version.split(".")[:2])  # "2.4"
+
+    # CUDA version (e.g., "cu124" from "12.4")
+    cuda_version = torch.version.cuda
+    if cuda_version is None:
+        print("No CUDA detected, skipping flash-attn installation.")
+        return
+    cuda_major_minor = cuda_version.replace(".", "")  # "124"
+    # flash-attn wheels use shortened CUDA tags like "cu12" (just major) or "cu121", "cu124"
+    # Check available tags: most wheels use "cu12" for any 12.x
+    cuda_tag_short = f"cu{cuda_version.split('.')[0]}"  # "cu12"
+
+    # CXX11 ABI
+    cxx11_abi = torch._C._GLIBCXX_USE_CXX11_ABI
+    abi_tag = "cxx11abiTRUE" if cxx11_abi else "cxx11abiFALSE"
+
+    # flash-attn version to install
+    flash_attn_version = "2.7.4.post1"
+
+    # Construct the wheel filename
+    # Example: flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
+    wheel_name = (
+        f"flash_attn-{flash_attn_version}+"
+        f"{cuda_tag_short}torch{torch_major_minor}{abi_tag}-"
+        f"{cp_tag}-{cp_tag}-linux_x86_64.whl"
+    )
+
+    base_url = f"https://github.com/Dao-AILab/flash-attention/releases/download/v{flash_attn_version}"
+    wheel_url = f"{base_url}/{wheel_name}"
+
+    print(f"Detected environment:")
+    print(f"  Python:    {py_major}.{py_minor} ({cp_tag})")
+    print(f"  PyTorch:   {torch_version} (torch{torch_major_minor})")
+    print(f"  CUDA:      {cuda_version} ({cuda_tag_short})")
+    print(f"  CXX11 ABI: {cxx11_abi} ({abi_tag})")
+    print(f"  Wheel URL: {wheel_url}")
+
+    result = subprocess.run(
+        [sys.executable, "-m", "pip", "install", wheel_url],
+        capture_output=True,
+        text=True,
+    )
+
+    if result.returncode != 0:
+        print(f"Pre-built wheel failed:\n{result.stderr}")
+        print("Falling back to building flash-attn from source (this may take a while)...")
+        subprocess.run(
+            [sys.executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
+            check=True,
+        )
+    else:
+        print("flash-attn installed successfully from pre-built wheel.")
+        print(result.stdout)
+
 
+install_flash_attn()
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import copy