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

@@ -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
 )

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

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)

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]