diff --git a/modules/devices.py b/modules/devices.py index 846d017dc..7a241c3cd 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -447,7 +447,7 @@ def set_sdpa_params(): if 'Triton Flash attention' in opts.sdp_options: try: - if sys.platform == "win32" and (backend == "zluda" or backend == "rocm"): + if backend in {"zluda", "rocm"}: from modules.flash_attn_triton_amd import interface_fa sdpa_pre_triton_flash_atten = torch.nn.functional.scaled_dot_product_attention @wraps(sdpa_pre_triton_flash_atten) diff --git a/modules/shared_defaults.py b/modules/shared_defaults.py index 504f0c239..d3993d2e9 100644 --- a/modules/shared_defaults.py +++ b/modules/shared_defaults.py @@ -40,15 +40,17 @@ def get_default_modes(cmd_opts, mem_stat): default_cross_attention = "Scaled-Dot-Product" - default_sdp_choices = ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'CK Flash attention', 'Sage attention'] if devices.backend == "zluda": - default_sdp_choices.insert(4, 'Triton Flash attention') # insert after Dynamic attention default_sdp_options = ['Math attention', 'Dynamic attention'] elif devices.backend in {"rocm", "directml", "cpu", "mps"}: default_sdp_options = ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention'] else: default_sdp_options = ['Flash attention', 'Memory attention', 'Math attention'] + default_sdp_choices = ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'CK Flash attention', 'Sage attention'] + if devices.backend in {"rocm", "zluda"}: + default_sdp_choices.insert(4, 'Triton Flash attention') # insert after Dynamic attention + return ( default_offload_mode, default_diffusers_offload_min_gpu_memory,