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

        Parameters
        ----------
        qkv : float16
            QKV tensor
            shape: (b, seqlen, num_heads, Kv)

        Returns
        ----------
            Tensor
        """

        head_size_v = v._attrs["shape"][3]._attrs["values"][0]
        self._attrs["head_size"] = head_size_v

        self._attrs["inputs"] = [q, k, v]
        if self._attrs["variable_seq_length_kv"]:
            assert lengths_kv is not None
            self._attrs["inputs"].append(lengths_kv)
        if self._attrs["variable_seq_length_q"]:
            assert lengths_q is not None
            self._attrs["inputs"].append(lengths_q)
        self._set_depth()
        self._extract_exec_path(q)
        output_shape = self._infer_shapes(q, v)

        required_workspace_size = self._compute_required_workspace(
            output_shape, q._attrs["shape"], k._attrs["shape"]
        )
        self._attrs["workspace"] = required_workspace_size
        _LOGGER.debug(f"Required workspace size: {required_workspace_size}")
        output = Tensor(
            output_shape,
            src_ops={self},
            dtype=self._attrs["inputs"][0]._attrs["dtype"],
        )
        self._attrs["outputs"] = [output]
        return output
call the op

Parameters
----------
qkv : float16
QKV tensor
shape: (b, seqlen, num_heads, Kv)

Returns
----------
Tensor
python/aitemplate/compiler/ops/attention/mem_eff_attention.py