diff --git a/modules/attention/backends/flex.py b/modules/attention/backends/flex.py index 4ca29be3e..e9b203786 100644 --- a/modules/attention/backends/flex.py +++ b/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'}), ) diff --git a/modules/attention/context.py b/modules/attention/context.py index 24aa8300a..84ca259aa 100644 --- a/modules/attention/context.py +++ b/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) diff --git a/modules/attention/registry.py b/modules/attention/registry.py index 23ca24d1a..96879101d 100644 --- a/modules/attention/registry.py +++ b/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() diff --git a/modules/attention/router.py b/modules/attention/router.py index c9ae07a48..8b8e6a432 100644 --- a/modules/attention/router.py +++ b/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}, } diff --git a/modules/attention/sparse/layout.py b/modules/attention/sparse/layout.py index 772e49b65..7362c96d6 100644 --- a/modules/attention/sparse/layout.py +++ b/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: diff --git a/modules/attention/sparse/stage.py b/modules/attention/sparse/stage.py new file mode 100644 index 000000000..6542e6968 --- /dev/null +++ b/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 diff --git a/modules/sd_hijack_modular.py b/modules/sd_hijack_modular.py index 3c151e2ab..528e4a0f2 100644 --- a/modules/sd_hijack_modular.py +++ b/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) diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index a873ea68d..1b0e05d14 100644 --- a/modules/ui_definitions.py +++ b/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("