Skip to content

[Task32] feat: fla mcore chunkwise cp stage 1 #251

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jambow0320 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[Task32] feat: fla mcore chunkwise cp stage 1 #251

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
79d3559
feat(megatron): backport FLA chunkwise CP for GDN (Task 32 v1)
jambow0320 Aug 4, 2026
31b38db
version v2
jambow0320 Aug 7, 2026
88c6a90
fix headwise resolve cu_seqlens check
jambow0320 Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion docker/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,13 @@ RUN pip install nvidia-cudnn-cu12==9.16.0.29

FROM base as train

# flash-linear-attention 0.4.2 (pulls fla-core 0.4.2) is the first release carrying
# `fla.ops.cp`, which GDN chunkwise context parallel needs. NOTE: `fla-core` depends
# on an unpinned `torch`, so pip re-resolves torch's own looser `nvidia-cudnn-cu12` pin here;
# the explicit `nvidia-cudnn-cu12==9.16.0.29` install further down runs after this layer
# and is what the final image keeps.
RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation --no-cache-dir && \
pip install --no-cache-dir flash-linear-attention==0.4.1 && \
pip install --no-cache-dir flash-linear-attention==0.4.2 && \
pip install --no-cache-dir tilelang -f https://tile-ai.github.io/whl/nightly/cu128/

# FA3 (Hopper flash-attention), built from source. This commit's _flash_attn_forward carries
Expand Down
2 changes: 1 addition & 1 deletion docker/patch/latest/megatron.patch
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
../megatron/20260506-85bced0ae.patch
../megatron/20260805-85bced0ae.patch
Loading
Loading