Reuse shared transform passes in MLX (pytorch#22858)

Use shared RMSNorm and GQA fusion passes in MLX, and normalize SDPA
input ranks before decomposition checks. Replace the SDPA pattern with a
direct op handler while preserving MLX support checks and causal
alignment. Extend shared GQA fusion to handle repeat_interleave and
update test coverage.

Author: metascroy

backends/mlx/ops.py

@@ -133,6 +133,7 @@
     RoundNode,
     RsqrtNode,
     ScatterAddNode,
+    SdpaNode,
     SigmoidNode,
     SignNode,
     SiluNode,
@@ -178,6 +179,7 @@
 from executorch.extension.llm.cache import (  # noqa: F401
     update_and_attend as _kvcache_op,
 )
+from torch.fx.experimental.symbolic_shapes import statically_known_true
 from torch.fx.node import Node
 
 _LEAKY_RELU_DEFAULT_NEGATIVE_SLOPE = 0.01
@@ -3074,6 +3076,108 @@ def _arange_start_step_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return out
 
 
+@REGISTRY.register(target=[torch.ops.aten.scaled_dot_product_attention.default])
+def _aten_sdpa_handler(P: MLXProgramBuilder, n: Node) -> Slot:  # noqa: C901
+    op_name = "aten.scaled_dot_product_attention"
+    require_args(n.args, 0, 6, op_name)
+    require_kwargs(
+        n.kwargs,
+        {
+            "query",
+            "key",
+            "value",
+            "attn_mask",
+            "dropout_p",
+            "is_causal",
+            "scale",
+            "enable_gqa",
+        },
+        op_name,
+    )
+
+    def arg(index, name, default=None):
+        return n.args[index] if len(n.args) > index else n.kwargs.get(name, default)
+
+    q_node, k_node, v_node = (
+        arg(i, name) for i, name in enumerate(("query", "key", "value"))
+    )
+    attn_mask = arg(3, "attn_mask")
+    dropout_p = arg(4, "dropout_p", 0.0)
+    is_causal = arg(5, "is_causal", False)
+    scale = n.kwargs.get("scale")
+    enable_gqa = n.kwargs.get("enable_gqa", False)
+    operand_vals = [
+        node.meta.get("val") if isinstance(node, Node) else None
+        for node in (q_node, k_node, v_node)
+    ]
+    if any(not isinstance(val, torch.Tensor) or val.dim() != 4 for val in operand_vals):
+        raise NotImplementedError("SDPA requires normalized rank-4 Q/K/V")
+    q_val, k_val, v_val = operand_vals
+    if any(
+        not statically_known_true(val.shape[0] == q_val.shape[0])
+        for val in (k_val, v_val)
+    ):
+        raise NotImplementedError("SDPA batch broadcasting is not supported")
+
+    q_heads, k_heads, v_heads = (val.shape[1] for val in operand_vals)
+    if not statically_known_true(k_heads > 0):
+        raise NotImplementedError("SDPA requires positive KV heads")
+    if not statically_known_true(k_heads == v_heads) or not statically_known_true(
+        q_heads % k_heads == 0
+    ):
+        raise NotImplementedError("SDPA head counts are not supported")
+    if attn_mask is not None:
+        mask_val = attn_mask.meta.get("val") if isinstance(attn_mask, Node) else None
+        if not isinstance(mask_val, torch.Tensor) or mask_val.dim() > 4:
+            raise NotImplementedError("SDPA mask rank is not supported")
+
+    require_static_float(dropout_p, "dropout_p", op_name)
+    if dropout_p != 0.0:
+        raise NotImplementedError("SDPA with dropout is not supported")
+    if not isinstance(is_causal, bool) or not isinstance(enable_gqa, bool):
+        raise NotImplementedError("SDPA flags must be static booleans")
+    q_len, k_len = q_val.shape[-2], k_val.shape[-2]
+    if is_causal and not statically_known_true(q_len <= k_len):
+        raise NotImplementedError("Causal SDPA requires query length <= key length")
+    if scale is None:
+        scale = q_val.shape[-1] ** -0.5
+    require_static_float(scale, "scale", op_name)
+
+    q, k, v, mask = P.slot_map([q_node, k_node, v_node, attn_mask])
+    inputs = [q, k, v]
+    # Slicing K/V makes causal attention square, reconciling PyTorch's top-left
+    # mask with MLX's bottom-right mask. Keys past the last query are never read.
+    if is_causal and not statically_known_true(q_len == k_len):
+        rows = emit_shape(P, q_node, q)[-2]
+        for i in (1, 2):
+            _, sliced = P.make_tmp_slot()
+            P.emit(
+                SliceNode(
+                    x=P.slot_to_tid(inputs[i]),
+                    out=P.slot_to_tid(sliced),
+                    axis=IntOrVid.from_literal(2),
+                    start=IntOrVid.from_literal(0),
+                    stop=rows,
+                )
+            )
+            inputs[i] = sliced
+
+    out = P.make_or_get_slot(n)
+    # MLX infers GQA from the Q/KV head counts; no separate flag is serialized.
+    P.emit(
+        SdpaNode(
+            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(out),
+            scale=scale,
+            mask=P.slot_to_tid(mask) if mask is not None else None,
+            causal=is_causal,
+        )
+    )
+    return out
+
+
 @REGISTRY.register(target=[torch.ops.aten.rms_norm.default])
 def _aten_rms_norm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     args = P.args(n)

backends/mlx/partitioner.py

@@ -23,6 +23,9 @@
 from executorch.backends.mlx._logging import logger
 from executorch.backends.mlx._memprofile import mem_phase
 from executorch.backends.mlx.preprocess import MLXBackend
+from executorch.backends.transforms.normalize_sdpa_input_rank import (
+    NormalizeSDPAInputRankPass,
+)
 from executorch.exir.backend.backend_details import CompileSpec
 from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import (
     generate_partitions_from_list_of_nodes,
@@ -33,6 +36,7 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
+from executorch.exir.pass_manager import ExportedProgramPassManager
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -110,6 +114,14 @@ def __init__(self, compile_specs: List[CompileSpec] | None = None) -> None:
             | None
         ) = None
 
+    def transform_for_pre_decomposition(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgram:
+        """Normalize SDPA ranks before deciding which ops to preserve."""
+        return ExportedProgramPassManager([NormalizeSDPAInputRankPass()])(
+            exported_program
+        ).exported_program
+
     def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> tuple[list[torch._ops.OpOverload], Callable[[torch.fx.Node], bool] | None]:

backends/mlx/passes.py

@@ -8,35 +8,31 @@
 Graph transformation passes for the MLX backend.
 """
 
-from dataclasses import dataclass
-from typing import List, Optional
+from typing import List
 
 import torch
-from executorch.backends.mlx.pattern_utils import (
-    extract_lifted_tensor_constant,
-    match_target,
-    OpStep,
-    PatternMatch,
-    walk_back,
-)
 from executorch.backends.transforms.collapse_view_copy import CollapseViewCopyPass
+from executorch.backends.transforms.fuse_gqa_with_sdpa import FuseGQAWithSDPAPass
+from executorch.backends.transforms.fuse_rms_norm import FuseRMSNormPass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import (
     ExportedProgramPassBase,
     ExportedProgramPassResult,
     ExportPass,
     PassResult,
 )
+from executorch.exir.pass_manager import PassType
 from executorch.exir.passes.cse_pass import CSEPass
 from torch.fx import GraphModule, Node
 
 
-def get_default_passes() -> List[ExportPass]:
+def get_default_passes() -> List[PassType]:
     """
     Returns a list of passes that are enabled by default for the MLX backend.
     """
     return [
-        FuseRMSNormPass(),
+        FuseRMSNormPass(fold_dtype_casts=True, allow_lossy_weight_casts=True),
+        FuseGQAWithSDPAPass(),
         CanonicalizePermutePass(),
         CollapseViewCopyPass(),
         CollapsePermutePass(),
@@ -103,175 +99,6 @@ def call(self, exported_program) -> ExportedProgramPassResult:
         return ExportedProgramPassResult(exported_program, True)
 
 
-@dataclass
-class RMSNormMatch(PatternMatch):
-    """
-    Matched RMSNorm pattern.
-
-    HuggingFace Llama's RMSNorm decomposes into:
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + eps)
-        return weight * hidden_states.to(input_dtype)
-
-    Graph pattern:
-        _to_copy (to f32) [optional]
-        pow(x, 2)
-        mean_dim(pow_out, [-1], keepdim=True)
-        add(mean_out, eps_tensor)
-        rsqrt(add_out)
-        mul(to_copy_out, rsqrt_out)
-        _to_copy (back to original dtype) [optional]
-        mul(weight, to_copy_out)
-    """
-
-    input_node: Node = None  # type: ignore[assignment]
-    weight_node: Node = None  # type: ignore[assignment]
-    eps: float = 0.0
-
-    @classmethod
-    def maybe_create(cls, head: Node, **context) -> Optional["RMSNormMatch"]:
-        """Match RMSNorm pattern starting from final mul(weight, normalized)."""
-        # Head must be mul
-        if not match_target(head, torch.ops.aten.mul.Tensor):
-            return None
-
-        if len(head.args) < 2:
-            return None
-
-        # Try both orderings: mul(weight, normalized) or mul(normalized, weight)
-        for weight_idx, norm_idx in [(0, 1), (1, 0)]:
-            weight_node = head.args[weight_idx]
-            norm_node = head.args[norm_idx]
-
-            if not isinstance(norm_node, Node):
-                continue
-
-            # Match entire chain with single walk_back:
-            #   [_to_copy] -> mul(input, rsqrt) -> rsqrt -> add -> mean -> pow -> [_to_copy]
-            # The mul follows arg_index=1 to get rsqrt (not input)
-            result = walk_back(
-                norm_node,
-                [
-                    OpStep(
-                        op=torch.ops.aten._to_copy.default,
-                        optional=True,
-                        kwargs={
-                            "dtype",
-                            "layout",
-                            "device",
-                            "pin_memory",
-                            "non_blocking",
-                            "memory_format",
-                        },
-                    ),
-                    OpStep(op=torch.ops.aten.mul.Tensor, nargs=2, arg_index=1),
-                    OpStep(op=torch.ops.aten.rsqrt.default),
-                    OpStep(op=torch.ops.aten.add.Tensor, nargs=2),
-                    OpStep(op=torch.ops.aten.mean.dim, nargs=(2, 3), kwargs={"dtype"}),
-                    OpStep(op=torch.ops.aten.pow.Tensor_Scalar, nargs=2),
-                    OpStep(
-                        op=torch.ops.aten._to_copy.default,
-                        optional=True,
-                        require_single_user=False,  # _to_copy output used by both pow and mul
-                        kwargs={
-                            "dtype",
-                            "layout",
-                            "device",
-                            "pin_memory",
-                            "non_blocking",
-                            "memory_format",
-                        },
-                    ),
-                ],
-            )
-            if result is None:
-                continue
-
-            original_input, entries = result
-            to_copy_out, mul, rsqrt, add, mean, pow, to_copy_in = entries
-
-            # If input _to_copy matched, verify it has exactly 2 users: pow and mul
-            if to_copy_in is not None:
-                users = set(to_copy_in.users.keys())
-                expected_users = {pow, mul}
-                if users != expected_users:
-                    continue
-
-            # Validate pow exponent is 2
-            if pow.args[1] != 2:
-                continue
-
-            # Extract epsilon from add node (it's a lifted tensor constant)
-            eps_value = None
-            for arg in add.args:
-                eps_value = extract_lifted_tensor_constant(arg)
-                if eps_value is not None:
-                    break
-
-            if eps_value is None:
-                continue
-
-            # Build body from non-None entries
-            body = [n for n in entries if n is not None]
-
-            return cls(
-                head=head,
-                body=body,
-                input_node=original_input,
-                weight_node=weight_node,
-                eps=eps_value,
-            )
-
-        return None
-
-
-class FuseRMSNormPass(ExportPass):
-    """
-    Fuses decomposed RMSNorm operations into aten.rms_norm.
-
-    This reduces ~7 ops to 1 fused op per RMSNorm layer.
-    """
-
-    def call(self, graph_module: GraphModule) -> PassResult:
-        graph = graph_module.graph
-        modified = False
-
-        for node in list(graph.nodes):
-            match = RMSNormMatch.maybe_create(node)
-            if match is None:
-                continue
-
-            # Get input shape for normalized_shape
-            input_meta = match.input_node.meta.get("val")
-            if input_meta is None:
-                continue
-
-            # Create fused rms_norm node
-            with graph.inserting_before(node):
-                normalized_shape = [input_meta.shape[-1]]
-                rms_norm_node = graph.call_function(
-                    torch.ops.aten.rms_norm.default,
-                    args=(
-                        match.input_node,
-                        normalized_shape,
-                        match.weight_node,
-                        match.eps,
-                    ),
-                )
-                rms_norm_node.meta = node.meta.copy()
-
-            node.replace_all_uses_with(rms_norm_node)
-            match.remove_body_nodes(graph)
-            graph.erase_node(node)
-            modified = True
-
-        if modified:
-            graph.eliminate_dead_code()
-            graph.lint()
-
-        return PassResult(graph_module, modified)
-
-
 class CanonicalizePermutePass(ExportPass):
     """
     Converts transpose_copy to permute_copy in the edge dialect graph.