From b2a6d721589ffce80bb9722e9ac9294b0c82e390 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 23 Aug 2026 22:16:46 +0100 Subject: [PATCH] feat(attention): compose the block selection with a token mask The sparse stage declined every call carrying an attention mask, so a model that passes a padding mask stayed dense at every length. It now serves those calls when the serving backend declares `masked_block`: sdnq takes attn_mask and block_mask together, flex would need a mask_mod and does not. A backend that cannot compose says so once instead of skipping quietly. --- modules/attention/backends/sdnq.py | 2 +- modules/attention/registry.py | 2 +- modules/attention/router.py | 2 +- modules/attention/sparse/stage.py | 12 +++++++++--- 4 files changed, 12 insertions(+), 6 deletions(-) diff --git a/modules/attention/backends/sdnq.py b/modules/attention/backends/sdnq.py index 5a56835e1..3769da964 100644 --- a/modules/attention/backends/sdnq.py +++ b/modules/attention/backends/sdnq.py @@ -41,5 +41,5 @@ backend = AttentionBackend( name='sdnq', label='SDNQ attention', priority=60, prepare=prepare, constraints=Constraints(min_tokens=32, min_long_side=512, min_heads=2), # sequences of 512 or fewer are text encoders, single-head calls the vae options=('sdnq_attention_matmul_type', 'sdnq_attention_pv_matmul_type', 'sdnq_attention_smooth_k', 'sdnq_attention_use_hadamard', 'sdnq_attention_hadamard_group_size', 'sdnq_attention_quantize_fp32', 'sdnq_attention_use_fp16_accum'), - caps=frozenset({'block_mask'}), + caps=frozenset({'block_mask', 'masked_block'}), # the kernel takes attn_mask and block_mask together ) diff --git a/modules/attention/registry.py b/modules/attention/registry.py index 96879101d..69f9cc231 100644 --- a/modules/attention/registry.py +++ b/modules/attention/registry.py @@ -64,7 +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' + caps: frozenset[str] = frozenset() # what the call can consume beyond plain sdpa arguments: 'block_mask', and 'masked_block' when it composes one with a token mask def available_on(self, platform: Platform) -> bool: return self.platforms is None or platform.backend in self.platforms diff --git a/modules/attention/router.py b/modules/attention/router.py index c1400de0a..6a58aa011 100644 --- a/modules/attention/router.py +++ b/modules/attention/router.py @@ -71,7 +71,7 @@ def make_router(plan: Plan, observer: Callable | None = None, stage: Callable | for entry in entries: if entry.backend.constraints.accepts(query, key, value, attn_mask): if stage is not None and 'block_mask' in entry.caps: - selection = stage(query, key, value, attn_mask, is_causal) + selection = stage(query, key, value, attn_mask, is_causal, entry.caps) if selection is not None: if observer is not None: observer(f'{entry.backend.name}+sparse', query, key, attn_mask) diff --git a/modules/attention/sparse/stage.py b/modules/attention/sparse/stage.py index 98eb85cea..d182e65d5 100644 --- a/modules/attention/sparse/stage.py +++ b/modules/attention/sparse/stage.py @@ -65,6 +65,7 @@ def make_stage(options: StageOptions): return None reported: set = set() inactive: set = set() + notified: set = set() cache: dict = {} static: dict = {} @@ -106,12 +107,17 @@ def make_stage(options: StageOptions): stage.last_skip = reason return None - def stage(query, key, value, attn_mask, is_causal): # pylint: disable=unused-argument + def stage(query, key, value, attn_mask, is_causal, caps=frozenset()): # pylint: disable=unused-argument state = context.current if state.role != 'transformer' or not state.active: return decline('not the denoiser') - 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 decline('masked or causal') + if is_causal: # the selection keeps the diagonal but encodes no causality + return decline('causal') + if attn_mask is not None and 'masked_block' not in caps: # flex would need a mask_mod to combine the two + if 'masked' not in notified: # an enabled setting that cannot act says so rather than doing nothing quietly + notified.add('masked') + log.info('Sparse attention: this model passes an attention mask and the serving backend cannot combine it with a block selection; attention stays dense') + return decline('masked') if query.device.type == 'cpu' or query.dim() != 4: return decline('unsupported tensor') seq_q, seq_kv = query.shape[-2], key.shape[-2]