Repository: Project-MONAI/MONAI
File: monai/networks/blocks/crossattention.py (Python)

Code before:
        """
        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        b, t, c = x.size()  # batch size, sequence length, embedding dimensionality (hidden_size)

        q = self.to_q(x)
        kv = context if context is not None else x
        _, kv_t, _ = kv.size()
        k = self.to_k(kv)
        v = self.to_v(kv)

        if self.attention_dtype is not None:
            q = q.to(self.attention_dtype)
            k = k.to(self.attention_dtype)

        q = q.view(b, t, self.num_heads, c // self.num_heads).transpose(1, 2)  # (b, nh, t,  hs)
        k = k.view(b, kv_t, self.num_heads, c // self.num_heads).transpose(1, 2)  # (b, nh, kv_t, hs)
        v = v.view(b, kv_t, self.num_heads, c // self.num_heads).transpose(1, 2)  # (b, nh, kv_t, hs)

        if self.use_flash_attention:
            x = torch.nn.functional.scaled_dot_product_attention(q, k, v).contiguous()
        else:

            att_mat = torch.einsum("blxd,blyd->blxy", q, k) * self.scale
            # apply relative positional embedding if defined
            att_mat = (
                self.rel_positional_embedding(x, att_mat, q) if self.rel_positional_embedding is not None else att_mat
            )

            if self.causal:
                att_mat = att_mat.masked_fill(self.causal_mask[:, :, :t, :kv_t] == 0, float("-inf"))

            att_mat = att_mat.softmax(dim=-1)

            if self.save_attn:
                # no gradients and new tensor;
                # https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html
                self.att_mat = att_mat.detach()

            att_mat = self.drop_weights(att_mat)
            x = torch.einsum("bhxy,bhyd->bhxd", att_mat, v)
        x = self.out_rearrange(x)
        x = self.out_proj(x)
        x = self.drop_output(x)
        return x


Review comment (suggestion):
Looks we can only add `rel_pos_embedding` when not using flash attention. I'm not sure whether it's same with the version in the generative repo, I didn't find the reference.
If that's the case, we may also need add a warning if `self.rel_positional_embedding is not None` and `use_flash_attention=True`.

Code after:
        Args:
            x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C
            context (torch.Tensor, optional): context tensor. B x (s_dim_1 * ... * s_dim_n) x C

        Return:
            torch.Tensor: B x (s_dim_1 * ... * s_dim_n) x C
        """
        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        b, t, c = x.size()  # batch size, sequence length, embedding dimensionality (hidden_size)

        q = self.to_q(x)
        kv = context if context is not None else x
        _, kv_t, _ = kv.size()
        k = self.to_k(kv)
        v = self.to_v(kv)

        if self.attention_dtype is not None:
            q = q.to(self.attention_dtype)
            k = k.to(self.attention_dtype)

        q = q.view(b, t, self.num_heads, c // self.num_heads).transpose(1, 2)  # (b, nh, t,  hs) #
        k = k.view(b, kv_t, self.num_heads, c // self.num_heads).transpose(1, 2)  # (b, nh, kv_t, hs)
        v = v.view(b, kv_t, self.num_heads, c // self.num_heads).transpose(1, 2)  # (b, nh, kv_t, hs)

        if self.use_flash_attention:
            x = torch.nn.functional.scaled_dot_product_attention(
                query=q.transpose(1, 2),
                key=k.transpose(1, 2),
                value=v.transpose(1, 2),
                scale=self.scale,
                dropout_p=self.dropout_rate,
                is_causal=self.causal,
            ).transpose(
                1, 2
            )  # Back to (b, nh, t, hs)
        else:
            att_mat = torch.einsum("blxd,blyd->blxy", q, k) * self.scale
            # apply relative positional embedding if defined
            if self.rel_positional_embedding is not None:
                att_mat = self.rel_positional_embedding(x, att_mat, q)

            if self.causal:
                att_mat = att_mat.masked_fill(self.causal_mask[:, :, :t, :kv_t] == 0, float("-inf"))

            att_mat = att_mat.softmax(dim=-1)

            if self.save_attn:
                # no gradients and new tensor;
                # https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html
                self.att_mat = att_mat.detach()