From 1395f5bf9ef42c9fb1ebbee61259255437ab9000 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Tue, 24 Sep 2024 12:53:18 +0900 Subject: [PATCH] add triton backend flash-attn (experimental) --- modules/rocm.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modules/rocm.py b/modules/rocm.py index 6f1782688..dd975f323 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -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"