magi attention support

[jordane95]

MagiAttention has a good support for flexible attention masking designs, and support almost all modern features in long context training. Should we add it as another attention backend?

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!

[guyueh1]

@jordane95 thanks for starting this topic. For my own education, how is Magi attention different from meta's FlexAttention in terms of support matrix and performance? NeMo-RL has NVIDIA-AutoModel as one of the training backend and AutoModel supports flexattention as attention backend, so in theory we have flex attention

[jordane95]

Hi @guyueh1 , I have tried flex attention in fsdp. It works in simple setting but the main issue is that it currently doesn't support cp in case we want to train large model on long context. It's hard to implement it in mcore as well since mcore delegate cp handling to transformer engine.

Your understanding is correct. Magi attention is basically flex-flash-attention + context parallelism with custom load-balancing technique for irregular attention mask. It has native cp support so we do not need to implement cp from scratch.

[guyueh1]

Ok, having cp + flex-flash-attention in automodel makes sense to me, but we will probably need AutoModel to drive this integration, @snowmanwwg can you tag AutoModel folks here to (1) explain if they already have a roadmap for CP + flex attention, and (2) give their opinions on integrating magiattention as an attention backend in AutoModel

[terrykong]

@akoumpa