    def forward(self, x: Tensor, thw_shape: List[int]) -> Tuple[Tensor, List[int]]:
        """
        Args:
            x (Tensor): Input tensor.
            thw_shape (List): The shape of the input tensor (before flattening).
        """

        B, N, C = get_shape(x)
        if self.pool_first:
            x = ops.reshape()(x, [B, N, self.num_heads, C // self.num_heads])
            x = ops.permute()(x, [0, 2, 1, 3])
            q = k = v = x
            pass
            q, q_shape, k, k_shape, v, v_shape = self._qkv_pool(q, k, v, thw_shape)
            q_N, k_N, v_N = self._get_qkv_length(q_shape, k_shape, v_shape)
            q, k, v = self._reshape_qkv_to_seq(q, k, v, q_N, v_N, k_N, B, C)
            q, k, v = self._qkv_proj(q, q_N, k, k_N, v, v_N, B, C)
        else:
            if self.separate_qkv:
                q = k = v = x
                pass
                # TODO: implement when separate_qkv
                # q, k, v = self._qkv_proj(q, N, k, N, v, N, B, C)
            else:
                # compute q, k, v and perform pooling
                qkv = ops.permute()(
                    ops.reshape()(self.qkv(x), [B, N, 3, self.num_heads, -1]),
                    [2, 0, 3, 1, 4],
                )
                # input shape: 3, B, num_heads, seqlen, head_dim
                shape = get_shape(qkv)
                # obtain q, k, v from qkv
                qkv = ops.reshape()(qkv, [3 * B, self.num_heads, N, shape[-1]])
                (q, k, v) = ops.split()(qkv, B, dim=0)
            q, q_thw_shape, k, k_thw_shape, v, v_thw_shape = self._qkv_pool(
                q, k, v, thw_shape
            )

        # attention
        q_shape = get_shape(q)
        B, num_heads, seqlen, head_dim = get_shape(q)
        score = ops.transpose()(ops.mem_eff_attention(causal=False)(q, k, v), 1, 2)

        if self.residual_pool:
            score = ops.elementwise(FuncEnum.ADD)(score, q)

        score = ops.reshape()(ops.transpose()(score, 1, 2), [B, -1, self.dim])

        score = self.proj(score)
        assert self.dropout_rate == 0.0
        if self.dropout_rate > 0.0:
            score = self.proj_drop(score)

        return score, q_thw_shape
Args:
x (Tensor): Input tensor.
thw_shape (List): The shape of the input tensor (before flattening).
python/aitemplate/frontend/nn/multiscale_attention.py