How to implement a Custom Block Diagonal Mask?

[TamirFriedman-RecoLabs]

Hi folks, I saw the implementation (or what I thought is the implementation) of BlockDiagonalCausalMask in which it just inherit BlockDiagonalMask and overrides the _create_block_mask method with a LowerTriangularMask blocks.
I thought I could also do it myself and implement a BlockDiagoNULLMask (where the diagonal itself is also masked out).
But as I saw, it doesn't really work out. all the implementations refuse to get it.
How can I implement it using the building blocks in xformers? with an efficient non-materialized block matrix like this one? there is an easy way or I'll need to implement a low-level kernel in CUDA/triton or some other scary interface? (although if that what is needed, I'm brave to deep dive)

class BlockDiagoNULLMask(fmha.attn_bias.BlockDiagonalMask):
    """
    Modification of `BlockDiagonalMask` where blocks are inner-connected
    except for the diagonal elements, which are masked from themselves.
    """

    def _create_block_mask(
        self: Self,
        shape: Tuple[int, ...],
        dtype: torch.dtype,
        device: str | torch.device,
    ) -> torch.Tensor:
        # Create a matrix filled with `-inf` on the diagonal and `0` elsewhere
        return torch.zeros(shape, dtype=dtype, device=device).fill_diagonal_(-torch.inf)

[danthe3rd]

Hi @TamirFriedman-RecoLabs
So the masks have 2 parts:
(1) The mask class itself + a way to materialize it - this is useful for testing for instance. This is the code you were looking at
(2) In the kernels, the mask type is passed as just a flag (eg causal=True for instance), because the trick to be efficient is to not materialize the full mask tensor

So you have a few options:
(a) Fully materialize your mask as a torch.Tensor, and pass it like that directly to memory_efficient_attention. It won't be fast, but it will work
(b) Write your own kernel that does it. It might be easier to start with the implementation in triton: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html