fix(attention): flex joins the chain instead of ending it

The flex backend never called the sdpa it replaced, so any backend
stacked before it was unreachable and every call it could not serve,
cpu or 3d inputs included, failed inside flex_attention. It is now an
ordinary entry gated on what flex_attention accepts: 4d tensors on one
non-cpu device. The mask path drops the 2d special case, which indexed
attn_mask.size and reshaped the mask onto the wrong axis; expanding to
(batch, heads, q, kv) already follows sdpa broadcast semantics.

modules/attention/backends/flex.py

@@ -1,6 +1,6 @@
 import torch
 from modules.logger import log
-from modules.attention.registry import AttentionBackend, Platform
+from modules.attention.registry import AttentionBackend, Constraints, Platform
 
 
 def prepare(platform: Platform, original): # pylint: disable=unused-argument
@@ -9,16 +9,14 @@ 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=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs): # pylint: disable=unused-argument
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
         score_mod = None
         block_mask = None
         if attn_mask is not None:
             batch_size, num_heads = query.shape[:2]
             seq_len_q = query.shape[-2]
             seq_len_kv = key.shape[-2]
-            if attn_mask.ndim == 2:
-                attn_mask = attn_mask.view(attn_mask.shape[0], 1, attn_mask.size[1], 1)
-            attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv)
+            attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv) # sdpa masks broadcast over the trailing dims
             if attn_mask.dtype == torch.bool:
                 def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
                     return attn_mask[batch_idx, head_idx, q_idx, kv_idx]
@@ -35,4 +33,7 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument
     return call
 
 
-backend = AttentionBackend(name='flex', label='Flex attention', priority=20, prepare=prepare, terminal=True)
+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
+)

modules/attention/registry.py

@@ -26,10 +26,13 @@ class Constraints:
     min_tokens: int = 0 # query and key sequences both at least this long
     min_long_side: int = 0 # query or key sequence longer than this
     min_heads: int = 0
+    min_ndim: int = 0
 
     def accepts(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: torch.Tensor | None) -> bool:
         if not self.allow_cpu and query.device.type == 'cpu':
             return False
+        if self.min_ndim and query.ndim < self.min_ndim:
+            return False
         if not self.allow_mask and attn_mask is not None:
             return False
         if not self.allow_float32 and query.dtype == torch.float32: