NVIDIA / NVIDIA/Megatron-LM

Feature Request: Dual Chunk Attention for Long Context

Open
#2,797 13 comments 0 reactions 1 assignee Claimed by @Ternura143 View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

## 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](https://arxiv.org/abs/2407.10671) - DCA for 128K context

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.