flash attention upgraded for cuda 12.8

app.py

@@ -24,20 +24,31 @@ importlib.invalidate_caches()
 
 def sh(cmd): subprocess.check_call(cmd, shell=True)
 
-flash_attention_wheel = hf_hub_download(
-        repo_id="alexnasa/flash-attn-3",
-        repo_type="model",
-        filename="flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl",
-    )
+flash_attention_installed = False
+
+try:
+    print("Attempting to download and install FlashAttention wheel...")
+    flash_attention_wheel = hf_hub_download(
+            repo_id="alexnasa/flash-attn-3",
+            repo_type="model",
+            filename="128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl",
+        )
 
-sh(f"pip install {flash_attention_wheel}")
+    sh(f"pip install {flash_attention_wheel}")
 
-# tell Python to re-scan site-packages now that the egg-link exists
-import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()
+    # tell Python to re-scan site-packages now that the egg-link exists
+    import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()
 
-import torch
+    flash_attention_installed = True
+    print("FlashAttention installed successfully.")
 
-print(f'torch version:{torch.__version__}')
+except Exception as e:
+    print(f"⚠️ Could not install FlashAttention: {e}")
+    print("Continuing without FlashAttention...")
+
+import torch
+print(f"Torch version: {torch.__version__}")
+print(f"FlashAttention available: {flash_attention_installed}")
 
 
 import torch.nn as nn

requirements.txt

@@ -1,4 +1,3 @@
-torch==2.7.1
 tqdm
 librosa==0.10.2.post1
 peft==0.15.1