mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
Enable triton flash atten option for rocm linux too
This commit is contained in:
+1
-1
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user