refactor attention handling

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -5,12 +5,12 @@
 - **Features**  
   - allow recursive inline wildcards using curly braces syntax  
   - simplify SDNQ pre-quantization saved config  
+  - refactor settings and improve handling of attention mechanisms  
 - **Fixes**
   - hires strength save/load in metadata, thanks @awsr  
   - fix imgi2img initial scale tab, thanks @awsr  
   - fix pony-v7 text-encoder  
   - detailer with face-restorers  
-  - fix sage-attention checks on sm86  
 
 ## Update for 2025-11-06
 

modules/attention.py

@@ -16,41 +16,40 @@ def set_dynamic_attention():
         log.error(f'Torch attention: type="dynamic attention" {err}')
         return None
 
-def set_triton_flash_attention(backend: str):
+def set_triton_flash_attention():
     try:
-        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)
-            def sdpa_triton_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: Optional[torch.FloatTensor] = None, dropout_p: float = 0.0, is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
-                if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
-                    if scale is None:
-                        scale = query.shape[-1] ** (-0.5)
-                    head_size_og = query.size(3)
-                    if head_size_og % 8 != 0:
-                        query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
-                        key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
-                        value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
-                    query = query.transpose(1, 2)
-                    key = key.transpose(1, 2)
-                    value = value.transpose(1, 2)
-                    out_padded = torch.zeros_like(query)
-                    interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal)
-                    return out_padded[..., :head_size_og].transpose(1, 2)
-                else:
-                    if enable_gqa:
-                        kwargs["enable_gqa"] = enable_gqa
-                    return sdpa_pre_triton_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
-            torch.nn.functional.scaled_dot_product_attention = sdpa_triton_flash_atten
-            log.debug('Torch attention: type="triton flash attention"')
+        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)
+        def sdpa_triton_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: Optional[torch.FloatTensor] = None, dropout_p: float = 0.0, is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
+            if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
+                if scale is None:
+                    scale = query.shape[-1] ** (-0.5)
+                head_size_og = query.size(3)
+                if head_size_og % 8 != 0:
+                    query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
+                    key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
+                    value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
+                query = query.transpose(1, 2)
+                key = key.transpose(1, 2)
+                value = value.transpose(1, 2)
+                out_padded = torch.zeros_like(query)
+                interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal)
+                return out_padded[..., :head_size_og].transpose(1, 2)
+            else:
+                if enable_gqa:
+                    kwargs["enable_gqa"] = enable_gqa
+                return sdpa_pre_triton_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
+        torch.nn.functional.scaled_dot_product_attention = sdpa_triton_flash_atten
+        log.debug('Torch attention: type="Triton Flash attention"')
     except Exception as err:
-        log.error(f'Torch attention: type="triton flash attention" {err}')
+        log.error(f'Torch attention: type="Triton Flash attention" {err}')
 
 def set_ck_flash_attention(backend: str, device: torch.device):
     try:
         if backend == "rocm":
             if not installed('flash-attn'):
-                log.info('Building CK Flash attention...')
+                log.info('Torch attention: type="CK Flash" building...')
                 agent = rocm.Agent(getattr(torch.cuda.get_device_properties(device), "gcnArchName", "gfx0000"))
                 install(rocm.get_flash_attention_command(agent), reinstall=True)
         else:
@@ -83,9 +82,9 @@ def set_ck_flash_attention(backend: str, device: torch.device):
                     kwargs["enable_gqa"] = enable_gqa
                 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, **kwargs)
         torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
-        log.debug('Torch attention: type="ck flash attention"')
+        log.debug('Torch attention: type="CK Flash attention"')
     except Exception as err:
-        log.error(f'Torch attention: type="ck flash attention" {err}')
+        log.error(f'Torch attention: type="CK Flash attention" {err}')
 
 def set_sage_attention(backend: str, device: torch.device):
     try:
@@ -100,7 +99,6 @@ def set_sage_attention(backend: str, device: torch.device):
                 use_cuda_backend = False
 
         if use_cuda_backend:
-            log.debug('Torch attention: type=SageAttention backend=cuda')
             from sageattention import sageattn_qk_int8_pv_fp16_cuda
             def sage_attn_impl(query, key, value, is_causal, scale):
                 return sageattn_qk_int8_pv_fp16_cuda(
@@ -112,7 +110,6 @@ def set_sage_attention(backend: str, device: torch.device):
                     pv_accum_dtype="fp32",
                 )
         else:
-            log.debug('Torch attention: type=SageAttention backend=auto')
             from sageattention import sageattn
             def sage_attn_impl(query, key, value, is_causal, scale):
                 return sageattn(
@@ -138,6 +135,60 @@ def set_sage_attention(backend: str, device: torch.device):
                     kwargs["enable_gqa"] = enable_gqa
                 return sdpa_pre_sage_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
         torch.nn.functional.scaled_dot_product_attention = sdpa_sage_atten
-        log.debug('Torch attention: type="sage attention"')
+        log.debug(f'Torch attention: type="Sage attention" backend={"cuda" if use_cuda_backend else "auto"}')
     except Exception as err:
-        log.error(f'Torch attention: type="sage attention" {err}')
+        log.error(f'Torch attention: type="Sage attention" {err}')
+
+
+def set_diffusers_attention(pipe, quiet:bool=False):
+    from modules import shared
+    import diffusers.models.attention_processor as p
+
+    def set_attn(pipe, attention, name:str=None):
+        if attention is None:
+            return
+        # other models uses their own attention processor
+        if getattr(pipe, "unet", None) is not None and hasattr(pipe.unet, "set_attn_processor"):
+            try:
+                pipe.unet.set_attn_processor(attention)
+            except Exception as e:
+                if 'Nunchaku' in pipe.unet.__class__.__name__:
+                    pass
+                else:
+                    shared.log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}')
+        """ # each transformer typically has its own attention processor
+        if getattr(pipe, "transformer", None) is not None and hasattr(pipe.transformer, "set_attn_processor"):
+            try:
+                pipe.transformer.set_attn_processor(attention)
+            except Exception as e:
+                if 'Nunchaku' in pipe.transformer.__class__.__name__:
+                    pass
+                else:
+                    shared.log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}')
+        """
+
+    shared.log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
+    if shared.opts.cross_attention_optimization == "Disabled":
+        pass # do nothing
+    elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers
+        # set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product")
+        pass
+    elif shared.opts.cross_attention_optimization == "xFormers":
+        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
+            pipe.enable_xformers_memory_efficient_attention()
+        else:
+            shared.log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
+    elif shared.opts.cross_attention_optimization == "Batch matrix-matrix":
+        set_attn(pipe, p.AttnProcessor(), name="Batch matrix-matrix")
+    elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM":
+        from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM
+        set_attn(pipe, DynamicAttnProcessorBMM(), name="Dynamic Attention BMM")
+
+    if shared.opts.attention_slicing != "Default" and hasattr(pipe, "enable_attention_slicing") and hasattr(pipe, "disable_attention_slicing"):
+        if shared.opts.attention_slicing:
+            pipe.enable_attention_slicing()
+        else:
+            pipe.disable_attention_slicing()
+        shared.log.debug(f"Torch attention: slicing={shared.opts.attention_slicing}")
+
+    pipe.current_attn_name = shared.opts.cross_attention_optimization

modules/devices.py

@@ -458,29 +458,29 @@ def set_sdpa_params():
             log.warning(f'Torch attention: type="sdpa" {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)
-            torch.backends.cuda.enable_math_sdp('Math attention' in opts.sdp_options)
+            torch.backends.cuda.enable_flash_sdp('Flash' in opts.sdp_options)
+            torch.backends.cuda.enable_mem_efficient_sdp('Memory' in opts.sdp_options)
+            torch.backends.cuda.enable_math_sdp('Math' in opts.sdp_options)
             if hasattr(torch.backends.cuda, "allow_fp16_bf16_reduction_math_sdp"): # only valid for torch >= 2.5
                 torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
-            log.debug(f'Torch attention: type="sdpa" opts={opts.sdp_options}')
+            log.debug(f'Torch attention: type="sdpa" kernels={opts.sdp_options} overrides={opts.sdp_overrides}')
         except Exception as err:
             log.warning(f'Torch attention: type="sdpa" {err}')
 
         # Stack hijcaks in reverse order. This gives priority to the last added hijack.
         # If the last hijack is not compatible, it will use the one before it and so on.
 
-        if 'Dynamic attention' in opts.sdp_options:
+        if 'Dynamic attention' in opts.sdp_overrides:
             global sdpa_pre_dyanmic_atten # pylint: disable=global-statement
             sdpa_pre_dyanmic_atten = attention.set_dynamic_attention()
 
-        if 'Triton Flash attention' in opts.sdp_options:
-            attention.set_triton_flash_attention(backend)
+        if 'Triton Flash attention' in opts.sdp_overrides:
+            attention.set_triton_flash_attention()
 
-        if 'CK Flash attention' in opts.sdp_options:
+        if 'CK Flash attention' in opts.sdp_overrides:
             attention.set_ck_flash_attention(backend, device)
 
-        if 'Sage attention' in opts.sdp_options:
+        if 'Sage attention' in opts.sdp_overrides:
             attention.set_sage_attention(backend, device)
 
         from importlib.metadata import version

modules/processing_diffusers.py

@@ -5,7 +5,7 @@ import numpy as np
 import torch
 import torchvision.transforms.functional as TF
 from PIL import Image
-from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks
+from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks, attention
 from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled, get_job_name
 from modules.processing_args import set_pipeline_args
 from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed
@@ -497,7 +497,7 @@ def update_pipeline(sd_model, p: processing.StableDiffusionProcessing):
         orig_pipeline = sd_model # processed ONNX pipeline should not be replaced with original pipeline.
     if getattr(sd_model, "current_attn_name", None) != shared.opts.cross_attention_optimization:
         shared.log.info(f"Setting attention optimization: {shared.opts.cross_attention_optimization}")
-        sd_models.set_diffusers_attention(sd_model)
+        attention.set_diffusers_attention(sd_model)
     return sd_model
 
 

modules/sd_models.py

@@ -10,7 +10,7 @@ import diffusers.loaders.single_file_utils
 import torch
 import huggingface_hub as hf
 from installer import log
-from modules import timer, paths, shared, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_detect, model_quant, sd_hijack_te, sd_hijack_accelerate, sd_hijack_safetensors
+from modules import timer, paths, shared, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_detect, model_quant, sd_hijack_te, sd_hijack_accelerate, sd_hijack_safetensors, attention
 from modules.memstats import memory_stats
 from modules.modeldata import model_data
 from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closest_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import
@@ -130,7 +130,7 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True,
 
     clear_caches()
     set_vae_options(sd_model, vae, op, quiet)
-    set_diffusers_attention(sd_model, quiet)
+    attention.set_diffusers_attention(sd_model, quiet)
 
     if shared.opts.diffusers_fuse_projections and hasattr(sd_model, 'fuse_qkv_projections'):
         try:
@@ -1157,60 +1157,6 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     return pipe
 
 
-def set_diffusers_attention(pipe, quiet:bool=False):
-    import diffusers.models.attention_processor as p
-
-    def set_attn(pipe, attention, name:str=None, quiet:bool=False):
-        if attention is None:
-            return
-        # other models uses their own attention processor
-        if pipe.__class__.__name__.startswith("StableDiffusion") and getattr(pipe, "unet", None) is not None and hasattr(pipe.unet, "set_attn_processor"):
-            try:
-                pipe.unet.set_attn_processor(attention)
-            except Exception as e:
-                if 'Nunchaku' in pipe.unet.__class__.__name__:
-                    pass
-                else:
-                    shared.log.error(f"Attention: {name if name is not None else attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}")
-        elif not quiet:
-            shared.log.warning(f"Attention: {name if name is not None else attention.__class__.__name__} is not compatible with {pipe.__class__.__name__}")
-
-    # if hasattr(pipe, 'pipe'):
-    #    set_diffusers_attention(pipe.pipe)
-
-    if 'Control' in pipe.__class__.__name__ or 'Adapter' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")):
-        if shared.opts.cross_attention_optimization not in {"Scaled-Dot-Product", "Disabled"}:
-            shared.log.warning(f"Attention: {shared.opts.cross_attention_optimization} is not compatible with {pipe.__class__.__name__}")
-        else:
-            pipe.current_attn_name = shared.opts.cross_attention_optimization
-        return
-
-    shared.log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
-    if shared.opts.cross_attention_optimization == "Disabled":
-        pass # do nothing
-    elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers
-        set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product", quiet=True)
-    elif shared.opts.cross_attention_optimization == "xFormers":
-        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
-            pipe.enable_xformers_memory_efficient_attention()
-        else:
-            shared.log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
-    elif shared.opts.cross_attention_optimization == "Batch matrix-matrix":
-        set_attn(pipe, p.AttnProcessor(), name="Batch matrix-matrix")
-    elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM":
-        from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM
-        set_attn(pipe, DynamicAttnProcessorBMM(), name="Dynamic Attention BMM")
-
-    if shared.opts.attention_slicing != "Default" and hasattr(pipe, "enable_attention_slicing") and hasattr(pipe, "disable_attention_slicing"):
-        if shared.opts.attention_slicing:
-            pipe.enable_attention_slicing()
-        else:
-            pipe.disable_attention_slicing()
-        shared.log.debug(f"Attention: slicing={shared.opts.attention_slicing}")
-
-    pipe.current_attn_name = shared.opts.cross_attention_optimization
-
-
 def add_noise_pred_to_diffusers_callback(pipe):
     if not hasattr(pipe, "_callback_tensor_inputs"):
         return pipe

modules/shared.py

@@ -136,7 +136,7 @@ def list_samplers():
     return modules.sd_samplers.all_samplers
 
 
-startup_offload_mode, startup_offload_min_gpu, startup_offload_max_gpu, startup_cross_attention, startup_sdp_options, startup_sdp_choices, startup_offload_always, startup_offload_never = get_default_modes(cmd_opts=cmd_opts, mem_stat=mem_stat)
+startup_offload_mode, startup_offload_min_gpu, startup_offload_max_gpu, startup_cross_attention, startup_sdp_options, startup_sdp_choices, startup_sdp_override_options, startup_sdp_override_choices, startup_offload_always, startup_offload_never = get_default_modes(cmd_opts=cmd_opts, mem_stat=mem_stat)
 
 options_templates.update(options_section(('sd', "Model Loading"), {
     "sd_backend": OptionInfo('diffusers', "Execution backend", gr.Radio, {"choices": ['diffusers', 'original'], "visible": False }),
@@ -296,13 +296,13 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "diffusers_generator_device": OptionInfo("GPU", "Generator device", gr.Radio, {"choices": ["GPU", "CPU", "Unset"]}),
 
     "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()}),
-    "attention_": OptionInfo("<h2>Cross Attention</h2>", "", gr.HTML),
+    "cross_attention_optimization": OptionInfo(startup_cross_attention, "Attention method", gr.Radio, lambda: {"choices": shared_items.list_crossattention()}),
+    "sdp_options": OptionInfo(startup_sdp_options, "SDP kernels", gr.CheckboxGroup, {"choices": startup_sdp_choices}),
+    "sdp_overrides": OptionInfo(startup_sdp_override_options, "SDP overrides", gr.CheckboxGroup, {"choices": startup_sdp_override_choices}),
     "attention_slicing": OptionInfo('Default', "Attention slicing", gr.Radio, {"choices": ['Default', 'Enabled', 'Disabled']}),
-    "sdp_options": OptionInfo(startup_sdp_options, "SDP options", gr.CheckboxGroup, {"choices": startup_sdp_choices}),
     "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}),
-    "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}),
+    "dynamic_attention_slice_rate": OptionInfo(0.5, "Dynamic Attention slicing rate", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4), "step": 0.01}),
+    "dynamic_attention_trigger_rate": OptionInfo(1, "Dynamic Attention trigger rate", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4)*2, "step": 0.01}),
 }))
 
 options_templates.update(options_section(('backends', "Backend Settings"), {

modules/shared_defaults.py

@@ -40,16 +40,17 @@ def get_default_modes(cmd_opts, mem_stat):
 
     default_cross_attention = "Scaled-Dot-Product"
 
+    default_sdp_override_options = []
     if devices.backend == "zluda":
         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']
+        default_sdp_options = ['Flash', 'Memory', 'Math']
+        default_sdp_override_options = ['Dynamic attention']
     else:
-        default_sdp_options = ['Flash attention', 'Memory attention', 'Math attention']
+        default_sdp_options = ['Flash', 'Memory', 'Math']
 
-    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
+    default_sdp_choices = ['Flash', 'Memory', 'Math']
+    default_sdp_override_choices = ['Dynamic attention', 'CK Flash attention', 'Triton Flash attention', 'Sage attention']
 
     return (
         default_offload_mode,
@@ -58,6 +59,8 @@ def get_default_modes(cmd_opts, mem_stat):
         default_cross_attention,
         default_sdp_options,
         default_sdp_choices,
+        default_sdp_override_options,
+        default_sdp_override_choices,
         default_diffusers_offload_always,
-        default_diffusers_offload_never
+        default_diffusers_offload_never,
     )