Skip to content

[Pytorch] MOE Sequential Block - #3350

Draft
vthumbe1503 wants to merge 47 commits into
NVIDIA:mainfrom
vthumbe1503:dispatch_combine_op
Draft

[Pytorch] MOE Sequential Block#3350
vthumbe1503 wants to merge 47 commits into
NVIDIA:mainfrom
vthumbe1503:dispatch_combine_op

Conversation

@vthumbe1503

@vthumbe1503 vthumbe1503 commented Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

vthumbe1503 and others added 30 commits July 28, 2026 23:05
@vthumbe1503
produce/consume extra output
b1574f0 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
allow for fusions with producer/consumer being part of same fuser wit…
63192ab 
…h error handling tests

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
cleanup
3b4b523 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
minor cleanup
de38ed8 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
dispatch combine impl
385b0d5 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
fusible ops test
ad3b044 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
5fb0d3a 
for more information, see https://pre-commit.ci
@vthumbe1503
Merge remote-tracking branch 'nvidia_origin/main' into enable_extra_o…
2ba4f6a 
…ut_consumption
@vthumbe1503
keep just ops infra changes
3af2ecc 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
cleanup with residual tests
d7d6380 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
74f563a 
for more information, see https://pre-commit.ci
@vthumbe1503
Merge branch 'main' into enable_extra_out_consumption
29d23f2
@vthumbe1503
address review comment
87e2b36 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
update to cleaner documentation
80601dc 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
address review comments
5070e34 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
ae41ad3 
for more information, see https://pre-commit.ci
@vthumbe1503
some cleanup
f82cbed 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
update docs
5a4e1ec 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
pin channels through channel version
0a479c7 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
unecessary handling removal
d679998 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
simplify
8f7ba95 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
doc update + extra_grad = None case
c62bb15 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
a93b820 
for more information, see https://pre-commit.ci
@vthumbe1503
test cleanup
35b73b1 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
no need to check staleness in every forward call
6801a6d 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
remove redundant tests
6688e8a 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
12430c2 
for more information, see https://pre-commit.ci
@vthumbe1503
revert from bad names
b189550 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
keep simple
a4cc112 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
Merge branch 'enable_extra_out_consumption' of github.com:vthumbe1503…
7edaf89 
…/TransformerEngine into enable_extra_out_consumption
vthumbe1503 and others added 17 commits August 9, 2026 23:37
@vthumbe1503
unecessary checks
5ba6055 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
minor doc
76826dc 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
Merge branch 'main' into enable_extra_out_consumption
827f8e9
@vthumbe1503
basic op and refrence implementation
60ca030 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
compare fused and unfused ops and they are matching
8ec3a39 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
fix lint
63a4ea3 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503 @timmoon10
Update transformer_engine/pytorch/ops/fuser.py
87c1cf6 
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
@vthumbe1503 @timmoon10
Update docs/examples/op_fuser/op_fuser.rst
307ab15 
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
@vthumbe1503 @timmoon10
Update transformer_engine/pytorch/ops/fuser.py
97a91cf 
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
@vthumbe1503
address review comments + extra output being configurable to be outpu…
6468a14 
…tted

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
cleanup
2eb21ed 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
Merge remote-tracking branch 'origin/enable_extra_out_consumption' in…
1b0a4b1 
…to dispatch_combine_op
@vthumbe1503
Restore ep_reference to its originally created content.
212a460
@vthumbe1503
Support BF16 grouped MLP in ep_reference and keep MOE extras internal.
c7dacdf 
Use compute_dtype for apples-to-apples TE comparisons, and set
output_to_caller=False so fused MOE does not materialize internal
dispatch extras for the Sequential caller.

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
update test to have internal extra_out
00b48cf 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@pre-commit-ci
[pre-commit.ci] auto fixes from pre-commit.com hooks
db24495 
for more information, see https://pre-commit.ci
@vthumbe1503
extra output grad can be None, equivalent to zero
3c37643 
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

address review comment

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

address review comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@timmoon10 timmoon10 Awaiting requested review from timmoon10 timmoon10 will be requested when the pull request is marked ready for review timmoon10 is a code owner
@ksivaman ksivaman Awaiting requested review from ksivaman ksivaman will be requested when the pull request is marked ready for review ksivaman is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vthumbe1503