def forward(self, qkv, encoder_kv=None, mask=None):
        """
        Apply QKV attention.

        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
        :return: an [N x (H * C) x T] tensor after attention.
        """
        bs, width, length = qkv.shape
        assert width % (3 * self.n_heads) == 0
        ch = width // (3 * self.n_heads)
        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
        if encoder_kv is not None:
            assert encoder_kv.shape[1] == self.n_heads * ch * 2
            ek, ev = encoder_kv.reshape(bs * self.n_heads, ch * 2, -1).split(ch, dim=1)
            k = th.cat([ek, k], dim=-1)
            v = th.cat([ev, v], dim=-1)
        scale = 1 / math.sqrt(math.sqrt(ch))
        weight = th.einsum("bct,bcs->bts", q * scale, k * scale)
        if mask is not None:
            mask = F.pad(mask, (0, length), value=0.0)
            mask = (
                mask.unsqueeze(1)
                .expand(-1, self.n_heads, -1)
                .reshape(bs * self.n_heads, 1, -1)
            )
            weight = weight + mask
        weight = th.softmax(weight, dim=-1)
        a = th.einsum("bts,bcs->bct", weight, v)
        return a.reshape(bs, -1, length)

        Apply QKV attention.

        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
        :return: an [N x (H * C) x T] tensor after attention.
        
ldm/modules/karlo/kakao/modules/unet.py