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

Convert Dynamic Attention SDP to a global SDP option

Browse Source
This commit is contained in:
Disty0
2024-08-18 01:44:27 +03:00
parent 5c857e8aaf
commit a3f26c9df0
5 changed files with 22 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/devices.py
+7 -7
View File
@@ -4,6 +4,7 @@ import sys
import time
import contextlib
import torch
from functools import wraps
from modules.errors import log
from modules import cmd_args, shared, memstats, errors
@@ -12,7 +13,6 @@ if sys.platform == "darwin":
previous_oom = 0
backup_sdpa = None
debug = os.environ.get('SD_DEVICE_DEBUG', None) is not None
@@ -250,18 +250,17 @@ def set_cuda_params():
except Exception:
pass
try:
if shared.opts.cross_attention_optimization == "Scaled-Dot-Product" or shared.opts.cross_attention_optimization == "Dynamic Attention SDP":
if shared.opts.cross_attention_optimization == "Scaled-Dot-Product":
torch.backends.cuda.enable_flash_sdp('Flash attention' in shared.opts.sdp_options)
torch.backends.cuda.enable_mem_efficient_sdp('Memory attention' in shared.opts.sdp_options)
torch.backends.cuda.enable_math_sdp('Math attention' in shared.opts.sdp_options)
if backend == "rocm":
global backup_sdpa # pylint: disable=global-statement
if 'Flash attention' in shared.opts.sdp_options:
try:
# https://github.com/huggingface/diffusers/discussions/7172
from flash_attn import flash_attn_func
if backup_sdpa is None:
backup_sdpa = torch.nn.functional.scaled_dot_product_attention
backup_sdpa = torch.nn.functional.scaled_dot_product_attention
@wraps(torch.nn.functional.scaled_dot_product_attention)
def sdpa_hijack(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
if query.shape[3] <= 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)
@@ -271,8 +270,9 @@ def set_cuda_params():
shared.log.debug('ROCm Flash Attention Hijacked')
except Exception as err:
log.error(f'ROCm Flash Attention failed: {err}')
elif backup_sdpa is not None: # Restore original SDPA
torch.nn.functional.scaled_dot_product_attention = backup_sdpa
if 'Dynamic attention' in shared.opts.sdp_options:
from modules.sd_hijack_dynamic_atten import sliced_scaled_dot_product_attention
torch.nn.functional.scaled_dot_product_attention = sliced_scaled_dot_product_attention
except Exception:
pass
if shared.cmd_opts.profile:
modules/sd_hijack_dynamic_atten.py
+7 -89
View File
@@ -1,7 +1,6 @@
from functools import cache
from functools import cache, wraps
import torch
import torch.nn.functional as F
from diffusers.utils import USE_PEFT_BACKEND
from modules import shared, devices
@@ -49,6 +48,8 @@ def find_slice_sizes(query_shape, query_element_size, slice_rate=4):
return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size
backup_sdpa = torch.nn.functional.scaled_dot_product_attention
@wraps(torch.nn.functional.scaled_dot_product_attention)
def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs):
do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate)
@@ -67,7 +68,7 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name
start_idx_3 = i3 * split_3_slice_size
end_idx_3 = (i3 + 1) * split_3_slice_size
hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = F.scaled_dot_product_attention(
hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = backup_sdpa(
query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
@@ -75,7 +76,7 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
dropout_p=dropout_p, is_causal=is_causal, **kwargs
)
else:
hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = F.scaled_dot_product_attention(
hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = backup_sdpa(
query[start_idx:end_idx, start_idx_2:end_idx_2],
key[start_idx:end_idx, start_idx_2:end_idx_2],
value[start_idx:end_idx, start_idx_2:end_idx_2],
@@ -83,7 +84,7 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
dropout_p=dropout_p, is_causal=is_causal, **kwargs
)
else:
hidden_states[start_idx:end_idx] = F.scaled_dot_product_attention(
hidden_states[start_idx:end_idx] = backup_sdpa(
query[start_idx:end_idx],
key[start_idx:end_idx],
value[start_idx:end_idx],
@@ -93,93 +94,10 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
if devices.backend != "directml":
getattr(torch, query.device.type).synchronize()
else:
return F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
return backup_sdpa(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
return hidden_states
class DynamicAttnProcessorSDP:
r"""
dynamically slices attention queries in order to keep them under the slice rate
slicing will not get triggered if the query size is smaller than the slice rate to gain performance
slice rate is in GB
based on AttnProcessor V2
"""
def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
def __call__(self, attn, hidden_states: torch.Tensor, encoder_hidden_states=None, attention_mask=None, temb=None, *args, **kwargs) -> torch.Tensor:
residual = hidden_states
if attn.spatial_norm is not None:
hidden_states = attn.spatial_norm(hidden_states, temb)
input_ndim = hidden_states.ndim
if input_ndim == 4:
batch_size, channel, height, width = hidden_states.shape
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
batch_size, sequence_length, _ = (
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
)
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
# scaled_dot_product_attention expects attention_mask shape to be
# (batch, heads, source_length, target_length)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
if attn.group_norm is not None:
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
query = attn.to_q(hidden_states)
if encoder_hidden_states is None:
encoder_hidden_states = hidden_states
elif attn.norm_cross:
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
key = attn.to_k(encoder_hidden_states)
value = attn.to_v(encoder_hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
# the output of sdp = (batch, num_heads, seq_len, head_dim)
# -: add support for attn.scale when we move to Torch 2.1
####################################################################
# Slicing part:
hidden_states = sliced_scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)
####################################################################
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
hidden_states = hidden_states.to(query.dtype)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
if input_ndim == 4:
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
if attn.residual_connection:
hidden_states = hidden_states + residual
hidden_states = hidden_states / attn.rescale_output_factor
return hidden_states
class DynamicAttnProcessorBMM:
r"""
dynamically slices attention queries in order to keep them under the slice rate
modules/sd_models.py
-3
View File
@@ -1525,9 +1525,6 @@ def set_diffusers_attention(pipe):
elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM":
from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM
set_attn(pipe, DynamicAttnProcessorBMM())
elif shared.opts.cross_attention_optimization == "Dynamic Attention SDP":
from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorSDP
set_attn(pipe, DynamicAttnProcessorSDP())
pipe.current_attn_name = shared.opts.cross_attention_optimization
modules/shared.py
+7 -6
View File
@@ -359,7 +359,7 @@ cpu_memory = psutil.virtual_memory().total / 1024 / 1024 / 1024
mem_stat = memory_stats()
if "gpu" in mem_stat:
gpu_memory = mem_stat['gpu']['total']
gpu_memory = mem_stat['gpu']['total'] + 0.1
if not (cmd_opts.lowvram or cmd_opts.medvram):
if "gpu" in mem_stat:
@@ -382,13 +382,11 @@ elif cmd_opts.lowvram:
if devices.backend == "directml": # Force BMM for DirectML instead of SDP
cross_attention_optimization_default = "Dynamic Attention BMM" if native else "Sub-quadratic"
elif native and (cmd_opts.lowvram or cmd_opts.medvram):
cross_attention_optimization_default = "Dynamic Attention SDP"
elif devices.backend == "cpu":
cross_attention_optimization_default = "Scaled-Dot-Product" if native else "Doggettx's"
elif devices.backend == "mps":
cross_attention_optimization_default = "Scaled-Dot-Product" if native else "Doggettx's"
else: # cuda, rocm, ipex
else: # cuda, rocm, ipex, openvino
cross_attention_optimization_default ="Scaled-Dot-Product"
@@ -399,6 +397,9 @@ if devices.backend == "rocm":
else:
sdp_options_default = ['Flash attention', 'Memory attention', 'Math attention']
if (cmd_opts.lowvram or cmd_opts.medvram) and 'Flash attention' not in sdp_options_default:
sdp_options_default.append('Dynamic attention')
options_templates.update(options_section(('sd', "Execution & Models"), {
"sd_backend": OptionInfo(default_backend, "Execution backend", gr.Radio, {"choices": ["diffusers", "original"] }),
"sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),
@@ -437,9 +438,9 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
"cross_attention_sep": OptionInfo("<h2>Cross Attention</h2>", "", gr.HTML),
"cross_attention_optimization": OptionInfo(cross_attention_optimization_default, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention(native) }),
"sdp_options": OptionInfo(sdp_options_default, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention'] }),
"sdp_options": OptionInfo(sdp_options_default, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention'] }),
"xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }),
"dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": 24, "step": 0.1, "visible": native}),
"dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory, "step": 0.1, "visible": native}),
"sub_quad_sep": OptionInfo("<h3>Sub-quadratic options</h3>", "", gr.HTML, {"visible": not native}),
"sub_quad_q_chunk_size": OptionInfo(512, "Attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8, "visible": not native}),
"sub_quad_kv_chunk_size": OptionInfo(512, "Attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8, "visible": not native}),
modules/shared_items.py
+1 -2
View File
@@ -31,8 +31,7 @@ def list_crossattention(diffusers=False):
"xFormers",
"Batch matrix-matrix",
"Split attention",
"Dynamic Attention BMM",
"Dynamic Attention SDP"
"Dynamic Attention BMM"
]
else:
return [