How to implement 4D attention mask within Flashattn?

[Yxxxb]

As shown in the figure, I would like to know how to implement such a 4D attention mask, or interleaved attention (i.e. some tokens use bidirectional attention and some tokens use casual attention) in Flashattn?

[Oasis-Git]

Hi,

Thanks for your interest in our job. This attention mask is mainly used for training and visualization. During inference the engine will split the request into sub request and for each sub request their attention mask is lower triangular(dense). So we do not need to design flash attention kernel for inference.