Feature Request: Dual Chunk Attention for Long Context

[sbhavani]

Summary

Request implementation of Dual Chunk Attention (DCA), the technique used by Qwen 2/2.5 for efficient long-context training and inference on 100K+ token sequences.

Motivation

Training on sequences >100K tokens is essential for document understanding, code generation, and video models. Qwen 2.5 demonstrates that DCA enables 128K context with sub-quadratic memory:

• Quadratic memory growth - Standard attention scales O(n²), making 128K+ sequences impractical even with FlashAttention
• CP doesn't reduce attention complexity - Context Parallelism helps with parallelism but still computes full attention
• YaRN/ABF are position-encoding only - They extend positional range but don't address attention memory

DCA reduces memory from O(n²) to O(n·c) where c = chunk size, by combining local intra-chunk attention with global inter-chunk attention.

Current State

Megatron has strong long-context parallelism (CP with p2p, a2a, allgather, hierarchical CP, YaRN, FlashAttention) but lacks algorithmic attention optimizations like chunked or sparse attention patterns.

Ask

1. New attention module - DualChunkAttention with intra-chunk (local) + inter-chunk (global) attention
2. Integration with Context Parallelism, FlashAttention, and GQA/MQA

References

• Qwen2 Technical Report - DCA for 128K context

[Ternura143]

Could you please assign this issue to me? I’ll start with the DCA implementation and open a draft PR soon.

[fengxy-03]

Could you please assign this issue to me? I’ll start with the DCA implementation and open a draft PR soon.
@Ternura143 Hi, Have you successfully integrated DCA into Megatron-LM?

[Ternura143]

Could you please assign this issue to me? I’ll start with the DCA implementation and open a draft PR soon.
@Ternura143 Hi, Have you successfully integrated DCA into Megatron-LM?

Hi @fengxy-03 , apologies for the delay recently due to other commitments. I'm currently working on it and will endeavour to upload the complete PR by the end of March.

[Ternura143]

Hi @fengxy-03, the draft PR is up: #4048. It includes the core DCA module with intra-chunk, successive-chunk, and inter-chunk attention, plus FlashAttention backend and YARN mscale support. Context Parallelism support is planned for a follow-up iteration.

[sbhavani]

@Ternura143 thanks for creating the PR!

Have you taken a look at Megatron Bridge? We aim to keep new model architectures and experimental features there. If it's possible to add the DCA module in Megatron Bridge it would simplify the review process.

[Ternura143]

Hi @sbhavani , thanks for the suggestion! That makes sense. I’ll check Megatron Bridge and evaluate moving the DCA implementation there. Happy to adapt if that’s the preferred direction.

[Ternura143]

Hi @sbhavani, thanks for the suggestion! I went and took a look at the Megatron Bridge repo to understand how it fits in.

From what I can see, Bridge sits on top of Megatron Core and is focused on things like HF checkpoint conversion, training recipes, and model providers. DCA on the other hand is modifying the core attention math itself — it changes how RoPE is applied, splits sequences into chunks, and merges outputs with log-sum-exp renormalization. This feels like it belongs in megatron/core, similar to how DSA and GatedDeltaNet are both implemented under megatron/core/transformer/experimental_attention_variant/.

The integration is also pretty lightweight — only 8 lines changed in attention.py (just a RoPE bypass), and the rest is self-contained within the existing experimental_attention_variant framework.

That said, I'd be happy to follow up with a Bridge-side contribution later — e.g., adding a DCA training recipe or HF conversion support for models that use DCA. Would that work as a plan?

If you had a different interpretation in mind though, I'm open to discussing! Just want to make sure we're on the same page about what moving to Bridge would look like.

[jagmarques]

Relevant addition to the DCA discussion: memory at long context is a two-layer problem — attention computation (O(n²)) and KV cache storage. DCA addresses the former; the latter keeps growing linearly with sequence length and becomes the dominant bottleneck beyond ~200K tokens even with chunked attention.

For the storage side, KV cache compression can be applied orthogonally. NexusQuant uses E8 lattice vector quantization to compress KV entries 7–10x training-free, validated on Mistral-7B and Llama-3-8B. Combined with DCA's O(n·c) attention memory, you'd have both dimensions covered.

The integration point in Megatron would be post-prefill: compress KV blocks before storing them in the distributed KV buffer, decompress at attention time. Since DCA already splits KV into intra/inter-chunk pools, NexusQuant could apply different compression ratios to each (inter-chunk tokens are accessed less frequently and tolerate higher compression).

Worth considering for the long-context roadmap — the DCA implementation you're proposing would be a natural place to plug this in.

[Ternura143]

@jagmarques Great point about the KV cache bottleneck at very long sequences. There's been some promising work on this front — NVIDIA's KVTC (ICLR 2026) achieves 20-40x compression via PCA decorrelation + adaptive quantization, and RocketKV (NVlabs, ICML 2025) pushes up to 400x with two-stage coarse-to-fine compression. The idea of applying different compression ratios to intra vs. inter-chunk KV pools is interesting and could complement either approach nicely.

For this PR, I'd like to keep the scope focused on the core DCA attention mechanism. KV cache compression is a great complementary direction for a future follow-up — and DCA's explicit chunk structure could indeed provide natural integration points. Thanks for flagging this!

[Ternura143]

Hi @sbhavani, thanks for the suggestion! I went and took a look at the Megatron Bridge repo to understand how it fits in.

From what I can see, Bridge sits on top of Megatron Core and is focused on things like HF checkpoint conversion, training recipes, and model providers. DCA on the other hand is modifying the core attention math itself — it changes how RoPE is applied, splits sequences into chunks, and merges outputs with log-sum-exp renormalization. This feels like it belongs in megatron/core, similar to how DSA and GatedDeltaNet are both implemented under megatron/core/transformer/experimental_attention_variant/.

The integration is also pretty lightweight — only 8 lines changed in attention.py (just a RoPE bypass), and the rest is self-contained within the existing experimental_attention_variant framework.

That said, I'd be happy to follow up with a Bridge-side contribution later — e.g., adding a DCA training recipe or HF conversion support for models that use DCA. Would that work as a plan?

If you had a different interpretation in mind though, I'm open to discussing! Just want to make sure we're on the same page about what moving to Bridge would look like.

Hi @sbhavani, just wanted to follow up on the Bridge discussion — I left a detailed analysis above about how DCA fits within megatron/core vs. Bridge. Would love your thoughts when you get a chance so I can move forward accordingly. Thanks!

[jagmarques]

Makes sense to keep this PR focused on DCA core. The intra vs inter-chunk compression idea can live in a follow-up. Good luck with the review.