Key pad mask support for flash attention

[MathematicalMagus]

❓ Questions and Help

Hello!
I noticed that FlashAttention forward supports the key_padding_mask argument, while your implementation allows only for the Causal mask. Can the key_padding_mask also (easily) be implemented in xformers?

[danthe3rd]

Hi,

Yes it's something we plan on adding. As you mentioned, the kernels already support it, so it's more a matter of figuring out the right API for this - we're looking into pytorch's nested tensors to see if it's the right abstraction

[danthe3rd]

We just merged this support - through a special attention bias.
You can learn more on the doc website.
Note that the backward is only supported by Flash, and hence won't work on V100 or older devices or even fp32.

[zen-d]

@danthe3rd But I think it means attn_mask in https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html#torch.nn.MultiheadAttention, not corresponding to key_padding_mask? correct me if I miss something.