Distributed Attention Methods

[tsrikris]

Do you have plans to enable distributed attention methods (Ring Attention, Striped Attention, Tree Attention, etc.) via a FlexAttention-like interface? Most of the functional implementations today are on JAX.

[drisspg]

It should be possible to build to use FlexAttention as a component in many of these algorithms, we have an arg to "return_lse" which is typically required to merge different attention chunks. We would accept + add an example showing how to do this