    def __call__(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        bias: Optional[Tensor] = None,
    ) -> Tensor:
        """call the op

        Parameters
        ----------
        q: Tensor, shape(B, M0, H, K0)
        k: Tensor, shape(B, N0, H, K0)
        v: Tensor, shape(B, N0, H, N1)
        bias: Tensor, shape(B, H, M0, N0), optional

        Returns
        ----------
        Tensor, shape(B, H, M0, N1)
        """

        if bias is not None:
            self._attrs["inputs"] = [q, k, v, bias]
        else:
            self._attrs["inputs"] = [q, k, v]
        self._set_depth()
        output_shape = self._infer_shapes()
        self._check_alignment()
        output = Tensor(
            output_shape,
            src_ops={self},
            dtype=self._attrs["inputs"][0]._attrs["dtype"],
        )
        self._attrs["outputs"] = [output]
        o_shape = [var.upper_bound() for var in output_shape]
        if o_shape[-1] > 128:
            self._attrs["workspace"] = 4 * np.prod(o_shape)

        return output
call the op

Parameters
----------
q: Tensor, shape(B, M0, H, K0)
k: Tensor, shape(B, N0, H, K0)
v: Tensor, shape(B, N0, H, N1)
bias: Tensor, shape(B, H, M0, N0), optional

Returns
----------
Tensor, shape(B, H, M0, N1)
python/aitemplate/compiler/ops/b2b_bmm/fmha_style_b2b_bmm.py