Fuse shared expert MLP with grouped GEMM by sraman-rgb · Pull Request #5604 · NVIDIA/Megatron-LM

sraman-rgb · 2026-07-01T18:23:59Z

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

What does this PR do?

This PR is for enabling usage of groupedGEMM based shared expert MLP consisting of dense GEMMs. This is needed for fusion of shared expert GEMM+swiglu+quant fusion

⚠️ 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 @NVIDIA/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 @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.

copy-pr-bot · 2026-07-01T18:24:02Z

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-07-01T18:24:06Z

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-07-01T18:28:22Z

/ok to test b687d5b

sraman-rgb · 2026-07-01T18:37:35Z

/ok to test b687d5b

copy-pr-bot · 2026-07-01T18:37:38Z

/ok to test b687d5b

@sraman-rgb, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

sraman-rgb · 2026-07-01T19:13:07Z

/ok to test a05f54d

Phlip79

Please add a description to your PR

sraman-rgb · 2026-07-02T17:35:32Z

Please add a description to your PR

done

kvareddy · 2026-07-06T09:52:52Z

@jiemingz @fanshiqing can you please take a look at this MR?

Signed-off-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>

sraman-rgb · 2026-07-07T05:21:45Z

/ok to test 30edd6c

jiemingz · 2026-07-07T14:07:36Z

+
+        glu_interleave_size = self.config.moe_shared_expert_glu_interleave_size
+        fc1_weight = self.linear_fc1.weight
+        op = te.pytorch.ops.GroupedLinear(


LGTM
The te op fuser code is starting to proliferate around the code base, we might want to consider collecting it in one location, probably under the extensions/transformer_engine directory

Signed-off-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>

sraman-rgb · 2026-07-07T15:45:20Z

/ok to test 5fe4e6b

svcnvidia-nemo-ci · 2026-07-07T17:46:52Z

🔄 Merge queue validation started!

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

sraman-rgb requested review from a team as code owners July 1, 2026 18:24

svcnvidia-nemo-ci marked this pull request as draft July 1, 2026 18:24

sraman-rgb marked this pull request as ready for review July 1, 2026 18:24

svcnvidia-nemo-ci added the complexity: medium label Jul 1, 2026

copy-pr-bot Bot temporarily deployed to public July 1, 2026 18:28 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 18:32 Inactive

sraman-rgb force-pushed the shared-expert-grouped-gemm branch from b687d5b to a05f54d Compare July 1, 2026 18:35

sraman-rgb enabled auto-merge July 1, 2026 19:13

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:17 Inactive

copy-pr-bot Bot temporarily deployed to test July 1, 2026 19:23 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:25 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:36 Inactive

copy-pr-bot Bot temporarily deployed to public July 1, 2026 19:40 Inactive

copy-pr-bot Bot temporarily deployed to public July 2, 2026 15:59 Inactive

copy-pr-bot Bot temporarily deployed to public July 2, 2026 16:08 Inactive

santhnm2 approved these changes Jul 2, 2026

View reviewed changes

Comment thread megatron/core/transformer/moe/shared_experts.py

Phlip79 reviewed Jul 2, 2026

View reviewed changes

yaoyu-33 approved these changes Jul 2, 2026

View reviewed changes

Merge branch 'main' into shared-expert-grouped-gemm

32e45c3

yaox12 approved these changes Jul 3, 2026

View reviewed changes

Comment thread megatron/core/transformer/moe/shared_experts.py Outdated

kvareddy requested review from fanshiqing and jiemingz July 6, 2026 09:52

sraman-rgb added 2 commits July 6, 2026 21:03

Merge branch 'main' into shared-expert-grouped-gemm

e821665

Use transformer config for fused shared expert recipe

30edd6c

Signed-off-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>

copy-pr-bot Bot temporarily deployed to public July 7, 2026 05:22 Inactive

copy-pr-bot Bot temporarily deployed to test July 7, 2026 05:22 Inactive

copy-pr-bot Bot temporarily deployed to public July 7, 2026 05:25 Inactive

fanshiqing approved these changes Jul 7, 2026

View reviewed changes

jiemingz reviewed Jul 7, 2026

View reviewed changes

jiemingz approved these changes Jul 7, 2026

View reviewed changes

kvareddy approved these changes Jul 7, 2026

View reviewed changes

sraman-rgb added 2 commits July 7, 2026 08:40

Fix fused shared expert recipe test config

c0ae967

Signed-off-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>

Merge branch 'main' into shared-expert-grouped-gemm

5fe4e6b

lhb8125 mentioned this pull request Aug 18, 2026

Megatron-LM Monthly Activity Report — July 2026 (Final, through July 31, 2026) #6607

Open

sudhu2k mentioned this pull request Oct 27, 2026

Megatron core_r0.19_0 IFU ROCm/Megatron-LM#159

Open

FeSens mentioned this pull request Oct 28, 2026

fix(MoE): apply the SwiGLU clamp and reject SiTU-GLU in TE fused MLPs #8640

Draft

3 of 7 tasks

Conversation

sraman-rgb commented Jul 1, 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 Jul 1, 2026

Uh oh!

github-actions Bot commented Jul 1, 2026

Uh oh!

sraman-rgb commented Jul 1, 2026

Uh oh!

sraman-rgb commented Jul 1, 2026

Uh oh!

copy-pr-bot Bot commented Jul 1, 2026

Uh oh!

sraman-rgb commented Jul 1, 2026

Uh oh!

Uh oh!

Phlip79 left a comment

Choose a reason for hiding this comment

Uh oh!

sraman-rgb commented Jul 2, 2026

Uh oh!

Uh oh!

kvareddy commented Jul 6, 2026

Uh oh!

sraman-rgb commented Jul 7, 2026

Uh oh!

jiemingz Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

sraman-rgb commented Jul 7, 2026

Uh oh!

svcnvidia-nemo-ci commented Jul 7, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

sraman-rgb commented Jul 1, 2026 •

edited

Loading