Enable triton flash atten option for rocm linux too

This commit is contained in:
Disty0
2025-09-03 16:26:00 +03:00
parent c42e0e0b37
commit c51552af90
2 changed files with 5 additions and 3 deletions
+1 -1
View File
@@ -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)
+4 -2
View File
@@ -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,