Method _parse_sdpa_args_and_kwargs

backends/mlx/patterns.py:415–422

(cls, sdpa_node: Node)

    @classmethod
    def _parse_sdpa_args_and_kwargs(cls, sdpa_node: Node):
        q, k, v = sdpa_node.args[0:3]
        attn_mask = sdpa_node.args[3] if len(sdpa_node.args) > 3 else None
        dropout_p = sdpa_node.args[4] if len(sdpa_node.args) > 4 else 0.0
        is_causal = sdpa_node.args[5] if len(sdpa_node.args) > 5 else False
        enable_gqa = sdpa_node.args[6] if len(sdpa_node.args) > 6 else False
        scale = sdpa_node.kwargs.get("scale", None)
        return q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa

    @classmethod
    def _try_unwrap_repeat_kv(cls, node: Node) -> Optional[Tuple[Node, List[Node]]]: