[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 <[email protected]>

…h error handling tests Signed-off-by: Varun Thumbe <[email protected]>

Signed-off-by: Varun Thumbe <[email protected]>

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

…ut_consumption

Signed-off-by: Varun Thumbe <[email protected]>

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

Signed-off-by: Varun Thumbe <[email protected]>

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

Signed-off-by: Varun Thumbe <[email protected]>

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

Signed-off-by: Varun Thumbe <[email protected]>

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

Signed-off-by: Varun Thumbe <[email protected]>

…/TransformerEngine into enable_extra_out_consumption

Signed-off-by: Varun Thumbe <[email protected]>

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

timmoon10 · 2026-08-26T18:08:45Z

-        # Make sure inputs are in correct dtype
+        # Prepare the activation input in the compute dtype.
        input_ = maybe_dequantize(input_, dtype)
-        scales = maybe_dequantize(extra_input, dtype)


This looks suspicious to me, but I think it's fine.

I see that the TE scaled activation kernel can handle different dtypes for input, scales, and output:

TransformerEngine/transformer_engine/common/activation/scaled_activation.cu

Lines 403 to 405 in d7340a4

TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, {

TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, {

TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(output->data.dtype, OutputT, {

I suppose there is some weirdness since we could be ignoring the dtype from torch.autocast. However, this is treated as an auxilliary tensor that behaves more like the MoE splits than like a real data tensor.

timmoon10 · 2026-08-26T18:17:30Z

+            raise TypeError(
+                "NCCL EP Dispatch supports MXFP8Quantizer only, got "
+                f"{type(input_quantizer).__name__}."
+            )


We want basic ops to be a safe fallback, so they should handle any quantizer. If we get a non-MXFP8 quantizer, then the correct behavior is to dequantize and perform the communication in high precision.

Handling fallback via BF16 comms

Signed-off-by: Varun Thumbe <[email protected]>

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

…ffer) and allow mxfp8 combine through env variable Signed-off-by: Varun Thumbe <[email protected]>

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

Signed-off-by: Varun Thumbe <[email protected]>

…merEngine into dispatch_combine_op

Signed-off-by: Varun Thumbe <[email protected]>

…merEngine into dispatch_combine_op

Signed-off-by: Varun Thumbe <[email protected]>

vthumbe1503 · 2026-08-28T19:18:52Z

/te-ci L1 pytorch

phu0ngng · 2026-08-28T20:25:44Z

+"""Pure PyTorch semantic reference for a SwiGLU MoE with expert parallelism.
+
+The implementation deliberately favors readable semantics over performance.  It
+supports a one-rank execution path and a variable-size ``all_to_all_single`` EP


Why do we need this ref impl btw?

phu0ngng · 2026-08-28T20:29:13Z

+            # Only BF16 dispatch is supported for now.
+            input_ = maybe_dequantize(input_, torch.bfloat16)


Can we error out instead of falling back to dequantize?

phu0ngng · 2026-08-28T20:44:39Z

+                device=device,
+            )
+        else:
+            grad_out = validate_buffer(
+                "MXFP8 grad_out storage",
+                grad_out,
+                device=device,
+                contiguous=True,


Why don't we need the grad_out to be contiguous when not MXFP8?

phu0ngng · 2026-08-28T21:14:04Z

+        expert_out = input_
+        if zero_copy:
+            expert_out = _alloc_io(tuple(input_.shape), torch.bfloat16, input_.device, True)
+            expert_out.copy_(input_)


Why do we need this .copy_(input_)? Is it a deep or shallow copy?

phu0ngng · 2026-08-28T21:23:10Z

+        zero_copy = buffer.zero_copy
+        expert_out = input_
+        if zero_copy:
+            expert_out = _alloc_io(tuple(input_.shape), torch.bfloat16, input_.device, True)


We should only allocate this when callers do not provide grad_expert_out.
Besides, it seems not to handle allocation for grad_expert_out_scales when using MXFP8 right now.

Signed-off-by: Varun Thumbe <[email protected]>

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

Signed-off-by: Varun Thumbe <[email protected]>

vthumbe1503 · 2026-08-29T08:03:26Z

/te-ci L1 pytorch

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

produce/consume extra output

b1574f0

Signed-off-by: Varun Thumbe <[email protected]>

allow for fusions with producer/consumer being part of same fuser wit…

63192ab

…h error handling tests Signed-off-by: Varun Thumbe <[email protected]>

cleanup

3b4b523

Signed-off-by: Varun Thumbe <[email protected]>

minor cleanup

de38ed8

Signed-off-by: Varun Thumbe <[email protected]>

dispatch combine impl

385b0d5

Signed-off-by: Varun Thumbe <[email protected]>

fusible ops test

ad3b044

Signed-off-by: Varun Thumbe <[email protected]>

[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 <[email protected]>

cleanup with residual tests

d7d6380

Signed-off-by: Varun Thumbe <[email protected]>

[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 <[email protected]>

update to cleaner documentation

80601dc

Signed-off-by: Varun Thumbe <[email protected]>

address review comments

5070e34

Signed-off-by: Varun Thumbe <[email protected]>

[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 <[email protected]>

update docs

5a4e1ec

Signed-off-by: Varun Thumbe <[email protected]>

pin channels through channel version

0a479c7

Signed-off-by: Varun Thumbe <[email protected]>

unecessary handling removal

d679998

Signed-off-by: Varun Thumbe <[email protected]>

simplify

8f7ba95

Signed-off-by: Varun Thumbe <[email protected]>

doc update + extra_grad = None case

c62bb15

Signed-off-by: Varun Thumbe <[email protected]>

[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 <[email protected]>

no need to check staleness in every forward call

6801a6d

Signed-off-by: Varun Thumbe <[email protected]>

remove redundant tests

6688e8a

Signed-off-by: Varun Thumbe <[email protected]>

[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 <[email protected]>

keep simple

a4cc112

Signed-off-by: Varun Thumbe <[email protected]>

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

7edaf89

…/TransformerEngine into enable_extra_out_consumption

vthumbe1503 added the 2.19 label Aug 24, 2026

vthumbe1503 and others added 3 commits August 25, 2026 14:32

Merge branch 'main' into dispatch_combine_op

290cb8d

refactor based on ep.py + reduce cpu overheads

b882b35

Signed-off-by: Varun Thumbe <[email protected]>

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

be9a0dc

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

timmoon10 reviewed Aug 26, 2026

View reviewed changes

vthumbe1503 and others added 2 commits August 26, 2026 20:17

limit the recv tokens per rank

f6c6306

Signed-off-by: Varun Thumbe <[email protected]>

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

38fd365

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

buptzyb mentioned this pull request Aug 28, 2026

[ROADMAP][2026 Q3] Megatron Core MoE Roadmap NVIDIA/Megatron-LM#6757

Open

vthumbe1503 and others added 14 commits August 28, 2026 18:27

address review comments + rework design for ops(epconfig + kwarg epbu…

fedcda0

…ffer) and allow mxfp8 combine through env variable Signed-off-by: Varun Thumbe <[email protected]>

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

a584be8

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

address review comments

bf9dd3f

Signed-off-by: Varun Thumbe <[email protected]>

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

0722ec6

…merEngine into dispatch_combine_op

t_flat was a better name

1d1e11a

Signed-off-by: Varun Thumbe <[email protected]>

simplify

5674825

Signed-off-by: Varun Thumbe <[email protected]>

clean cursor mess

6646ca3

Signed-off-by: Varun Thumbe <[email protected]>

restore comments

62d9b06

Signed-off-by: Varun Thumbe <[email protected]>

bf16 hardcoding for combine as well

e8d1bb7

Signed-off-by: Varun Thumbe <[email protected]>

scale type should be autocast type if autocast enabled

cb4c57e

Signed-off-by: Varun Thumbe <[email protected]>

Merge branch 'main' into dispatch_combine_op

b518f6e

update copyrights

ebed56a

Signed-off-by: Varun Thumbe <[email protected]>

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

3e9812b

…merEngine into dispatch_combine_op

match the right syntax

a44091e

Signed-off-by: Varun Thumbe <[email protected]>

phu0ngng reviewed Aug 28, 2026

View reviewed changes

vthumbe1503 and others added 5 commits August 28, 2026 14:36

Merge branch 'main' into dispatch_combine_op

4956554

better changes

ae6dc49

Signed-off-by: Varun Thumbe <[email protected]>

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

dc59ae3

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

minor review comments pending

f83257e

Signed-off-by: Varun Thumbe <[email protected]>

Merge branch 'main' into dispatch_combine_op

a18254c

	TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(input->data.dtype, InputT, {
	TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(act_scales->data.dtype, ScaleT, {
	TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(output->data.dtype, OutputT, {

		# Only BF16 dispatch is supported for now.
		input_ = maybe_dequantize(input_, torch.bfloat16)

Uh oh!

Conversation

vthumbe1503 commented Aug 11, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

timmoon10 Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

timmoon10 Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vthumbe1503 commented Aug 28, 2026

Uh oh!

phu0ngng Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 commented Aug 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vthumbe1503 Aug 28, 2026 •

edited

Loading