Query about batch dimension in MultiHeadCrossAttention code

[Yanjia0]

Here is the code:
https://github.com/PixArt-alpha/PixArt-sigma/blob/1b9027e84e37b7efb6d02f00cc71826c20a2c0cb/diffusion/model/nets/PixArt_blocks.py#L28

def forward(self, x, cond, mask=None):
        # query: img tokens; key/value: condition; mask: if padding tokens
        B, N, C = x.shape

        q = self.q_linear(x).view(1, -1, self.num_heads, self.head_dim)
        kv = self.kv_linear(cond).view(1, -1, 2, self.num_heads, self.head_dim)
        k, v = kv.unbind(2)
        attn_bias = None
        if mask is not None:
            attn_bias = xformers.ops.fmha.BlockDiagonalMask.from_seqlens([N] * B, mask)
        x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, attn_bias=attn_bias)
        x = x.view(B, -1, C)
        x = self.proj(x)
        x = self.proj_drop(x)

The first parameter of q, k is 1. We are wondering why you fix this number as 1 instead of using B. Won't this mess up the information of different videos and images within a batch? Or during the whole training process, you use batchsize=1 for training?

[wangqiang9]

I have the same confusion.