    def __call__(self, x: Tensor, cu_seqlens: Tensor) -> Tensor:
        """call the op

        Parameters
        ----------
        x : float16
            QKV tensor
            shape: (batch*seqlen, 3, num_heads, head_size)
        cu_seqlens : int
            seq lens tensor
            shape (batch_size + 1)

        Returns
        ----------
            Tensor
        """
        self._attrs["inputs"] = [x, cu_seqlens]
        self._set_depth()
        self._extract_exec_path(x)
        output_shape = self._infer_shapes(x, cu_seqlens)
        output = Tensor(output_shape, src_ops={self})

        batch_size = self._attrs["batch_size"]
        max_seq_len = self._attrs["max_seq_len"]
        total = x._attrs["shape"][0]._attrs["values"][0]
        num_heads = x._attrs["shape"][2]._attrs["values"][0]
        head_size = x._attrs["shape"][3]._attrs["values"][0]
        assert head_size in [8, 16, 32, 64, 128]
        self._attrs["head_size"] = head_size

        base_N = 256  # SM80
        if max_seq_len <= 128:
            seq_len = 128
        elif max_seq_len <= 256:
            seq_len = 256
        else:
            seq_len = ((max_seq_len + base_N - 1) // base_N) * base_N
        self._attrs["seq_len"] = seq_len

        self._attrs["workspace"] = (
            4 * num_heads * (total * head_size + batch_size * seq_len)
        )
        self._attrs["outputs"] = [output]
        return output
call the op

Parameters
----------
x : float16
QKV tensor
shape: (batch*seqlen, 3, num_heads, head_size)
cu_seqlens : int
seq lens tensor
shape (batch_size + 1)

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