feat(attention): sparse attention settings and router stage

Sparse attention is a stage over the chain rather than a member of it:
one switch, and the router hands the selection to whichever active
backend advertises that it consumes a block mask, currently flex. A
backend declares that through a capability set, so the quantized kernel
joins later without touching the router.

The stage gates on the component role, self attention, a minimum
sequence length defaulting to the measured 8192 token crossover, and the
absence of a token mask or causal flag, which flex cannot combine with a
block only mask. Budgets follow a precomputed per step schedule with at
most two distinct values. Enabling the feature with no capable backend
in the chain warns and leaves attention dense rather than doing nothing
quietly.

The modular pre-forward hook now receives kwargs and publishes whatever
token layout the pipeline passes by name, so a packed sequence gets its
conditioning pinned without any model specific code. Without a layout
the whole sequence is sparsified and that is logged once per length.

modules/attention/backends/flex.py

@@ -9,7 +9,10 @@ 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, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=None): # pylint: disable=unused-argument
+        if selection is not None:
+            from modules.attention.sparse import flex as sparse_flex
+            return sparse_flex.attend(query, key, value, selection, scale=scale, enable_gqa=enable_gqa)
         score_mod = None
         block_mask = None
         if attn_mask is not None:
@@ -36,4 +39,5 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument
 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
+    caps=frozenset({'block_mask'}),
 )

modules/attention/context.py

@@ -13,6 +13,7 @@ class GenerationContext:
     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
+    layout: object | None = None # TokenLayout published by whoever knows the packing, None until something does
 
 
 current = GenerationContext()
@@ -30,6 +31,7 @@ def begin(pipe, steps: int = 0) -> None:
     from modules import devices
     current.active = True
     current.role = 'transformer'
+    current.layout = None
     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:
@@ -56,10 +58,16 @@ def tick(step: int | None = None) -> None:
     current.forwards = current.step + 1
 
 
+def set_layout(layout) -> None:
+    """Publish what the packed sequence holds; callers that know the packing set this per forward."""
+    current.layout = layout
+
+
 def end() -> None:
     current.active = False
     current.role = None
     current.model_key = None
+    current.layout = None
     new_pass(0)
 
 

modules/attention/registry.py

@@ -64,6 +64,7 @@ class AttentionBackend:
     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
+    caps: frozenset[str] = frozenset() # what the call can consume beyond plain sdpa arguments, currently 'block_mask'
 
     def available_on(self, platform: Platform) -> bool:
         return self.platforms is None or platform.backend in self.platforms
@@ -94,5 +95,8 @@ class Registry:
     def options(self) -> list[str]:
         return sorted({name for backend in self.backends.values() for name in backend.options})
 
+    def with_cap(self, cap: str) -> list[AttentionBackend]:
+        return [backend for backend in self.ordered() if cap in backend.caps]
+
 
 registry = Registry()

modules/attention/router.py

@@ -59,7 +59,7 @@ def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registr
     return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels))
 
 
-def make_router(plan: Plan, observer: Callable | None = None) -> AttentionCall:
+def make_router(plan: Plan, observer: Callable | None = None, stage: 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'
@@ -69,10 +69,16 @@ def make_router(plan: Plan, observer: Callable | None = None) -> AttentionCall:
     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 stage is not None and 'block_mask' in entry.backend.caps:
+                    selection = stage(query, key, value, attn_mask, is_causal)
+                    if selection is not None:
+                        if observer is not None:
+                            observer(f'{entry.backend.name}+sparse', query, key, attn_mask)
+                        return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=selection)
                 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:
+        if observer is not None: # pylint: disable=duplicate-code
             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)
@@ -89,21 +95,43 @@ def install_router(labels, platform: Platform, original: AttentionCall, reg: Reg
     plan = build_plan(labels, platform, original, reg)
     debug.reset()
     observer = debug.observe if debug.enabled else None
-    torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer) if (plan.entries or plan.terminal is not None) else original
+    stage = build_sparse_stage(plan)
+    torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer, stage) 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}')
+    log.debug(f'Torch attention: chain={">".join(plan.chain())} overrides={list(labels)} backend={platform.backend} sparse={stage is not None}')
     return plan
 
 
+def build_sparse_stage(plan: Plan):
+    """Sparse attention is a stage over the chain rather than a chain member, so it needs a backend in the chain that consumes a block mask."""
+    from modules.attention.sparse import stage as sparse_stage
+    try:
+        options = sparse_stage.read_options()
+    except Exception:
+        return None
+    if not options.enabled:
+        return None
+    capable = [entry.backend.name for entry in plan.entries if 'block_mask' in entry.backend.caps]
+    if not capable:
+        names = [backend.label for backend in default_registry.with_cap('block_mask')]
+        log.warning(f'Sparse attention: enabled but no active backend consumes a block mask, enable one of {names} in sdp overrides; attention stays dense')
+        return None
+    built = sparse_stage.make_stage(options)
+    if built is not None:
+        log.info(f'Sparse attention: backend={capable[0]} budget={options.budget:.0%} gate={options.gate} schedule={options.schedule_steps}x+{options.schedule_bump:.0%}')
+    return built
+
+
 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."""
+    """Settings whose change rebuilds the chain: the override set, the torch kernel flags, every option a backend captures, and the sparse stage."""
+    from modules.attention.sparse import stage as sparse_stage
     reg = reg if reg is not None else default_registry
-    return ['sdp_options', 'sdp_overrides', *reg.options()]
+    return ['sdp_options', 'sdp_overrides', *reg.options(), *sparse_stage.OPTION_NAMES]
 
 
 def reapply() -> None:
@@ -117,12 +145,22 @@ def reapply() -> None:
 
 
 def report() -> dict:
-    """The active chain and generation context, for the api and the debug log."""
+    """The active chain, sparse stage and generation context, for the api and the debug log."""
+    from modules.attention.sparse import stage as sparse_stage
     plan = current_plan
     state = context.current
+    options = sparse_stage.read_options()
+    layout = state.layout
     return {
         'chain': plan.chain() if plan is not None else ['sdpa'],
         'overrides': list(plan.labels) if plan is not None else [],
+        'sparse': {
+            'enabled': options.enabled,
+            'budget': options.budget,
+            'gate': options.gate,
+            'capable': [entry.backend.name for entry in plan.entries if 'block_mask' in entry.backend.caps] if plan is not None else [],
+            'layout': {'source': layout.source, 'kinds': list(layout.kinds()), 'length': layout.length} if layout is not None else None,
+        },
         'backend': plan.platform.backend if plan is not None else None,
         'context': {'active': state.active, 'role': state.role, 'step': state.step, 'steps': state.steps, 'model': state.model_key},
     }

modules/attention/sparse/layout.py

@@ -53,7 +53,7 @@ def runs(indices: torch.Tensor) -> list[tuple[int, int]]:
     return [(int(values[bounds[i]].item()), int(values[bounds[i + 1] - 1].item()) + 1) for i in range(len(bounds) - 1)]
 
 
-def layout_from_index_kwargs(kwargs: dict, length: int) -> TokenLayout | None:
+def layout_from_index_kwargs(kwargs: dict, length: int | None = None) -> TokenLayout | None:
     """Read a layout off the *_indices tensors a pipeline passes its transformer by name."""
     spans: list[Span] = []
     for name, value in kwargs.items():
@@ -68,7 +68,7 @@ def layout_from_index_kwargs(kwargs: dict, length: int) -> TokenLayout | None:
     if not spans:
         return None
     spans.sort(key=lambda s: s.start)
-    return TokenLayout(spans=tuple(spans), length=length, source='indices')
+    return TokenLayout(spans=tuple(spans), length=length if length is not None else spans[-1].end, source='indices')
 
 
 def layout_from_segments(segments, length: int | None = None, source: str = 'segments') -> TokenLayout:

modules/attention/sparse/stage.py

@@ -0,0 +1,98 @@
+"""The router stage that turns settings plus a published layout into a per call block selection."""
+from dataclasses import dataclass
+from modules.logger import log
+from modules.attention import context
+from modules.attention.sparse import layout as layout_mod
+from modules.attention.sparse.selector import SparseSpec, block_count, schedule, select_blocks
+
+
+# measured on a 3090: below roughly this length a 30 percent budget caps under 1.25x per block,
+# so the selector cannot pay for itself; see docs/sparse-attention-tracker.md
+AUTO_MIN_TOKENS = 8192
+
+# settings the stage reads, so a change to any of them rebuilds the chain
+OPTION_NAMES = ('sparse_attention_enabled', 'sparse_attention_budget', 'sparse_attention_min_tokens', 'sparse_attention_schedule_steps', 'sparse_attention_schedule_bump', 'sparse_attention_head_shared')
+
+
+@dataclass(frozen=True)
+class StageOptions:
+    enabled: bool = False
+    budget: float = 0.30
+    min_tokens: int = 0 # 0 selects AUTO_MIN_TOKENS
+    schedule_steps: int = 0
+    schedule_bump: float = 0.0
+    head_shared: bool = False
+
+    @property
+    def gate(self) -> int:
+        return self.min_tokens if self.min_tokens > 0 else AUTO_MIN_TOKENS
+
+
+def read_options() -> StageOptions:
+    from modules import shared
+    opts = shared.opts
+    return StageOptions(
+        enabled=bool(getattr(opts, 'sparse_attention_enabled', False)),
+        budget=float(getattr(opts, 'sparse_attention_budget', 30)) / 100.0,
+        min_tokens=int(getattr(opts, 'sparse_attention_min_tokens', 0)),
+        schedule_steps=int(getattr(opts, 'sparse_attention_schedule_steps', 0)),
+        schedule_bump=float(getattr(opts, 'sparse_attention_schedule_bump', 0)) / 100.0,
+        head_shared=bool(getattr(opts, 'sparse_attention_head_shared', False)),
+    )
+
+
+def resolve_layout(seq: int, reported: set) -> layout_mod.TokenLayout:
+    """The published layout when there is one, otherwise sparsify the whole sequence and say so once."""
+    published = context.current.layout
+    if isinstance(published, layout_mod.TokenLayout) and published.length == seq:
+        return published
+    if seq not in reported:
+        reported.add(seq)
+        detail = 'none published' if published is None else f'published length {getattr(published, "length", None)} does not match {seq}'
+        log.info(f'Sparse attention: no token layout ({detail}), sparsifying the whole sequence at tokens={seq}')
+    return layout_mod.layout_from_prefix(seq, 0)
+
+
+def make_stage(options: StageOptions):
+    """Return the per call selector, or None when the feature is off."""
+    if not options.enabled or options.budget >= 1.0:
+        return None
+    reported: set = set()
+    cache: dict = {}
+
+    def budget_for_step() -> float:
+        state = context.current
+        if options.schedule_steps <= 0 or options.schedule_bump <= 0 or state.steps <= 0:
+            return options.budget
+        key = (state.steps, options.budget, options.schedule_bump, options.schedule_steps)
+        table = cache.get(key)
+        if table is None:
+            table = schedule(state.steps, options.budget, options.schedule_bump, options.schedule_steps)
+            cache.clear()
+            cache[key] = table
+        return table[min(state.step, len(table) - 1)] if table else options.budget
+
+    def stage(query, key, value, attn_mask, is_causal): # pylint: disable=unused-argument
+        state = context.current
+        if state.role != 'transformer' or not state.active:
+            return None
+        if attn_mask is not None or is_causal: # flex would need a mask_mod to combine these; the quantized kernel composes them in R2
+            return None
+        if query.device.type == 'cpu' or query.dim() != 4:
+            return None
+        seq_q, seq_kv = query.shape[-2], key.shape[-2]
+        if seq_q != seq_kv or seq_q < options.gate: # cross attention is short and already cheap
+            return None
+        budget = budget_for_step()
+        if budget >= 1.0:
+            return None
+        spec = SparseSpec(budget=budget, head_shared=options.head_shared)
+        token_layout = resolve_layout(seq_q, reported)
+        nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
+        pins, drops = layout_mod.block_pins(token_layout, seq_q, seq_kv, spec.block_q, spec.block_kv, query.device)
+        if pins.shape[-2:] != (nq, nk):
+            return None
+        return select_blocks(query, key, spec, pins=pins, drops=drops)
+
+    stage.options = options # pylint: disable=attribute-defined-outside-init
+    return stage

modules/sd_hijack_modular.py

@@ -75,6 +75,15 @@ class InterruptLogFilter(logging.Filter):
         return 'Interrupted...' not in record.msg
 
 
+def publish_layout(kwargs):
+    """Hand the attention router whatever the pipeline says about its packed sequence, keyed on the *_indices tensors rather than the model."""
+    try:
+        from modules.attention.sparse import layout as sparse_layout
+        attention_context.set_layout(sparse_layout.layout_from_index_kwargs(kwargs or {}))
+    except Exception as e:
+        log.debug(f'Pipeline: token layout {e}')
+
+
 def install_state_hook(pipe):
     runner_log = logging.getLogger('diffusers.modular_pipelines.modular_pipeline')
     if not any(isinstance(f, InterruptLogFilter) for f in runner_log.filters):
@@ -91,9 +100,10 @@ def install_state_hook(pipe):
             return True
         return False
 
-    def _pre_transformer_hook(module, args): # pylint: disable=unused-argument
+    def _pre_transformer_hook(module, args, kwargs): # pylint: disable=unused-argument
         new_phase = set_phase('Generate', module)
         attention_context.set_role('transformer')
+        publish_layout(kwargs)
         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:
@@ -138,7 +148,8 @@ def install_state_hook(pipe):
         if module is not None:
             target = getattr(module, 'model', module) # conditioning calls the inner model directly
             if isinstance(target, torch.nn.Module) and getattr(target, 'sdnext_state_hook', None) is None:
-                target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook)
+                # with_kwargs, because the blocks call the transformer entirely by keyword and the token layout rides in those kwargs
+                target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook, with_kwargs=True)
 
     for name in ('text_encoder', 'text_encoder_2'):
         module = getattr(pipe, name, None)

modules/ui_definitions.py

@@ -262,6 +262,14 @@ def create_settings(cmd_opts):
         "sdnq_attention_pv_matmul_type": OptionInfo("disabled", "SDNQ Attention PV MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}),
         "sdnq_attention_hadamard_group_size": OptionInfo(256, "SDNQ Attention Hadamard Group Size", gr.Slider, {"minimum": 4, "maximum": 1024, "step": 1}),
 
+        "sparse_attention_sep": OptionInfo("<h2>Sparse Attention</h2>", "", gr.HTML),
+        "sparse_attention_enabled": OptionInfo(False, "Sparse Attention", gr.Checkbox),
+        "sparse_attention_budget": OptionInfo(30, "Sparse Attention KV budget", gr.Slider, {"minimum": 5, "maximum": 100, "step": 5}),
+        "sparse_attention_min_tokens": OptionInfo(0, "Sparse Attention minimum sequence", gr.Slider, {"minimum": 0, "maximum": 65536, "step": 1024}),
+        "sparse_attention_schedule_steps": OptionInfo(0, "Sparse Attention dense steps", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
+        "sparse_attention_schedule_bump": OptionInfo(30, "Sparse Attention dense step bonus", gr.Slider, {"minimum": 0, "maximum": 70, "step": 5}),
+        "sparse_attention_head_shared": OptionInfo(False, "Sparse Attention share selection across heads", gr.Checkbox),
+
         "hf_attention_sep": OptionInfo("<h2>Attention Dispatcher</h2>", "", gr.HTML),
         "hf_attention": OptionInfo('', "Attention dispatcher kernel", gr.Textbox),
     }))