Mask for flash attention

[hokiyoshi]

Can you explain to me why Swin Flash attention didn't use the mask for attention?

[zzd1992]

I suppose "mask" means the causal mask for language models.
In this project, we deal with images. We suppose a pixel can access to all other pixels within a window.
Thus, currently there is no mask for Swin Flash attention.

However, if you deal with auto-regressive image processing tasks, you should add mask for attention.

[hokiyoshi]

Hi, I mean the mask when shift windows attention.

This code:

if self.shift_size > 0:
            # calculate attention mask for SW-MSA
            H, W = self.input_resolution
            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
            h_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            w_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            cnt = 0
            for h in h_slices:
                for w in w_slices:
                    img_mask[:, h, w, :] = cnt
                    cnt += 1

            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))`

Anyway, I want to learn triton. Do you have source ?
Besides, I want apply Swin Transformer in to my model but I realize that it requires at least C = 16 after split heads. Do you have any solutions for lower channel ?