Rename ROCm Flash atten hijack to CK Flash atten and enable AOTriton memory and flash atten by default

2026-09-18 16:54:33 +02:00 · 2025-02-13 22:01:06 +03:00
parent 5ee6768eba
commit f94196bcd1
5 changed files with 36 additions and 50 deletions
@@ -6,7 +6,7 @@ from functools import wraps
 import torch
 from modules import rocm
 from modules.errors import log, display, install as install_traceback
-from installer import install
+from installer import install, installed


 debug = os.environ.get('SD_DEVICE_DEBUG', None) is not None
@@ -417,22 +417,7 @@ def set_sdpa_params():
                sdpa_original = torch.nn.functional.scaled_dot_product_attention
        except Exception as err:
            log.warning(f'Torch attention: type="sdpa" {err}')
-        if backend == "rocm":
-            if 'Flash attention' in opts.sdp_options:
-                try:
-                    # https://github.com/huggingface/diffusers/discussions/7172
-                    from flash_attn import flash_attn_func
-                    sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention
-                    @wraps(sdpa_pre_flash_atten)
-                    def sdpa_flash_atten(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
-                        if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
-                            return flash_attn_func(q=query.transpose(1, 2), k=key.transpose(1, 2), v=value.transpose(1, 2), dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
-                        else:
-                            return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
-                    torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
-                    log.debug('Torch attention: type="rocm flash attention"')
-                except Exception as err:
-                    log.error(f'Torch attention: type="rocm flash attention" {err}')
+
        try:
            torch.backends.cuda.enable_flash_sdp('Flash attention' in opts.sdp_options)
            torch.backends.cuda.enable_mem_efficient_sdp('Memory attention' in opts.sdp_options)
@@ -440,6 +425,10 @@ def set_sdpa_params():
            log.debug(f'Torch attention: type="sdpa" flash={"Flash attention" in opts.sdp_options} memory={"Memory attention" in opts.sdp_options} math={"Math attention" in opts.sdp_options}')
        except Exception as err:
            log.warning(f'Torch attention: type="sdpa" {err}')
+
+        # stack hijcaks with order: Sage -> CK Flash -> Dynamic
+        # if the first is not compatible, uses the second and so on
+
        if 'Sage attention' in opts.sdp_options:
            try:
                install('sageattention')
@@ -455,7 +444,29 @@ def set_sdpa_params():
                log.debug('Torch attention: type="sage attention"')
            except Exception as err:
                log.error(f'Torch attention: type="sage attention" {err}')
-        elif 'Dynamic attention' in opts.sdp_options:
+
+        if 'CK Flash attention' in opts.sdp_options:
+            try:
+                if backend == "rocm":
+                    if not installed('flash-attn'):
+                        agent = rocm.Agent(getattr(torch.cuda.get_device_properties(device), "gcnArchName", "gfx0000"))
+                        install(rocm.get_flash_attention_command(agent), reinstall=True)
+                else:
+                    install('flash-attn')
+                from flash_attn import flash_attn_func
+                sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention
+                @wraps(sdpa_pre_flash_atten)
+                def sdpa_flash_atten(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
+                    if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
+                        return flash_attn_func(q=query.transpose(1, 2), k=key.transpose(1, 2), v=value.transpose(1, 2), dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
+                    else:
+                        return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
+                torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
+                log.debug('Torch attention: type="ck flash attention"')
+            except Exception as err:
+                log.error(f'Torch attention: type="ck flash attention" {err}')
+
+        if 'Dynamic attention' in opts.sdp_options:
            try:
                global sdpa_pre_dyanmic_atten # pylint: disable=global-statement
                sdpa_pre_dyanmic_atten = torch.nn.functional.scaled_dot_product_attention
@@ -464,24 +475,6 @@ def set_sdpa_params():
                log.debug('Torch attention: type="dynamic attention"')
            except Exception as err:
                log.error(f'Torch attention: type="dynamic attention" {err}')
-        """
-        elif 'Flash2 attention' in opts.sdp_options:
-            try:
-                install('flash-attn')
-                from flash_attn import flash_attn_func
-                sdpa_pre_flash2_atten = torch.nn.functional.scaled_dot_product_attention
-                @wraps(sdpa_pre_flash2_atten)
-                def sdpa_flash2_atten(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
-                    if (query.shape[-1] in {128, 96, 64}) and (attn_mask is None) and (query.dtype != torch.float32) and (query.shape[-1] % key.shape[-1] == 0) and (query.shape[-1] % value.shape[-1] == 0):
-                        print('HERE', query.shape, key.shape, value.shape)
-                        return flash_attn_func(q=query, k=key, v=value, causal=is_causal)
-                    else:
-                        return sdpa_pre_flash2_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
-                torch.nn.functional.scaled_dot_product_attention = sdpa_flash2_atten
-                log.debug('Torch attention: type="flash2 attention"')
-            except Exception as err:
-                log.error(f'Torch attention: type="flash2 attention" {err}')
-        """
    except Exception as e:
        log.warning(f'Torch SDPA: {e}')

@@ -202,7 +202,7 @@ else:
        return version == version_torch and 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":
+        if os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "false").lower() == "true":
            return "pytest git+https://github.com/ROCm/flash-attention@micmelesse/upstream_pr"
        default = "git+https://github.com/ROCm/flash-attention"
        if agent.gfx_version >= 0x1100:
@@ -452,14 +452,10 @@ def get_default_modes():
    else: # cuda, rocm, zluda, ipex, openvino
        default_cross_attention = "Scaled-Dot-Product"

-    if devices.backend == "rocm":
-        default_sdp_options =  ['Memory attention', 'Math attention']
-    elif devices.backend == "zluda":
+    if devices.backend == "zluda":
        default_sdp_options =  ['Math attention', 'Dynamic attention']
    else:
        default_sdp_options = ['Flash attention', 'Memory attention', 'Math attention']
-    if (cmd_opts.lowvram or cmd_opts.medvram) and ('Flash attention' not in default_sdp_options and 'Dynamic attention' not in default_sdp_options):
-        default_sdp_options.append('Dynamic attention')

    return default_offload_mode, default_diffusers_offload_min_gpu_memory, default_cross_attention, default_sdp_options

@@ -537,7 +533,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {

    "cross_attention_sep": OptionInfo("<h2>Cross Attention</h2>", "", gr.HTML),
    "cross_attention_optimization": OptionInfo(startup_cross_attention, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention(native)}),
-    "sdp_options": OptionInfo(startup_sdp_options, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'Sage attention'], "visible": native}),
+    "sdp_options": OptionInfo(startup_sdp_options, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'CK Flash attention', 'Sage attention'], "visible": native}),
    "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }),
    "dynamic_attention_slice_rate": OptionInfo(0.5, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4), "step": 0.01, "visible": native}),
    "dynamic_attention_trigger_rate": OptionInfo(1, "Dynamic Attention trigger rate in GB", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4)*2, "step": 0.01, "visible": native}),