Combine GEMM + SwiGLU fused MLP PRs (3890, 4071, 4095, 4219, 4311, 4324) → main - #4636

Merged

Connor-XY merged 33 commits into

NVIDIA:mainfrom

Connor-XY:yxu1/pr3971-rebase-fix-checks

May 16, 2026

Commits on May 13, 2026

Support GEMM + Swiglu fused MLP (#3890)

ksivaman
authored and
Connor-XY
committed
2e618b3
Update megatron/core/transformer/moe/experts.py

authored and
Connor-XY
committed
d3fcd28
fix lint

ksivaman
authored and
Connor-XY
committed
a78c09d
small fix

yaox12
authored and
Connor-XY
committed
24d04f0
Address PR 3971 rebase follow-ups
Connor-XY
committed
f8515fa
Fix test_fused_forward assertions

Connor-XY
and
claude
committed
97a0a53
Add tests for uncovered PR 3971 paths

Connor-XY
and
claude
committed
c6275d2
Address strict-review findings on _make_fused_ops

Connor-XY
and
claude
committed
71b2913
Rename single_grouped_parameter to single_grouped_weight for current TE

Connor-XY
and
claude
committed
2b33d05
Test normalize_grouped_parameter_keys hook orchestration

Connor-XY
and
claude
committed
b6a13ab
Fix lint: project black formatting + pylint line length

Connor-XY
and
claude
committed
6512a1d
TE fused grouped mlp with grouped bias and delayed wgrad (#4095)

authored and
Connor-XY
committed
af6a39e
Enabled fused grouped MLP for `quick_gelu` and add config for grouped params (#4219)

authored and
Connor-XY
committed
4cd3e8e
Fix checkpoint loading with `load_main_params_from_ckpt=True` for grouped weight (#4324)

authored and
Connor-XY
committed
81d4d95
test: align grouped-MLP fakes with PR 4095/4219 fields

Connor-XY
and
claude
committed
c824c66
Use wrapper's combined delay_wgrad mode in fused path; add coverage

Connor-XY
and
claude
committed
4a38f4f
test: cover _normalize_state_dict_for_grouped_params; fix GeGLU typo

Connor-XY
and
claude
committed
29c90ad
Fix fused grouped MLP wgrad hooks for DDP reduce-scatter (#4311)

authored and
Connor-XY
committed
f3e955d
test: cover wrapper hook trigger in fused backward_dw

Connor-XY
and
claude
committed
0a9d86c
[Dev] Skip routed expert padding for graph-safe MoE (#4071)

zhongbozhu
authored and
Connor-XY
committed
2a4e41a
test: align fused_forward fake config with PR 4071 padding helper

Connor-XY
and
claude
committed
4812c96
test: unwrap MlpBuilder partial via get_submodules in real-TE fused-ops test

Connor-XY
and
claude
committed
23186ec
Reject HybridEP dispatcher in fused TE op-fuser support check

Connor-XY
and
claude
committed
e9704da
test: robust fused-path convergence test (Adam + target-driven loss)

Connor-XY
and
claude
committed
1a99f7e
test: add functional MoE recipe exercising the fused TE op-fuser path

Connor-XY
and
claude
committed
093caf1
lint: collapse pytest.skip onto a single line

Connor-XY
and
claude
committed
5d66788
Revert HybridEP guard and duplicate CLI flag

Connor-XY
and
claude
committed
20cfe67
test: add golden values for fused TE op-fuser MoE functional test

Connor-XY
and
claude
committed
0909634
test: set NVTE_CUTEDSL_FUSED_GROUPED_MLP in functional recipe; clarify unit-test scope

Connor-XY
and
claude
committed
3f4cd68
test: clarify cuDSL also requires SM100 (Blackwell)

Connor-XY
and
claude
committed
4868335
test: export NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 for MoE unit-test buckets

Connor-XY
and
claude
committed
46be5d1
Fail fast: reject moe_single_grouped_weight outside fp8+mxfp8

Connor-XY
and
claude
committed
8d3920d

Commits on May 14, 2026

Drop stale grad-mode tracking from #3890's param_and_grad_buffer hunk

Connor-XY
and
claude
committed
d49ad27

Navigation Menu

Uh oh!

Combine GEMM + SwiGLU fused MLP PRs (3890, 4071, 4095, 4219, 4311, 4324) → main - #4636

Combine GEMM + SwiGLU fused MLP PRs (3890, 4071, 4095, 4219, 4311, 4324) → main#4636
Connor-XY merged 33 commits into
NVIDIA:mainfrom
Connor-XY:yxu1/pr3971-rebase-fix-checks

Commits on May 13, 2026

Support GEMM + Swiglu fused MLP (#3890)

Update megatron/core/transformer/moe/experts.py

fix lint

small fix

Address PR 3971 rebase follow-ups

Fix test_fused_forward assertions

Add tests for uncovered PR 3971 paths

Address strict-review findings on _make_fused_ops

Rename single_grouped_parameter to single_grouped_weight for current TE

Test normalize_grouped_parameter_keys hook orchestration

Fix lint: project black formatting + pylint line length

TE fused grouped mlp with grouped bias and delayed wgrad (#4095)

Enabled fused grouped MLP for `quick_gelu` and add config for grouped params (#4219)

Fix checkpoint loading with `load_main_params_from_ckpt=True` for grouped weight (#4324)

test: align grouped-MLP fakes with PR 4095/4219 fields

Use wrapper's combined delay_wgrad mode in fused path; add coverage

test: cover _normalize_state_dict_for_grouped_params; fix GeGLU typo

Fix fused grouped MLP wgrad hooks for DDP reduce-scatter (#4311)

test: cover wrapper hook trigger in fused backward_dw

[Dev] Skip routed expert padding for graph-safe MoE (#4071)

test: align fused_forward fake config with PR 4071 padding helper

test: unwrap MlpBuilder partial via get_submodules in real-TE fused-ops test

Reject HybridEP dispatcher in fused TE op-fuser support check

test: robust fused-path convergence test (Adam + target-driven loss)

test: add functional MoE recipe exercising the fused TE op-fuser path

lint: collapse pytest.skip onto a single line

Revert HybridEP guard and duplicate CLI flag

test: add golden values for fused TE op-fuser MoE functional test

test: set NVTE_CUTEDSL_FUSED_GROUPED_MLP in functional recipe; clarify unit-test scope

test: clarify cuDSL also requires SM100 (Blackwell)

test: export NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 for MoE unit-test buckets

Fail fast: reject moe_single_grouped_weight outside fp8+mxfp8

Commits on May 14, 2026

Drop stale grad-mode tracking from #3890's param_and_grad_buffer hunk