    def _infer_shapes(self):
        """infer the output shape for fmha_style_b2b_bmm."""
        q, k, v = self._attrs["inputs"][0:3]
        q_shape = q._attrs["shape"]
        k_shape = k._attrs["shape"]
        v_shape = v._attrs["shape"]
        if len(q_shape) != len(k_shape) or len(q_shape) != len(v_shape):
            raise RuntimeError(
                f"QKV ranks must be the same! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        if len(q_shape) != 4:
            raise RuntimeError(
                f"QKV must have rank == 4! Current rank: {len(q_shape)}, QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )

        if q_shape[0] != k_shape[0] or q_shape[0] != v_shape[0]:
            raise RuntimeError(
                f"QKV must have same batch size! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        if q_shape[2] != k_shape[2] or q_shape[2] != v_shape[2]:
            raise RuntimeError(
                f"QKV must have same head size! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        batch_size = q_shape[0]
        M0 = q_shape[1]
        K0 = q_shape[3]
        if K0 != k_shape[3]:
            raise RuntimeError(
                f"Q K shapes are not compatible! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        N0 = k_shape[1]
        if N0 != v_shape[1]:
            raise RuntimeError(
                f"K V shapes are not compatible! QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}."
            )
        N1 = v_shape[3]

        if self._attrs["causal_type"] != CausalType.NO_CAUSAL:
            if M0 != N0:
                raise RuntimeError(
                    f"When causal_type is enabled, M0 must be equal to N0. Current {M0=}, {N0=}."
                )

        head_size = q_shape[2]
        output_shape = [batch_size, M0, head_size, N1]

        if len(self._attrs["inputs"]) == 4:
            bias = self._attrs["inputs"][3]
            bias_shape = bias._attrs["shape"]
            bias_expected_shape = [batch_size, head_size, M0, N0]
            broadcastable, _ = shape_utils.get_broadcast_max_shape(
                bias_shape, bias_expected_shape
            )
            if len(bias_shape) != 4:
                raise RuntimeError(
                    f"Expected bias rank 4. Current bias rank: {len(bias_shape)}."
                )
            if not broadcastable:
                raise RuntimeError(
                    f"bias shape is not compatible with Q K! "
                    f"QKV shapes: {q_shape=}, {k_shape=}, {v_shape=}, "
                    f"bias shapes: {bias_shape=}, {bias_expected_shape=}."
                )
            if bias_shape[-1] != N0:
                raise RuntimeError(
                    f"Bias last dim is not broadcastable! Expected shape: {N0}, current bias shape: {bias_shape}"
                )
        return output_shape
infer the output shape for fmha_style_b2b_bmm.
python/aitemplate/compiler/ops/b2b_bmm/fmha_style_b2b_bmm.py