From 9eed27d758bfbf5f9acf9b932766a57e2f78d461 Mon Sep 17 00:00:00 2001 From: PyTorch Bot Date: Tue, 1 Sep 2026 17:32:39 -0700 Subject: [PATCH] Run attention on MLX wherever the fused kernel can compute it Attention on rank-3 tensors exports without complaint and then fails when you run the model: MLX execute failed: [scaled_dot_product_attention] input with shape (2,16,64) expected to be rank 4 PyTorch accepts rank 2, 3, 4 and 5 here. The fused kernel takes rank 4 only, and has other requirements besides. The handler matched the operator by name and checked none of them, so it claimed calls the kernel cannot run. Rank 2 and rank 3 gain the missing leading dimensions, run on the fused kernel, and the result is squeezed back. The dimensions go at the front, because for a rank-3 input the first dimension is already the head one, and inserting in the middle moves it into the batch slot, which misaligns masks and grouped heads. Admitting a lower rank also reaches code that was only ever given rank 4. The grouped key and value unwrap looks for a repeat on dimension 1, which is the head dimension at rank 4 and the key sequence below it, so at rank 3 it would absorb a repeat that carries real keys. Without a mask the result still matches, because a duplicated key and its value give the same weighted sum, so this hides. With a causal mask it is wrong by whole units. Both unwraps now run only at rank 4. Three shapes are left for the lowering step to decompose into primitives, which still run on this backend: Rank 5 and above, because folding the leading dimensions pairs the wrong operands as soon as one of them broadcasts a batch the others do not. Unequal batch sizes at rank 4, because the kernel requires them to match and rejects the call, so it fails at execute today. Causal attention at rank 2 or rank 3 whose query and key lengths differ. Torch anchors a causal mask at the top left and MLX at the bottom right, so the two disagree there and the disagreement is silent. Rank 4 already reaches the kernel today and keeps its current behaviour, since changing it is a separate question. Declining a call turns out not to be free, which is the other half of this change. Preservation from decomposition is requested per operator, so one claimed call would keep the operator whole for a declined one in the same graph, leaving that call neither lowered nor decomposed and stopping export. Give the whole operator back when any of its calls is unsupported. The framework does offer a per-node filter that would keep the other calls fused, and it is deliberately not used: it puts the program on a path that fails on an ordinary attention block that reshapes its output. The docstring says so. That half is not specific to attention. Two calls to torch.roll in one graph, one supported and one not, already fail at execute today for the same reason. Test Plan: Fifteen partitioner tests. Eleven fail before this change and all fifteen pass after. They assert on the serialized nodes, so they check which path a call took rather than only that it answered, and the rejection cases also assert the work stayed on this backend, so they cannot pass when nothing is delegated at all. Added a rank-3 case to the operator suite too, which runs the compiled runtime. Ran every combination on an Apple Silicon Mac against eager: plain, causal, explicit scale, float mask, boolean mask, grouped-query attention, batch size 1, float16 and bfloat16, at ranks 2 through 5. Rank 3 matches to 3.6e-07 and keeps the fused kernel in all seven of its variants, rank 2 to 3.6e-07, rank 5 decomposes and matches to 3.6e-07 where it previously failed, and unequal batch goes from failing to 2.4e-07. The grouped key case was measured both ways: at rank 4 the repeat is absorbed and the result matches, and at rank 3 with a causal mask absorbing it differs from what the model asked for by 5.1, so the guard is what keeps that correct. A zero key head count used to divide by zero and abort the whole export; it now declines that node. A graph holding one supported and one unsupported call exports and runs. Two rolls in one graph go from failing at execute to matching eager exactly. Exported the speech example and compared it against the same export without this change: identical serialized node counts and an identical greedy token sequence over twelve decode steps, so that model is unaffected. Ran the neighbouring backend test files, 71 tests, all passing. --- backends/mlx/partitioner.py | 17 ++ backends/mlx/patterns.py | 131 ++++++++++++++- backends/mlx/test/test_ops.py | 24 +++ backends/mlx/test/test_partitioner.py | 231 +++++++++++++++++++++++++- 4 files changed, 395 insertions(+), 8 deletions(-) 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()