[PyTorch] Deprecate grouped linear and MLP environment gates by zhongbozhu · Pull Request #3584 · NVIDIA/TransformerEngine

zhongbozhu · 2026-09-29T16:26:03Z

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

greptile-apps · 2026-09-29T16:31:50Z

Confidence Score: 4/5

[Medium impact] Removes environment variable gates for grouped linear and MLP features.

The PR is not yet safe to merge because the outstanding CUDA-free grouped-MLP support check can raise instead of returning unsupported.

Findings

CUDA-free support checks raise ▶

Summary

The PR removes environment gates for grouped parameters and grouped-MLP fusion, updates EP alignment defaults and examples, and expands grouped-operation coverage. Changes since the previous review also add compiled attention and operation tests, FP8 recipe parameter restoration, and native kernel synchronization safeguards.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[GroupedLinear → activation → GroupedLinear] --> B{Recipe and pattern eligible?}
  B -->|Yes| C{Device and frontend supported?}
  C -->|Yes| D[Fused grouped MLP]
  B -->|No| E[Separate operations]
  C -->|No| E

_{Reviews (10) · Last reviewed commit: "attempt to fix ci issue" · Reviewed by Greptile}

zhongbozhu · 2026-09-29T16:39:32Z

/te-ci pytorch L1

zhongbozhu · 2026-09-30T17:39:08Z

/te-ci pytorch L1

zhongbozhu · 2026-10-01T23:09:33Z

/te-ci pytorch L1

vthumbe1503 · 2026-10-05T01:41:42Z

LGTM. Thanks for cleaning these env variables. All @timmoon10's comments seem to be resolved as well. @zhongbozhu I resolved the merge conflicts.

vthumbe1503 · 2026-10-05T01:41:50Z

/te-ci pytorch

zhongbozhu · 2026-10-05T23:40:05Z

/te-ci pytorch

greptile-apps · 2026-10-05T23:43:26Z

        """Whether this fused operation is supported on the current system."""
-        if int(os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "0")) <= 0:
-            return False
        if get_device_compute_capability()[0] != 10:


CUDA-free support checks raise If CUDA is unavailable, is_supported() now calls get_device_compute_capability() instead of returning False. That helper calls torch.cuda.current_device(), so the support check raises rather than letting the caller use the unfused path. The unary-fusion callback can hit this check while building a plan, before checking whether the recipe or operations are candidates for fusion.

Knowledge Base Used: PyTorch tensor and operator abstractions

timmoon10

LGTM

pggPL

We need to investigate #3640 before merging that PR

zhongbozhu · 2026-10-06T22:35:27Z

/te-ci pytorch

zhongbozhu · 2026-10-07T21:18:56Z

/te-ci pytorch L1

zhongbozhu · 2026-10-08T08:34:30Z

/te-ci pytorch L1

Signed-off-by: zhongboz <zhongboz@nvidia.com>

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

Signed-off-by: zhongboz <zhongboz@nvidia.com>

zhongbozhu · 2026-10-09T02:20:41Z

/te-ci pytorch L1

zhongbozhu requested review from ksivaman and timmoon10 as code owners September 29, 2026 16:26

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Sep 29, 2026

greptile-apps Bot reviewed Sep 29, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py

Comment thread docs/examples/te_mixtral/run_finetune_ep.py

Comment thread tests/pytorch/test_grouped_mlp.py Outdated

ptrendx added the 2.21 label Sep 30, 2026

timmoon10 reviewed Sep 30, 2026

View reviewed changes

zhongbozhu mentioned this pull request Oct 1, 2026

Relax TE Op Fuser Env Var Assert NVTE_CUTEDSL_FUSED_GROUPED_MLP NVIDIA/Megatron-LM#7499

Merged

1 of 7 tasks

zhongbozhu force-pushed the cleanup/grouped-mlp-env-gates branch from 6409749 to 1ac8a2c Compare October 1, 2026 23:09

greptile-apps Bot reviewed Oct 1, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py Outdated

zhongbozhu force-pushed the cleanup/grouped-mlp-env-gates branch from 1ac8a2c to 91032db Compare October 2, 2026 23:57

vthumbe1503 previously approved these changes Oct 5, 2026

View reviewed changes

ghadiaravi13 mentioned this pull request Oct 5, 2026

[PyTorch] Add NVTE_CUTEDSL_FUSED_GROUPED_MLP_WARN_FALLBACK to warn on grouped MLP fusion fallback #3631

Draft

6 of 13 tasks

zhongbozhu dismissed vthumbe1503’s stale review via 7a47536 October 5, 2026 23:39

zhongbozhu force-pushed the cleanup/grouped-mlp-env-gates branch from a87a37a to 7a47536 Compare October 5, 2026 23:39

greptile-apps Bot reviewed Oct 5, 2026

View reviewed changes

vthumbe1503 previously approved these changes Oct 6, 2026

View reviewed changes

timmoon10 previously approved these changes Oct 6, 2026

View reviewed changes

pggPL requested changes Oct 6, 2026

View reviewed changes

zhongbozhu dismissed stale reviews from timmoon10 and vthumbe1503 via accb0b6 October 6, 2026 22:34

greptile-apps Bot reviewed Oct 7, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ep.py

Comment thread transformer_engine/pytorch/ep.py

zhongbozhu added 7 commits October 8, 2026 18:20

[PyTorch] Deprecate grouped linear and MLP environment gates

b7e669d

Signed-off-by: zhongboz <zhongboz@nvidia.com>

resolve comments

8da121d

Signed-off-by: zhongboz <zhongboz@nvidia.com>

no need to check if cuda exists

a3f6eab

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

fix torch compile CI

68f106a

Signed-off-by: zhongboz <zhongboz@nvidia.com>

fix alignment config

fe51c82

Signed-off-by: zhongboz <zhongboz@nvidia.com>

add more tests

adb3ad5

Signed-off-by: zhongboz <zhongboz@nvidia.com>

attempt to fix ci issue

b2ebc26

Signed-off-by: zhongboz <zhongboz@nvidia.com>

zhongbozhu force-pushed the cleanup/grouped-mlp-env-gates branch from 6b78903 to b2ebc26 Compare October 9, 2026 02:19

Conversation

zhongbozhu commented Sep 29, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Sep 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 4/5

Findings

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhongbozhu commented Sep 29, 2026

Uh oh!

zhongbozhu commented Sep 30, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhongbozhu commented Oct 1, 2026

Uh oh!

Uh oh!

vthumbe1503 commented Oct 5, 2026

Uh oh!

vthumbe1503 commented Oct 5, 2026

Uh oh!

zhongbozhu commented Oct 5, 2026

Uh oh!

greptile-apps Bot Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

pggPL left a comment

Choose a reason for hiding this comment

Uh oh!

zhongbozhu commented Oct 6, 2026

Uh oh!

Uh oh!

Uh oh!

zhongbozhu commented Oct 7, 2026

Uh oh!

zhongbozhu commented Oct 8, 2026

Uh oh!

zhongbozhu commented Oct 9, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

greptile-apps Bot commented Sep 29, 2026 •

edited

Loading