Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Disty0
2025-02-13 22:01:06 +03:00
parent 5ee6768eba
commit f94196bcd1
5 changed files with 36 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/locale-sanitize-override.py
Regular → Executable
View File
installer.py
+4 -7
View File
@@ -655,6 +655,10 @@ def install_rocm_zluda():
torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
else:
check_python(supported_minors=[9, 10, 11, 12], reason='ROCm backend requires a Python version between 3.9 and 3.12')
if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", None) is None:
os.environ.setdefault('TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL', '1')
if args.use_nightly:
if rocm.version is None or float(rocm.version) >= 6.3: # assume the latest if version check fails
torch_command = os.environ.get('TORCH_COMMAND', '--pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.3')
@@ -679,13 +683,6 @@ def install_rocm_zluda():
# older rocm (5.7) uses torch 2.3 or older
torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm{rocm.version}')
if installed("torch") and device is not None:
if 'Flash attention' in opts.get('sdp_options', ''):
if not installed('flash-attn'):
install(rocm.get_flash_attention_command(device), reinstall=True)
#elif not args.experimental:
# uninstall('flash-attn')
if device is not None and rocm.version != "6.2" and rocm.get_blaslt_enabled():
log.debug(f'ROCm hipBLASLt: arch={device.name} available={device.blaslt_supported}')
rocm.set_blaslt_enabled(device.blaslt_supported)
modules/devices.py
+29 -36
View File
@@ -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}')
modules/rocm.py
+1 -1
View File
@@ -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:
modules/shared.py
+2 -6
View File
@@ -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}),