Add Magi Attention for varlen + Context Parallelism

[wangbinluo]

Summary

The existing _ContextParallel with FlexAttention uses AllGather to collect the full K/V sequence, then relies on BlockMask to skip cross-document computation. This works, but has limitations for varlen packed sequence training (SFT, DPO):

1. No varlen input support — upstream CP requires (B, H, S, D) padded tensors. Real SFT/DPO pipelines produce packed tokens with cu_seqlens (variable-length format used by Flash Attention). Using upstream CP requires converting back to padded format, wasting memory and compute on padding tokens.

2. Per-document length constraints — _PerDocumentHeadTailLoadBalancer requires every document to be divisible by 2 × cp_world_size. With CP=8, each document must be a multiple of 16 tokens. Real training data has arbitrary document lengths (137, 2041, 89, ...), making this impractical.

3. Communication inefficiency for packed sequences — AllGather sends the entire K/V sequence to every rank, including tokens from documents that a rank has no Q tokens for. With many short documents packed into a long sequence, most communicated K/V is unused.

This PR implements MagiAttention to address these limitations, enabling varlen packed sequences with Context Parallelism.

Closes #2536

What changed

New module: torchtitan/distributed/varlen_cp/ (~2900 lines)

• magi_attention.py — Custom autograd.Function implementing the Magi Attention forward/backward. Accepts cu_seqlens + packed tokens directly. Uses AllToAll-V to communicate only needed K/V per document (zero redundancy). LPT dispatch redistributes Q sub-chunks across ranks for load balancing.
• magi_dispatch.py — AllToAll-V communication helpers, Q redistribution, result gathering.
• dispatch_solver.py — LPT greedy solver that assigns Q sub-chunks to ranks based on actual attention workload (trapezoidal area estimation from cu_seqlens).
• flex_attn_kernels.py — Range-based flex_attention wrappers. Converts (Q_range, K_range, attn_type) descriptors into BlockMask for PyTorch-native flex_attention. Both forward and backward use flex_attention exclusively.
• mask_primitives.py — Utilities for converting cu_seqlens to attention slices and ranges.
• dispatch_ops.py — Dispatch/undispatch operations for Q redistribution.
• ring_attention.py — Entry point (varlen_ring_attention) and LSE merge utilities.

Modified files:

• context_parallel.py — "varlen" case sets _cp_mesh on VarlenAttentionWrapper modules.
• attention.py — VarlenAttentionWrapper._forward_cp() creates dispatch plan from cu_seqlens and calls Magi Attention when CP is active.

Tests (tests/unit_tests/test_varlen_cp/):

• Forward/backward correctness against single-GPU reference (5 configs × fwd + bwd)
• Ring-pass / all-gather reference equivalence
• Dispatch solver, mask primitives, dispatch ops unit tests

Design decisions

• Zero external dependencies — Uses only torch.nn.attention.flex_attention and standard torch.distributed collectives. No external kernel libraries.
• flex_attention for both forward and backward — No fallback paths. The compiled Triton kernel handles all attention computation.
• Cross-layer metadata cache — Dispatch plan, AllToAll-V split sizes, and BlockMask are cached across transformer layers within one forward pass (same cu_seqlens → same plan for all 96+ layers).
• Backward K/V reuse — Forward saves packed K/V via save_for_backward; backward reuses them instead of re-gathering.

Comparison with FlexAttention + CP

	FlexAttention + CP (upstream)	Varlen + Magi Attention (this PR)
Input format	(B, H, S, D) padded tensor	(total_tokens, H, D) + cu_seqlens
Document length constraint	Each doc divisible by 2 × cp_world_size	Only total sequence divisible by cp_world_size
Communication	AllGather full K/V — O(seq_len)	Per-doc AllToAll-V — O(needed_KV)
Document isolation	BlockMask mask_mod (computation-level)	cu_seqlens + per-doc packed K/V (communication + computation)
Load balancing	PTRR index reordering	LPT Q sub-chunk redistribution across ranks
Attention kernel	flex_attention	flex_attention (same)

Test plan

• Unit tests: dispatch solver, mask primitives, dispatch ops
• Unit tests: Magi Attention forward (5 configs) and backward (3 configs)
• Ring-pass / all-gather reference equivalence
• Integration test: varlen + CP + TP in training loop
• Multi-node benchmark

🤖 Generated with Claude Code

[meta-cla]

Hi @wangbinluo!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!