[Pytorch] MOE Sequential Block by vthumbe1503 · Pull Request #3350 · NVIDIA/TransformerEngine

vthumbe1503 · 2026-08-11T21:33:21Z

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

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

…h error handling tests Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

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

for more information, see https://pre-commit.ci

…ut_consumption

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

for more information, see https://pre-commit.ci

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

for more information, see https://pre-commit.ci

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

for more information, see https://pre-commit.ci

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

for more information, see https://pre-commit.ci

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

…/TransformerEngine into enable_extra_out_consumption

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

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>

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

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

…to dispatch_combine_op

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>

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

for more information, see https://pre-commit.ci

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>

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

for more information, see https://pre-commit.ci

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

…merEngine into dispatch_combine_op

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

for more information, see https://pre-commit.ci

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

for more information, see https://pre-commit.ci

phu0ngng · 2026-08-20T17:06:50Z

+
+
+class Dispatch(BasicOperation):
+    """Dispatch BF16 tokens to local experts with NCCL EP.


Tokens can be in FP8 or other precision as well, right?

phu0ngng · 2026-08-20T17:08:36Z

+        input_: torch.Tensor,
+        *,
+        basic_op_extra_inputs: list[tuple[torch.Tensor, ...]],
+        prev_op_grad_output_quantizer: Optional[Quantizer],


What do we plan to do with prev_op_grad_output_quantizer?

phu0ngng · 2026-08-20T17:10:35Z

+        kwargs = basic_op_kwargs[0]
+
+        if input_.dtype is not torch.bfloat16:
+            raise NotImplementedError(f"NCCL EP requires BF16 dispatch input, got {input_.dtype}.")


I think the input of TE Dispatch op should be BF16, and we do internal quantization if requested, but NCCL EP Dispatch does not require BF16 input.

phu0ngng · 2026-08-20T17:21:29Z

+            raise RuntimeError("NCCL EP dispatch receive size is unavailable.")
+        rows = int(rows)
+        recv_shape = (rows, self.buffer.hidden_dim)
+        recv_tokens = _validate_output_buffer(


I think this validation may not work when callers provide the recv_tokens which packs quantized data and scale for MXFP8.

phu0ngng · 2026-08-20T17:24:43Z

+                device=grad_output.device,
+            )
+        else:
+            grad_recv_weights = grad_recv_weights.to(dtype=torch.float32).contiguous()


why do we need to(dtype=torch.float32).contiguous() here?

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

for more information, see https://pre-commit.ci

vthumbe1503 and others added 30 commits July 28, 2026 23:05

produce/consume extra output

b1574f0

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

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>

cleanup

3b4b523

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

minor cleanup

de38ed8

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

dispatch combine impl

385b0d5

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

fusible ops test

ad3b044

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

5fb0d3a

for more information, see https://pre-commit.ci

Merge remote-tracking branch 'nvidia_origin/main' into enable_extra_o…

2ba4f6a

…ut_consumption

keep just ops infra changes

3af2ecc

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

cleanup with residual tests

d7d6380

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

74f563a

for more information, see https://pre-commit.ci

Merge branch 'main' into enable_extra_out_consumption

29d23f2

address review comment

87e2b36

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

update to cleaner documentation

80601dc

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

address review comments

5070e34

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

ae41ad3

for more information, see https://pre-commit.ci

some cleanup

f82cbed

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

update docs

5a4e1ec

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

pin channels through channel version

0a479c7

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

unecessary handling removal

d679998

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

simplify

8f7ba95

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

doc update + extra_grad = None case

c62bb15

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

a93b820

for more information, see https://pre-commit.ci

test cleanup

35b73b1

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

no need to check staleness in every forward call

6801a6d

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

remove redundant tests

6688e8a

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

12430c2

for more information, see https://pre-commit.ci

revert from bad names

b189550

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

keep simple

a4cc112

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

Merge branch 'enable_extra_out_consumption' of github.com:vthumbe1503…

7edaf89

…/TransformerEngine into enable_extra_out_consumption

vthumbe1503 and others added 13 commits August 10, 2026 07:22

compare fused and unfused ops and they are matching

8ec3a39

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

fix lint

63a4ea3

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

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>

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>

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>

address review comments + extra output being configurable to be outpu…

6468a14

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

cleanup

2eb21ed

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

Merge remote-tracking branch 'origin/enable_extra_out_consumption' in…

1b0a4b1

…to dispatch_combine_op

Restore ep_reference to its originally created content.

212a460

update test to have internal extra_out

00b48cf

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

db24495

for more information, see https://pre-commit.ci

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>

vthumbe1503 force-pushed the dispatch_combine_op branch from 517d0ff to 3c37643 Compare August 12, 2026 05:20

vthumbe1503 and others added 5 commits August 17, 2026 16:42

fix merge conflict

e10a13b

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

revert commit for nccl

5b555ce

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

14d9ef9

for more information, see https://pre-commit.ci

more fix

42c622e

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

Merge branch 'dispatch_combine_op' of github.com:vthumbe1503/Transfor…

1172c72

…merEngine into dispatch_combine_op

vthumbe1503 requested a review from phu0ngng August 17, 2026 22:41

vthumbe1503 and others added 7 commits August 19, 2026 09:49

cudnn kernel

e40eb74

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

c7eceaa

for more information, see https://pre-commit.ci

remove the previous artifact

fb8fc83

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

minor change

76d7fdd

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

minor bug fix

c946162

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

single_groupe_weight + remove copies

80f6063

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

a779c6c

for more information, see https://pre-commit.ci

phu0ngng reviewed Aug 20, 2026

View reviewed changes

vthumbe1503 and others added 2 commits August 21, 2026 09:27

cpu init for megatron needs

f197433

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

075fc44

for more information, see https://pre-commit.ci

Navigation Menu

Uh oh!

[Pytorch] MOE Sequential Block - #3350

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

vthumbe1503 commented Aug 11, 2026

Uh oh!

phu0ngng Aug 20, 2026

Uh oh!

phu0ngng Aug 20, 2026

Uh oh!

phu0ngng Aug 20, 2026

Uh oh!

phu0ngng Aug 20, 2026

Uh oh!

phu0ngng Aug 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		class Dispatch(BasicOperation):
		"""Dispatch BF16 tokens to local experts with NCCL EP.

Uh oh!

Conversation

vthumbe1503 commented Aug 11, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

phu0ngng Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants