Method __call__

backends/mlx/patterns.py:512–541

(self, P: MLXProgramBuilder, n: Node)

        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
        q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa = (
            SDPAHandler._parse_sdpa_args_and_kwargs(n)
        )
        head_dim = q.meta["val"].shape[-1]
        if scale is None:
            scale = head_dim**-0.5

        q = self.q_node
        k = self.k_node
        v = self.v_node

        assert dropout_p == 0.0, "SDPA with dropout is not supported"

        q, k, v, attn_mask = P.slot_map([q, k, v, attn_mask])
        out = P.make_or_get_slot(n)

        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),
                scale=scale,
                mask=P.slot_to_tid(attn_mask) if attn_mask else None,
                causal=is_causal,
            )
        )
        return out


@REGISTRY.register_pattern(name="NVFP4_QUANTIZED_EMBEDDING")