From 6ed1b99aaab43bcd4df096b76a04d35d9c9f0403 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 21:12:21 +0100 Subject: [PATCH 01/12] refactor(attention): move into a package modules/attention.py becomes modules/attention/: hijacks.py keeps the six sdpa monkeypatch setters, dispatcher.py the diffusers-side processor and dispatcher setup with the kernels hub hijack, and the package facade re-exports every public name so call sites are unchanged. The devices import moves inside set_diffusers_attention, which removes the devices <-> attention import cycle. --- modules/attention/__init__.py | 8 ++ modules/attention/dispatcher.py | 115 ++++++++++++++++++ .../{attention.py => attention/hijacks.py} | 114 +---------------- 3 files changed, 124 insertions(+), 113 deletions(-) create mode 100644 modules/attention/__init__.py create mode 100644 modules/attention/dispatcher.py rename modules/{attention.py => attention/hijacks.py} (70%) diff --git a/modules/attention/__init__.py b/modules/attention/__init__.py new file mode 100644 index 000000000..a0eae843f --- /dev/null +++ b/modules/attention/__init__.py @@ -0,0 +1,8 @@ +"""Attention backends: the SDPA hijacks stacked by devices.set_sdpa_params, plus the diffusers-side processor and dispatcher setup.""" +from modules.attention.hijacks import set_dynamic_attention, set_sdnq_attention, set_triton_flash_attention, set_flex_attention, set_ck_flash_attention, set_sage_attention +from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, hijack_kernels, get_kernel_hijack, get_hf_api_hijack + +__all__ = [ + 'set_dynamic_attention', 'set_sdnq_attention', 'set_triton_flash_attention', 'set_flex_attention', 'set_ck_flash_attention', 'set_sage_attention', + 'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack', +] diff --git a/modules/attention/dispatcher.py b/modules/attention/dispatcher.py new file mode 100644 index 000000000..4cf232022 --- /dev/null +++ b/modules/attention/dispatcher.py @@ -0,0 +1,115 @@ +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 + 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}') diff --git a/modules/attention.py b/modules/attention/hijacks.py similarity index 70% rename from modules/attention.py rename to modules/attention/hijacks.py index 1aad29692..e837f79cf 100644 --- a/modules/attention.py +++ b/modules/attention/hijacks.py @@ -1,6 +1,6 @@ from functools import wraps import torch -from modules import rocm, errors, devices +from modules import rocm from modules.logger import log from installer import install, installed, torch_info @@ -240,115 +240,3 @@ def set_sage_attention(backend: str, device: torch.device): 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}') From 3302e78af631f81acffb37b42febe0debc53867c Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 21:51:29 +0100 Subject: [PATCH 02/12] refactor(attention): backend registry and a single sdpa router Replace the six closure hijacks stacked in devices.set_sdpa_params with a registry of declarative backends and one router installed in their place. Each backend declares the constraints its closure carried as a predicate, a priority matching its old stacking position, and a prepare step that imports and configures the implementation; the router walks the prepared entries by priority and hands declined calls to the terminal backend (dynamic, flex) or the original sdpa, so fallback is the router's job rather than each closure's. - parity held: gates transcribed literally, the same kernel kwargs, enable_gqa passed to the original only when set, torch_info keeps the last prepared backend, the dynamic pin still set - a backend enabled on a platform without it warns instead of silently doing nothing - the legacy set_* entry points are gone; devices.py installs the router - test/test-attention-router.py checks every override subset against the old stacking order, gate parity over 16,000 shape cases, dispatch, terminal handoff and prepare isolation, offline --- modules/attention/__init__.py | 10 +- modules/attention/backends/__init__.py | 10 + modules/attention/backends/dynamic.py | 11 + modules/attention/backends/flash_ck.py | 44 +++ modules/attention/backends/flex.py | 38 +++ modules/attention/backends/sage.py | 53 ++++ modules/attention/backends/sdnq.py | 28 ++ modules/attention/backends/triton_amd.py | 32 +++ modules/attention/hijacks.py | 242 ---------------- modules/attention/registry.py | 91 ++++++ modules/attention/router.py | 90 ++++++ modules/devices.py | 22 +- test/test-attention-router.py | 343 +++++++++++++++++++++++ 13 files changed, 748 insertions(+), 266 deletions(-) create mode 100644 modules/attention/backends/__init__.py create mode 100644 modules/attention/backends/dynamic.py create mode 100644 modules/attention/backends/flash_ck.py create mode 100644 modules/attention/backends/flex.py create mode 100644 modules/attention/backends/sage.py create mode 100644 modules/attention/backends/sdnq.py create mode 100644 modules/attention/backends/triton_amd.py delete mode 100644 modules/attention/hijacks.py create mode 100644 modules/attention/registry.py create mode 100644 modules/attention/router.py create mode 100644 test/test-attention-router.py diff --git a/modules/attention/__init__.py b/modules/attention/__init__.py index a0eae843f..7b923ca6b 100644 --- a/modules/attention/__init__.py +++ b/modules/attention/__init__.py @@ -1,8 +1,12 @@ -"""Attention backends: the SDPA hijacks stacked by devices.set_sdpa_params, plus the diffusers-side processor and dispatcher setup.""" -from modules.attention.hijacks import set_dynamic_attention, set_sdnq_attention, set_triton_flash_attention, set_flex_attention, set_ck_flash_attention, set_sage_attention +"""Attention backends: one scaled_dot_product_attention router over the registered backends, plus 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 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 __all__ = [ - 'set_dynamic_attention', 'set_sdnq_attention', 'set_triton_flash_attention', 'set_flex_attention', 'set_ck_flash_attention', 'set_sage_attention', + 'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry', + 'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack', + 'backends', ] diff --git a/modules/attention/backends/__init__.py b/modules/attention/backends/__init__.py new file mode 100644 index 000000000..49d166e9c --- /dev/null +++ b/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) diff --git a/modules/attention/backends/dynamic.py b/modules/attention/backends/dynamic.py new file mode 100644 index 000000000..dcbdd8a5a --- /dev/null +++ b/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) diff --git a/modules/attention/backends/flash_ck.py b/modules/attention/backends/flash_ck.py new file mode 100644 index 000000000..09bd86e75 --- /dev/null +++ b/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), +) diff --git a/modules/attention/backends/flex.py b/modules/attention/backends/flex.py new file mode 100644 index 000000000..7997d54fd --- /dev/null +++ b/modules/attention/backends/flex.py @@ -0,0 +1,38 @@ +import torch +from modules.logger import log +from modules.attention.registry import AttentionBackend, 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=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs): # 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(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, terminal=True) diff --git a/modules/attention/backends/sage.py b/modules/attention/backends/sage.py new file mode 100644 index 000000000..d51055970 --- /dev/null +++ b/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), +) diff --git a/modules/attention/backends/sdnq.py b/modules/attention/backends/sdnq.py new file mode 100644 index 000000000..c0c49da67 --- /dev/null +++ b/modules/attention/backends/sdnq.py @@ -0,0 +1,28 @@ +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 + + 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, + 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, + use_fp16_accum=shared.opts.sdnq_attention_use_fp16_accum, + ) + + 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}') + 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 +) diff --git a/modules/attention/backends/triton_amd.py b/modules/attention/backends/triton_amd.py new file mode 100644 index 000000000..a1f3be7df --- /dev/null +++ b/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'}), +) diff --git a/modules/attention/hijacks.py b/modules/attention/hijacks.py deleted file mode 100644 index e837f79cf..000000000 --- a/modules/attention/hijacks.py +++ /dev/null @@ -1,242 +0,0 @@ -from functools import wraps -import torch -from modules import rocm -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, - 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}') diff --git a/modules/attention/registry.py b/modules/attention/registry.py new file mode 100644 index 000000000..ef33d009e --- /dev/null +++ b/modules/attention/registry.py @@ -0,0 +1,91 @@ +"""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 + + 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 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 + + 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()] + + +registry = Registry() diff --git a/modules/attention/router.py b/modules/attention/router.py new file mode 100644 index 000000000..f917424cd --- /dev/null +++ b/modules/attention/router.py @@ -0,0 +1,90 @@ +"""The single scaled_dot_product_attention entry point over the prepared backends.""" +from dataclasses import dataclass +from functools import wraps +import torch +from installer import torch_info +from modules.logger import log +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 and owns the torch_info record + 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) + torch_info.set(attention=backend.name) + entries.reverse() + return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels)) + + +def make_router(plan: Plan) -> AttentionCall: + entries = plan.entries + terminal = plan.terminal.call if plan.terminal is not None else None + 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): + return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa) + 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) + torch.nn.functional.scaled_dot_product_attention = make_router(plan) if (plan.entries or plan.terminal is not None) else original + current_plan = plan + return plan + + +def get_plan() -> Plan | None: + return current_plan diff --git a/modules/devices.py b/modules/devices.py index 901014505..b29b0213c 100644 --- a/modules/devices.py +++ b/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: diff --git a/test/test-attention-router.py b/test/test-attention-router.py new file mode 100644 index 000000000..c51f4b3f5 --- /dev/null +++ b/test/test-attention-router.py @@ -0,0 +1,343 @@ +#!/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 + +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 and flex replaced the chain end instead of wrapping it. +OLD_ORDER = ['Dynamic attention', 'Flex attention', 'Triton Flash attention', 'Flash attention', 'Sage attention', 'SDNQ attention'] +OLD_TERMINALS = {'Dynamic attention', 'Flex 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'} + +OLD_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, +} + + +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 in itertools.product(('cpu', 'meta'), ('cpu', 'meta'), (torch.float16, torch.float32), (1, 8), lengths, lengths, (40, 64, 96, 128, 256), (False, True)): + q = shaped((1, heads, q_len, head_dim), dtype, q_device) + k = shaped((1, heads, k_len, head_dim), dtype, kv_device) + v = shaped((1, heads, k_len, head_dim), dtype, kv_device) + m = shaped((1, 1, q_len, k_len), torch.bool, q_device) if masked else None + for name, gate in OLD_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_terminals_carry_no_gate(): + for name in ('dynamic', 'flex'): + backend = attention.registry.backends[name] + assert backend.terminal, name + assert backend.constraints == attention.Constraints(), name + 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 run_all(): + log.warning('=== attention router ===') + cat = category('router') + for fn in [ + test_plan_matches_stacking_oracle, + test_gates_match_transcribed_predicates, + test_terminals_carry_no_gate, + 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('=== 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) From bb0cc5328ea2bff208674a5eb9486b00f65bd17e Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 21:59:54 +0100 Subject: [PATCH 03/12] fix(attention): flex joins the chain instead of ending it The flex backend never called the sdpa it replaced, so any backend stacked before it was unreachable and every call it could not serve, cpu or 3d inputs included, failed inside flex_attention. It is now an ordinary entry gated on what flex_attention accepts: 4d tensors on one non-cpu device. The mask path drops the 2d special case, which indexed attn_mask.size and reshaped the mask onto the wrong axis; expanding to (batch, heads, q, kv) already follows sdpa broadcast semantics. --- modules/attention/backends/flex.py | 13 ++++++------ modules/attention/registry.py | 3 +++ test/test-attention-router.py | 33 ++++++++++++++++-------------- 3 files changed, 28 insertions(+), 21 deletions(-) diff --git a/modules/attention/backends/flex.py b/modules/attention/backends/flex.py index 7997d54fd..4ca29be3e 100644 --- a/modules/attention/backends/flex.py +++ b/modules/attention/backends/flex.py @@ -1,6 +1,6 @@ import torch from modules.logger import log -from modules.attention.registry import AttentionBackend, Platform +from modules.attention.registry import AttentionBackend, Constraints, Platform def prepare(platform: Platform, original): # pylint: disable=unused-argument @@ -9,16 +9,14 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument 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=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs): # pylint: disable=unused-argument + 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] - 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) + 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] @@ -35,4 +33,7 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument return call -backend = AttentionBackend(name='flex', label='Flex attention', priority=20, prepare=prepare, terminal=True) +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 +) diff --git a/modules/attention/registry.py b/modules/attention/registry.py index ef33d009e..820dc77f8 100644 --- a/modules/attention/registry.py +++ b/modules/attention/registry.py @@ -26,10 +26,13 @@ class Constraints: 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: diff --git a/test/test-attention-router.py b/test/test-attention-router.py index c51f4b3f5..44a95c150 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -102,9 +102,10 @@ def run_test(cat: str, fn): # ============================================================ # devices.set_sdpa_params applied the hijacks in this order; each wrapped the previous, so the -# last applied was tried first. Dynamic and flex replaced the chain end instead of wrapping it. +# 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', 'Flex attention'} +OLD_TERMINALS = {'Dynamic attention'} OLD_NAMES = { 'Dynamic attention': 'dynamic', 'Flex attention': 'flex', @@ -117,11 +118,13 @@ OLD_NAMES = { CHOICES = OLD_ORDER TRITON_PLATFORMS = {'rocm', 'zluda'} -OLD_GATES = { +# 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, } @@ -194,12 +197,13 @@ def test_plan_matches_stacking_oracle(): 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 in itertools.product(('cpu', 'meta'), ('cpu', 'meta'), (torch.float16, torch.float32), (1, 8), lengths, lengths, (40, 64, 96, 128, 256), (False, True)): - q = shaped((1, heads, q_len, head_dim), dtype, q_device) - k = shaped((1, heads, k_len, head_dim), dtype, kv_device) - v = shaped((1, heads, k_len, head_dim), dtype, kv_device) - m = shaped((1, 1, q_len, k_len), torch.bool, q_device) if masked else None - for name, gate in OLD_GATES.items(): + 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}' @@ -208,11 +212,10 @@ def test_gates_match_transcribed_predicates(): return True -def test_terminals_carry_no_gate(): - for name in ('dynamic', 'flex'): - backend = attention.registry.backends[name] - assert backend.terminal, name - assert backend.constraints == attention.Constraints(), name +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 @@ -312,7 +315,7 @@ def run_all(): for fn in [ test_plan_matches_stacking_oracle, test_gates_match_transcribed_predicates, - test_terminals_carry_no_gate, + test_only_dynamic_is_terminal, test_choices_match_backends, test_router_dispatch_prefers_priority_then_terminal_then_original, test_prepare_failure_skips_backend, From 9d1d7c839afb1e1c10b6e314149a430213677155 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 22:06:36 +0100 Subject: [PATCH 04/12] fix(processing): re-apply attention when the overrides change The generate-time gate compared the stored processor name against the sdp_overrides list, which can never be equal, so the check reduced to the processor name alone and a changed override set was never applied until the next model load. set_diffusers_attention now stamps the override set it applied beside the processor name, the gate compares both, and pipe switches carry the new attribute with the old one. --- modules/attention/dispatcher.py | 1 + modules/processing_diffusers.py | 5 +++-- modules/sd_models.py | 3 +++ 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/modules/attention/dispatcher.py b/modules/attention/dispatcher.py index 4cf232022..e7d31e477 100644 --- a/modules/attention/dispatcher.py +++ b/modules/attention/dispatcher.py @@ -48,6 +48,7 @@ def set_diffusers_attention(pipe, quiet = False): 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 diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 0efbc4904..9e71050e3 100644 --- a/modules/processing_diffusers.py +++ b/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 diff --git a/modules/sd_models.py b/modules/sd_models.py index a21877dc8..6140571e9 100644 --- a/modules/sd_models.py +++ b/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: From 4509b145cd6c9dd7a0ebe48c0a6f181e466c7147 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 22:16:26 +0100 Subject: [PATCH 05/12] feat(attention): generation context A module-level context tells attention consumers what is running: the component role (transformer, text encoder, vae), the index of the denoiser forward about to run, the pass length, and the model. It is opened and closed around process_images, reset per denoising pass beside the callback setup, and advanced by both step sources: the classic callback passes the completed step plus one, the modular pre-forward hook counts forwards. Roles come from the existing text encoder and vae hijacks and the modular phase hooks. The step also lives in a device scalar updated in place, so a compiled reader keeps its graph across steps. --- modules/attention/__init__.py | 6 +-- modules/attention/context.py | 77 +++++++++++++++++++++++++++++++++ modules/modular_load.py | 6 +++ modules/processing.py | 3 ++ modules/processing_args.py | 2 + modules/processing_callbacks.py | 2 + modules/sd_hijack_te.py | 10 +++-- modules/sd_hijack_vae.py | 25 ++++++----- test/test-attention-router.py | 62 ++++++++++++++++++++++++++ 9 files changed, 176 insertions(+), 17 deletions(-) create mode 100644 modules/attention/context.py diff --git a/modules/attention/__init__.py b/modules/attention/__init__.py index 7b923ca6b..359ab81ca 100644 --- a/modules/attention/__init__.py +++ b/modules/attention/__init__.py @@ -1,12 +1,12 @@ -"""Attention backends: one scaled_dot_product_attention router over the registered backends, plus the diffusers-side processor and dispatcher setup.""" +"""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 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 +from modules.attention import backends, context __all__ = [ 'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry', 'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack', - 'backends', + 'backends', 'context', ] diff --git a/modules/attention/context.py b/modules/attention/context.py new file mode 100644 index 000000000..24aa8300a --- /dev/null +++ b/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 diff --git a/modules/modular_load.py b/modules/modular_load.py index b8e8b8cc0..b5c784c39 100644 --- a/modules/modular_load.py +++ b/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): @@ -41,6 +42,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: @@ -52,11 +54,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: @@ -64,6 +68,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 @@ -71,6 +76,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 diff --git a/modules/processing.py b/modules/processing.py index b72c8e151..0e919cd23 100644 --- a/modules/processing.py +++ b/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 @@ -199,6 +200,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 @@ -232,6 +234,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 diff --git a/modules/processing_args.py b/modules/processing_args.py index 94a66153f..115bc6c61 100644 --- a/modules/processing_args.py +++ b/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: diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 48ac6e75b..9b8728695 100644 --- a/modules/processing_callbacks.py +++ b/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: diff --git a/modules/sd_hijack_te.py b/modules/sd_hijack_te.py index 0e9111c09..06826cd33 100644 --- a/modules/sd_hijack_te.py +++ b/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'): diff --git a/modules/sd_hijack_vae.py b/modules/sd_hijack_vae.py index cc6b83919..990495782 100644 --- a/modules/sd_hijack_vae.py +++ b/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') diff --git a/test/test-attention-router.py b/test/test-attention-router.py index 44a95c150..91c56145b 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -15,6 +15,8 @@ Covers: - 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 No running server required. Nothing is moved to the accelerator. @@ -309,6 +311,57 @@ def test_dynamic_backend_pins_pre_dynamic_sdpa(): 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 run_all(): log.warning('=== attention router ===') cat = category('router') @@ -324,6 +377,15 @@ def run_all(): ]: 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('=== Results ===') total_passed = 0 total_failed = 0 From 125ae8e2bf3737b7a1d673d02cf178057f81b599 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 22:23:47 +0100 Subject: [PATCH 06/12] feat(attention): route tracing and chain telemetry SD_ATTN_DEBUG logs each distinct route once: backend, component role, step, shapes, dtype and mask presence. The router takes an optional observer for it, so the clean path carries one pointer check. report() returns the active chain and generation context, and torch_info records the whole chain as one string instead of the last prepared backend. --- modules/attention/__init__.py | 8 ++--- modules/attention/debug.py | 21 +++++++++++++ modules/attention/router.py | 30 +++++++++++++++--- test/test-attention-router.py | 58 +++++++++++++++++++++++++++++++++++ 4 files changed, 109 insertions(+), 8 deletions(-) create mode 100644 modules/attention/debug.py diff --git a/modules/attention/__init__.py b/modules/attention/__init__.py index 359ab81ca..f34edf81e 100644 --- a/modules/attention/__init__.py +++ b/modules/attention/__init__.py @@ -1,12 +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 +from modules.attention.router import Plan, PlanEntry, build_plan, get_plan, install_router, 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 +from modules.attention import backends, context, debug __all__ = [ 'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry', - 'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', + 'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'report', 'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack', - 'backends', 'context', + 'backends', 'context', 'debug', ] diff --git a/modules/attention/debug.py b/modules/attention/debug.py new file mode 100644 index 000000000..0b69ed76b --- /dev/null +++ b/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() diff --git a/modules/attention/router.py b/modules/attention/router.py index f917424cd..46539427f 100644 --- a/modules/attention/router.py +++ b/modules/attention/router.py @@ -1,9 +1,11 @@ """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 @@ -35,7 +37,7 @@ def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registr 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 and owns the torch_info record + 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): @@ -53,21 +55,25 @@ def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registr terminal = entry else: entries.append(entry) - torch_info.set(attention=backend.name) entries.reverse() return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels)) -def make_router(plan: Plan) -> AttentionCall: +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 @@ -81,10 +87,26 @@ def install_router(labels, platform: Platform, original: AttentionCall, reg: Reg """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) - torch.nn.functional.scaled_dot_product_attention = make_router(plan) if (plan.entries or plan.terminal is not None) else original + 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 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}, + } diff --git a/test/test-attention-router.py b/test/test-attention-router.py index 91c56145b..1a69ff757 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -17,6 +17,8 @@ Covers: - 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 No running server required. Nothing is moved to the accelerator. @@ -362,6 +364,53 @@ def test_context_roles_nest_and_stick(): 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 run_all(): log.warning('=== attention router ===') cat = category('router') @@ -386,6 +435,15 @@ def run_all(): ]: 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, + ]: + run_test(cat, fn) + log.warning('=== Results ===') total_passed = 0 total_failed = 0 From b57b8c6a64c3b71c2f936ef1766e0c21f56b65de Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 22:31:14 +0100 Subject: [PATCH 07/12] perf(attention): capture backend options once per chain The sdnq backend read six settings on every call; it now captures them when the chain is built. Each backend declares the settings its call captures, and webui registers one onchange over those names plus the override set and the torch kernel flags, so a change rebuilds the chain between jobs. When a compiled model is resident the rebuild also resets dynamo, since its graphs hold the previous router. --- modules/attention/__init__.py | 4 ++-- modules/attention/backends/sdnq.py | 22 +++++++++++----------- modules/attention/registry.py | 4 ++++ modules/attention/router.py | 16 ++++++++++++++++ test/test-attention-router.py | 13 +++++++++++++ webui.py | 3 +++ 6 files changed, 49 insertions(+), 13 deletions(-) diff --git a/modules/attention/__init__.py b/modules/attention/__init__.py index f34edf81e..78994780c 100644 --- a/modules/attention/__init__.py +++ b/modules/attention/__init__.py @@ -1,12 +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, report +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', 'report', + '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', ] diff --git a/modules/attention/backends/sdnq.py b/modules/attention/backends/sdnq.py index c0c49da67..61fc14bbf 100644 --- a/modules/attention/backends/sdnq.py +++ b/modules/attention/backends/sdnq.py @@ -5,24 +5,24 @@ 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, + '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, - 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, - use_fp16_accum=shared.opts.sdnq_attention_use_fp16_accum, - ) + 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={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}') + log.debug(f'Torch attention: type="SDNQ attention" matmul={options["matmul_dtype"]}:{options["pv_matmul_dtype"]} smooth={options["smooth_k"]} hadamard={options["use_hadamard"]} 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_use_fp16_accum'), ) diff --git a/modules/attention/registry.py b/modules/attention/registry.py index 820dc77f8..23ca24d1a 100644 --- a/modules/attention/registry.py +++ b/modules/attention/registry.py @@ -63,6 +63,7 @@ class AttentionBackend: 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 @@ -90,5 +91,8 @@ class Registry: 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() diff --git a/modules/attention/router.py b/modules/attention/router.py index 46539427f..c9ae07a48 100644 --- a/modules/attention/router.py +++ b/modules/attention/router.py @@ -100,6 +100,22 @@ 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 diff --git a/test/test-attention-router.py b/test/test-attention-router.py index 1a69ff757..7438e36ea 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -411,6 +411,18 @@ def test_debug_observe_logs_each_route_once(): 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') @@ -441,6 +453,7 @@ def run_all(): 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, ]: run_test(cat, fn) diff --git a/webui.py b/webui.py index 911d35ccf..378cfc184 100644 --- a/webui.py +++ b/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") From e0c3e00af2938eb88950a3923c174acd85511869 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 22:40:09 +0100 Subject: [PATCH 08/12] chore(settings): remove dead attention options xformers_options had no reader, and Sub-quadratic has not been an attention choice for a long time, so the hypertile branches keyed on it never ran. Configs that still store xformers_options load without the unknown-setting warning. --- modules/options_handler.py | 2 +- modules/sd_hijack_hypertile.py | 7 ------- modules/ui_definitions.py | 1 - 3 files changed, 1 insertion(+), 9 deletions(-) diff --git a/modules/options_handler.py b/modules/options_handler.py index 9006673f6..d0bb8d454 100644 --- a/modules/options_handler.py +++ b/modules/options_handler.py @@ -17,7 +17,7 @@ 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'] secrets_pattern = ['_version', '_token', '_key', '_secret', '_password'] diff --git a/modules/sd_hijack_hypertile.py b/modules/sd_hijack_hypertile.py index eead36266..0cfe4e220 100644 --- a/modules/sd_hijack_hypertile.py +++ b/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) diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index f77ec04d1..552ca6d97 100644 --- a/modules/ui_definitions.py +++ b/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}), From b0f004bf83c882fcfb16033fb36b7cf208ecfa53 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sat, 22 Aug 2026 23:05:07 +0100 Subject: [PATCH 09/12] test(attention): cover the sdpa escape hatches bypass_sdpa_hijacks and llm_context restore the pinned original over the router and put the router back, including when the body raises. Captioners and detailers run inside them, so the router must be removable. --- test/test-attention-router.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/test/test-attention-router.py b/test/test-attention-router.py index 7438e36ea..70154f140 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -19,6 +19,8 @@ Covers: 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. @@ -411,6 +413,35 @@ def test_debug_observe_logs_each_route_once(): 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() @@ -454,6 +485,7 @@ def run_all(): test_install_router_records_the_chain, test_debug_observe_logs_each_route_once, test_reapply_options_cover_declared_backend_options, + test_escape_hatch_bypasses_the_router, ]: run_test(cat, fn) From 26d922371db551655b0f3cb1e36f44e413e92aec Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 24 Aug 2026 19:00:10 +0100 Subject: [PATCH 10/12] fix(attention): honor the disabled choice for attention slicing attention_slicing holds one of Default, Enabled or Disabled, so testing the string for truth sent Disabled down the enable branch and left the disable call unreachable, while the log line below it reported the choice rather than the action taken. --- modules/attention/dispatcher.py | 2 +- test/test-attention-router.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/modules/attention/dispatcher.py b/modules/attention/dispatcher.py index e7d31e477..82cbc323c 100644 --- a/modules/attention/dispatcher.py +++ b/modules/attention/dispatcher.py @@ -41,7 +41,7 @@ def set_diffusers_attention(pipe, quiet = False): 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: + if shared.opts.attention_slicing == "Enabled": pipe.enable_attention_slicing() else: pipe.disable_attention_slicing() diff --git a/test/test-attention-router.py b/test/test-attention-router.py index 70154f140..245467326 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -413,6 +413,36 @@ def test_debug_observe_logs_each_route_once(): 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_escape_hatch_bypasses_the_router(): from modules import devices saved_sdpa = torch.nn.functional.scaled_dot_product_attention @@ -485,6 +515,7 @@ def run_all(): 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_escape_hatch_bypasses_the_router, ]: run_test(cat, fn) From 4d6f2b65c8ce0bdcbda847382e90b1f2bb982752 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 24 Aug 2026 19:59:41 +0100 Subject: [PATCH 11/12] chore(settings): remove the bmm attention methods Batch matrix-matrix and Dynamic Attention BMM applied a legacy Attention processor to pipe.unet, which a diffusion transformer does not have, so they served unet models alone and said nothing elsewhere. The choices, the processor and its slice helper are removed, an unrecognized method now warns rather than selecting nothing, and a stored value is rewritten to Scaled-Dot-Product on load. --- modules/attention/dispatcher.py | 24 +---- modules/options_handler.py | 16 +++ modules/sd_hijack_dynamic_atten.py | 157 ----------------------------- modules/shared_items.py | 2 - test/test-attention-router.py | 22 ++++ ui/locale/locale_en.json | 2 - 6 files changed, 40 insertions(+), 183 deletions(-) diff --git a/modules/attention/dispatcher.py b/modules/attention/dispatcher.py index 82cbc323c..19d58e99b 100644 --- a/modules/attention/dispatcher.py +++ b/modules/attention/dispatcher.py @@ -5,40 +5,20 @@ from installer import install, torch_info def set_diffusers_attention(pipe, quiet = False): from modules import shared, devices - 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") + 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": diff --git a/modules/options_handler.py b/modules/options_handler.py index d0bb8d454..2400a2069 100644 --- a/modules/options_handler.py +++ b/modules/options_handler.py @@ -18,9 +18,22 @@ if TYPE_CHECKING: cmd_opts = cmd_args.parse_args() 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) diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index 0c0c4771d..cf932496d 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/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 diff --git a/modules/shared_items.py b/modules/shared_items.py index e7d8dab8e..ad37125db 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -142,8 +142,6 @@ def list_crossattention(): "Disabled", "Scaled-Dot-Product", "xFormers", - "Batch matrix-matrix", - "Dynamic Attention BMM" ] diff --git a/test/test-attention-router.py b/test/test-attention-router.py index 245467326..d100e6168 100644 --- a/test/test-attention-router.py +++ b/test/test-attention-router.py @@ -443,6 +443,27 @@ def test_attention_slicing_follows_the_choice(): 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 @@ -516,6 +537,7 @@ def run_all(): 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) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index 089b239b6..d9260dc01 100644 --- a/ui/locale/locale_en.json +++ b/ui/locale/locale_en.json @@ -189,7 +189,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"}, @@ -397,7 +396,6 @@ {"id":"","label":"Dequantize using torch.compile","localized":"","hint":"Compiles the dequantization step with torch.compile for faster inference. Requires Triton.

Changing this needs a full restart to take effect.

Enabled by default when Triton is available.","reload":"server","ui":"settings_quantization"}, {"id":"","label":"Dequantize using full precision","localized":"","hint":"Uses FP32 for the dequantization step for better numerical accuracy, at a small speed cost.

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"}, From fcadb4fa9b0fdb9f12342f0bd15429eb905caa65 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 24 Aug 2026 21:59:58 +0100 Subject: [PATCH 12/12] docs(settings): hint the attention settings Everything under Cross Attention, SDNQ Attention and Attention Dispatcher shipped with no hint, which left the layering invisible: the SDP kernel boxes are candidates torch chooses from per call rather than a selection, and the Flash box is torch's own build of the kernel rather than the flash-attn package that the Flash attention override installs. - cover attention method, sdp kernels, sdp overrides and attention slicing - give the dynamic attention rates their unit and the estimate they compare against - cover the sdnq attention kernel settings, including the head dimension clamp that the hadamard group size resolves through - cover the diffusers attention dispatcher and the backend names it takes - record the constraints each override serves, leaving the throughput comparison to the workload rather than naming a winner --- ui/locale/locale_en.json | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/ui/locale/locale_en.json b/ui/locale/locale_en.json index d9260dc01..3019dfb1f 100644 --- a/ui/locale/locale_en.json +++ b/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.

Scaled-Dot-Product uses torch attention and is the path the rest of this section builds on; SDP kernels and SDP overrides apply to it and to nothing else.
xFormers uses the xFormers library on the modules that accept it.
Disabled leaves whatever the pipeline was built with.

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 Attention method is active.

Default leaves the pipeline's own choice untouched, which for most models means off.
Enabled and Disabled override that choice.

Unrelated to Dynamic attention, which slices by a memory estimate; this one is diffusers' fixed slicing.

Default Default.","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.

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.

Accepted names come from the diffusers build in use and currently include native, flash, flash_varlen, flex, sage, sage_varlen and xformers, alongside underscore-prefixed variants that pin one specific kernel.
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.
Names ending in _hub fetch the kernel from the Hugging Face Hub on first use, which needs a download and the kernels package.

Empty by default.","ui":"settings_cuda"} ], "b": [ {"id":"","label":"Batch","localized":"","hint":"Batch processing settings","ui":"img2img"}, @@ -233,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.

Three layers apply in order. Attention method picks the attention processor the pipeline is loaded with. SDP kernels limits which kernels torch may choose from inside its own attention. SDP overrides replaces torch attention with a different implementation for the calls that implementation accepts.
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,8 +401,8 @@ {"id":"","label":"Dequantize using full precision","localized":"","hint":"Uses FP32 for the dequantization step for better numerical accuracy, at a small speed cost.

Enabled by default.","reload":"model","ui":"settings_quantization"}, {"id":"","label":"Disabled","localized":"","hint":"","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 Dynamic attention starts slicing. Smaller slices hold the peak lower and add more per-slice overhead.
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.

Applies while Dynamic attention is enabled in SDP overrides.
Default 0.5.","ui":"settings_cuda"}, + {"id":"","label":"Dynamic Attention trigger rate","localized":"","hint":"Estimated attention matrix size in GB above which Dynamic attention begins slicing. Below it the call runs in one pass.
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.

Applies while Dynamic attention is enabled in SDP overrides.
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"}, @@ -1452,7 +1456,16 @@ {"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 Flash, dropping to Memory for calls flash cannot serve such as those carrying an arbitrary attention mask, and to Math when neither fits. Clearing a box removes a candidate; it never pins the remaining one to every call.

Flash is torch's own build of the FlashAttention kernel. It is not the same thing as the Flash attention entry in SDP overrides, which calls the separately installed flash-attn package and bypasses torch entirely.
Memory is the memory-efficient kernel, which accepts arbitrary masks that flash does not.
Math 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.

Applies while Attention method is Scaled-Dot-Product, and continues to govern the calls that an enabled override declines.

All three by default. ZLUDA starts with Math 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.

Flash attention 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.
Sage attention computes attention with quantized matmuls, for head dimensions of 64, 96 or 128 and no attention mask.
SDNQ attention is SD.Next's own quantized Triton kernel, configured in the section below. It takes attention masks, which the other quantized backends do not.
Flex attention uses torch's compiled flex_attention.
Dynamic attention slices attention to fit available memory and serves every call the others decline, standing in for the torch fallback.
Triton Flash attention is a Triton implementation for ROCm and ZLUDA, listed only on those backends.

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.

None by default on CUDA. ZLUDA, CPU and MPS start with Dynamic attention, as do ROCm GPUs older than RDNA3.","ui":"settings_cuda"}, + {"id":"","label":"SDNQ Attention","localized":"","hint":"Settings for the SDNQ attention entry in SDP overrides. They do nothing until that override is enabled.

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.
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.
Softmax ignores a constant shift applied to every score in a row, so removing that offset changes the quantization error and not the attention.

Costs one mean and one subtraction per call.
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.
With SDNQ Attention PV MatMul type also set, the values are rotated as well and the output is rotated back.

Costs a rotation pass on every attention call, so it is worth enabling where a model shows quantization artifacts without it.
Idle while SDNQ Attention MatMul type is disabled, since nothing is quantized then.

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.
Operands are pre-scaled to keep products inside the fp16 range, which covers ordinary activations with less headroom than fp32 leaves.

Reaches the parts of the kernel that run in floating point. An int8 matmul accumulates in int32 and is unaffected, so at the default SDNQ Attention MatMul type this applies to the probability-value matmul alone.

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.

enabled selects int8, and int8 and uint8 reach the same kernel.
float16 and float8_e4m3fn take the floating point path. fp8 needs a GPU with fp8 tensor cores and fails on hardware without them rather than falling back.
disabled leaves queries and keys in the model's own precision, which also idles SDNQ Attention use Smooth K and SDNQ Attention use Hadamard.

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 SDNQ Attention MatMul type.

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.
disabled keeps this matmul in the model's own precision.

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.

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.
Applies while SDNQ Attention use Hadamard is enabled.

Default 256.","ui":"settings_cuda"} ], "t": [ {"id":"txt2img_nav","label":"T2I","localized":"","hint":"Create image from text
Legacy interface that mimics original text-to-image interface and behavior"},