Skip to content

Re-add fused GDN TransformerEngine integration - #7583

Merged
gautham-kollu merged 29 commits into
NVIDIA:mainfrom
ksivaman:gdn_attn
Oct 10, 2026
Merged

gautham-kollu merged 29 commits into
NVIDIA:mainfrom
ksivaman:gdn_attn

Conversation

@ksivaman

Copy link
Copy Markdown
Member

Re-add the GDN TransformerEngine integration from upstream #6645 which was reverted in #7326 as part of a combined revert for unrelated CI failures.

  • I, the PR author, have personally reviewed every line of this PR.

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests
  • If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
  • 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"

  1. When your PR is ready, click Ready for Review.
  2. 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.

Re-add the GDN TransformerEngine integration from upstream PR NVIDIA#6645
(54c62df), which was reverted in PR NVIDIA#7326 (1e137dd) as part of a
combined revert for unrelated CI failures. Adds a TransformerEngine
GatedDeltaNetAttention kernel backend option (gdn_kernel_backend:
transformer_engine|fla|torch) alongside the existing FLA/torch paths,
resolved against the current gdn.py two-stage forward split introduced
by the ShortcutMoE change (NVIDIA#6959) that landed after the revert.

Fixes a golden-config mismatch in the original commit: the
gdn_kernel_backend default is "fla", but
tests/unit_tests/models/test_hybrid_moe_model.py's GOLDEN_CONFIG had
recorded "transformer_engine", which would have failed
test_constructor's config-drift check.

Co-authored-by: gautham-kollu <gkollu@nvidia.com>
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 22, 2026

Copy link
Copy Markdown

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.

@ksivaman
ksivaman requested review from a team as code owners September 22, 2026 15:13
@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft September 22, 2026 15:13
@github-actions

Copy link
Copy Markdown
Contributor

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.

@ksivaman
ksivaman marked this pull request as ready for review September 22, 2026 15:16
@gautham-kollu

Copy link
Copy Markdown
Contributor

/ok to test 5a3f5f3

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Sep 22, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Sep 22, 2026
@gautham-kollu

Copy link
Copy Markdown
Contributor

Deterministic mode is not supported with TE backend. Hence adding "exempt" label

@gautham-kollu

Copy link
Copy Markdown
Contributor

/claude fix

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

❌ Claude fix stopped because a workflow step failed. Inspect the run.

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Oct 7, 2026

Copy link
Copy Markdown
Member Author

/ok to test 4fbd64f

@ksivaman

ksivaman commented Oct 8, 2026

Copy link
Copy Markdown
Member Author

/ok to test bf3c4fb

@balasaajay

Copy link
Copy Markdown
Contributor

/ok to test 6f31b04

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
@ksivaman

ksivaman commented Oct 9, 2026

Copy link
Copy Markdown
Member Author

/ok to test 6078a33

@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

This branch was previously deployed

2 inactive deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants