diff --git a/backends/mlx/partitioner.py b/backends/mlx/partitioner.py
index 7814e883588..a82f54cda00 100644
--- a/backends/mlx/partitioner.py
+++ b/backends/mlx/partitioner.py
@@ -133,6 +133,18 @@ def ops_to_not_decompose(
         handler that rejects the 6-arg edge form, for instance). Preserving an op
         the handler then rejects is worse than not preserving it, because the op
         neither decomposes into something delegatable nor lowers itself.
+
+        A target is only preserved when every node carrying it is supported. One
+        unsupported node is enough to give the whole operator back to decomposition,
+        because keeping it would leave that node neither lowered nor decomposed and
+        export would stop.
+
+        The second return value is a per-node filter, which would keep the supported
+        calls fused and decompose only the rest. It is deliberately not used: it puts
+        the program on exir's EDGE_DO_NOT_DECOMP path, which fails on an ordinary
+        attention block that reshapes its output, a shape this backend has to lower.
+        The cost of the coarser choice is that one declined call also unfuses the
+        operator's other calls in that graph.
         """
         from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder
 
@@ -157,6 +169,7 @@ def ops_to_not_decompose(
 
         # Collect ops for nodes that are actually supported
         do_not_decompose: list[torch._ops.OpOverload] = []
+        declined: set[torch._ops.OpOverload] = set()
 
         for node in ep.graph.nodes:
             if node.op == "call_function" and isinstance(
@@ -166,6 +179,10 @@ def ops_to_not_decompose(
                 if info is not None and info.supported:
                     if node.target not in do_not_decompose:
                         do_not_decompose.append(node.target)
+                else:
+                    declined.add(node.target)
+
+        do_not_decompose = [op for op in do_not_decompose if op not in declined]
 
         self._not_decompose_cache = (weakref.ref(ep), do_not_decompose)
 
diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py
index 1760dc63b9a..5a8525eacce 100644
--- a/backends/mlx/patterns.py
+++ b/backends/mlx/patterns.py
@@ -44,6 +44,7 @@
     AddIntNode,
     AddNode,
     AsTypeNode,
+    ExpandDimsNode,
     IndexCopyNode,
     IntOrVid,
     ModIntNode,
@@ -52,10 +53,12 @@
     SdpaNode,
     SliceNode,
     SliceUpdateNode,
+    SqueezeNode,
     SubtractIntNode,
     SymSizeNode,
 )
 from torch.export.exported_program import ExportedProgram
+from torch.fx.experimental.symbolic_shapes import statically_known_true
 from torch.fx.node import Node
 
 
@@ -527,6 +530,61 @@ def _try_unwrap_repeat_kv(cls, node: Node) -> Optional[Tuple[Node, List[Node]]]:
         body = [e for e in entries if e is not None]
         return base, body
 
+    @classmethod
+    def _kernel_can_compute(cls, sdpa_node: Node) -> bool:
+        """Whether the fused kernel can compute this call faithfully.
+
+        Its preconditions are narrower than what PyTorch accepts, and a claimed call
+        that the kernel cannot compute fails loudly at execute rather than falling
+        back, so declining it here is what sends it to decomposition instead.
+        """
+        q, k, v, attn_mask, _, is_causal, _, _ = cls._parse_sdpa_args_and_kwargs(
+            sdpa_node
+        )
+        operand_vals = [
+            operand.meta.get("val") if isinstance(operand, Node) else None
+            for operand in (q, k, v)
+        ]
+        if any(val is None for val in operand_vals):
+            return False
+
+        # Ranks 2 and 3 are lifted to 4 on emission. Beyond 4 the leading dimensions
+        # would have to fold together, and a fold pairs the wrong operands as soon as
+        # one of them broadcasts a batch the others do not.
+        ranks = {val.dim() for val in operand_vals}
+        if len(ranks) != 1 or not 2 <= next(iter(ranks)) <= 4:
+            return False
+
+        # Torch anchors a causal mask at the top left and MLX at the bottom right, so
+        # the two agree only when the query and key lengths are equal. Rank 4 already
+        # reaches the kernel today and is left alone; the lower ranks are newly lifted
+        # here, so do not open a path that returns wrong values with no error.
+        if (
+            next(iter(ranks)) < 4
+            and is_causal
+            and not statically_known_true(
+                operand_vals[0].shape[-2] == operand_vals[1].shape[-2]
+            )
+        ):
+            return False
+
+        # The kernel requires the batch sizes to match and rejects the call otherwise,
+        # so a broadcast batch has to be decomposed rather than fused.
+        if next(iter(ranks)) == 4 and any(
+            not statically_known_true(val.shape[0] == operand_vals[0].shape[0])
+            for val in operand_vals[1:]
+        ):
+            return False
+
+        if attn_mask is not None:
+            mask_val = (
+                attn_mask.meta.get("val") if isinstance(attn_mask, Node) else None
+            )
+            if mask_val is None or mask_val.dim() > 4:
+                return False
+
+        return True
+
     @classmethod
     def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"]:
         sdpa_node = head
@@ -535,15 +593,23 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"
         ):
             return None
 
+        if not cls._kernel_can_compute(sdpa_node):
+            return None
+
         q, k, v, _, _, _, _, _ = cls._parse_sdpa_args_and_kwargs(sdpa_node)
 
-        # Detect grouped kv attention pattern with repeat_interleave before SDPA
+        # Detect grouped kv attention pattern with repeat_interleave before SDPA.
+        # Both unwraps below key on dim 1, which is the head dimension only at rank 4.
+        # At a lower rank dim 1 is the key sequence, and absorbing a repeat there
+        # drops keys, which a causal mask then turns into a wrong answer.
+        is_rank4 = q.meta["val"].dim() == 4 if isinstance(q, Node) else False
         is_grouped_kv = False
         k_base = k
         v_base = v
         body: List[Node] = []
         if (
-            match_target(k, torch.ops.aten.repeat_interleave.self_int)
+            is_rank4
+            and match_target(k, torch.ops.aten.repeat_interleave.self_int)
             and has_single_user(k)
             and (len(k.args) == 3)
             and (len(k.kwargs) == 0)
@@ -563,7 +629,7 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"
 
         # Detect HuggingFace repeat_kv pattern:
         # unsqueeze(dim=2) → expand → clone → view
-        if not is_grouped_kv:
+        if is_rank4 and not is_grouped_kv:
             k_unwrap = cls._try_unwrap_repeat_kv(k)
             v_unwrap = cls._try_unwrap_repeat_kv(v)
             if k_unwrap is not None and v_unwrap is not None:
@@ -572,6 +638,27 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"
                 is_grouped_kv = True
                 body = k_body + v_body
 
+        # Checked after the unwrapping above, because grouped-query attention reaches
+        # the kernel with its original head counts. MLX pairs heads only when the key
+        # and value agree and the query is a whole multiple of them.
+        kernel_vals = [
+            node.meta.get("val") if isinstance(node, Node) else None
+            for node in (q, k_base, v_base)
+        ]
+        if any(val is None for val in kernel_vals):
+            return None
+        q_heads, k_heads, v_heads = (
+            1 if val.dim() == 2 else val.shape[-3] for val in kernel_vals
+        )
+        # A zero head count would make the multiple test below divide by zero, and a
+        # raise here aborts the whole export rather than declining this one node.
+        if not statically_known_true(k_heads > 0):
+            return None
+        if not statically_known_true(k_heads == v_heads):
+            return None
+        if not statically_known_true(q_heads % k_heads == 0):
+            return None
+
         head = sdpa_node
         if not is_grouped_kv:
             body = []
@@ -593,19 +680,49 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
         assert dropout_p == 0.0, "SDPA with dropout is not supported"
 
         q, k, v, attn_mask = P.slot_map([q, k, v, attn_mask])
+        # Add the dimensions the kernel is missing at the front, never in the middle.
+        # For a rank-3 input the first dimension is already the head one, so inserting
+        # there would move it into the batch slot and misalign masks and grouped heads.
+        input_nodes = (self.q_node, self.k_node, self.v_node)
+        inputs = [q, k, v]
+        for i, input_node in enumerate(input_nodes):
+            for _ in range(4 - input_node.meta["val"].dim()):
+                _, expanded = P.make_tmp_slot()
+                P.emit(
+                    ExpandDimsNode(
+                        x=P.slot_to_tid(inputs[i]),
+                        out=P.slot_to_tid(expanded),
+                        axis=0,
+                    )
+                )
+                inputs[i] = expanded
+
+        output_rank = n.meta["val"].dim()
+
         out = P.make_or_get_slot(n)
+        sdpa_out = out
+        if output_rank < 4:
+            _, sdpa_out = P.make_tmp_slot()
 
         P.emit(
             SdpaNode(
-                q=P.slot_to_tid(q),
-                k=P.slot_to_tid(k),
-                v=P.slot_to_tid(v),
-                out=P.slot_to_tid(out),
+                q=P.slot_to_tid(inputs[0]),
+                k=P.slot_to_tid(inputs[1]),
+                v=P.slot_to_tid(inputs[2]),
+                out=P.slot_to_tid(sdpa_out),
                 scale=scale,
                 mask=P.slot_to_tid(attn_mask) if attn_mask else None,
                 causal=is_causal,
             )
         )
+        if output_rank < 4:
+            P.emit(
+                SqueezeNode(
+                    x=P.slot_to_tid(sdpa_out),
+                    out=P.slot_to_tid(out),
+                    dims=list(range(4 - output_rank)),
+                )
+            )
         return out
 
 
diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py
index 452ae37c40b..9ecdf5a8e58 100644
--- a/backends/mlx/test/test_ops.py
+++ b/backends/mlx/test/test_ops.py
@@ -6275,6 +6275,30 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (q, k, v)
 
 
+@register_test
+class SDPARank3Test(OpTestCase):
+    """Attention on rank-3 tensors, which PyTorch accepts and the fused kernel does not.
+
+    The node counts are the point of the test: they assert the fused kernel is still
+    used, rather than the operator having been decomposed into primitives.
+    """
+
+    name = "sdpa_rank3"
+    rtol = 1e-3
+    atol = 1e-3
+    expected_node_counts = {
+        "SdpaNode": 1,
+        "ExpandDimsNode": 3,
+        "SqueezeNode": 1,
+    }
+
+    def create_model(self) -> nn.Module:
+        return SDPAModel()
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        return tuple(torch.randn(2, 16, 64) for _ in range(3))
+
+
 class CustomSDPAModel(nn.Module):
     """
     Test model for mlx::custom_sdpa with KVCache.
diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py
index 4a5833aa656..3b82e306b1f 100644
--- a/backends/mlx/test/test_partitioner.py
+++ b/backends/mlx/test/test_partitioner.py
@@ -9,12 +9,16 @@
 Tests for the MLX partitioner.
 """
 
+import tempfile
 import unittest
+from pathlib import Path
 
 import torch
 import torch.nn as nn
 from executorch.backends.mlx.partitioner import MLXPartitioner
-from executorch.exir import EdgeCompileConfig, to_edge
+from executorch.backends.mlx.test.test_utils import get_mlx_node_counts
+from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower
+from executorch.runtime import Runtime
 from torch.export import export
 
 
@@ -41,5 +45,230 @@ def forward(self, x):
         self.assertIn("to_edge_transform_and_lower", str(ctx.exception))
 
 
+def _lower(model, inputs):
+    return to_edge_transform_and_lower(
+        export(model, inputs, strict=False),
+        partitioner=[MLXPartitioner()],
+    ).to_executorch()
+
+
+def _delegate_count(program) -> int:
+    return sum(
+        1
+        for node in program.exported_program().graph_module.graph.nodes
+        if node.op == "call_function" and "executorch_call_delegate" in str(node.target)
+    )
+
+
+def _run(model, inputs):
+    """Lower, execute, and return the node counts, the delegate count and the error.
+
+    The delegate count is returned so a test can tell "decomposed onto this backend"
+    apart from "not lowered here at all", which a node count alone cannot show.
+    """
+    with torch.no_grad():
+        ref = model(*inputs)
+    program = _lower(model, inputs)
+    delegates = _delegate_count(program)
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "model.pte"
+        path.write_bytes(program.buffer)
+        counts = get_mlx_node_counts(path)
+        method = Runtime.get().load_program(path).load_method("forward")
+        out = method.execute(list(inputs))[0]
+    return counts, delegates, (out - ref).abs().max().item()
+
+
+class Sdpa(nn.Module):
+    def __init__(self, is_causal: bool = False):
+        super().__init__()
+        self.is_causal = is_causal
+
+    def forward(self, q, k, v):
+        return torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, is_causal=self.is_causal
+        )
+
+
+class GroupedSdpa(nn.Module):
+    """Grouped key/value attention, where the repeat is unwrapped before the kernel."""
+
+    def __init__(self, dim: int, is_causal: bool = False):
+        super().__init__()
+        self.dim = dim
+        self.is_causal = is_causal
+
+    def forward(self, q, k, v):
+        k = k.repeat_interleave(2, dim=self.dim)
+        v = v.repeat_interleave(2, dim=self.dim)
+        return torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, is_causal=self.is_causal
+        )
+
+
+class TestMLXPartitionerSdpaShapes(unittest.TestCase):
+    """The fused kernel takes rank 4, so other ranks are adapted or left alone."""
+
+    def test_rank4_is_unchanged(self):
+        counts, _, err = _run(
+            Sdpa(), tuple(torch.randn(2, 4, 16, 64) for _ in range(3))
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 1)
+        self.assertEqual(counts.get("ExpandDimsNode", 0), 0)
+        self.assertEqual(counts.get("SqueezeNode", 0), 0)
+        self.assertLess(err, 1e-4)
+
+    def test_rank3_is_lifted_once(self):
+        counts, _, err = _run(Sdpa(), tuple(torch.randn(2, 16, 64) for _ in range(3)))
+        self.assertEqual(counts.get("SdpaNode", 0), 1)
+        self.assertEqual(counts.get("ExpandDimsNode", 0), 3)
+        self.assertEqual(counts.get("SqueezeNode", 0), 1)
+        self.assertLess(err, 1e-4)
+
+    def test_rank2_is_lifted_twice(self):
+        counts, _, err = _run(Sdpa(), tuple(torch.randn(16, 64) for _ in range(3)))
+        self.assertEqual(counts.get("SdpaNode", 0), 1)
+        self.assertEqual(counts.get("ExpandDimsNode", 0), 6)
+        self.assertEqual(counts.get("SqueezeNode", 0), 1)
+        self.assertLess(err, 1e-4)
+
+    def test_rank5_is_decomposed_on_this_backend(self):
+        # Folding the leading dimensions pairs the wrong operands once one of them
+        # broadcasts a batch, so this decomposes rather than fusing.
+        counts, delegates, err = _run(
+            Sdpa(), tuple(torch.randn(2, 2, 4, 16, 64) for _ in range(3))
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 0)
+        self.assertGreater(delegates, 0)
+        self.assertLess(err, 1e-4)
+
+    def test_unequal_batch_is_decomposed_on_this_backend(self):
+        counts, delegates, err = _run(
+            Sdpa(),
+            (
+                torch.randn(2, 4, 16, 64),
+                torch.randn(1, 4, 16, 64),
+                torch.randn(1, 4, 16, 64),
+            ),
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 0)
+        self.assertGreater(delegates, 0)
+        self.assertLess(err, 1e-4)
+
+    def test_zero_head_count_declines_instead_of_raising(self):
+        # The head multiple test would divide by zero here, and raising from the
+        # matcher aborts the whole export rather than declining this one node. Only
+        # lowering is checked: a zero-size operand is not executable either way.
+        program = _lower(
+            Sdpa(),
+            (
+                torch.randn(1, 4, 8, 16),
+                torch.randn(1, 0, 8, 16),
+                torch.randn(1, 0, 8, 16),
+            ),
+        )
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "model.pte"
+            path.write_bytes(program.buffer)
+            self.assertEqual(get_mlx_node_counts(path).get("SdpaNode", 0), 0)
+
+
+class TestMLXPartitionerGroupedKeys(unittest.TestCase):
+    """The grouped key/value unwrap reads dim 1 as the head, which holds at rank 4."""
+
+    def test_rank4_head_repeat_is_absorbed(self):
+        counts, _, err = _run(
+            GroupedSdpa(dim=1),
+            (
+                torch.randn(2, 4, 16, 64),
+                torch.randn(2, 2, 16, 64),
+                torch.randn(2, 2, 16, 64),
+            ),
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 1)
+        self.assertEqual(counts.get("RepeatNode", 0), 0)
+        self.assertLess(err, 1e-4)
+
+    def test_rank3_sequence_repeat_is_kept(self):
+        # At rank 3 dim 1 is the key sequence, so absorbing the repeat would drop
+        # half the keys. Without a mask that still sums correctly, which is what
+        # makes it easy to miss; with a causal mask it is wrong by whole units.
+        counts, _, err = _run(
+            GroupedSdpa(dim=1, is_causal=True),
+            (torch.randn(2, 16, 64), torch.randn(2, 8, 64), torch.randn(2, 8, 64)),
+        )
+        self.assertEqual(counts.get("RepeatNode", 0), 2)
+        self.assertLess(err, 1e-4)
+
+
+class TestMLXPartitionerSdpaCausal(unittest.TestCase):
+    """MLX anchors a causal mask at the bottom right and torch at the top left."""
+
+    def test_equal_lengths_stay_fused(self):
+        counts, _, err = _run(
+            Sdpa(is_causal=True), tuple(torch.randn(1, 4, 16, 64) for _ in range(3))
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 1)
+        self.assertLess(err, 1e-4)
+
+    def test_rank3_equal_lengths_are_lifted(self):
+        counts, _, err = _run(
+            Sdpa(is_causal=True), tuple(torch.randn(2, 16, 64) for _ in range(3))
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 1)
+        self.assertEqual(counts.get("ExpandDimsNode", 0), 3)
+        self.assertLess(err, 1e-4)
+
+    def test_rank3_unequal_lengths_are_not_lifted(self):
+        # The two conventions disagree here and the disagreement is silent, so a
+        # shape this backend could not previously reach is not opened up.
+        counts, delegates, err = _run(
+            Sdpa(is_causal=True),
+            (torch.randn(2, 6, 64), torch.randn(2, 16, 64), torch.randn(2, 16, 64)),
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 0)
+        self.assertGreater(delegates, 0)
+        self.assertLess(err, 1e-4)
+
+    def test_rank2_unequal_lengths_are_not_lifted(self):
+        counts, _, err = _run(
+            Sdpa(is_causal=True),
+            (torch.randn(6, 64), torch.randn(16, 64), torch.randn(16, 64)),
+        )
+        self.assertEqual(counts.get("SdpaNode", 0), 0)
+        self.assertLess(err, 1e-4)
+
+
+class TestMLXPartitionerMixedSupport(unittest.TestCase):
+    """An operator is preserved from decomposition per operator, not per call."""
+
+    def test_supported_and_unsupported_calls_in_one_graph(self):
+        class Mixed(nn.Module):
+            def forward(self, a, b):
+                x = torch.nn.functional.scaled_dot_product_attention(a, a, a)
+                y = torch.nn.functional.scaled_dot_product_attention(b, b, b)
+                return x.sum() + y.sum()
+
+        # Without giving the whole operator back, the rank-5 call would be neither
+        # lowered nor decomposed and this would raise a missing out variant.
+        counts, delegates, err = _run(
+            Mixed().eval(),
+            (torch.randn(1, 4, 16, 64), torch.randn(2, 2, 4, 16, 64)),
+        )
+        # The cost of the coarse choice: the supported call is unfused as well.
+        self.assertEqual(counts.get("SdpaNode", 0), 0)
+        self.assertGreater(delegates, 0)
+        self.assertLess(err, 1e-3)
+
+    def test_mixed_support_outside_attention(self):
+        class TwoRolls(nn.Module):
+            def forward(self, x):
+                return torch.roll(x, 1, dims=0).sum() + torch.roll(x, 1).sum()
+
+        _, delegates, err = _run(TwoRolls().eval(), (torch.randn(4, 8),))
+        self.assertGreater(delegates, 0)
+        self.assertLess(err, 1e-4)
+
+
 if __name__ == "__main__":
     unittest.main()