Fuse per-sequence AlltoAll into a unified one in GDN forward by xuantengh · Pull Request #4913 · NVIDIA/Megatron-LM

xuantengh · 2026-05-21T06:21:44Z

What does this PR do ?

PR #2645 adds the THD support for GDN, but it requires per-sequence, per-section (qkv, alpha, beta, gate) AlltoAll converting activation between CP <-> HP. This PR fuses the per-sequence AlltoAll into a unified one, with local indexing mappings to reorder at the sequence dimension and feature dimension.

wandb link for running a Qwen3 Next proxy model with 16K sequence length and CP size = 4, the iteration time speedup is around 7% - 10%.

Contribution process

Pre-checks

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

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

copy-pr-bot · 2026-05-21T06:21:47Z

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.

xuantengh · 2026-05-21T06:21:59Z

/claude review

claude

LGTM

xuantengh · 2026-05-21T06:26:41Z

/ok to test 8dcfd7d

xuantengh · 2026-05-22T14:55:59Z

/ok to test 6d28256

xuantengh · 2026-05-22T15:00:29Z

/claude review

xuantengh · 2026-05-23T03:55:18Z

/claude review

claude

Light review: the core logic for fusing per-sequence AlltoAll into a single batched a2a + local permutations looks correct. The mamba_context_parallel.py cleanup is a clean replacement with existing utility functions. Test coverage in TestFusedThdAllToAll is solid — it validates both directions and a round-trip identity check across multiple cu_seqlens configurations and cp_sizes.

Only nit: unused import os and a commented-out os.environ.setdefault line in the test file (see inline comments).

xuantengh · 2026-06-09T04:45:51Z

/ok to test 5d261d1

xuantengh · 2026-06-09T10:49:54Z

/ok to test 4137f93

svcnvidia-nemo-ci · 2026-06-09T12:07:37Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/27205020791

claude Bot approved these changes May 21, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public May 21, 2026 06:27 Inactive

copy-pr-bot Bot temporarily deployed to test May 21, 2026 06:27 Inactive

copy-pr-bot Bot temporarily deployed to public May 21, 2026 06:30 Inactive

copy-pr-bot Bot temporarily deployed to public May 21, 2026 06:38 Inactive

xuantengh force-pushed the xuantengh/gdn_a2a branch from 2c59663 to 6d28256 Compare May 22, 2026 10:34

copy-pr-bot Bot temporarily deployed to public May 22, 2026 14:56 Inactive

copy-pr-bot Bot temporarily deployed to test May 22, 2026 14:57 Inactive

copy-pr-bot Bot temporarily deployed to public May 22, 2026 15:00 Inactive

claude Bot reviewed May 22, 2026

View reviewed changes

Comment thread tests/unit_tests/ssm/test_gated_delta_net.py Outdated

copy-pr-bot Bot temporarily deployed to public May 22, 2026 15:08 Inactive

xuantengh force-pushed the xuantengh/gdn_a2a branch from f5714af to efb43a2 Compare May 23, 2026 03:54

claude Bot reviewed May 23, 2026

View reviewed changes

Comment thread tests/unit_tests/ssm/test_gated_delta_net.py

claude Bot reviewed May 23, 2026

View reviewed changes

Comment thread tests/unit_tests/ssm/test_gated_delta_net.py Outdated

claude Bot reviewed May 23, 2026

View reviewed changes

xuantengh force-pushed the xuantengh/gdn_a2a branch from efb43a2 to de01425 Compare May 23, 2026 04:01

xuantengh added 10 commits June 8, 2026 20:00

use qwen3 model config for testing

0dd7928

add head perm

5ad0a44

fix test

a188b24

move unused function to test file

aa3fbf9

update comments

757612e

move head_perm to forward

2a80dcf

add lru cache

a2bf804

type annotation

41a848f

fix tuple

859a6dc

fix test nvls

5d261d1

xuantengh force-pushed the xuantengh/gdn_a2a branch from 7f7b668 to 5d261d1 Compare June 9, 2026 04:32

copy-pr-bot Bot temporarily deployed to public June 9, 2026 04:46 Inactive

copy-pr-bot Bot temporarily deployed to test June 9, 2026 04:46 Inactive

copy-pr-bot Bot temporarily deployed to public June 9, 2026 04:49 Inactive

copy-pr-bot Bot temporarily deployed to public June 9, 2026 04:50 Inactive

Merge branch 'main' into xuantengh/gdn_a2a

4137f93

copy-pr-bot Bot temporarily deployed to public June 9, 2026 10:50 Inactive

copy-pr-bot Bot temporarily deployed to test June 9, 2026 10:50 Inactive

copy-pr-bot Bot temporarily deployed to public June 9, 2026 10:53 Inactive

Victarry mentioned this pull request Jun 10, 2026

[ROADMAP][2026 Q2] Megatron Core MoE Roadmap #4815

Open

71 tasks

svcnvidia-nemo-ci mentioned this pull request Jun 12, 2026

chore: nightly sync main into dev (12_06_2026) #5314

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fuse per-sequence AlltoAll into a unified one in GDN forward#4913

Fuse per-sequence AlltoAll into a unified one in GDN forward#4913
xuantengh merged 13 commits into
NVIDIA:mainfrom
xuantengh:xuantengh/gdn_a2a

xuantengh commented May 21, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented May 21, 2026

Uh oh!

xuantengh commented May 21, 2026

Uh oh!

claude Bot left a comment

Uh oh!

xuantengh commented May 21, 2026

Uh oh!

xuantengh commented May 22, 2026

Uh oh!

xuantengh commented May 22, 2026

Uh oh!

Uh oh!

xuantengh commented May 23, 2026

Uh oh!

Uh oh!

Uh oh!

claude Bot left a comment

Uh oh!

xuantengh commented Jun 9, 2026

Uh oh!

xuantengh commented Jun 9, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Conversation

xuantengh commented May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented May 21, 2026

Uh oh!

xuantengh commented May 21, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

xuantengh commented May 21, 2026

Uh oh!

xuantengh commented May 22, 2026

Uh oh!

xuantengh commented May 22, 2026

Uh oh!

Uh oh!

xuantengh commented May 23, 2026

Uh oh!

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

xuantengh commented Jun 9, 2026

Uh oh!

xuantengh commented Jun 9, 2026

Uh oh!

svcnvidia-nemo-ci commented Jun 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

xuantengh commented May 21, 2026 •

edited

Loading