Merge pull request #5052 from vladmandic/feat/attention-core

Feat/attention core

modules/attention.py

@@ -1,355 +0,0 @@
-from functools import wraps
-import torch
-from modules import rocm, errors, devices
-from modules.logger import log
-from installer import install, installed, torch_info
-
-
-def set_dynamic_attention():
-    try:
-        sdpa_pre_dyanmic_atten = torch.nn.functional.scaled_dot_product_attention
-        from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
-        torch.nn.functional.scaled_dot_product_attention = dynamic_scaled_dot_product_attention
-        torch_info.set(attention='dynamic')
-        return sdpa_pre_dyanmic_atten
-    except Exception as err:
-        log.error(f'Torch attention: type="dynamic attention" {err}')
-        return None
-
-
-def set_sdnq_attention():
-    try:
-        from modules import shared
-        from sdnq.kernels.triton_atten import sdnq_triton_atten
-        sdpa_pre_sdnq_atten = torch.nn.functional.scaled_dot_product_attention
-        @wraps(sdpa_pre_sdnq_atten)
-        def sdpa_sdnq_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
-            if (
-                query.device.type != "cpu"
-                and (query.shape[-2] >= 32 and key.shape[-2] >= 32)
-                and (query.shape[-2] > 512 or key.shape[-2] > 512) # Skip TE
-                and query.shape[-3] > 1 # Skip VAE
-            ):
-                return sdnq_triton_atten(
-                    query=query, key=key, value=value, attn_mask=attn_mask,
-                    is_causal=is_causal, scale=scale, enable_gqa=enable_gqa,
-                    matmul_dtype=shared.opts.sdnq_attention_matmul_type,
-                    pv_matmul_dtype=shared.opts.sdnq_attention_pv_matmul_type,
-                    smooth_k=shared.opts.sdnq_attention_smooth_k,
-                    use_hadamard=shared.opts.sdnq_attention_use_hadamard,
-                    hadamard_group_size=shared.opts.sdnq_attention_hadamard_group_size,
-                    quantize_fp32=shared.opts.sdnq_attention_quantize_fp32,
-                    use_fp16_accum=shared.opts.sdnq_attention_use_fp16_accum,
-                )
-            else:
-                if enable_gqa:
-                    kwargs["enable_gqa"] = enable_gqa
-                return sdpa_pre_sdnq_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_sdnq_atten
-        torch_info.set(attention='sdnq')
-        log.debug(f'Torch attention: type="SDNQ attention" matmul={shared.opts.sdnq_attention_matmul_type}:{shared.opts.sdnq_attention_pv_matmul_type} smooth={shared.opts.sdnq_attention_smooth_k} hadamard={shared.opts.sdnq_attention_use_hadamard} fp16_accum={shared.opts.sdnq_attention_use_fp16_accum}')
-    except Exception as err:
-        log.error(f'Torch attention: type="SDNQ attention" {err}')
-
-
-def set_triton_flash_attention(backend: str):
-    try:
-        if backend in {"rocm", "zluda"}: # flash_attn_triton_amd only works with AMD
-            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: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
-                use_triton = (
-                    query.shape[-1] <= 128
-                    and attn_mask is None
-                    and query.device.type != "cpu"
-                    and key.device == query.device
-                    and value.device == query.device
-                )
-                if use_triton:
-                    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
-            torch_info.set(attention='triton')
-            log.debug('Torch attention: type="Triton Flash attention"')
-    except Exception as err:
-        log.error(f'Torch attention: type="Triton Flash attention" {err}')
-
-
-def set_flex_attention():
-    try:
-        from torch.nn.attention.flex_attention import flex_attention, create_block_mask
-        def flex_attention_causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument
-            return q_idx >= kv_idx
-
-        sdpa_pre_flex_atten = torch.nn.functional.scaled_dot_product_attention
-        @wraps(sdpa_pre_flex_atten)
-        def sdpa_flex_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor: # pylint: disable=unused-argument
-            score_mod = None
-            block_mask = None
-            if attn_mask is not None:
-                batch_size, num_heads = query.shape[:2]
-                seq_len_q = query.shape[-2]
-                seq_len_kv = key.shape[-2]
-                if attn_mask.ndim == 2:
-                    attn_mask = attn_mask.view(attn_mask.shape[0], 1, attn_mask.size[1], 1)
-                attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv)
-                if attn_mask.dtype == torch.bool:
-                    def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
-                        return attn_mask[batch_idx, head_idx, q_idx, kv_idx]
-                    block_mask = create_block_mask(mask_mod, batch_size, None, seq_len_q, seq_len_kv, device=query.device)
-                else:
-                    def score_mod_fn(score, batch_idx, head_idx, q_idx, kv_idx):
-                        return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx]
-                    score_mod = score_mod_fn
-            elif is_causal:
-                block_mask = create_block_mask(flex_attention_causal_mask, query.shape[0], query.shape[1], query.shape[-2], key.shape[-2], device=query.device)
-            return flex_attention(query, key, value, score_mod=score_mod, block_mask=block_mask, scale=scale, enable_gqa=enable_gqa)
-
-        torch.nn.functional.scaled_dot_product_attention = sdpa_flex_atten
-        torch_info.set(attention="flex")
-        log.debug('Torch attention: type="Flex attention"')
-    except Exception as err:
-        log.error(f'Torch attention: type="Flex attention" {err}')
-
-
-def set_ck_flash_attention(backend: str, device: torch.device):
-    try:
-        if backend == "rocm":
-            if not installed('flash-attn'):
-                log.info('Torch attention: type="Flash attention" building...')
-                agent = rocm.Agent(device)
-                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: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
-            use_flash = (
-                query.shape[-1] <= 128
-                and attn_mask is None
-                and query.dtype != torch.float32
-                and query.device.type != "cpu"
-                and key.device == query.device
-                and value.device == query.device
-            )
-            if use_flash:
-                is_unsqueezed = False
-                if query.dim() == 3:
-                    query = query.unsqueeze(0)
-                    is_unsqueezed = True
-                    if key.dim() == 3:
-                        key = key.unsqueeze(0)
-                    if value.dim() == 3:
-                        value = value.unsqueeze(0)
-                if enable_gqa:
-                    key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
-                    value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
-                query = query.transpose(1, 2)
-                key = key.transpose(1, 2)
-                value = value.transpose(1, 2)
-                attn_output = flash_attn_func(q=query, k=key, v=value, dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
-                if is_unsqueezed:
-                    attn_output = attn_output.squeeze(0)
-                return attn_output
-            else:
-                if enable_gqa:
-                    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
-        torch_info.set(attention="flash")
-        log.debug('Torch attention: type="Flash attention"')
-    except Exception as err:
-        log.error(f'Torch attention: type="Flash attention" {err}')
-
-
-def set_sage_attention(backend: str, device: torch.device):
-    try:
-        install('sageattention')
-
-        use_cuda_backend = False
-        if (backend == "cuda") and (torch.cuda.get_device_capability(device) == (8, 6)):
-            use_cuda_backend = True # Detect GPU architecture - sm86 confirmed to need CUDA backend workaround as Sage Attention + Triton causes NaNs
-            try:
-                from sageattention import sageattn_qk_int8_pv_fp16_cuda
-            except Exception:
-                use_cuda_backend = False
-
-        if use_cuda_backend:
-            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(
-                    q=query, k=key, v=value,
-                    tensor_layout="HND",
-                    is_causal=is_causal,
-                    sm_scale=scale,
-                    return_lse=False,
-                    pv_accum_dtype="fp32",
-                )
-        else:
-            from sageattention import sageattn
-            def sage_attn_impl(query, key, value, is_causal, scale):
-                return sageattn(
-                    q=query, k=key, v=value,
-                    attn_mask=None,
-                    dropout_p=0.0,
-                    is_causal=is_causal,
-                    scale=scale,
-                )
-
-        sdpa_pre_sage_atten = torch.nn.functional.scaled_dot_product_attention
-        @wraps(sdpa_pre_sage_atten)
-        def sdpa_sage_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
-            use_sage = (
-                query.shape[-1] in {128, 96, 64}
-                and attn_mask is None
-                and query.device.type != "cpu"
-                and key.device == query.device
-                and value.device == query.device
-            )
-            if use_sage:
-                if enable_gqa:
-                    key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
-                    value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
-
-                # Call preselected sage attention implementation
-                return sage_attn_impl(query, key, value, is_causal, scale)
-            else:
-                if enable_gqa:
-                    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
-        torch_info.set(attention="sage")
-        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}')
-
-
-def set_diffusers_attention(pipe, quiet = False):
-    from modules import shared
-    import diffusers.models.attention_processor as p
-
-    def set_attn(pipe, attention, name: str | None = 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:
-                    log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}')
-
-    log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
-    if shared.opts.cross_attention_optimization == "Disabled":
-        torch_info.set(attention="disabled")
-    elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product":  # The default set by Diffusers
-        devices.set_sdpa_params()
-        # set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product")
-    elif shared.opts.cross_attention_optimization == "xFormers":
-        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
-            torch_info.set(attention="xformers")
-            pipe.enable_xformers_memory_efficient_attention()
-        else:
-            log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
-    elif shared.opts.cross_attention_optimization == "Batch matrix-matrix":
-        torch_info.set(attention="bmm")
-        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
-        torch_info.set(attention="dynamic_bmm")
-        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()
-        log.debug(f"Torch attention: slicing={shared.opts.attention_slicing}")
-
-    pipe.current_attn_name = shared.opts.cross_attention_optimization
-
-
-orig_get_kernel = None
-def get_kernel_hijack(repo_id, revision=None, version=None, backend=None, user_agent=None, trust_remote_code: bool | list[str] = False): # pylint: disable=unused-argument
-    log.debug(f'Attention dispatcher hub: repo="{repo_id}" revision={revision} version={version} backend={backend}')
-    user_agent = 'kernels/0.16.0'
-    module = None
-    try:
-        module = orig_get_kernel(repo_id, revision=revision, version=version, backend=backend, user_agent=user_agent, trust_remote_code=True)
-    except Exception as e:
-        log.error(f'Attention dispatcher hub: {e}')
-        errors.display(e, 'kernels')
-    return module
-
-
-def get_hf_api_hijack(user_agent = None): # pylint: disable=unused-argument
-    from huggingface_hub import HfApi
-    return HfApi(library_name="kernels", user_agent="donottrack")
-
-
-def hijack_kernels():
-    global orig_get_kernel # pylint: disable=global-statement
-    try:
-        install('kernels==0.16.0')
-        import kernels
-        import kernels.utils
-        log.debug(f'Attention dispatcher: kernels={kernels.__version__}')
-        if orig_get_kernel is None:
-            orig_get_kernel = kernels.get_kernel
-            kernels.get_kernel = get_kernel_hijack
-            kernels.utils._get_hf_api = get_hf_api_hijack # pylint: disable=protected-access
-        from diffusers.utils import import_utils
-        import_utils._kernels_available = True # pylint: disable=protected-access
-        import_utils._kernels_version = kernels.__version__ # pylint: disable=protected-access
-    except Exception as e:
-        log.error(f'Attention dispatcher kernels: {e}')
-        return
-
-
-def set_attention_dispatcher(pipe):
-    from modules import shared
-    attn = shared.opts.hf_attention.strip().lower()
-    if pipe is None or not hasattr(pipe, 'transformer') or not hasattr(pipe.transformer, 'set_attention_backend'):
-        return
-
-    from diffusers.models import attention_dispatch as a
-    backends = [b.value for b in a._AttentionBackendRegistry.list_backends()] # pylint: disable=protected-access
-    # https://huggingface.co/docs/kernels/index
-    # https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends
-
-    if 'hub' in attn:
-        hijack_kernels()
-
-    prev = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
-    if attn in backends:
-        try:
-            pipe.transformer.set_attention_backend(attn)
-        except Exception as e:
-            log.error(f'Attention dispatcher: target={attn} {e}')
-        current = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
-        log.debug(f'Attention dispatcher: target={attn} previous={prev[0].value} active={current[0]} list={backends}')
-    elif len(attn) > 0:
-        log.warning(f'Attention dispatcher: active={prev[0].value} list={backends} target={attn} not found')
-    else:
-        log.debug(f'Attention dispatcher: active={prev[0].value} list={backends}')

modules/attention/__init__.py

@@ -0,0 +1,12 @@
+"""Attention backends: one scaled_dot_product_attention router over the registered backends, the per-generation context, and the diffusers-side processor and dispatcher setup."""
+from modules.attention.registry import AttentionBackend, AttentionCall, Constraints, Platform, Registry, registry
+from modules.attention.router import Plan, PlanEntry, build_plan, get_plan, install_router, reapply, reapply_options, report
+from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, hijack_kernels, get_kernel_hijack, get_hf_api_hijack
+from modules.attention import backends, context, debug
+
+__all__ = [
+    'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry',
+    'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'reapply', 'reapply_options', 'report',
+    'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack',
+    'backends', 'context', 'debug',
+]

modules/attention/backends/__init__.py

@@ -0,0 +1,10 @@
+"""Built-in backends, registered in ascending priority."""
+from modules.attention.registry import registry
+from modules.attention.backends import dynamic, flex, triton_amd, flash_ck, sage, sdnq
+
+registry.register(dynamic.backend)
+registry.register(flex.backend)
+registry.register(triton_amd.backend)
+registry.register(flash_ck.backend)
+registry.register(sage.backend)
+registry.register(sdnq.backend)

modules/attention/backends/dynamic.py

@@ -0,0 +1,11 @@
+from modules.attention.registry import AttentionBackend, Platform
+
+
+def prepare(platform: Platform, original): # pylint: disable=unused-argument
+    from modules import devices
+    devices.sdpa_pre_dyanmic_atten = original # the sliced path calls this pin for every slice
+    from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
+    return dynamic_scaled_dot_product_attention
+
+
+backend = AttentionBackend(name='dynamic', label='Dynamic attention', priority=10, prepare=prepare, terminal=True)

modules/attention/backends/flash_ck.py

@@ -0,0 +1,44 @@
+from installer import install, installed
+from modules import rocm
+from modules.logger import log
+from modules.attention.registry import AttentionBackend, Constraints, Platform
+
+
+def prepare(platform: Platform, original): # pylint: disable=unused-argument
+    if platform.backend == 'rocm':
+        if not installed('flash-attn'):
+            log.info('Torch attention: type="Flash attention" building...')
+            agent = rocm.Agent(platform.device)
+            install(rocm.get_flash_attention_command(agent), reinstall=True)
+    else:
+        install('flash-attn')
+    from flash_attn import flash_attn_func
+
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+        is_unsqueezed = False
+        if query.dim() == 3:
+            query = query.unsqueeze(0)
+            is_unsqueezed = True
+            if key.dim() == 3:
+                key = key.unsqueeze(0)
+            if value.dim() == 3:
+                value = value.unsqueeze(0)
+        if enable_gqa:
+            key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
+            value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+        attn_output = flash_attn_func(q=query, k=key, v=value, dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
+        if is_unsqueezed:
+            attn_output = attn_output.squeeze(0)
+        return attn_output
+
+    log.debug('Torch attention: type="Flash attention"')
+    return call
+
+
+backend = AttentionBackend(
+    name='flash', label='Flash attention', priority=40, prepare=prepare,
+    constraints=Constraints(max_head_dim=128, allow_mask=False, allow_float32=False, same_device=True),
+)

modules/attention/backends/flex.py

@@ -0,0 +1,39 @@
+import torch
+from modules.logger import log
+from modules.attention.registry import AttentionBackend, Constraints, Platform
+
+
+def prepare(platform: Platform, original): # pylint: disable=unused-argument
+    from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+
+    def causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument
+        return q_idx >= kv_idx
+
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+        score_mod = None
+        block_mask = None
+        if attn_mask is not None:
+            batch_size, num_heads = query.shape[:2]
+            seq_len_q = query.shape[-2]
+            seq_len_kv = key.shape[-2]
+            attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv) # sdpa masks broadcast over the trailing dims
+            if attn_mask.dtype == torch.bool:
+                def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
+                    return attn_mask[batch_idx, head_idx, q_idx, kv_idx]
+                block_mask = create_block_mask(mask_mod, batch_size, None, seq_len_q, seq_len_kv, device=query.device)
+            else:
+                def score_mod_fn(score, batch_idx, head_idx, q_idx, kv_idx):
+                    return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx]
+                score_mod = score_mod_fn
+        elif is_causal:
+            block_mask = create_block_mask(causal_mask, query.shape[0], query.shape[1], query.shape[-2], key.shape[-2], device=query.device)
+        return flex_attention(query, key, value, score_mod=score_mod, block_mask=block_mask, scale=scale, enable_gqa=enable_gqa)
+
+    log.debug('Torch attention: type="Flex attention"')
+    return call
+
+
+backend = AttentionBackend(
+    name='flex', label='Flex attention', priority=20, prepare=prepare,
+    constraints=Constraints(min_ndim=4, same_device=True), # flex_attention takes 4d tensors on one device and compiles on cpu
+)

modules/attention/backends/sage.py

@@ -0,0 +1,53 @@
+import torch
+from installer import install
+from modules.logger import log
+from modules.attention.registry import AttentionBackend, Constraints, Platform
+
+
+def prepare(platform: Platform, original): # pylint: disable=unused-argument
+    install('sageattention')
+
+    use_cuda_backend = False
+    if platform.backend == 'cuda' and torch.cuda.get_device_capability(platform.device) == (8, 6):
+        use_cuda_backend = True # sm86 needs the cuda backend, sage attention over triton produces NaNs there
+        try:
+            from sageattention import sageattn_qk_int8_pv_fp16_cuda
+        except Exception:
+            use_cuda_backend = False
+
+    if use_cuda_backend:
+        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(
+                q=query, k=key, v=value,
+                tensor_layout="HND",
+                is_causal=is_causal,
+                sm_scale=scale,
+                return_lse=False,
+                pv_accum_dtype="fp32",
+            )
+    else:
+        from sageattention import sageattn
+        def sage_attn_impl(query, key, value, is_causal, scale):
+            return sageattn(
+                q=query, k=key, v=value,
+                attn_mask=None,
+                dropout_p=0.0,
+                is_causal=is_causal,
+                scale=scale,
+            )
+
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+        if enable_gqa:
+            key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
+            value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
+        return sage_attn_impl(query, key, value, is_causal, scale)
+
+    log.debug(f'Torch attention: type="Sage attention" backend={"cuda" if use_cuda_backend else "auto"}')
+    return call
+
+
+backend = AttentionBackend(
+    name='sage', label='Sage attention', priority=50, prepare=prepare,
+    constraints=Constraints(head_dims=frozenset({64, 96, 128}), allow_mask=False, same_device=True),
+)

modules/attention/backends/sdnq.py

@@ -0,0 +1,29 @@
+from modules.logger import log
+from modules.attention.registry import AttentionBackend, Constraints, Platform
+
+
+def prepare(platform: Platform, original): # pylint: disable=unused-argument
+    from modules import shared
+    from sdnq.kernels.triton_atten import sdnq_triton_atten
+    options = {
+        'matmul_dtype': shared.opts.sdnq_attention_matmul_type,
+        'pv_matmul_dtype': shared.opts.sdnq_attention_pv_matmul_type,
+        'smooth_k': shared.opts.sdnq_attention_smooth_k,
+        'use_hadamard': shared.opts.sdnq_attention_use_hadamard,
+        'hadamard_group_size': shared.opts.sdnq_attention_hadamard_group_size,
+        'quantize_fp32': shared.opts.sdnq_attention_quantize_fp32,
+        'use_fp16_accum': shared.opts.sdnq_attention_use_fp16_accum,
+    }
+
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+        return sdnq_triton_atten(query=query, key=key, value=value, attn_mask=attn_mask, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **options)
+
+    log.debug(f'Torch attention: type="SDNQ attention" matmul={options["matmul_dtype"]}:{options["pv_matmul_dtype"]} smooth={options["smooth_k"]} hadamard={options["use_hadamard"]} quantize_fp32={options["quantize_fp32"]} fp16_accum={options["use_fp16_accum"]}')
+    return call
+
+
+backend = AttentionBackend(
+    name='sdnq', label='SDNQ attention', priority=60, prepare=prepare,
+    constraints=Constraints(min_tokens=32, min_long_side=512, min_heads=2), # sequences of 512 or fewer are text encoders, single-head calls the vae
+    options=('sdnq_attention_matmul_type', 'sdnq_attention_pv_matmul_type', 'sdnq_attention_smooth_k', 'sdnq_attention_use_hadamard', 'sdnq_attention_hadamard_group_size', 'sdnq_attention_quantize_fp32', 'sdnq_attention_use_fp16_accum'),
+)

modules/attention/backends/triton_amd.py

@@ -0,0 +1,32 @@
+import torch
+from modules.logger import log
+from modules.attention.registry import AttentionBackend, Constraints, Platform
+
+
+def prepare(platform: Platform, original): # pylint: disable=unused-argument
+    from modules.flash_attn_triton_amd import interface_fa
+
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+        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)
+
+    log.debug('Torch attention: type="Triton Flash attention"')
+    return call
+
+
+backend = AttentionBackend(
+    name='triton', label='Triton Flash attention', priority=30, prepare=prepare,
+    constraints=Constraints(max_head_dim=128, allow_mask=False, same_device=True),
+    platforms=frozenset({'rocm', 'zluda'}),
+)

modules/attention/context.py

@@ -0,0 +1,77 @@
+"""Per-generation state for attention consumers: the component running, the denoiser forward about to run, and the model."""
+from contextlib import contextmanager
+from dataclasses import dataclass
+import torch
+
+
+@dataclass
+class GenerationContext:
+    active: bool = False
+    role: str | None = None # 'transformer', 'te' or 'vae' while a generation runs, None outside one
+    step: int = 0 # index of the denoiser forward about to run
+    steps: int = 0 # forwards in the current pass
+    forwards: int = 0
+    model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, telemetry only
+    step_buffer: torch.Tensor | None = None # the step as a device scalar updated in place, so compiled readers keep their graph
+
+
+current = GenerationContext()
+
+
+def denoiser_name(pipe) -> str | None:
+    for name in ('transformer', 'unet'):
+        module = getattr(pipe, name, None)
+        if module is not None:
+            return module.__class__.__name__
+    return None
+
+
+def begin(pipe, steps: int = 0) -> None:
+    from modules import devices
+    current.active = True
+    current.role = 'transformer'
+    current.model_key = (pipe.__class__.__name__, denoiser_name(pipe)) if pipe is not None else None
+    device = devices.device if devices.device is not None else torch.device('cpu')
+    if current.step_buffer is None or current.step_buffer.device != device:
+        current.step_buffer = torch.zeros((), dtype=torch.int64, device=device)
+    new_pass(steps)
+
+
+def new_pass(steps: int = 0) -> None:
+    """Restart the step count for a denoising pass: base, hires or refiner."""
+    current.steps = int(steps or 0)
+    current.forwards = 0
+    set_step(0)
+
+
+def set_step(step: int) -> None:
+    current.step = int(step)
+    if current.step_buffer is not None:
+        current.step_buffer.fill_(current.step)
+
+
+def tick(step: int | None = None) -> None:
+    """Advance to the next forward: the classic callback passes the completed step plus one, the modular pre-hook passes nothing and counts forwards."""
+    set_step(current.forwards if step is None else step)
+    current.forwards = current.step + 1
+
+
+def end() -> None:
+    current.active = False
+    current.role = None
+    current.model_key = None
+    new_pass(0)
+
+
+def set_role(name: str | None) -> None:
+    current.role = name
+
+
+@contextmanager
+def role(name: str):
+    previous = current.role
+    current.role = name
+    try:
+        yield
+    finally:
+        current.role = previous

modules/attention/debug.py

@@ -0,0 +1,21 @@
+"""Opt-in route tracing for the sdpa router, enabled by SD_ATTN_DEBUG."""
+import os
+import torch
+from modules.logger import log
+from modules.attention import context
+
+enabled = os.environ.get('SD_ATTN_DEBUG', None) is not None
+seen: set[tuple] = set()
+
+
+def observe(name: str, query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None) -> None:
+    """Log each distinct route once: backend, component role, step, shapes, dtype and mask presence."""
+    signature = (name, context.current.role, tuple(query.shape), tuple(key.shape), str(query.dtype), attn_mask is not None)
+    if signature in seen:
+        return
+    seen.add(signature)
+    log.debug(f'Attention route: backend={name} role={context.current.role} step={context.current.step} q={list(query.shape)} k={list(key.shape)} dtype={query.dtype} mask={attn_mask is not None}')
+
+
+def reset() -> None:
+    seen.clear()

modules/attention/dispatcher.py

@@ -0,0 +1,96 @@
+from modules import errors
+from modules.logger import log
+from installer import install, torch_info
+
+
+def set_diffusers_attention(pipe, quiet = False):
+    from modules import shared, devices
+
+    log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
+    if shared.opts.cross_attention_optimization == "Disabled":
+        torch_info.set(attention="disabled")
+    elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product":  # The default set by Diffusers
+        devices.set_sdpa_params()
+    elif shared.opts.cross_attention_optimization == "xFormers":
+        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
+            torch_info.set(attention="xformers")
+            pipe.enable_xformers_memory_efficient_attention()
+        else:
+            log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
+    else:
+        log.warning(f'Torch attention: method="{shared.opts.cross_attention_optimization}" unknown, pipe={pipe.__class__.__name__} keeps its own attention processor')
+
+    if shared.opts.attention_slicing != "Default" and hasattr(pipe, "enable_attention_slicing") and hasattr(pipe, "disable_attention_slicing"):
+        if shared.opts.attention_slicing == "Enabled":
+            pipe.enable_attention_slicing()
+        else:
+            pipe.disable_attention_slicing()
+        log.debug(f"Torch attention: slicing={shared.opts.attention_slicing}")
+
+    pipe.current_attn_name = shared.opts.cross_attention_optimization
+    pipe.current_attn_overrides = list(shared.opts.sdp_overrides)
+
+
+orig_get_kernel = None
+def get_kernel_hijack(repo_id, revision=None, version=None, backend=None, user_agent=None, trust_remote_code: bool | list[str] = False): # pylint: disable=unused-argument
+    log.debug(f'Attention dispatcher hub: repo="{repo_id}" revision={revision} version={version} backend={backend}')
+    user_agent = 'kernels/0.16.0'
+    module = None
+    try:
+        module = orig_get_kernel(repo_id, revision=revision, version=version, backend=backend, user_agent=user_agent, trust_remote_code=True)
+    except Exception as e:
+        log.error(f'Attention dispatcher hub: {e}')
+        errors.display(e, 'kernels')
+    return module
+
+
+def get_hf_api_hijack(user_agent = None): # pylint: disable=unused-argument
+    from huggingface_hub import HfApi
+    return HfApi(library_name="kernels", user_agent="donottrack")
+
+
+def hijack_kernels():
+    global orig_get_kernel # pylint: disable=global-statement
+    try:
+        install('kernels==0.16.0')
+        import kernels
+        import kernels.utils
+        log.debug(f'Attention dispatcher: kernels={kernels.__version__}')
+        if orig_get_kernel is None:
+            orig_get_kernel = kernels.get_kernel
+            kernels.get_kernel = get_kernel_hijack
+            kernels.utils._get_hf_api = get_hf_api_hijack # pylint: disable=protected-access
+        from diffusers.utils import import_utils
+        import_utils._kernels_available = True # pylint: disable=protected-access
+        import_utils._kernels_version = kernels.__version__ # pylint: disable=protected-access
+    except Exception as e:
+        log.error(f'Attention dispatcher kernels: {e}')
+        return
+
+
+def set_attention_dispatcher(pipe):
+    from modules import shared
+    attn = shared.opts.hf_attention.strip().lower()
+    if pipe is None or not hasattr(pipe, 'transformer') or not hasattr(pipe.transformer, 'set_attention_backend'):
+        return
+
+    from diffusers.models import attention_dispatch as a
+    backends = [b.value for b in a._AttentionBackendRegistry.list_backends()] # pylint: disable=protected-access
+    # https://huggingface.co/docs/kernels/index
+    # https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends
+
+    if 'hub' in attn:
+        hijack_kernels()
+
+    prev = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
+    if attn in backends:
+        try:
+            pipe.transformer.set_attention_backend(attn)
+        except Exception as e:
+            log.error(f'Attention dispatcher: target={attn} {e}')
+        current = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
+        log.debug(f'Attention dispatcher: target={attn} previous={prev[0].value} active={current[0]} list={backends}')
+    elif len(attn) > 0:
+        log.warning(f'Attention dispatcher: active={prev[0].value} list={backends} target={attn} not found')
+    else:
+        log.debug(f'Attention dispatcher: active={prev[0].value} list={backends}')

modules/attention/registry.py

@@ -0,0 +1,98 @@
+"""Declarative backend registry behind the scaled_dot_product_attention router."""
+from dataclasses import dataclass, field
+from typing import Callable
+import torch
+
+
+AttentionCall = Callable[..., torch.Tensor]
+
+
+@dataclass(frozen=True)
+class Platform:
+    """Where the router runs: the devices backend name and the selected device."""
+    backend: str
+    device: torch.device | None = None
+
+
+@dataclass(frozen=True)
+class Constraints:
+    """Shape, dtype and device conditions a backend serves; a call failing any of them moves on to the next entry."""
+    allow_cpu: bool = False
+    allow_mask: bool = True
+    allow_float32: bool = True
+    same_device: bool = False
+    head_dims: frozenset[int] | None = None
+    max_head_dim: int | None = None
+    min_tokens: int = 0 # query and key sequences both at least this long
+    min_long_side: int = 0 # query or key sequence longer than this
+    min_heads: int = 0
+    min_ndim: int = 0
+
+    def accepts(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: torch.Tensor | None) -> bool:
+        if not self.allow_cpu and query.device.type == 'cpu':
+            return False
+        if self.min_ndim and query.ndim < self.min_ndim:
+            return False
+        if not self.allow_mask and attn_mask is not None:
+            return False
+        if not self.allow_float32 and query.dtype == torch.float32:
+            return False
+        if self.same_device and (key.device != query.device or value.device != query.device):
+            return False
+        head_dim = query.shape[-1]
+        if self.head_dims is not None and head_dim not in self.head_dims:
+            return False
+        if self.max_head_dim is not None and head_dim > self.max_head_dim:
+            return False
+        if self.min_tokens and (query.shape[-2] < self.min_tokens or key.shape[-2] < self.min_tokens):
+            return False
+        if self.min_long_side and query.shape[-2] <= self.min_long_side and key.shape[-2] <= self.min_long_side:
+            return False
+        if self.min_heads and query.shape[-3] < self.min_heads:
+            return False
+        return True
+
+
+@dataclass(frozen=True)
+class AttentionBackend:
+    """One attention implementation: how to prepare it once and which calls it serves."""
+    name: str
+    label: str # the sdp_overrides choice that enables it
+    priority: int # higher priority entries are tried first
+    prepare: Callable[[Platform, AttentionCall], AttentionCall | None] # imports and configures the implementation, returns its call or None
+    constraints: Constraints = field(default_factory=Constraints)
+    terminal: bool = False # serves every call the entries decline, in place of the original sdpa
+    platforms: frozenset[str] | None = None # devices backends the implementation exists for, None for all
+    options: tuple[str, ...] = () # settings the prepared call captures; a change to one rebuilds the chain
+
+    def available_on(self, platform: Platform) -> bool:
+        return self.platforms is None or platform.backend in self.platforms
+
+
+class Registry:
+    def __init__(self):
+        self.backends: dict[str, AttentionBackend] = {}
+
+    def register(self, backend: AttentionBackend) -> AttentionBackend:
+        if backend.name in self.backends:
+            raise ValueError(f'attention backend registered twice: name={backend.name}')
+        if self.by_label(backend.label) is not None:
+            raise ValueError(f'attention backend label registered twice: label="{backend.label}"')
+        self.backends[backend.name] = backend
+        return backend
+
+    def by_label(self, label: str) -> AttentionBackend | None:
+        return next((backend for backend in self.backends.values() if backend.label == label), None)
+
+    def ordered(self) -> list[AttentionBackend]:
+        """Backends by ascending priority, the order they are prepared in."""
+        return sorted(self.backends.values(), key=lambda backend: backend.priority)
+
+    def labels(self) -> list[str]:
+        return [backend.label for backend in self.ordered()]
+
+    def options(self) -> list[str]:
+        return sorted({name for backend in self.backends.values() for name in backend.options})
+
+
+registry = Registry()

modules/attention/router.py

@@ -0,0 +1,128 @@
+"""The single scaled_dot_product_attention entry point over the prepared backends."""
+from dataclasses import dataclass
+from functools import wraps
+from typing import Callable
+import torch
+from installer import torch_info
+from modules.logger import log
+from modules.attention import context, debug
+from modules.attention.registry import AttentionBackend, AttentionCall, Platform, Registry, registry as default_registry
+
+
+@dataclass(frozen=True)
+class PlanEntry:
+    backend: AttentionBackend
+    call: AttentionCall
+
+
+@dataclass(frozen=True)
+class Plan:
+    """The prepared chain for one set of overrides: entries by descending priority, then the terminal or the original sdpa."""
+    entries: tuple[PlanEntry, ...]
+    terminal: PlanEntry | None
+    original: AttentionCall
+    platform: Platform
+    labels: tuple[str, ...]
+
+    def chain(self) -> list[str]:
+        names = [entry.backend.name for entry in self.entries]
+        names.append(self.terminal.backend.name if self.terminal is not None else 'sdpa')
+        return names
+
+
+current_plan: Plan | None = None
+
+
+def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registry | None = None) -> Plan:
+    reg = reg if reg is not None else default_registry
+    entries: list[PlanEntry] = []
+    terminal: PlanEntry | None = None
+    for backend in reg.ordered(): # ascending priority: the last prepared backend is tried first
+        if backend.label not in labels:
+            continue
+        if not backend.available_on(platform):
+            log.warning(f'Torch attention: type="{backend.label}" not available on backend={platform.backend}')
+            continue
+        try:
+            call = backend.prepare(platform, original)
+        except Exception as err:
+            log.error(f'Torch attention: type="{backend.label}" {err}')
+            continue
+        if call is None:
+            continue
+        entry = PlanEntry(backend=backend, call=call)
+        if backend.terminal:
+            terminal = entry
+        else:
+            entries.append(entry)
+    entries.reverse()
+    return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels))
+
+
+def make_router(plan: Plan, observer: Callable | None = None) -> AttentionCall:
+    entries = plan.entries
+    terminal = plan.terminal.call if plan.terminal is not None else None
+    terminal_name = plan.terminal.backend.name if plan.terminal is not None else 'sdpa'
+    original = plan.original
+
+    @wraps(original)
+    def sdpa_router(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs):
+        for entry in entries:
+            if entry.backend.constraints.accepts(query, key, value, attn_mask):
+                if observer is not None:
+                    observer(entry.backend.name, query, key, attn_mask)
+                return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa)
+        if observer is not None:
+            observer(terminal_name, query, key, attn_mask)
+        if terminal is not None:
+            return terminal(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **kwargs)
+        if enable_gqa: # older sdpa signatures and platform wrappers reject the keyword, so it only travels when set
+            kwargs['enable_gqa'] = enable_gqa
+        return original(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
+
+    return sdpa_router
+
+
+def install_router(labels, platform: Platform, original: AttentionCall, reg: Registry | None = None) -> Plan:
+    """Prepare the enabled backends and install the router; an empty plan leaves the original sdpa in place."""
+    global current_plan # pylint: disable=global-statement
+    plan = build_plan(labels, platform, original, reg)
+    debug.reset()
+    observer = debug.observe if debug.enabled else None
+    torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer) if (plan.entries or plan.terminal is not None) else original
+    current_plan = plan
+    torch_info.set(attention='>'.join(plan.chain()))
+    log.debug(f'Torch attention: chain={">".join(plan.chain())} overrides={list(labels)} backend={platform.backend}')
+    return plan
+
+
+def get_plan() -> Plan | None:
+    return current_plan
+
+
+def reapply_options(reg: Registry | None = None) -> list[str]:
+    """Settings whose change rebuilds the chain: the override set, the torch kernel flags, and every option a backend captures."""
+    reg = reg if reg is not None else default_registry
+    return ['sdp_options', 'sdp_overrides', *reg.options()]
+
+
+def reapply() -> None:
+    """Rebuild the chain from the current settings; a resident compiled model is reset so its graphs trace the new router."""
+    from modules import devices, shared
+    devices.set_sdpa_params()
+    compiled = getattr(shared, 'compiled_model_state', None)
+    if compiled is not None and getattr(compiled, 'is_compiled', False):
+        torch._dynamo.reset() # pylint: disable=protected-access
+        log.debug('Torch attention: dynamo reset, compiled model resident')
+
+
+def report() -> dict:
+    """The active chain and generation context, for the api and the debug log."""
+    plan = current_plan
+    state = context.current
+    return {
+        'chain': plan.chain() if plan is not None else ['sdpa'],
+        'overrides': list(plan.labels) if plan is not None else [],
+        'backend': plan.platform.backend if plan is not None else None,
+        'context': {'active': state.active, 'role': state.role, 'step': state.step, 'steps': state.steps, 'model': state.model_key},
+    }

modules/devices.py

@@ -515,27 +515,7 @@ def set_sdpa_params():
         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_overrides:
-            global sdpa_pre_dyanmic_atten # pylint: disable=global-statement
-            sdpa_pre_dyanmic_atten = attention.set_dynamic_attention()
-
-        if 'Flex attention' in opts.sdp_overrides:
-            attention.set_flex_attention()
-
-        if 'Triton Flash attention' in opts.sdp_overrides:
-            attention.set_triton_flash_attention(backend)
-
-        if 'Flash attention' in opts.sdp_overrides:
-            attention.set_ck_flash_attention(backend, device)
-
-        if 'Sage attention' in opts.sdp_overrides:
-            attention.set_sage_attention(backend, device)
-
-        if 'SDNQ attention' in opts.sdp_overrides:
-            attention.set_sdnq_attention()
+        attention.install_router(opts.sdp_overrides, attention.Platform(backend=backend, device=device), sdpa_original)
 
         from importlib.metadata import version
         try:

modules/modular_load.py

@@ -3,6 +3,7 @@ import logging
 import torch
 from modules import shared, errors, devices, sd_offload
 from modules.logger import log
+from modules.attention import context as attention_context
 
 
 class InterruptLogFilter(logging.Filter):
@@ -49,6 +50,7 @@ def install_state_hook(pipe):
 
     def _pre_transformer_hook(module, args): # pylint: disable=unused-argument
         new_phase = set_phase('Generate', module)
+        attention_context.set_role('transformer')
         if new_phase:
             sd_offload.offload_ondemand(pipe, exclude=['transformer', 'transformer_ref'], reason='generate', force=hasattr(pipe, 'sdnext_force_offload'))
         if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0:
@@ -60,11 +62,13 @@ def install_state_hook(pipe):
                     raise AssertionError('Interrupted...')
                 time.sleep(0.1)
         shared.state.step()
+        attention_context.tick()
         if shared.state.interrupted or shared.state.skipped:
             raise AssertionError('Interrupted...')
 
     def _pre_text_encode_hook(module, args): # pylint: disable=unused-argument
         new_phase = set_phase('Text Encode', module)
+        attention_context.set_role('te')
         if new_phase:
             sd_offload.offload_ondemand(pipe, exclude=['text_encoder'], reason='text encode', force=hasattr(pipe, 'sdnext_force_offload'))
         if shared.state.interrupted or shared.state.skipped:
@@ -72,6 +76,7 @@ def install_state_hook(pipe):
 
     def _pre_vae_decode_hook(module, args): # pylint: disable=unused-argument
         new_phase = set_phase('Decode', module)
+        attention_context.set_role('vae')
         if new_phase:
             sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae decode', force=hasattr(pipe, 'sdnext_force_offload'))
         if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled decodes abort promptly
@@ -79,6 +84,7 @@ def install_state_hook(pipe):
 
     def _pre_vae_encode_hook(module, args): # pylint: disable=unused-argument
         new_phase = set_phase('Encode', module)
+        attention_context.set_role('vae')
         if new_phase:
             sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae encode', force=hasattr(pipe, 'sdnext_force_offload'))
         if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled encodes abort promptly

modules/options_handler.py

@@ -17,10 +17,23 @@ if TYPE_CHECKING:
     import builtins
 
 cmd_opts = cmd_args.parse_args()
-compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order']
+compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order', 'xformers_options']
+removed_values = { # a stored choice that no longer exists is kept by validate, so it has to be rewritten or it selects nothing
+    'cross_attention_optimization': (['Batch matrix-matrix', 'Dynamic Attention BMM'], 'Scaled-Dot-Product'),
+}
 secrets_pattern = ['_version', '_token', '_key', '_secret', '_password']
 
 
+def migrate_removed_values(data: dict) -> list:
+    """Rewrite stored settings whose choice was removed, returning what changed."""
+    migrated = []
+    for key, (removed, replacement) in removed_values.items():
+        if data.get(key, None) in removed:
+            migrated.append(f'{key}={data[key]} replaced={replacement}')
+            data[key] = replacement
+    return migrated
+
+
 class Options:
     data_labels: dict[str, OptionInfo | LegacyOption]
     data: dict[str, Any]
@@ -203,6 +216,9 @@ class Options:
         self.secrets = readfile(secretsfn, lock=True, as_type="dict")
         if self.data.get('quicksettings') is not None and self.data.get('quicksettings_list') is None:
             self.data['quicksettings_list'] = [i.strip() for i in self.data.get('quicksettings', '').split(',')]
+        migrated = migrate_removed_values(self.data)
+        if len(migrated) > 0:
+            log.warning(f"Setting migration: {migrated}")
         unknown_settings = []
         for k, v in self.data.items():
             info = self.data_labels.get(k, None)

modules/processing.py

@@ -5,6 +5,7 @@ import numpy as np
 from PIL import Image, ImageOps
 from modules import shared, devices, errors, images, scripts_manager, memstats, script_callbacks, extra_networks, sd_models, sd_checkpoint, sd_vae, processing_helpers, processing_grading, timer, masking
 from modules.logger import log
+from modules.attention import context as attention_context
 from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet
 from modules.processing_info import create_infotext
 from modules.processing_class import ( # pylint: disable=unused-import
@@ -198,6 +199,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed | None:
 
         script_callbacks.before_process_callback(p)
         timer.process.record('pre')
+        attention_context.begin(shared.sd_model, p.steps)
 
         if shared.cmd_opts.profile:
             timer.startup.profile = True
@@ -231,6 +233,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed | None:
                 results = process_images_inner(p)
 
     finally:
+        attention_context.end()
         script_callbacks.after_process_callback(p)
 
         if p.override_settings_restore_afterwards: # restore opts to original state

modules/processing_args.py

@@ -8,6 +8,7 @@ import numpy as np
 from PIL import Image
 from modules import shared, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, sd_vae
 from modules.logger import log
+from modules.attention import context as attention_context
 from modules.processing_callbacks import diffusers_callback_legacy, diffusers_callback, set_callbacks_p
 from modules.processing_helpers import get_generator, apply_circular # pylint: disable=unused-import
 from modules.processing_prompt import set_prompt
@@ -366,6 +367,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:l
         args['callback_steps'] = 1
 
     set_callbacks_p(p)
+    attention_context.new_pass(steps)
     if 'prior_callback_on_step_end' in possible: # Wuerstchen / Cascade
         args['prior_callback_on_step_end'] = diffusers_callback
         if 'prior_callback_on_step_end_tensor_inputs' in possible:

modules/processing_callbacks.py

@@ -4,6 +4,7 @@ import torch
 import numpy as np
 from modules import shared, devices, processing_correction, timer, prompt_parser_diffusers
 from modules.logger import log
+from modules.attention import context as attention_context
 
 
 p = None
@@ -87,6 +88,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
     if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0:
         shared.state.sampling_steps = pipe.num_timesteps
     shared.state.step()
+    attention_context.tick(step + 1)
     if shared.state.interrupted or shared.state.skipped:
         raise AssertionError('Interrupted...')
     if latents is None or p is None:

modules/processing_diffusers.py

@@ -546,8 +546,9 @@ def update_pipeline(sd_model, p: processing.StableDiffusionProcessing):
         global orig_pipeline # pylint: disable=global-statement
         orig_pipeline = updated_model # processed ONNX pipeline should not be replaced with original pipeline.
     current_attn = getattr(updated_model, "current_attn_name", None)
-    if (current_attn != shared.opts.cross_attention_optimization) and (current_attn != shared.opts.sdp_overrides):
-        log.info(f"Setting attention optimization: {shared.opts.cross_attention_optimization}")
+    current_overrides = getattr(updated_model, "current_attn_overrides", None)
+    if current_attn != shared.opts.cross_attention_optimization or current_overrides != list(shared.opts.sdp_overrides):
+        log.info(f"Setting attention optimization: {shared.opts.cross_attention_optimization} overrides={shared.opts.sdp_overrides}")
         attention.set_diffusers_attention(updated_model)
     return updated_model
 

modules/sd_hijack_dynamic_atten.py

@@ -113,160 +113,3 @@ def dynamic_scaled_dot_product_attention(query: torch.FloatTensor, key: torch.Fl
     if is_unsqueezed:
         hidden_states = hidden_states.squeeze(0)
     return hidden_states
-
-
-@cache
-def find_bmm_slice_sizes(query_shape, query_element_size, slice_rate=2, trigger_rate=4):
-    if len(query_shape) == 3:
-        batch_size_attention, query_tokens, shape_three = query_shape
-        shape_four = 1
-    else:
-        batch_size_attention, query_tokens, shape_three, shape_four = query_shape
-
-    slice_block_size = query_tokens * shape_three * shape_four / 1024 / 1024 * query_element_size
-    block_size = batch_size_attention * slice_block_size
-
-    split_slice_size = batch_size_attention
-    split_2_slice_size = query_tokens
-    split_3_slice_size = shape_three
-
-    do_split = False
-    do_split_2 = False
-    do_split_3 = False
-
-    if block_size > trigger_rate:
-        do_split = True
-        split_slice_size = find_split_size(split_slice_size, slice_block_size, slice_rate=slice_rate)
-        if split_slice_size * slice_block_size > slice_rate:
-            slice_2_block_size = split_slice_size * shape_three * shape_four / 1024 / 1024 * query_element_size
-            do_split_2 = True
-            split_2_slice_size = find_split_size(split_2_slice_size, slice_2_block_size, slice_rate=slice_rate)
-            if split_2_slice_size * slice_2_block_size > slice_rate:
-                slice_3_block_size = split_slice_size * split_2_slice_size * shape_four / 1024 / 1024 * query_element_size
-                do_split_3 = True
-                split_3_slice_size = find_split_size(split_3_slice_size, slice_3_block_size, slice_rate=slice_rate)
-
-    return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size
-
-
-class DynamicAttnProcessorBMM:
-    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 V1
-    """
-
-    def __call__(self, attn, hidden_states: torch.Tensor, encoder_hidden_states=None, attention_mask=None, temb=None, *args, **kwargs) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches, keyword-arg-before-vararg
-
-        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
-        )
-        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-
-        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)
-
-        query = attn.head_to_batch_dim(query)
-        key = attn.head_to_batch_dim(key)
-        value = attn.head_to_batch_dim(value)
-
-        ####################################################################
-        # Slicing parts:
-        batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2]
-        hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype)
-        do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_bmm_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate*4, trigger_rate=shared.opts.dynamic_attention_trigger_rate*4)
-
-        if do_split:
-            for i in range(batch_size_attention // split_slice_size):
-                start_idx = i * split_slice_size
-                end_idx = (i + 1) * split_slice_size
-                if do_split_2:
-                    for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name
-                        start_idx_2 = i2 * split_2_slice_size
-                        end_idx_2 = (i2 + 1) * split_2_slice_size
-                        if do_split_3:
-                            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
-
-                                query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]
-                                key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]
-                                attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attention_mask is not None else None
-
-                                attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-                                del query_slice
-                                del key_slice
-                                del attn_mask_slice
-                                attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3])
-
-                                hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = attn_slice
-                                del attn_slice
-                        else:
-                            query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2]
-                            key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2]
-                            attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None
-
-                            attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-                            del query_slice
-                            del key_slice
-                            del attn_mask_slice
-                            attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2])
-
-                            hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice
-                            del attn_slice
-                else:
-                    query_slice = query[start_idx:end_idx]
-                    key_slice = key[start_idx:end_idx]
-                    attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None
-
-                    attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
-                    del query_slice
-                    del key_slice
-                    del attn_mask_slice
-                    attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])
-
-                    hidden_states[start_idx:end_idx] = attn_slice
-                    del attn_slice
-        else:
-            attention_probs = attn.get_attention_scores(query, key, attention_mask)
-            hidden_states = torch.bmm(attention_probs, value)
-        ####################################################################
-        hidden_states = attn.batch_to_head_dim(hidden_states)
-
-        # 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

modules/sd_hijack_hypertile.py

@@ -186,12 +186,8 @@ def context_hypertile_vae(p):
     from modules import shared
     if shared.sd_model is None or not get_opt(p, 'hypertile_vae_enabled'):
         return nullcontext()
-    if shared.opts.cross_attention_optimization == 'Sub-quadratic':
-        log.warning('Hypertile UNet is not compatible with Sub-quadratic cross-attention optimization')
-        return nullcontext()
     global max_h, max_w, error_reported # pylint: disable=global-statement
     error_reported = False
-    error_reported = False
     set_resolution(p)
     max_h, max_w = 0, 0
     vae = getattr(shared.sd_model, "vae", None)
@@ -217,9 +213,6 @@ def context_hypertile_unet(p):
     from modules import shared
     if shared.sd_model is None or not get_opt(p, 'hypertile_unet_enabled'):
         return nullcontext()
-    if shared.opts.cross_attention_optimization == 'Sub-quadratic' and not shared.cmd_opts.experimental:
-        log.warning('Hypertile UNet is not compatible with Sub-quadratic cross-attention optimization')
-        return nullcontext()
     global max_h, max_w, error_reported # pylint: disable=global-statement
     error_reported = False
     set_resolution(p)

modules/sd_hijack_te.py

@@ -2,6 +2,7 @@ import os
 import time
 from modules import shared, errors, timer, sd_models
 from modules.logger import log
+from modules.attention import context as attention_context
 
 
 class PromptCache:
@@ -65,10 +66,11 @@ def hijack_encode_prompt(*args, **kwargs):
             res = cached
         else:
             log.debug(f'Encode: prompt="{prompt}" hijack=True')
-            if hasattr(shared.sd_model, 'orig_encode_prompt'):
-                res = shared.sd_model.orig_encode_prompt(*args_copy, **kwargs)
-            else:
-                res = shared.sd_model.encode_prompt(*args_copy, **kwargs)
+            with attention_context.role('te'):
+                if hasattr(shared.sd_model, 'orig_encode_prompt'):
+                    res = shared.sd_model.orig_encode_prompt(*args_copy, **kwargs)
+                else:
+                    res = shared.sd_model.encode_prompt(*args_copy, **kwargs)
             prompt_cache.set(prompt, res)
 
         if hasattr(shared.sd_model, 'after_prompt_encode'):

modules/sd_hijack_vae.py

@@ -3,6 +3,7 @@ import time
 import torch
 from modules import shared, sd_models, devices, timer, errors
 from modules.logger import log
+from modules.attention import context as attention_context
 
 
 debug = log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -25,20 +26,22 @@ def hijack_vae_decode(*args, **kwargs):
         sd_models.move_model(shared.sd_model.vae, devices.device)
         if torch.is_tensor(args[0]):
             latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype
-            if hasattr(shared.sd_model.vae, '_asymmetric_upscale_vae'):
-                res = hijack_vae_upscale(latents, *args[1:], **kwargs)
-            elif getattr(shared.sd_model, 'sdnext_vae_type', None) == 'Tiny':
-                from modules.video_models import video_vae
-                res = video_vae.vae_decode_tiny(latents) # None when the model has no tiny counterpart, and it says so
-            if res is None:
-                res = shared.sd_model.vae.orig_decode(latents, *args[1:], **kwargs)
+            with attention_context.role('vae'):
+                if hasattr(shared.sd_model.vae, '_asymmetric_upscale_vae'):
+                    res = hijack_vae_upscale(latents, *args[1:], **kwargs)
+                elif getattr(shared.sd_model, 'sdnext_vae_type', None) == 'Tiny':
+                    from modules.video_models import video_vae
+                    res = video_vae.vae_decode_tiny(latents) # None when the model has no tiny counterpart, and it says so
+                if res is None:
+                    res = shared.sd_model.vae.orig_decode(latents, *args[1:], **kwargs)
             t1 = time.time()
             try:
                 log.debug(f'Decode: vae={shared.sd_model.vae.__class__.__name__} dtype={latents.dtype} latents={list(latents.shape)}:{latents.device} decoded={list(res[0].shape)} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} time={t1-t0:.3f}')
             except Exception:
                 pass
         else:
-            res = shared.sd_model.vae.orig_decode(*args, **kwargs)
+            with attention_context.role('vae'):
+                res = shared.sd_model.vae.orig_decode(*args, **kwargs)
     except Exception as e:
         log.error(f'Decode: vae={shared.sd_model.vae.__class__.__name__} {e}')
         errors.display(e, 'vae')
@@ -58,11 +61,13 @@ def hijack_vae_encode(*args, **kwargs):
         sd_models.move_model(shared.sd_model.vae, devices.device)
         if torch.is_tensor(args[0]):
             latents = args[0].to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype
-            res = shared.sd_model.vae.orig_encode(latents, *args[1:], **kwargs)
+            with attention_context.role('vae'):
+                res = shared.sd_model.vae.orig_encode(latents, *args[1:], **kwargs)
             t1 = time.time()
             log.debug(f'Encode: vae={shared.sd_model.vae.__class__.__name__} slicing={getattr(shared.sd_model.vae, "use_slicing", None)} tiling={getattr(shared.sd_model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}')
         else:
-            res = shared.sd_model.vae.orig_encode(*args, **kwargs)
+            with attention_context.role('vae'):
+                res = shared.sd_model.vae.orig_encode(*args, **kwargs)
     except Exception as e:
         log.error(f'Encode: vae={shared.sd_model.vae.__class__.__name__} {e}')
         errors.display(e, 'vae')

modules/sd_models.py

@@ -1261,6 +1261,7 @@ def copy_diffuser_options(new_pipe, orig_pipe):
     new_pipe.sd_model_hash = getattr(orig_pipe, 'sd_model_hash', None)
     new_pipe.has_accelerate = getattr(orig_pipe, 'has_accelerate', False)
     new_pipe.current_attn_name = getattr(orig_pipe, 'current_attn_name', None)
+    new_pipe.current_attn_overrides = getattr(orig_pipe, 'current_attn_overrides', None)
     new_pipe.default_scheduler = getattr(orig_pipe, 'default_scheduler', None)
     new_pipe.image_encoder = getattr(orig_pipe, 'image_encoder', None)
     new_pipe.feature_extractor = getattr(orig_pipe, 'feature_extractor', None)
@@ -1288,6 +1289,7 @@ def backup_pipe_components(pipe):
         'sd_model_hash': getattr(pipe, "sd_model_hash", None),
         'has_accelerate': getattr(pipe, "has_accelerate", None),
         'current_attn_name': getattr(pipe, "current_attn_name", None),
+        'current_attn_overrides': getattr(pipe, "current_attn_overrides", None),
         'default_scheduler': getattr(pipe, "default_scheduler", None),
         'image_encoder': getattr(pipe, "image_encoder", None),
         'feature_extractor': getattr(pipe, "feature_extractor", None),
@@ -1309,6 +1311,7 @@ def restore_pipe_components(pipe, components):
     pipe.sd_model_hash = components['sd_model_hash']
     pipe.has_accelerate = components['has_accelerate']
     pipe.current_attn_name = components['current_attn_name']
+    pipe.current_attn_overrides = components.get('current_attn_overrides')
     pipe.default_scheduler = components['default_scheduler']
 
     if components['image_encoder'] is not None:

modules/shared_items.py

@@ -142,8 +142,6 @@ def list_crossattention():
         "Disabled",
         "Scaled-Dot-Product",
         "xFormers",
-        "Batch matrix-matrix",
-        "Dynamic Attention BMM"
     ]
 
 

modules/ui_definitions.py

@@ -250,7 +250,6 @@ def create_settings(cmd_opts):
         "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']}),
-        "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }),
         "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}),
 

test/test-attention-router.py

@@ -0,0 +1,564 @@
+#!/usr/bin/env python
+"""
+Offline unit tests for the attention router in modules.attention.
+
+Covers:
+
+- plan construction over every subset of the sdp_overrides choices on cuda, rocm, zluda and cpu
+  against an oracle of the stacking order the closure hijacks used: priority, terminal selection,
+  platform gating
+- gate parity: every backend's declared constraints against a literal transcription of the
+  predicate its closure carried, over a grid of shapes, dtypes, devices and masks
+- every sdp_overrides choice maps to a registered backend and every backend to a choice
+- router dispatch: the first accepting entry wins, the terminal receives declined calls, the
+  original sdpa only receives enable_gqa when it is set
+- a backend whose prepare raises is skipped without disturbing the rest
+- install_router leaves the original sdpa in place for an empty plan
+- the dynamic backend pins the pre-dynamic sdpa the sliced path reads
+- the generation context: step normalized to the forward about to run on both the classic
+  callback and the modular pre-hook, per-pass resets, the in-place step buffer, role scopes
+- telemetry: the route observer, the chain string recorded in torch_info, report(), and the
+  SD_ATTN_DEBUG route log deduplication
+- the escape hatches captioners and detailers depend on: bypass_sdpa_hijacks and llm_context
+  restore the original sdpa over the router, and put the router back even when the body raises
+
+No running server required. Nothing is moved to the accelerator.
+
+Usage:
+    python test/test-attention-router.py
+"""
+
+import itertools
+import logging
+import os
+import sys
+from dataclasses import replace
+
+import torch
+
+script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, script_dir)
+os.chdir(script_dir)
+
+os.environ['SD_INSTALL_QUIET'] = '1'
+
+# Bootstrap cmd_args before any module that pulls in shared.py.
+import modules.cmd_args  # pylint: disable=wrong-import-position
+import installer  # pylint: disable=wrong-import-position
+orig_argv = sys.argv
+sys.argv = [sys.argv[0]]
+try:
+    modules.cmd_args.parse_args()
+finally:
+    sys.argv = orig_argv
+installer.add_args(modules.cmd_args.parser)
+modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
+
+stock_sdpa = torch.nn.functional.scaled_dot_product_attention # importing shared installs the configured hijacks in-process
+
+from modules.errors import log                                   # pylint: disable=wrong-import-position
+from modules import attention                                    # pylint: disable=wrong-import-position
+from modules.attention import router as attention_router         # pylint: disable=wrong-import-position
+
+
+# ============================================================
+# Test infrastructure
+# ============================================================
+
+results: dict[str, dict] = {}
+
+
+def category(name: str):
+    if name not in results:
+        results[name] = {'passed': 0, 'failed': 0, 'tests': []}
+    return name
+
+
+def record(cat: str, passed: bool, name: str, detail: str = ''):
+    status = 'PASS' if passed else 'FAIL'
+    results[cat]['passed' if passed else 'failed'] += 1
+    results[cat]['tests'].append((status, name))
+    msg = f'  {status}: {name}'
+    if detail:
+        msg += f' ({detail})'
+    if passed:
+        log.info(msg)
+    else:
+        log.error(msg)
+
+
+def run_test(cat: str, fn):
+    name = fn.__name__
+    try:
+        ok = fn()
+        if ok is False:
+            record(cat, False, name)
+        else:
+            record(cat, True, name)
+    except AssertionError as e:
+        record(cat, False, name, str(e))
+    except Exception as e:  # pylint: disable=broad-except
+        record(cat, False, name, f'exception: {e}')
+        import traceback
+        traceback.print_exc()
+
+
+# ============================================================
+# The closure hijacks this router replaces, transcribed
+# ============================================================
+
+# devices.set_sdpa_params applied the hijacks in this order; each wrapped the previous, so the
+# last applied was tried first. Dynamic replaced the chain end instead of wrapping it; flex did
+# too, which left everything stacked before it unreachable, so it is an ordinary entry now.
+OLD_ORDER = ['Dynamic attention', 'Flex attention', 'Triton Flash attention', 'Flash attention', 'Sage attention', 'SDNQ attention']
+OLD_TERMINALS = {'Dynamic attention'}
+OLD_NAMES = {
+    'Dynamic attention': 'dynamic',
+    'Flex attention': 'flex',
+    'Triton Flash attention': 'triton',
+    'Flash attention': 'flash',
+    'Sage attention': 'sage',
+    'SDNQ attention': 'sdnq',
+}
+# mirrors shared_defaults.get_default_modes: five choices everywhere, Triton Flash attention added on rocm and zluda
+CHOICES = OLD_ORDER
+TRITON_PLATFORMS = {'rocm', 'zluda'}
+
+# the four closure predicates transcribed literally, plus the contract flex_attention itself enforces
+GATES = {
+    'sdnq': lambda q, k, v, m: q.device.type != "cpu" and (q.shape[-2] >= 32 and k.shape[-2] >= 32) and (q.shape[-2] > 512 or k.shape[-2] > 512) and q.shape[-3] > 1,
+    'triton': lambda q, k, v, m: q.shape[-1] <= 128 and m is None and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
+    'flash': lambda q, k, v, m: q.shape[-1] <= 128 and m is None and q.dtype != torch.float32 and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
+    'sage': lambda q, k, v, m: q.shape[-1] in {128, 96, 64} and m is None and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
+    'flex': lambda q, k, v, m: q.ndim == 4 and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
+}
+
+
+def oracle_chain(labels, platform_backend):
+    enabled = [label for label in OLD_ORDER if label in labels]
+    if platform_backend not in TRITON_PLATFORMS:
+        enabled = [label for label in enabled if label != 'Triton Flash attention']
+    terminal = None
+    entries = []
+    for label in enabled:
+        if label in OLD_TERMINALS:
+            terminal = label
+        else:
+            entries.append(label)
+    entries.reverse()
+    return [OLD_NAMES[label] for label in entries], (OLD_NAMES[terminal] if terminal else None)
+
+
+def stub_registry(failing=()):
+    """The registered backends with prepares that return a tagged call instead of importing anything."""
+    reg = attention.Registry()
+    for backend in attention.registry.ordered():
+        def prepare(platform, original, name=backend.name): # pylint: disable=unused-argument
+            if name in failing:
+                raise RuntimeError(f'{name} unavailable')
+            def call(*args, **kwargs): # pylint: disable=unused-argument
+                return name
+            return call
+        reg.register(replace(backend, prepare=prepare))
+    return reg
+
+
+def shaped(shape, dtype=torch.float16, device='meta'):
+    """A tensor of the given shape without allocating it."""
+    return torch.empty(1, dtype=dtype, device=device).expand(*shape)
+
+
+def sdpa_stub(**kwargs): # pylint: disable=unused-argument
+    return 'sdpa'
+
+
+# ============================================================
+# Tests
+# ============================================================
+
+def test_plan_matches_stacking_oracle():
+    level = log.level
+    log.setLevel(logging.ERROR) # platform gating warns per plan
+    try:
+        plans = 0
+        for platform_backend in ('cuda', 'rocm', 'zluda', 'cpu'):
+            reg = stub_registry()
+            platform = attention.Platform(backend=platform_backend)
+            for count in range(len(OLD_ORDER) + 1):
+                for labels in itertools.combinations(OLD_ORDER, count):
+                    plan = attention.build_plan(list(labels), platform, sdpa_stub, reg)
+                    expected_entries, expected_terminal = oracle_chain(labels, platform_backend)
+                    got_entries = [entry.backend.name for entry in plan.entries]
+                    got_terminal = plan.terminal.backend.name if plan.terminal is not None else None
+                    assert got_entries == expected_entries, f'{platform_backend} {labels}: entries {got_entries} != {expected_entries}'
+                    assert got_terminal == expected_terminal, f'{platform_backend} {labels}: terminal {got_terminal} != {expected_terminal}'
+                    assert plan.chain() == got_entries + [got_terminal or 'sdpa']
+                    plans += 1
+    finally:
+        log.setLevel(level)
+    log.info(f'  {plans} plans match the stacking oracle')
+    return True
+
+
+def test_gates_match_transcribed_predicates():
+    cases = 0
+    lengths = (16, 32, 512, 513, 4096)
+    for q_device, kv_device, dtype, heads, q_len, k_len, head_dim, masked, batched in itertools.product(('cpu', 'meta'), ('cpu', 'meta'), (torch.float16, torch.float32), (1, 8), lengths, lengths, (40, 64, 96, 128, 256), (False, True), (False, True)):
+        lead = (1,) if batched else ()
+        q = shaped((*lead, heads, q_len, head_dim), dtype, q_device)
+        k = shaped((*lead, heads, k_len, head_dim), dtype, kv_device)
+        v = shaped((*lead, heads, k_len, head_dim), dtype, kv_device)
+        m = shaped((*lead, 1, q_len, k_len), torch.bool, q_device) if masked else None
+        for name, gate in GATES.items():
+            expected = bool(gate(q, k, v, m))
+            got = attention.registry.backends[name].constraints.accepts(q, k, v, m)
+            assert got == expected, f'{name}: q={tuple(q.shape)} k={tuple(k.shape)} dtype={dtype} devices={q_device}/{kv_device} mask={masked} got={got} expected={expected}'
+            cases += 1
+    log.info(f'  {cases} gate cases match the transcribed predicates')
+    return True
+
+
+def test_only_dynamic_is_terminal():
+    for name, backend in attention.registry.backends.items():
+        assert backend.terminal == (name == 'dynamic'), name
+    assert attention.registry.backends['dynamic'].constraints == attention.Constraints()
+    return True
+
+
+def test_choices_match_backends():
+    labels = attention.registry.labels()
+    assert sorted(labels) == sorted(CHOICES), f'registered={labels} choices={CHOICES}'
+    for label in CHOICES:
+        assert attention.registry.by_label(label) is not None, label
+    triton = attention.registry.backends['triton']
+    assert triton.platforms == frozenset(TRITON_PLATFORMS), triton.platforms
+    for name, backend in attention.registry.backends.items():
+        if name != 'triton':
+            assert backend.platforms is None, name
+    return True
+
+
+def test_router_dispatch_prefers_priority_then_terminal_then_original():
+    calls = []
+
+    def original(**kwargs):
+        calls.append(('sdpa', kwargs))
+        return 'sdpa'
+
+    reg = attention.Registry()
+
+    def add(name, constraints, priority, terminal=False):
+        def prepare(platform, orig): # pylint: disable=unused-argument
+            def call(*args, **kwargs): # pylint: disable=unused-argument
+                calls.append((name, kwargs))
+                return name
+            return call
+        reg.register(attention.AttentionBackend(name=name, label=f'{name} attention', priority=priority, prepare=prepare, constraints=constraints, terminal=terminal))
+
+    add('narrow', attention.Constraints(head_dims=frozenset({64})), priority=20)
+    add('wide', attention.Constraints(), priority=10)
+    platform = attention.Platform(backend='cuda')
+    router = attention_router.make_router(attention.build_plan(['narrow attention', 'wide attention'], platform, original, reg))
+    q64 = shaped((1, 8, 128, 64))
+    q128 = shaped((1, 8, 128, 128))
+    cpu = shaped((1, 8, 128, 64), device='cpu')
+    assert router(q64, q64, q64) == 'narrow'
+    assert router(q128, q128, q128) == 'wide'
+    assert router(cpu, cpu, cpu) == 'sdpa'
+    assert 'enable_gqa' not in calls[-1][1], calls[-1]
+    assert router(cpu, cpu, cpu, enable_gqa=True) == 'sdpa'
+    assert calls[-1][1].get('enable_gqa') is True, calls[-1]
+
+    add('term', attention.Constraints(), priority=5, terminal=True)
+    router = attention_router.make_router(attention.build_plan(['narrow attention', 'term attention'], platform, original, reg))
+    assert router(q64, q64, q64) == 'narrow'
+    assert router(cpu, cpu, cpu, extra=1) == 'term'
+    assert calls[-1][1].get('extra') == 1 and calls[-1][1].get('enable_gqa') is False, calls[-1]
+    return True
+
+
+def test_prepare_failure_skips_backend():
+    reg = stub_registry(failing=('sage',))
+    plan = attention.build_plan(['Sage attention', 'SDNQ attention', 'Flash attention'], attention.Platform(backend='cuda'), sdpa_stub, reg)
+    assert [entry.backend.name for entry in plan.entries] == ['sdnq', 'flash'], plan.chain()
+    return True
+
+
+def test_install_router_keeps_original_for_empty_plan():
+    saved = torch.nn.functional.scaled_dot_product_attention
+    saved_plan = attention_router.current_plan
+    try:
+        platform = attention.Platform(backend='cuda')
+        plan = attention.install_router([], platform, sdpa_stub, stub_registry())
+        assert torch.nn.functional.scaled_dot_product_attention is sdpa_stub
+        assert plan.chain() == ['sdpa'], plan.chain()
+        plan = attention.install_router(['SDNQ attention', 'Sage attention'], platform, sdpa_stub, stub_registry())
+        assert torch.nn.functional.scaled_dot_product_attention is not sdpa_stub
+        assert plan.chain() == ['sdnq', 'sage', 'sdpa'], plan.chain()
+        assert attention.get_plan() is plan
+    finally:
+        torch.nn.functional.scaled_dot_product_attention = saved
+        attention_router.current_plan = saved_plan
+    return True
+
+
+def test_dynamic_backend_pins_pre_dynamic_sdpa():
+    from modules import devices
+    saved = devices.sdpa_pre_dyanmic_atten
+    try:
+        call = attention.registry.backends['dynamic'].prepare(attention.Platform(backend='cuda'), sdpa_stub)
+        from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
+        assert call is dynamic_scaled_dot_product_attention
+        assert devices.sdpa_pre_dyanmic_atten is sdpa_stub
+    finally:
+        devices.sdpa_pre_dyanmic_atten = saved
+    return True
+
+
+def test_context_classic_ticks_follow_the_callback():
+    ctx = attention.context
+
+    class Pipe:
+        transformer = object()
+
+    ctx.begin(Pipe(), steps=4)
+    assert ctx.current.active and ctx.current.role == 'transformer' and ctx.current.step == 0 and ctx.current.steps == 4
+    assert ctx.current.model_key == ('Pipe', 'object'), ctx.current.model_key
+    buffer = ctx.current.step_buffer
+    for completed in range(4):
+        ctx.tick(completed + 1) # the diffusers callback reports the step just completed
+        assert ctx.current.step == completed + 1
+        assert ctx.current.step_buffer is buffer and int(buffer.item()) == completed + 1
+    ctx.new_pass(2) # hires or refiner pass
+    assert ctx.current.step == 0 and ctx.current.steps == 2 and int(buffer.item()) == 0
+    ctx.end()
+    assert not ctx.current.active and ctx.current.role is None and ctx.current.model_key is None and ctx.current.step == 0
+    return True
+
+
+def test_context_modular_ticks_count_forwards():
+    ctx = attention.context
+    ctx.begin(None, steps=3)
+    assert ctx.current.model_key is None
+    for expected in range(3):
+        ctx.tick() # the modular pre-hook fires before each forward
+        assert ctx.current.step == expected, ctx.current.step
+    ctx.end()
+    return True
+
+
+def test_context_roles_nest_and_stick():
+    ctx = attention.context
+    ctx.begin(None)
+    with ctx.role('te'):
+        assert ctx.current.role == 'te'
+        with ctx.role('vae'):
+            assert ctx.current.role == 'vae'
+        assert ctx.current.role == 'te'
+    assert ctx.current.role == 'transformer'
+    ctx.set_role('vae')
+    assert ctx.current.role == 'vae'
+    ctx.end()
+    assert ctx.current.role is None
+    with ctx.role('te'): # outside a generation the scope still restores what it found
+        assert ctx.current.role == 'te'
+    assert ctx.current.role is None
+    return True
+
+
+def test_router_observer_sees_each_route():
+    routes = []
+    reg = attention.Registry()
+
+    def prepare(platform, original): # pylint: disable=unused-argument
+        return lambda *args, **kwargs: 'narrow'
+
+    reg.register(attention.AttentionBackend(name='narrow', label='narrow attention', priority=20, prepare=prepare, constraints=attention.Constraints(head_dims=frozenset({64}))))
+    plan = attention.build_plan(['narrow attention'], attention.Platform(backend='cuda'), sdpa_stub, reg)
+    router = attention_router.make_router(plan, observer=lambda name, q, k, m: routes.append(name))
+    q64 = shaped((1, 8, 128, 64))
+    q128 = shaped((1, 8, 128, 128))
+    router(q64, q64, q64)
+    router(q128, q128, q128)
+    assert routes == ['narrow', 'sdpa'], routes
+    return True
+
+
+def test_install_router_records_the_chain():
+    saved = torch.nn.functional.scaled_dot_product_attention
+    saved_plan = attention_router.current_plan
+    saved_info = installer.torch_info.get('attention')
+    try:
+        attention.install_router(['SDNQ attention', 'Dynamic attention'], attention.Platform(backend='cuda'), sdpa_stub, stub_registry())
+        assert installer.torch_info.get('attention') == 'sdnq>dynamic', installer.torch_info.get('attention')
+        info = attention.report()
+        assert info['chain'] == ['sdnq', 'dynamic'] and info['overrides'] == ['SDNQ attention', 'Dynamic attention'] and info['backend'] == 'cuda', info
+        assert info['context']['active'] is False and info['context']['role'] is None, info
+    finally:
+        torch.nn.functional.scaled_dot_product_attention = saved
+        attention_router.current_plan = saved_plan
+        installer.torch_info.set(attention=saved_info)
+    return True
+
+
+def test_debug_observe_logs_each_route_once():
+    attention.debug.reset()
+    q = shaped((1, 8, 128, 64))
+    attention.debug.observe('sdnq', q, q, None)
+    attention.debug.observe('sdnq', q, q, None)
+    attention.debug.observe('sdnq', q, q, shaped((1, 1, 128, 128), torch.bool))
+    assert len(attention.debug.seen) == 2, attention.debug.seen
+    attention.debug.reset()
+    assert not attention.debug.seen
+    return True
+
+
+def test_attention_slicing_follows_the_choice():
+    from modules import shared
+
+    class Pipe:
+        def __init__(self):
+            self.calls = []
+
+        def enable_attention_slicing(self):
+            self.calls.append('enable')
+
+        def disable_attention_slicing(self):
+            self.calls.append('disable')
+
+    saved = {key: shared.opts.data.get(key, None) for key in ['attention_slicing', 'cross_attention_optimization']}
+    try:
+        shared.opts.data['cross_attention_optimization'] = 'Disabled' # the branch under test is the only one that should act
+        for choice, expected in [('Default', []), ('Enabled', ['enable']), ('Disabled', ['disable'])]:
+            shared.opts.data['attention_slicing'] = choice
+            pipe = Pipe()
+            attention.set_diffusers_attention(pipe, quiet=True)
+            assert pipe.calls == expected, f'{choice} produced {pipe.calls}'
+    finally:
+        for key, value in saved.items():
+            if value is None:
+                shared.opts.data.pop(key, None)
+            else:
+                shared.opts.data[key] = value
+    return True
+
+
+def test_removed_attention_methods_are_gone():
+    from modules import shared_items
+    from modules import options_handler
+    from modules import sd_hijack_dynamic_atten
+
+    removed = ['Batch matrix-matrix', 'Dynamic Attention BMM']
+    choices = shared_items.list_crossattention()
+    assert not [name for name in removed if name in choices], choices
+    for name in removed:
+        data = {'cross_attention_optimization': name}
+        migrated = options_handler.migrate_removed_values(data)
+        assert data['cross_attention_optimization'] == 'Scaled-Dot-Product', data
+        assert len(migrated) == 1, migrated
+    kept = {'cross_attention_optimization': 'xFormers'}
+    assert options_handler.migrate_removed_values(kept) == [], 'a live choice is left alone'
+    assert kept['cross_attention_optimization'] == 'xFormers', kept
+    assert not hasattr(sd_hijack_dynamic_atten, 'DynamicAttnProcessorBMM'), 'the bmm processor is removed'
+    assert hasattr(sd_hijack_dynamic_atten, 'dynamic_scaled_dot_product_attention'), 'the sliced sdpa path stays'
+    return True
+
+
+def test_escape_hatch_bypasses_the_router():
+    from modules import devices
+    saved_sdpa = torch.nn.functional.scaled_dot_product_attention
+    saved_plan = attention_router.current_plan
+    saved_original = devices.sdpa_original
+    try:
+        devices.sdpa_original = sdpa_stub # what set_sdpa_params pinned before building the chain
+        attention.install_router(['SDNQ attention'], attention.Platform(backend='cuda'), sdpa_stub, stub_registry())
+        router = torch.nn.functional.scaled_dot_product_attention
+        assert router is not sdpa_stub
+        with devices.bypass_sdpa_hijacks():
+            assert torch.nn.functional.scaled_dot_product_attention is sdpa_stub # captioners and detailers run here
+        assert torch.nn.functional.scaled_dot_product_attention is router
+        with devices.llm_context():
+            assert torch.nn.functional.scaled_dot_product_attention is sdpa_stub
+        assert torch.nn.functional.scaled_dot_product_attention is router
+        try:
+            with devices.bypass_sdpa_hijacks():
+                raise RuntimeError('captioner failed')
+        except RuntimeError:
+            pass
+        assert torch.nn.functional.scaled_dot_product_attention is router, 'the chain must survive a failure inside the bypass'
+    finally:
+        devices.sdpa_original = saved_original
+        torch.nn.functional.scaled_dot_product_attention = saved_sdpa
+        attention_router.current_plan = saved_plan
+    return True
+
+
+def test_reapply_options_cover_declared_backend_options():
+    from modules import shared
+    names = attention.reapply_options()
+    assert names[:2] == ['sdp_options', 'sdp_overrides'], names
+    declared = attention.registry.options()
+    assert set(declared) <= set(names), (declared, names)
+    assert attention.registry.backends['sdnq'].options and set(attention.registry.backends['sdnq'].options) <= set(declared)
+    for name in names:
+        assert name in shared.opts.data_labels, name
+    return True
+
+
+def run_all():
+    log.warning('=== attention router ===')
+    cat = category('router')
+    for fn in [
+        test_plan_matches_stacking_oracle,
+        test_gates_match_transcribed_predicates,
+        test_only_dynamic_is_terminal,
+        test_choices_match_backends,
+        test_router_dispatch_prefers_priority_then_terminal_then_original,
+        test_prepare_failure_skips_backend,
+        test_install_router_keeps_original_for_empty_plan,
+        test_dynamic_backend_pins_pre_dynamic_sdpa,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== generation context ===')
+    cat = category('context')
+    for fn in [
+        test_context_classic_ticks_follow_the_callback,
+        test_context_modular_ticks_count_forwards,
+        test_context_roles_nest_and_stick,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== telemetry ===')
+    cat = category('telemetry')
+    for fn in [
+        test_router_observer_sees_each_route,
+        test_install_router_records_the_chain,
+        test_debug_observe_logs_each_route_once,
+        test_reapply_options_cover_declared_backend_options,
+        test_attention_slicing_follows_the_choice,
+        test_removed_attention_methods_are_gone,
+        test_escape_hatch_bypasses_the_router,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== Results ===')
+    total_passed = 0
+    total_failed = 0
+    for cat_name, info in results.items():
+        ok = info['failed'] == 0
+        status = 'PASS' if ok else 'FAIL'
+        log.info(f"  {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]")
+        total_passed += info['passed']
+        total_failed += info['failed']
+    log.warning(f'Total: {total_passed} passed, {total_failed} failed')
+    return total_failed == 0
+
+
+if __name__ == '__main__':
+    import time
+    t0 = time.time()
+    ok = run_all()
+    torch.nn.functional.scaled_dot_product_attention = stock_sdpa
+    log.warning(f'Total time: {time.time() - t0:.2f}s')
+    sys.exit(0 if ok else 1)

ui/locale/locale_en.json

@@ -152,7 +152,11 @@
     {"id":"","label":"Adapter 2","localized":"","hint":"","ui":"txt2img"},
     {"id":"","label":"Adapter 3","localized":"","hint":"","ui":"txt2img"},
     {"id":"","label":"Adapter 4","localized":"","hint":"","ui":"txt2img"},
-    {"id":"","label":"Audio","localized":"","hint":"","ui":"video"}
+    {"id":"","label":"Audio","localized":"","hint":"","ui":"video"},
+    {"id":"","label":"Attention method","localized":"","hint":"Attention processor the pipeline is loaded with.<br><br><b>Scaled-Dot-Product</b> uses torch attention and is the path the rest of this section builds on; <b><i>SDP kernels</i></b> and <b><i>SDP overrides</i></b> apply to it and to nothing else.<br><b>xFormers</b> uses the xFormers library on the modules that accept it.<br><b>Disabled</b> leaves whatever the pipeline was built with.<br><br>Default Scaled-Dot-Product.","ui":"settings_cuda"},
+    {"id":"","label":"Attention slicing","localized":"","hint":"Computes attention in slices instead of one pass, trading speed for a lower memory peak. Applied by the pipeline rather than by the kernel, so it stacks on whichever <b><i>Attention method</i></b> is active.<br><br><b>Default</b> leaves the pipeline's own choice untouched, which for most models means off.<br><b>Enabled</b> and <b>Disabled</b> override that choice.<br><br>Unrelated to <b>Dynamic attention</b>, which slices by a memory estimate; this one is diffusers' fixed slicing.<br><br>Default <b>Default</b>.","ui":"settings_cuda"},
+    {"id":"","label":"Attention Dispatcher","localized":"","hint":"Diffusers keeps its own attention backend registry, separate from the torch level settings above. A backend chosen here is installed on the model itself and changes what diffusers dispatches to inside its transformer blocks.<br><br>Applies to diffusion transformer models that expose the call. UNet models and pipelines without it are left alone.","ui":"settings_cuda"},
+    {"id":"","label":"Attention dispatcher kernel","localized":"","hint":"Name of the diffusers attention backend to install on the model. Empty leaves the diffusers default in place.<br><br>Accepted names come from the diffusers build in use and currently include <b>native</b>, <b>flash</b>, <b>flash_varlen</b>, <b>flex</b>, <b>sage</b>, <b>sage_varlen</b> and <b>xformers</b>, alongside underscore-prefixed variants that pin one specific kernel.<br>A name is checked against that list, and an unrecognized one is logged as a warning with the full list while the model keeps its current backend. A recognized name whose package is missing is logged as an error at the same point, with the same result.<br>Names ending in <b>_hub</b> fetch the kernel from the Hugging Face Hub on first use, which needs a download and the kernels package.<br><br>Empty by default.","ui":"settings_cuda"}
   ],
   "b": [
     {"id":"","label":"Batch","localized":"","hint":"Batch processing settings","ui":"img2img"},
@@ -189,7 +193,6 @@
     {"id":"","label":"block_level","localized":"","hint":"","ui":"settings_offload"},
     {"id":"","label":"Backend storage","localized":"","hint":"","ui":"settings_quantization"},
     {"id":"","label":"BF16","localized":"","hint":"Use modified 16-bit floating point precision for calculations","ui":"settings_cuda"},
-    {"id":"","label":"Batch matrix-matrix","localized":"","hint":"Standard batched matrix multiplication for attention. Reliable but not VRAM-efficient.","ui":"settings_cuda"},
     {"id":"","label":"BCFHW","localized":"","hint":"","ui":"settings_advanced"},
     {"id":"","label":"BFCHW","localized":"","hint":"","ui":"settings_advanced"},
     {"id":"","label":"BCHW","localized":"","hint":"","ui":"settings_advanced"},
@@ -234,7 +237,7 @@
     {"id":"","label":"Control Media","localized":"","hint":"Add input image as separate initialization image for control processing","ui":"control"},
     {"id":"","label":"Create Video","localized":"","hint":"","ui":"extras"},
     {"id":"","label":"ChronoEdit","localized":"","hint":"","ui":"settings_model_options"},
-    {"id":"","label":"Cross Attention","localized":"","hint":"","ui":"settings_cuda"},
+    {"id":"","label":"Cross Attention","localized":"","hint":"Selects how attention is computed during generation. Attention is where a diffusion model spends most of its time and most of its peak memory, so these settings move both.<br><br>Three layers apply in order. <b><i>Attention method</i></b> picks the attention processor the pipeline is loaded with. <b><i>SDP kernels</i></b> limits which kernels torch may choose from inside its own attention. <b><i>SDP overrides</i></b> replaces torch attention with a different implementation for the calls that implementation accepts.<br>Overrides are tried in priority order and each one declines the calls it cannot serve, so several can be enabled at once and whatever is left over falls back to torch.","ui":"settings_cuda"},
     {"id":"","label":"CLiP Skip","localized":"","hint":"Early stopping parameter for the CLiP text encoder; 1 is stop at last layer as usual, 2 is stop at penultimate layer, etc","ui":"settings_advanced"},
     {"id":"","label":"Cache-DiT","localized":"","hint":"","ui":"settings_advanced"},
     {"id":"","label":"CFG-Zero","localized":"","hint":"","ui":"settings_advanced"},
@@ -397,10 +400,9 @@
     {"id":"","label":"Dequantize using torch.compile","localized":"","hint":"Compiles the dequantization step with <i>torch.compile</i> for faster inference. Requires <i>Triton</i>.<br><br>Changing this needs a full restart to take effect.<br><br>Enabled by default when Triton is available.","reload":"server","ui":"settings_quantization"},
     {"id":"","label":"Dequantize using full precision","localized":"","hint":"Uses <b>FP32</b> for the dequantization step for better numerical accuracy, at a small speed cost.<br><br>Enabled by default.","reload":"model","ui":"settings_quantization"},
     {"id":"","label":"Disabled","localized":"","hint":"","ui":"settings_cuda"},
-    {"id":"","label":"Dynamic Attention BMM","localized":"","hint":"Performs attention computation in steps instead of all at once. Slower inference times, but greatly reduced memory usage","ui":"settings_cuda"},
     {"id":"","label":"Dynamic attention","localized":"","hint":"Adjusts attention computation dynamically per step. Saves VRAM but slows generation.","ui":"settings_cuda"},
-    {"id":"","label":"Dynamic Attention slicing rate","localized":"","hint":"","ui":"settings_cuda"},
-    {"id":"","label":"Dynamic Attention trigger rate","localized":"","hint":"","ui":"settings_cuda"},
+    {"id":"","label":"Dynamic Attention slicing rate","localized":"","hint":"Target size in GB for each attention slice once <b>Dynamic attention</b> starts slicing. Smaller slices hold the peak lower and add more per-slice overhead.<br>Slicing is applied across the batch first, then across attention heads, then across query tokens, going a level deeper whenever the level above is still over target.<br><br>Applies while <b>Dynamic attention</b> is enabled in <b><i>SDP overrides</i></b>.<br>Default 0.5.","ui":"settings_cuda"},
+    {"id":"","label":"Dynamic Attention trigger rate","localized":"","hint":"Estimated attention matrix size in GB above which <b>Dynamic attention</b> begins slicing. Below it the call runs in one pass.<br>The estimate is batch x heads x query length x key length x bytes per element, so it grows with the square of the sequence length and crosses the threshold at high resolution or on video long before it does anywhere else.<br><br>Applies while <b>Dynamic attention</b> is enabled in <b><i>SDP overrides</i></b>.<br>Default 1.","ui":"settings_cuda"},
     {"id":"","label":"Deterministic mode","localized":"","hint":"Forces deterministic output across runs. Useful for reproducibility, but may disable some optimizations.","ui":"settings_backends"},
     {"id":"","label":"DirectML retry ops for NaN","localized":"","hint":"","ui":"settings_backends"},
     {"id":"","label":"deep-cache","localized":"","hint":"","ui":"settings_compile"},
@@ -1457,7 +1459,17 @@
     {"id":"","label":"Specify model revision","localized":"","hint":"","ui":"models_huggingface_tab"},
     {"id":"","label":"SegmentAnything","localized":"","hint":"","ui":"control"},
     {"id":"","label":"Sections","localized":"","hint":"","ui":"video"},
-    {"id":"","label":"Samplers","localized":"","hint":"Samplers/schedulers advanced settings","ui":"tab_txt2img"}
+    {"id":"","label":"Samplers","localized":"","hint":"Samplers/schedulers advanced settings","ui":"tab_txt2img"},
+    {"id":"","label":"SDP kernels","localized":"","hint":"Which of torch's built-in attention kernels torch is allowed to choose from. These are permissions rather than a selection: torch picks one per call from whatever is left enabled, preferring <b>Flash</b>, dropping to <b>Memory</b> for calls flash cannot serve such as those carrying an arbitrary attention mask, and to <b>Math</b> when neither fits. Clearing a box removes a candidate; it never pins the remaining one to every call.<br><br><b>Flash</b> is torch's own build of the FlashAttention kernel. It is not the same thing as the <b>Flash attention</b> entry in <b><i>SDP overrides</i></b>, which calls the separately installed flash-attn package and bypasses torch entirely.<br><b>Memory</b> is the memory-efficient kernel, which accepts arbitrary masks that flash does not.<br><b>Math</b> is the unfused reference path, the widest in what it accepts and the least optimized. Leaving it enabled keeps a fallback for calls the other two decline.<br><br>Applies while <b><i>Attention method</i></b> is <b>Scaled-Dot-Product</b>, and continues to govern the calls that an enabled override declines.<br><br>All three by default. ZLUDA starts with <b>Math</b> alone.","ui":"settings_cuda"},
+    {"id":"","label":"SDP overrides","localized":"","hint":"Replaces torch attention with another implementation. Each entry declares the shapes, dtypes and mask conditions it can serve; a call that fails them moves to the next entry and finally back to torch, so several can be enabled together and the chain resolves per call.<br><br><b>Flash attention</b> installs and calls the flash-attn package directly, for calls with no attention mask, half precision inputs and a head dimension of 128 or less.<br><b>Sage attention</b> computes attention with quantized matmuls, for head dimensions of 64, 96 or 128 and no attention mask.<br><b>SDNQ attention</b> is SD.Next's own quantized Triton kernel, configured in the section below. It takes attention masks, which the other quantized backends do not.<br><b>Flex attention</b> uses torch's compiled flex_attention.<br><b>Dynamic attention</b> slices attention to fit available memory and serves every call the others decline, standing in for the torch fallback.<br><b>Triton Flash attention</b> is a Triton implementation for ROCm and ZLUDA, listed only on those backends.<br><br>The quantized and compiled backends trade some numerical accuracy for throughput. How much of each arrives depends on the model, the sequence length and the GPU, so comparing them on the actual workload settles it faster than picking by reputation.<br><br>None by default on CUDA. ZLUDA, CPU and MPS start with <b>Dynamic attention</b>, as do ROCm GPUs older than RDNA3.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention","localized":"","hint":"Settings for the <b>SDNQ attention</b> entry in <b><i>SDP overrides</i></b>. They do nothing until that override is enabled.<br><br>The kernel quantizes the two matmuls inside attention, computing them on lower precision operands and rescaling the result. It is written in Triton, so it needs a working Triton for the active device.<br>Short sequences and single-head calls are left to the rest of the chain, so text encoders and the VAE keep ordinary attention.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention use Smooth K","localized":"","hint":"Subtracts the mean of the keys before quantizing them. Keys carry a large offset that is shared across the sequence, which spends most of the quantized range representing a value identical for every key and leaves little of it for the differences that decide the attention.<br>Softmax ignores a constant shift applied to every score in a row, so removing that offset changes the quantization error and not the attention.<br><br>Costs one mean and one subtraction per call.<br>Enabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention use Hadamard","localized":"","hint":"Rotates queries and keys by a Hadamard transform before quantizing them. The rotation spreads a few oversized channels across all of them, which is the error shape quantization handles worst. The transform is orthogonal, so the scores it produces are the ones the unrotated tensors would produce, minus the quantization error it removes.<br>With <b><i>SDNQ Attention PV MatMul type</i></b> also set, the values are rotated as well and the output is rotated back.<br><br>Costs a rotation pass on every attention call, so it is worth enabling where a model shows quantization artifacts without it.<br>Idle while <b><i>SDNQ Attention MatMul type</i></b> is <b>disabled</b>, since nothing is quantized then.<br><br>Disabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention use FP16 Accumulation","localized":"","hint":"Accumulates the floating point matmuls in fp16 rather than fp32. Some tensor cores run fp16 accumulation at a higher rate than fp32, and on those the kernel is cheaper for it.<br>Operands are pre-scaled to keep products inside the fp16 range, which covers ordinary activations with less headroom than fp32 leaves.<br><br>Reaches the parts of the kernel that run in floating point. An int8 matmul accumulates in int32 and is unaffected, so at the default <b><i>SDNQ Attention MatMul type</i></b> this applies to the probability-value matmul alone.<br><br>Disabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention MatMul type","localized":"","hint":"Precision the query-key matmul is computed in, the first of the two matmuls in attention.<br><br><b>enabled</b> selects int8, and <b>int8</b> and <b>uint8</b> reach the same kernel.<br><b>float16</b> and <b>float8_e4m3fn</b> take the floating point path. fp8 needs a GPU with fp8 tensor cores and fails on hardware without them rather than falling back.<br><b>disabled</b> leaves queries and keys in the model's own precision, which also idles <b><i>SDNQ Attention use Smooth K</i></b> and <b><i>SDNQ Attention use Hadamard</i></b>.<br><br>Default enabled.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention PV MatMul type","localized":"","hint":"Precision the probability-value matmul is computed in, the second of the two matmuls in attention. Choices match <b><i>SDNQ Attention MatMul type</i></b>.<br><br>Quantizing this one as well takes out the floating point work the first setting leaves behind, and it is the more delicate of the two: its inputs are already normalized probabilities, and the small ones among them carry the fine detail.<br><b>disabled</b> keeps this matmul in the model's own precision.<br><br>Default disabled.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention Hadamard Group Size","localized":"","hint":"Width of the Hadamard rotation in channels. Wider groups mix more channels together and spread outliers further.<br><br>Clamped to the head dimension of the running model, rounded down to a power of two that divides it. On a model with 64 or 128 channels per head the upper part of this range resolves to that head dimension rather than to the number shown. Rotation is skipped below 4.<br>Applies while <b><i>SDNQ Attention use Hadamard</i></b> is enabled.<br><br>Default 256.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention Quantize FP32","localized":"","hint":"Upcasts queries, keys and values to fp32 for the quantization step, meaning the mean subtraction, scale and rounding that produce the low precision operands. The matmuls themselves are unaffected, and the kernel applies the scales in fp32 either way.<br>Turned off, that arithmetic runs in the model's own precision. bf16 carries eight mantissa bits, so a scale derived in it is coarser than one derived in fp32, and <b><i>SDNQ Attention use Smooth K</i></b> loses the most from it, since a mean across the whole sequence is exactly the kind of sum that wants the extra bits.<br><br>Whether the upcast costs anything depends on how the GPU runs fp32 vector work against fp16 and bf16. NVIDIA and AMD run them at the same rate here, so there is nothing to save; Intel runs fp32 slower and takes a noticeable hit.<br><br>Enabled by default.","ui":"settings_cuda"}
   ],
   "t": [
     {"id":"txt2img_nav","label":"T2I","localized":"","hint":"Create image from text<br>Legacy interface that mimics original text-to-image interface and behavior"},

webui.py

@@ -17,6 +17,7 @@ import modules.loader
 import modules.hashes
 import modules.paths
 import modules.devices
+import modules.attention
 import modules.migrate
 from modules import shared
 from modules import call_queue
@@ -205,6 +206,8 @@ def load_model():
     shared.opts.onchange("temp_dir", modules.gr_tempdir.on_tmpdir_changed)
     for opt in modules.sd_offload_state.offload_reapply_options:
         shared.opts.onchange(opt, call_queue.wrap_queued_call(modules.sd_models.reapply_offload), call=False)
+    for opt in modules.attention.reapply_options():
+        shared.opts.onchange(opt, call_queue.wrap_queued_call(modules.attention.reapply), call=False)
     timer.startup.record("onchange")