[RFC] Refactor and simplify scaled_dot_product_attention for easier extensibility

[fmassa]

🚀 Feature

scaled_dot_product_attention is a key component of most transformer architectures, and is used almost everywhere in xformers.
Over time, its API has been extended to support a few different cases (either performance improvements, or new functionality). With additional improvements planned down the road to enable even other cases, I think it would be a good moment to step back and re-write scaled_dot_product_attention so that it is extensible enough and generic enough for the use-cases we care about.

Here is a list of proposed changes:

1 - Let it take [batch, head, seq, dim] instead of 3d tensors

We originally made scaled_dot_product_attention take 3d tensors to simplify the implementation of some early sparse components. While this is ok in general, handling different types of masks types (same per batch, or per head only) means that we would need to materialize the full mask, instead of relying on broadcast for that (which would save memory). Plus, letting it take 4d tensors would make its API more aligned with what PyTorch expects as well.

2 - Remove AttentionMask support inside scaled_dot_product_attention

The use of AttentionMask inside scaled_dot_product_attention is mostly there as a performance optimization for the softmax in the causal case.
Instead, I propose to slightly modify its behavior so that we have this be an implementation detail of each specific Tensor subclass passed as a mask in scaled_dot_product_attention.

3 - Use __torch_function__ to handle the dispatch for different mask types

There is a lot of code in scaled_dot_product_attention which is basically performing the dispatching between different mask types (SparseCS / Tensor), but it doesn't handle block-sparse for example.

I think we should instead always use custom Tensor types and the dispatch mechanism from PyTorch to perform the dispatch under the hood. This would enable new Tensor subclasses to be added without having to change the underlying implementation in scaled_dot_product_attention. This is drafted in https://github.com/facebookresearch/xformers/tree/sparse_cleanup (but needs rebase)

4 - attn_mask becomes attn_bias

We currently have two behaviors which are (partly) enabled: masking and bias addition. I think we should only handle bias addition, and we would change the behavior of the sparse tensors to actually be masked tensors which are backed by sparse operators.
If users want the masking behavior, this can be done through a -inf in the bias.

Thoughts?

cc @blefaudeux @dianaml0 for thoughts

[dianaml0]

These seem like really nice improvements, I think 1 and 4 would help with integration with other codebases. For 2, would the functionality of AttentionMask just be encompassed by the Tensor classes instead?

[fmassa]

For 2, would the functionality of AttentionMask just be encompassed by the Tensor classes instead?

Yes, the idea would be that most of what is present in AttentionMask would be present instead inside the tensor subclasses. What would be left out are convenience functions to go from bool to float masks and vice-versa, but this can be left as a separate function.

[blefaudeux]

Thanks for the write up @fmassa, and sorry for the delay..

1 and 2 are kind of no brainers to me: 1 (4d tensors, aligned with pytorch) is side effect free as long as the repo is internally coherent (and the unit test will catch that), there are material benefits. 2 (AttentionMask) is completely fine by me, I think that it goes with 3 and 4 really but even just 2 would be ok. At the time AttentionMask was a way to factorize some redundant code and do the checks and transforms only once (not just causal actually, but accepting bool and float masks, I think that you'll see that in the tests).

3 (torch_function to dispatch the mask types) is also fine by me, I think that it makes it a little more complicated for outsiders to contribute, but it's not a part which gets people's attention typically (it's super useful, but I think that people mostly look for this new fancy thing instead), so that's all good. I agree that all the dispatch and size and type handling being nicely wrapped somewhere is nice. Top of head torch_function was not available until a couple of torch releases ?

4 (attn_mask -> attn_bias) is also fine by me, just an open question with respect to the mask size in that case, floats take 16 to 32 bits typically, which is a lot more than your typical bool, isn't that an issue ? Other than that it makes it very clear what it does, it's nice for models like Swin, and it handles both mask & bias purposes, so I just see wins