[PyTorch] Add ScaledTanhSReLU and wire it into the fused grouped MLP by wanyingw · Pull Request #3463 · NVIDIA/TransformerEngine

wanyingw · 2026-09-02T16:42:10Z

Description

Adds a new activation op ScaledTanhSReLU — a tanh soft-clamped squared ReLU with
per-row post-scaling — and wires it into the fused grouped MLP
(GroupedLinear + activation + GroupedLinear → cuDNN grouped GEMM):

y = (s * tanh(relu(x) / s))^2 * scales

ScaledSReLU is the s → ∞ limit and is left completely untouched.

Motivation: this activation is rolling out in Megatron-Core MoE models, where it currently runs as a standalone elementwise kernel between the two grouped GEMMs. The cuDNN frontend srelu/dsrelu grouped GEMM kernels now accept a tanh_clamp_scale parameter (cuDNN frontend PR: NVIDIA/cudnn-frontend#858); this PR exposes it through TE.

Dependencies: cuDNN frontend with tanh_clamp_scale in grouped_gemm_srelu_wrapper_sm100 / grouped_gemm_dsrelu_wrapper_sm100 (cuDNN FE PR: NVIDIA/cudnn-frontend#858). Feature detection is by wrapper signature, not version, so this PR is safe to merge before the frontend release: without clamp support the fuser simply declines ScaledTanhSReLU and construction raises, the clamp is never silently dropped. A min-version gate constant will follow once the frontend release ships. Follows the SiTU-GLU integration pattern (#3402).

Type of change

New feature (non-breaking change which adds functionality)

Changes

ops/basic/activation.py: new ScaledTanhSReLU(clamp_scale=...) — a sibling of ScaledSReLU under _ScaledUnary (not a subclass), mirroring how ScaledSiTUGLU sits beside the other gated ops. The unfused path is implemented in torch: the scaled-unary CUDA kernels are parameterless for SReLU, and this path only serves non-SM100 / unfused configurations. A CUDA functor pair can follow.
ops/fused/grouped_mlp.py: pass tanh_clamp_scale to the cuDNN srelu (forward) and dsrelu (backward) wrappers when the activation is ScaledTanhSReLU. The feature check requires both wrappers to accept the parameter — a frontend with only the forward clamp would train against an unclamped backward.
Sibling-class audit of every SReLU-family check in the fused grouped MLP:
- Extended: validate_grouped_mlp_dims, the single-group-runtime-offsets exclusion, the NVFP4 fc2 alpha in fuser_backward (dsrelu applies alpha once — a sibling falling into the gated-kernel sqrt branch would scale the gradient wrong), and fuse_srelu_ops' activation_op_types (gated on the feature check).
- Deliberately left ScaledSReLU-only, with reasoning in comments: the NVFP4 RHT hadamard gate (its kernel has no clamp support; the generic quantize path handles it) and activation recomputation (falls back to saving fc2_x — costs memory, stays correct).
Tests: clamped forward/backward vs the torch reference in test_grouped_mlp.py and test_fusible_ops.py; existing ScaledSReLU tests pass unmodified.

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 (op docstrings)
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

Testing

On GB300 (sm_103): released cuDNN-frontend wheel with the frontend branch's Python overlaid (the srelu/dsrelu kernels are CuTe DSL, JIT-compiled at runtime, so no frontend rebuild is needed), and the prebuilt TE wheel with this branch's pytorch/ops/ overlaid (this change is pure Python). pytest tests/pytorch/test_fusible_ops.py tests/pytorch/test_grouped_mlp.py -k "tanh_srelu or scaled_srelu": 49 passed / 66 skipped / 0 failed, including:

clamped cases checking forward and backward against a torch autograd reference;
all pre-existing ScaledSReLU cases passing unmodified.

Fixes # (issue)

Type of change

New feature (non-breaking change which adds functionality)

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

/te-ci pytorch

zhongbozhu

LGTM

greptile-apps · 2026-09-02T17:46:06Z

Greptile Summary

Adds the ScaledTanhSReLU PyTorch operation and conditionally integrates it with the fused grouped-MLP path when both cuDNN forward and backward wrappers support the clamp parameter.

Implements validated soft-clamp configuration and matching unfused forward/backward computations.
Extends grouped-MLP dimension validation, fusion selection, quantization scaling, and cuDNN argument forwarding.
Adds standalone and grouped-MLP numerical, gradient, saturation, validation, and quantization coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/basic/activation.py	Defines the validated soft-clamped activation and matching fp32-intermediate forward and backward formulas.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Feature-detects complete cuDNN support and wires the activation through grouped-MLP fusion, quantization, and backward paths.
tests/pytorch/test_fusible_ops.py	Adds standalone forward, gradient, saturation, parameter-validation, and configuration coverage.
tests/pytorch/test_grouped_mlp.py	Adds fused grouped-MLP reference comparisons and low-precision path coverage for the new activation.
transformer_engine/pytorch/ops/basic/init.py	Exposes the new activation through the public basic-operations import chain.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[GroupedLinear] --> B[ScaledTanhSReLU]
  B --> C[GroupedLinear]
  B --> D{Both cuDNN wrappers accept tanh_clamp_scale?}
  D -->|Yes| E[Fused grouped GEMM forward and backward]
  D -->|No| F[Unfused PyTorch activation path]

_{Reviews (4): Last reviewed commit: "[PyTorch] Add ScaledTanhSReLU and wire i..." | Re-trigger Greptile}

timmoon10

LGTM, pending CI

timmoon10 · 2026-09-02T18:00:54Z

    )


 def fuse_srelu_ops(


Nit: Beyond the scope of this PR, but we're seeing that the SReLU naming scheme was short-sighted. The underlying kernel should be able to support any unary activation (normal ReLU, sigmoid, whatever), and SReLU just happened to be the first one. At some point we should change to something like fuse_grouped_mlp_ops_with_unary_activation.

sraman-rgb · 2026-09-02T18:09:16Z

LGTM

Adds a tanh soft-clamped squared ReLU activation alongside ScaledSReLU: y = (s * tanh(relu(x) / s))^2 * scales ScaledTanhSReLU is a sibling of ScaledSReLU under _ScaledUnary, not a subclass, matching how ScaledSiTUGLU sits beside the other gated ops. That choice costs explicit handling at each place the fused grouped MLP tests for the SReLU family, so all six were audited: extended, because a sibling would otherwise misbehave: * _cudnn_frontend_supports_single_group_runtime_offsets -- the srelu and dsrelu wrappers take no use_single_group_runtime_offsets argument, so the whole family must be excluded or the single-group (shared expert) path passes an unexpected kwarg. * validate_grouped_mlp_dims -- otherwise falls through to the GLU branch and raises TypeError on a valid configuration. * the NVFP4 fc2 alpha in fuser_backward -- dsrelu applies alpha once and needs the full product; the gated kernels need sqrt(product). A sibling would silently take the sqrt branch and scale the gradient wrong. * fuse_srelu_ops' activation_op_types, gated on the cuDNN feature check. deliberately left as ScaledSReLU only, with the reasoning in comments: * the NVFP4 RHT hadamard gate -- the kernel it selects is the GLU hadamard wrapper, which has no soft-clamp support, so routing tanh-SReLU through it would compute an unclamped activation. The cost is that NVFP4 RHT gives up hadamard fusion here; the generic quantize path handles it. * activation recomputation -- falls back to saving fc2_x, which costs memory but stays correct. The unfused path is implemented in torch: the scaled-unary CUDA kernels are parameterless for SReLU, and this activation has no fused kernel yet. It only serves non-SM100 and unfused configurations; the fused path goes straight to the cuDNN epilogue. A CUDA functor pair can follow. Feature detection is by wrapper signature rather than version so this can be developed against an editable cuDNN frontend; both the forward and backward wrappers must accept tanh_clamp_scale, since a frontend with only the forward clamp would train against an unclamped backward. If the check fails the fuser declines and construction raises -- the clamp is never silently dropped. Signed-off-by: Wanying Wang <wanyingw@nvidia.com>

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

zhongbozhu approved these changes Sep 2, 2026

View reviewed changes

wanyingw marked this pull request as ready for review September 2, 2026 17:43

wanyingw requested a review from timmoon10 as a code owner September 2, 2026 17:43

timmoon10 approved these changes Sep 2, 2026

View reviewed changes

wanyingw force-pushed the srelu-tanh-clamp branch 2 times, most recently from 9247701 to effdfe8 Compare September 2, 2026 20:30

wanyingw force-pushed the srelu-tanh-clamp branch from effdfe8 to e28c0cc Compare September 2, 2026 20:34

Navigation Menu

Uh oh!

[PyTorch] Add ScaledTanhSReLU and wire it into the fused grouped MLP - #3463

[PyTorch] Add ScaledTanhSReLU and wire it into the fused grouped MLP#3463
wanyingw wants to merge 1 commit into
NVIDIA:mainfrom
wanyingw:srelu-tanh-clamp

wanyingw commented Sep 2, 2026 •

edited

Loading

Uh oh!

zhongbozhu commented Sep 2, 2026

Uh oh!

zhongbozhu left a comment

Uh oh!

greptile-apps Bot commented Sep 2, 2026 •

edited

Loading

Uh oh!

timmoon10 left a comment

Uh oh!

timmoon10 Sep 2, 2026

Uh oh!

sraman-rgb commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

wanyingw commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist

Testing

Type of change

Uh oh!

zhongbozhu commented Sep 2, 2026

Uh oh!

zhongbozhu left a comment

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

timmoon10 Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

sraman-rgb commented Sep 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wanyingw commented Sep 2, 2026 •

edited

Loading

greptile-apps Bot commented Sep 2, 2026 •

edited

Loading