feat(xyz): add attention and sparse attention axes

Fourteen axes drive the attention subsystem from the grid: the diffusers
method, the sdp override chain, the dispatcher kernel, five sdnq attention
knobs and six sparse settings. An axis writes shared.opts.data and rebuilds
the chain itself, since the onchange for these settings runs through the
queue lock a grid cell already holds. SharedSettingsStackHelper saves and
restores every attention setting around the grid, keys an axis introduced
included, and an axis for a setting a backend owns warns when that backend
is not in the active chain.

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, 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.dispatcher import set_diffusers_attention, set_attention_dispatcher, list_dispatcher_backends, hijack_kernels, get_kernel_hijack, get_hf_api_hijack
 from modules.attention import backends, context, debug
 
 __all__ = [
     'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry',
     'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'reapply', 'reapply_options', 'report',
-    'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack',
+    'set_diffusers_attention', 'set_attention_dispatcher', 'list_dispatcher_backends', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack',
     'backends', 'context', 'debug',
 ]

modules/attention/dispatcher.py

@@ -94,3 +94,13 @@ def set_attention_dispatcher(pipe):
         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}')
+
+
+def list_dispatcher_backends() -> list:
+    """The kernels diffusers can dispatch attention to, for anything that offers hf_attention as a choice."""
+    try:
+        from diffusers.models import attention_dispatch as a
+        return sorted(b.value for b in a._AttentionBackendRegistry.list_backends()) # pylint: disable=protected-access
+    except Exception as e:
+        log.error(f'Attention dispatcher: {e}')
+        return []

scripts/xyz/xyz_grid_classes.py

@@ -3,6 +3,12 @@ from scripts.xyz.xyz_grid_shared import ( # pylint: disable=no-name-in-module, u
     apply_task_arg,
     apply_task_args,
     apply_setting,
+    apply_attention,
+    apply_attention_overrides,
+    apply_attention_dispatcher,
+    list_sdp_overrides,
+    save_attention,
+    restore_attention,
     apply_prompt_primary,
     apply_prompt_refine,
     apply_prompt_detailer,
@@ -40,7 +46,7 @@ from scripts.xyz.xyz_grid_shared import ( # pylint: disable=no-name-in-module, u
     format_nothing,
     str_permutations,
  )
-from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet
+from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet, attention
 from modules.control.units import controlnet, t2iadapter
 from modules.control import processor
 
@@ -107,6 +113,7 @@ class SharedSettingsStackHelper():
     disable_apply_metadata = None
     disable_apply_params = None
     sdnq_quant_mode = None
+    attention_settings = None
 
     def __enter__(self):
         # Save overridden settings so they can be restored later
@@ -140,6 +147,7 @@ class SharedSettingsStackHelper():
         self.disable_apply_metadata = shared.opts.disable_apply_metadata
         self.disable_apply_params = shared.opts.disable_apply_params
         self.sdnq_quant_mode = shared.opts.sdnq_quantize_weights_mode
+        self.attention_settings = save_attention()
         shared.opts.data["disable_apply_metadata"] = []
         shared.opts.data["disable_apply_params"] = ''
 
@@ -188,6 +196,7 @@ class SharedSettingsStackHelper():
         if self.sdnq_quant_mode != shared.opts.sdnq_quantize_weights_mode:
             shared.opts.data["sdnq_quantize_weights_mode"] = self.sdnq_quant_mode
             sd_models.reload_model_weights(op='model')
+        restore_attention(self.attention_settings)
 
 
 axis_options = [
@@ -250,6 +259,20 @@ axis_options = [
     AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")),
     AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared_items.sdnq_quant_modes)),
     AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared_items.sdnq_quant_modes)),
+    AxisOption("[Attention] Method", str, apply_setting('cross_attention_optimization'), cost=0.2, choices=shared_items.list_crossattention),
+    AxisOption("[Attention] SDP override", str, apply_attention_overrides, cost=0.2, choices=list_sdp_overrides),
+    AxisOption("[Attention] Dispatcher", str, apply_attention_dispatcher, cost=0.2, choices=lambda: ['None'] + attention.list_dispatcher_backends()),
+    AxisOption("[Attention] SDNQ matmul", str, apply_attention('sdnq_attention_matmul_type'), cost=0.2, choices=lambda: list(shared_items.sdnq_matmul_modes)),
+    AxisOption("[Attention] SDNQ PV matmul", str, apply_attention('sdnq_attention_pv_matmul_type'), cost=0.2, choices=lambda: list(shared_items.sdnq_matmul_modes)),
+    AxisOption("[Attention] SDNQ smooth K", str, apply_attention('sdnq_attention_smooth_k'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Attention] SDNQ hadamard", str, apply_attention('sdnq_attention_use_hadamard'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Attention] SDNQ fp16 accumulation", str, apply_attention('sdnq_attention_use_fp16_accum'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Sparse] Enabled", str, apply_attention('sparse_attention_enabled'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Sparse] KV budget", int, apply_attention('sparse_attention_budget'), cost=0.2),
+    AxisOption("[Sparse] Minimum sequence", int, apply_attention('sparse_attention_min_tokens'), cost=0.2),
+    AxisOption("[Sparse] Dense steps", int, apply_attention('sparse_attention_schedule_steps'), cost=0.2),
+    AxisOption("[Sparse] Dense step bonus", int, apply_attention('sparse_attention_schedule_bump'), cost=0.2),
+    AxisOption("[Sparse] Shared heads", str, apply_attention('sparse_attention_head_shared'), cost=0.2, choices=lambda: ['False', 'True']),
     AxisOption("[HDR] Mode", int, apply_field("hdr_mode")),
     AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")),
     AxisOption("[HDR] Color", float, apply_field("hdr_color")),

scripts/xyz/xyz_grid_shared.py

@@ -81,6 +81,75 @@ def apply_setting(field):
     return fun
 
 
+def attention_options() -> list:
+    """Attention settings an axis can change; the stack helper restores exactly this set."""
+    from modules import attention
+    return ['cross_attention_optimization', 'hf_attention', *attention.reapply_options()]
+
+
+def list_sdp_overrides() -> list:
+    item = shared.opts.data_labels.get('sdp_overrides', None)
+    args = item.component_args if item is not None else None
+    args = args() if callable(args) else args
+    return ['None'] + list((args or {}).get('choices', None) or [])
+
+
+def apply_attention(field):
+    def fun(p, x, xs):
+        from modules import attention
+        apply_setting(field)(p, x, xs)
+        attention.reapply() # backends read their settings when the chain is built, so a write on its own changes nothing
+        owner = next((backend for backend in attention.registry.backends.values() if field in backend.options), None)
+        plan = attention.get_plan()
+        if owner is not None and plan is not None and owner.name not in plan.chain():
+            log.warning(f'XYZ grid apply attention: {field} is read by "{owner.label}" which is not in the active chain={plan.chain()}')
+    return fun
+
+
+def apply_attention_overrides(p, x, xs):
+    from modules import attention
+    labels = [label.strip() for label in str(x).split('+') if len(label.strip()) > 0 and label.strip().lower() != 'none']
+    unknown = [label for label in labels if attention.registry.by_label(label) is None]
+    if len(unknown) > 0:
+        log.warning(f'XYZ grid apply attention: unknown overrides={unknown} available={attention.registry.labels()}')
+    shared.opts.data['sdp_overrides'] = labels
+    attention.reapply()
+    log.debug(f'XYZ grid apply attention: overrides={labels}')
+
+
+def apply_attention_dispatcher(p, x, xs):
+    from modules import attention
+    value = '' if str(x).strip().lower() in ['none', 'default'] else str(x).strip()
+    shared.opts.data['hf_attention'] = value
+    if shared.sd_loaded:
+        attention.set_attention_dispatcher(shared.sd_model)
+    log.debug(f'XYZ grid apply attention: dispatcher="{value}"')
+
+
+def save_attention() -> dict:
+    return {field: shared.opts.data[field] for field in attention_options() if field in shared.opts.data}
+
+
+def restore_attention(saved: dict):
+    """Put back whatever an attention axis changed, keys it introduced included, then rebuild what reads them."""
+    from modules import attention
+    changed = []
+    for field in attention_options():
+        if (field in saved) == (field in shared.opts.data) and saved.get(field, None) == shared.opts.data.get(field, None):
+            continue
+        changed.append(field)
+        if field in saved:
+            shared.opts.data[field] = saved[field]
+        else:
+            shared.opts.data.pop(field, None)
+    if len(changed) == 0:
+        return
+    attention.reapply()
+    if 'hf_attention' in changed and shared.sd_loaded:
+        attention.set_attention_dispatcher(shared.sd_model)
+    log.debug(f'XYZ grid restore attention: {changed}')
+
+
 def apply_seed(p, x, xs):
     p.seed = x
     p.all_seeds = None

test/test-xyz-attention.py

@@ -0,0 +1,252 @@
+#!/usr/bin/env python
+"""
+Offline unit tests for the attention and sparse axes of the xyz grid.
+
+Covers:
+
+- every [Attention] and [Sparse] axis resolves its choices and targets a registered option
+- applying an axis and then leaving the grid restores shared.opts.data exactly, keys the axis
+  introduced included, so a grid never leaks its last cell into the session
+- the axes backed by boolean options take the string the dropdown hands them
+- the sdp override axis turns a label, a plus joined pair or None into the option's list
+- an override axis apply reaches the router: the rebuilt chain contains the backend it named
+- the restore set covers every setting the sparse stage reads
+
+No running server required. Nothing is moved to the accelerator, and no axis value that would
+install a package is used.
+
+Usage:
+    python test/test-xyz-attention.py
+"""
+
+import os
+import sys
+
+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, shared                             # pylint: disable=wrong-import-position
+from modules.attention.sparse import stage as sparse_stage        # pylint: disable=wrong-import-position
+from scripts.xyz import xyz_grid_shared as xyz                    # pylint: disable=wrong-import-position
+from scripts.xyz.xyz_grid_classes import axis_options             # pylint: disable=wrong-import-position
+
+
+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):
+    try:
+        ok = fn()
+        record(cat, ok is not False, fn.__name__)
+    except AssertionError as e:
+        record(cat, False, fn.__name__, str(e))
+    except Exception as e:  # pylint: disable=broad-except
+        record(cat, False, fn.__name__, f'exception: {e}')
+        import traceback
+        traceback.print_exc()
+
+
+def attention_axes():
+    return [axis for axis in axis_options if axis.label.startswith('[Attention]') or axis.label.startswith('[Sparse]')]
+
+
+def sample_value(axis):
+    """A value the axis accepts that is not the current one, avoiding backends whose prepare installs a package."""
+    if axis.choices is None:
+        return int(shared.opts.get(option_of(axis) or 'sparse_attention_budget') or 0) + 5
+    choices = [choice for choice in axis.choices() if choice not in ['Sage attention', 'Flash attention', 'Triton Flash attention']]
+    current = str(shared.opts.get(option_of(axis)) if option_of(axis) else '')
+    return next((choice for choice in choices if str(choice) != current), choices[0])
+
+
+def option_of(axis):
+    """The option an axis writes, read back from the closure the axis factory built."""
+    closure = getattr(axis.apply, '__closure__', None) or ()
+    for cell in closure:
+        if isinstance(cell.cell_contents, str) and cell.cell_contents in shared.opts.data_labels:
+            return cell.cell_contents
+    return {'[Attention] SDP override': 'sdp_overrides', '[Attention] Dispatcher': 'hf_attention'}.get(axis.label, None)
+
+
+def bool_axes():
+    return [axis for axis in attention_axes() if isinstance(shared.opts.get(option_of(axis) or ''), bool)]
+
+
+# ============================================================
+# Tests
+# ============================================================
+
+def test_axes_are_registered():
+    axes = attention_axes()
+    assert len(axes) >= 10, f'only {len(axes)} attention axes'
+    for axis in axes:
+        option = option_of(axis)
+        assert option is not None, f'{axis.label} names no option'
+        assert option in shared.opts.data_labels, f'{axis.label} targets unknown option {option}'
+    return True
+
+
+def test_axis_choices_resolve():
+    for axis in attention_axes():
+        if axis.choices is None:
+            assert axis.type is int, f'{axis.label} has no choices and is not numeric'
+            continue
+        choices = axis.choices()
+        assert isinstance(choices, list) and len(choices) > 0, f'{axis.label} resolved {choices}'
+    return True
+
+
+def test_axes_write_and_restore_exactly():
+    for axis in attention_axes():
+        before = dict(shared.opts.data)
+        saved = xyz.save_attention()
+        axis.apply(None, sample_value(axis), [])
+        assert dict(shared.opts.data) != before, f'{axis.label} wrote nothing'
+        xyz.restore_attention(saved)
+        after = dict(shared.opts.data)
+        leaked = {key for key in set(before) | set(after) if before.get(key, '<absent>') != after.get(key, '<absent>')}
+        assert not leaked, f'{axis.label} leaked {sorted(leaked)}'
+    return True
+
+
+def test_bool_axes_coerce_the_string_the_dropdown_sends():
+    axes = bool_axes()
+    assert len(axes) >= 3, f'only {len(axes)} boolean axes'
+    for axis in axes:
+        option = option_of(axis)
+        saved = xyz.save_attention()
+        axis.apply(None, 'True', [])
+        assert shared.opts.data[option] is True, f'{axis.label} took "True" as {shared.opts.data[option]!r}'
+        axis.apply(None, 'False', [])
+        assert shared.opts.data[option] is False, f'{axis.label} took "False" as {shared.opts.data[option]!r}'
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_override_axis_parses_labels():
+    saved = xyz.save_attention()
+    try:
+        xyz.apply_attention_overrides(None, 'None', [])
+        assert shared.opts.data['sdp_overrides'] == [], shared.opts.data['sdp_overrides']
+        xyz.apply_attention_overrides(None, 'Flex attention', [])
+        assert shared.opts.data['sdp_overrides'] == ['Flex attention'], shared.opts.data['sdp_overrides']
+        xyz.apply_attention_overrides(None, 'Flex attention+SDNQ attention', [])
+        assert shared.opts.data['sdp_overrides'] == ['Flex attention', 'SDNQ attention'], shared.opts.data['sdp_overrides']
+    finally:
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_override_axis_rebuilds_the_chain():
+    saved = xyz.save_attention()
+    try:
+        xyz.apply_attention_overrides(None, 'Flex attention', [])
+        assert 'flex' in attention.get_plan().chain(), attention.get_plan().chain()
+        xyz.apply_attention_overrides(None, 'None', [])
+        assert 'flex' not in attention.get_plan().chain(), attention.get_plan().chain()
+    finally:
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_dispatcher_axis_clears_on_none():
+    saved = xyz.save_attention()
+    try:
+        xyz.apply_attention_dispatcher(None, 'native', [])
+        assert shared.opts.data['hf_attention'] == 'native', shared.opts.data['hf_attention']
+        xyz.apply_attention_dispatcher(None, 'None', [])
+        assert shared.opts.data['hf_attention'] == '', repr(shared.opts.data['hf_attention'])
+    finally:
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_restore_set_covers_every_attention_setting():
+    covered = set(xyz.attention_options())
+    missing = set(sparse_stage.OPTION_NAMES) - covered
+    assert not missing, f'sparse settings outside the restore set: {sorted(missing)}'
+    for backend in attention.registry.backends.values():
+        outside = set(backend.options) - covered
+        assert not outside, f'{backend.label} settings outside the restore set: {sorted(outside)}'
+    return True
+
+
+def run_all():
+    log.warning('=== xyz attention axes ===')
+    cat = category('axes')
+    for fn in [
+        test_axes_are_registered,
+        test_axis_choices_resolve,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== apply and restore ===')
+    cat = category('apply')
+    for fn in [
+        test_axes_write_and_restore_exactly,
+        test_bool_axes_coerce_the_string_the_dropdown_sends,
+        test_override_axis_parses_labels,
+        test_override_axis_rebuilds_the_chain,
+        test_dispatcher_axis_clears_on_none,
+        test_restore_set_covers_every_attention_setting,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== Results ===')
+    total_passed = 0
+    total_failed = 0
+    for cat_name, info in results.items():
+        status = 'PASS' if info['failed'] == 0 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)