Add Dual Chunk Attention (DCA) for long-context training

[Ternura143]

What does this PR do ?

Implement Dual Chunk Attention (DCA) for efficient long-context training on 100K+ token sequences with sub-quadratic memory complexity.

Resolves #2797.

Changes

1. New: megatron/core/transformer/experimental_attention_variant/dca.py

   • DualChunkAttention module with three attention components:
     • Intra-chunk: standard causal attention within each chunk
     • Successive-chunk: locality-preserving attention to the immediately preceding chunk
     • Inter-chunk: fixed-distance attention to all earlier chunks
   • LSE-based output merging for correct softmax renormalization across chunks
   • FlashAttention backend with native GQA support (auto-fallback to unfused on CPU)
   • YARN mscale integration for RoPE concentration factor
   • Seamless fallback to standard attention for sequences shorter than chunk_len

2. Modified: transformer_config.py — Add dca_chunk_size (default: 8192), dca_local_size (default: 1024) config parameters with validation

3. Modified: attention.py — DCA integration: skip standard RoPE and pass rotary_pos_emb to DCA core_attention

4. Modified: experimental_attention_variant_module_specs.py — Add get_dca_module_spec() and register "dca" in the experimental attention variant framework

5. New: tests/unit_tests/transformer/test_attention_variant_dca.py — Unit tests for output shape, short-sequence equivalence, GQA, gradient flow, multi-chunk, causality, YARN mscale, FlashAttention

Usage

config = TransformerConfig(
    experimental_attention_variant="dca",
    dca_chunk_size=8192,
    dca_local_size=1024,
)

Status: Draft

Planned next steps:

• FlashAttention integration for memory-efficient chunk attention
• YARN mscale integration
• Context Parallelism support
• Packed sequence support
• Functional tests with end-to-end training

References

• Training-Free Long-Context Scaling of Large Language Models — Original DCA paper (ICML 2024)
• Qwen2 Technical Report — DCA + YARN integration

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[Ternura143]

Hi @ko3n1g , this is a draft PR implementing Dual Chunk Attention. Would appreciate any early feedback on the architecture direction before I proceed with Context Parallelism integration. Thank you!