Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

add triton backend flash-attn (experimental)

Browse Source
This commit is contained in:
Seunghoon Lee
2024-09-24 12:53:18 +09:00
parent e246e55734
commit 1395f5bf9e
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/rocm.py
+2
View File
@@ -185,6 +185,8 @@ else:
return bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1")))
def get_flash_attention_command(agent: Agent):
if os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "FALSE") == "TRUE":
return "pytest git+https://github.com/ROCm/flash-attention@micmelesse/upstream_pr"
default = "git+https://github.com/ROCm/flash-attention"
if agent.arch == MicroArchitecture.RDNA:
default = "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"