diff --git a/modules/attention/backends/flex.py b/modules/attention/backends/flex.py
index 7997d54fd..4ca29be3e 100644
--- a/modules/attention/backends/flex.py
+++ b/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
+)
diff --git a/modules/attention/registry.py b/modules/attention/registry.py
index ef33d009e..820dc77f8 100644
--- a/modules/attention/registry.py
+++ b/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:
diff --git a/test/test-attention-router.py b/test/test-attention-router.py
index c51f4b3f5..44a95c150 100644
--- a/test/test-attention-router.py
+++ b/test/test-attention-router.py
@@ -102,9 +102,10 @@ def run_test(cat: str, fn):
 # ============================================================
 
 # devices.set_sdpa_params applied the hijacks in this order; each wrapped the previous, so the
-# last applied was tried first. Dynamic and flex replaced the chain end instead of wrapping it.
+# last applied was tried first. Dynamic replaced the chain end instead of wrapping it; flex did
+# too, which left everything stacked before it unreachable, so it is an ordinary entry now.
 OLD_ORDER = ['Dynamic attention', 'Flex attention', 'Triton Flash attention', 'Flash attention', 'Sage attention', 'SDNQ attention']
-OLD_TERMINALS = {'Dynamic attention', 'Flex attention'}
+OLD_TERMINALS = {'Dynamic attention'}
 OLD_NAMES = {
     'Dynamic attention': 'dynamic',
     'Flex attention': 'flex',
@@ -117,11 +118,13 @@ OLD_NAMES = {
 CHOICES = OLD_ORDER
 TRITON_PLATFORMS = {'rocm', 'zluda'}
 
-OLD_GATES = {
+# the four closure predicates transcribed literally, plus the contract flex_attention itself enforces
+GATES = {
     'sdnq': lambda q, k, v, m: q.device.type != "cpu" and (q.shape[-2] >= 32 and k.shape[-2] >= 32) and (q.shape[-2] > 512 or k.shape[-2] > 512) and q.shape[-3] > 1,
     'triton': lambda q, k, v, m: q.shape[-1] <= 128 and m is None and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
     'flash': lambda q, k, v, m: q.shape[-1] <= 128 and m is None and q.dtype != torch.float32 and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
     'sage': lambda q, k, v, m: q.shape[-1] in {128, 96, 64} and m is None and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
+    'flex': lambda q, k, v, m: q.ndim == 4 and q.device.type != "cpu" and k.device == q.device and v.device == q.device,
 }
 
 
@@ -194,12 +197,13 @@ def test_plan_matches_stacking_oracle():
 def test_gates_match_transcribed_predicates():
     cases = 0
     lengths = (16, 32, 512, 513, 4096)
-    for q_device, kv_device, dtype, heads, q_len, k_len, head_dim, masked in itertools.product(('cpu', 'meta'), ('cpu', 'meta'), (torch.float16, torch.float32), (1, 8), lengths, lengths, (40, 64, 96, 128, 256), (False, True)):
-        q = shaped((1, heads, q_len, head_dim), dtype, q_device)
-        k = shaped((1, heads, k_len, head_dim), dtype, kv_device)
-        v = shaped((1, heads, k_len, head_dim), dtype, kv_device)
-        m = shaped((1, 1, q_len, k_len), torch.bool, q_device) if masked else None
-        for name, gate in OLD_GATES.items():
+    for q_device, kv_device, dtype, heads, q_len, k_len, head_dim, masked, batched in itertools.product(('cpu', 'meta'), ('cpu', 'meta'), (torch.float16, torch.float32), (1, 8), lengths, lengths, (40, 64, 96, 128, 256), (False, True), (False, True)):
+        lead = (1,) if batched else ()
+        q = shaped((*lead, heads, q_len, head_dim), dtype, q_device)
+        k = shaped((*lead, heads, k_len, head_dim), dtype, kv_device)
+        v = shaped((*lead, heads, k_len, head_dim), dtype, kv_device)
+        m = shaped((*lead, 1, q_len, k_len), torch.bool, q_device) if masked else None
+        for name, gate in GATES.items():
             expected = bool(gate(q, k, v, m))
             got = attention.registry.backends[name].constraints.accepts(q, k, v, m)
             assert got == expected, f'{name}: q={tuple(q.shape)} k={tuple(k.shape)} dtype={dtype} devices={q_device}/{kv_device} mask={masked} got={got} expected={expected}'
@@ -208,11 +212,10 @@ def test_gates_match_transcribed_predicates():
     return True
 
 
-def test_terminals_carry_no_gate():
-    for name in ('dynamic', 'flex'):
-        backend = attention.registry.backends[name]
-        assert backend.terminal, name
-        assert backend.constraints == attention.Constraints(), name
+def test_only_dynamic_is_terminal():
+    for name, backend in attention.registry.backends.items():
+        assert backend.terminal == (name == 'dynamic'), name
+    assert attention.registry.backends['dynamic'].constraints == attention.Constraints()
     return True
 
 
@@ -312,7 +315,7 @@ def run_all():
     for fn in [
         test_plan_matches_stacking_oracle,
         test_gates_match_transcribed_predicates,
-        test_terminals_carry_no_gate,
+        test_only_dynamic_is_terminal,
         test_choices_match_backends,
         test_router_dispatch_prefers_priority_then_terminal_then_original,
         test_prepare_failure_skips_backend,