Is there an efficient way to use memory_efficient_attention with a causal mask that has a small rectangle of zeros?

[arilato]

❓ Questions and Help

I understand I can pass a custom mask to memory_efficient_attention, but it is very inefficient for what I'm trying to do. Essentially, I'm adding a small rectangle of zeros (or -inf) to the attention mask near the lower right diagonal. Essentially, I want to formulate a sequence
(context, m1, m2)
s.t. m2 cannot attend to m1, each being a series of tokens.

Is there a memory-efficient way to do this in xformers without materializing the entire mask?