Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (#4318) by sraman-rgb · Pull Request #4786 · NVIDIA/Megatron-LM

sraman-rgb · 2026-05-13T19:21:18Z

What does this PR do ?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

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 the @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.

) Co-authored-by: Siddhartha Raman S <sraman@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Xin Yao <xiny@nvidia.com>

copy-pr-bot · 2026-05-13T19:21:22Z

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.

github-actions · 2026-05-13T19:21:31Z

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:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

sraman-rgb · 2026-05-19T15:39:16Z

/ok to test 94ed88f

gautham-kollu · 2026-05-19T23:49:18Z

/ok to test 93be030

sraman-rgb · 2026-05-21T23:01:51Z

/ok to test 30d2b64

svcnvidia-nemo-ci · 2026-05-22T19:53:29Z

🔄 Merge queue validation started!

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

) (NVIDIA#4786) Co-authored-by: Siddhartha Raman S <sraman@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: gautham-kollu <gkollu@nvidia.com> Co-authored-by: Siddhartha Raman S <sraman@login-lyris01.lyris.clusters.nvidia.com>

* origin/main: (50 commits) Drain predecessor reduce-scatter at dispatch time (NVIDIA#4940) ci: Add allow_failure flag to gpt and moe recipes that are failing in nightlies (NVIDIA#4905) fix(tests): initialize num_microbatches calculator in vision cudagraph tests (NVIDIA#4986) test: re-enable test_pp2_create_cudagraphs_first_stage on TE 2.15+ (NVIDIA#4985) ci: Add support for MBridge job gating based on PR labels (NVIDIA#4926) test(ci): re-enable 8experts2parallel_multi_dist_optimizer_instances_1node (NVIDIA#4984) test: re-enable paged stashing MoE tests (NVIDIA#4978) Fix elastification unwrap_model import (NVIDIA#4972) Avoid offsetting functional test master port (NVIDIA#4973) test: enable NVTE_CUTEDSL_FUSED_GROUPED_MLP via pytest fixture (NVIDIA#4931) chore(beep boop 🤖): Bump (main) (2026-05-25) test(release): add release goldens for deepseekv3/nemotron3 and set tp2pp2 exit-interval (NVIDIA#4932) Fix `get_batch` return order to ignore BlendedDataset provenance fields (NVIDIA#4952) ci: restore perf test torchrun logs (NVIDIA#4951) Various training utils (NVIDIA#4872) ci: Update training script paths in BERT and T5 (NVIDIA#4939) [MXFP8/FP4-param-gather] Post processing after forced param AG in eval (NVIDIA#4562) Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA#4800) Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (NVIDIA#4318) (NVIDIA#4786) Fix paged stashing test submodules lookup (NVIDIA#4925) ... # Conflicts: # megatron/training/training.py

) (NVIDIA#4786) Co-authored-by: Siddhartha Raman S <sraman@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: gautham-kollu <gkollu@nvidia.com> Co-authored-by: Siddhartha Raman S <sraman@login-lyris01.lyris.clusters.nvidia.com>

Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (NVIDIA#4318

b7168d8

) Co-authored-by: Siddhartha Raman S <sraman@login-lyris02.lyris.clusters.nvidia.com> Co-authored-by: Xin Yao <xiny@nvidia.com>

sraman-rgb requested review from a team as code owners May 13, 2026 19:21

svcnvidia-nemo-ci marked this pull request as draft May 13, 2026 19:21

sraman-rgb marked this pull request as ready for review May 13, 2026 19:21

svcnvidia-nemo-ci requested a review from a team May 13, 2026 19:22

svcnvidia-nemo-ci added the complexity: medium label May 13, 2026

Merge branch 'main' into Dense_Grouped_GEMM_main

94ed88f

copy-pr-bot Bot temporarily deployed to public May 19, 2026 15:40 Inactive

copy-pr-bot Bot temporarily deployed to test May 19, 2026 15:40 Inactive

copy-pr-bot Bot temporarily deployed to public May 19, 2026 15:43 Inactive

copy-pr-bot Bot temporarily deployed to public May 19, 2026 15:51 Inactive

gautham-kollu requested a review from santhnm2 May 19, 2026 17:04

santhnm2 approved these changes May 19, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label May 19, 2026

timmoon10 reviewed May 19, 2026

View reviewed changes

Merge branch 'main' into Dense_Grouped_GEMM_main

93be030

copy-pr-bot Bot temporarily deployed to public May 19, 2026 23:49 Inactive

copy-pr-bot Bot temporarily deployed to test May 19, 2026 23:50 Inactive

copy-pr-bot Bot temporarily deployed to public May 19, 2026 23:52 Inactive

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

copy-pr-bot Bot temporarily deployed to public May 21, 2026 17:31 Inactive

copy-pr-bot Bot temporarily deployed to public May 21, 2026 17:39 Inactive

ericharper approved these changes May 21, 2026

View reviewed changes

Merge branch 'main' into Dense_Grouped_GEMM_main

30d2b64

sraman-rgb enabled auto-merge May 21, 2026 23:01

copy-pr-bot Bot temporarily deployed to public May 21, 2026 23:02 Inactive

copy-pr-bot Bot temporarily deployed to test May 21, 2026 23:02 Inactive

copy-pr-bot Bot temporarily deployed to public May 21, 2026 23:05 Inactive

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

jaredcasper approved these changes May 22, 2026

View reviewed changes

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

sraman-rgb added this pull request to the merge queue May 22, 2026

Merged via the queue into NVIDIA:main with commit fa7a23b May 22, 2026
77 of 82 checks passed

sraman-rgb deleted the Dense_Grouped_GEMM_main branch May 22, 2026 20:31

svcnvidia-nemo-ci mentioned this pull request May 27, 2026

chore: nightly sync main into dev (27_05_2026) #5029

Closed

timmoon10 mentioned this pull request May 28, 2026

Enable NVFP4 fused grouped MLP NVIDIA/TransformerEngine#3048

Merged

13 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (#4318)#4786

Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (#4318)#4786
sraman-rgb merged 6 commits into
NVIDIA:mainfrom
sraman-rgb:Dense_Grouped_GEMM_main

sraman-rgb commented May 13, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented May 13, 2026

Uh oh!

github-actions Bot commented May 13, 2026

Uh oh!

sraman-rgb commented May 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gautham-kollu commented May 19, 2026

Uh oh!

sraman-rgb commented May 21, 2026

Uh oh!

svcnvidia-nemo-ci commented May 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

sraman-rgb commented May 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issue tracking

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 13, 2026

Uh oh!

github-actions Bot commented May 13, 2026

Uh oh!

sraman-rgb commented May 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gautham-kollu commented May 19, 2026

Uh oh!

sraman-rgb commented May 21, 2026

Uh oh!

svcnvidia-nemo-ci commented May 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

sraman-rgb commented May 13, 2026 •

edited

Loading