    def __call__(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        bias: Tensor,
    ) -> Tensor:
        """call the op

        Note: [H,] means optional num-heads,
        if it exists for one input tensor, all need to have it,
        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)

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

        self._attrs["inputs"] = [q, k, v, bias]
        self._set_depth()
        output_shape, max_seq_len = self._infer_shapes()
        self._check_alignment()
        output = Tensor(
            output_shape,
            src_ops={self},
            dtype=self._attrs["inputs"][0]._attrs["dtype"],
        )
        self._attrs["outputs"] = [output]
        self._attrs["max_seq_len"] = max_seq_len

        return output
call the op

Note: [H,] means optional num-heads,
if it exists for one input tensor, all need to have it,
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)

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