Support Ring Attention https://arxiv.org/abs/2310.01889 and Tree Attention https://arxiv.org/abs/2408.04093 using Flex Attention.
Support Ring Attention https://arxiv.org/abs/2310.01889 and Tree Attention https://arxiv.org/abs/2408.04093 using Flex Attention.