Does magi support zigzag dispatch pattern?

[jordane95]

Hi, I'm trying to use magi attention as backend in mcore to replace nvte. But it seems that a lot of adaptations is required: cp split, rope, undispatch. The complexity comes from the custom comm pattern adopted by magi designed for specific attn mask. I just want to use regular varlen + casual mask + swa, so simple zigzag would be efficient enough. I'm wondering if there is support for using zigzag dispatch mode so I won't need to change data dispatch and rope part.

[hanwen-sun]

@jordane95 Thank you for using MagiAttention. currently, we do not support the native zigzag strategy. Integrating MagiAttention as a third-party library into Megatron does require some modifications. You can refer to our previous example here: https://github.com/SandAI-org/Megatron-LM-MagiAttention/blob/d464966381f6eec0324473a35132e73931a1a8df/magiattention_example/README.md.

We are also actively working on merging MagiAttention into Megatron-LM: NVIDIA/Megatron-LM#2749. Once merged, using MagiAttention within Megatron will be much more convenient. If you need any further assistance, please feel free to continue the discussion in this issue.