Method maybe_create

backends/mlx/patterns.py:463–510

(cls, ep: ExportedProgram, head: Node)

    @classmethod
    def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"]:
        sdpa_node = head
        if not match_target(
            sdpa_node, torch.ops.aten.scaled_dot_product_attention.default
        ):
            return None

        q, k, v, _, _, _, _, _ = cls._parse_sdpa_args_and_kwargs(sdpa_node)

        # Detect grouped kv attention pattern with repeat_interleave before SDPA
        is_grouped_kv = False
        k_base = k
        v_base = v
        body: List[Node] = []
        if (
            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)
            and match_target(v, torch.ops.aten.repeat_interleave.self_int)
            and has_single_user(v)
            and (len(v.args) == 3)
            and (len(v.kwargs) == 0)
        ):
            k_unrepeated, k_reps, k_dim = k.args
            v_unrepeated, v_reps, v_dim = v.args

            if (k_dim == 1 and v_dim == 1) and (k_reps == v_reps):
                is_grouped_kv = True
                k_base = k_unrepeated
                v_base = v_unrepeated
                body = [k, v]

        # Detect HuggingFace repeat_kv pattern:
        # unsqueeze(dim=2) → expand → clone → view
        if 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:
                k_base, k_body = k_unwrap
                v_base, v_body = v_unwrap
                is_grouped_kv = True
                body = k_body + v_body

        head = sdpa_node
        if not is_grouped_kv:
            body = []
        return SDPAHandler(head, body, q_node=q, k_node=k_base, v_node=v_base)

    def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
        assert n == self.head