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

tests/pytorch/test_fusible_ops.py

@@ -3101,6 +3101,101 @@ def test_scaled_srelu(
         assert_close_grads(x_test, x_ref, **tols)
         assert_close_grads(scales_test, scales_ref, **tols)
 
+    @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128)))
+    @pytest.mark.parametrize("input_requires_grad", (False, True))
+    @pytest.mark.parametrize("scales_requires_grad", (False, True))
+    @pytest.mark.parametrize("tanh_clamp_scale", (0.5, 2.0))
+    def test_scaled_tanh_srelu(
+        self,
+        *,
+        in_shape: Iterable[int],
+        dtype: torch.dtype = torch.float32,
+        device: torch.device = "cuda",
+        input_requires_grad: bool,
+        scales_requires_grad: bool,
+        tanh_clamp_scale: float,
+    ) -> None:
+        """Tanh soft-clamped SReLU with post-scale.
+
+        Covers the unfused path specifically: the fused grouped-MLP op goes straight
+        to the cuDNN srelu_tanh epilogue and never runs this code. Small clamp scales
+        are used so tanh genuinely saturates -- with a large scale the result is
+        numerically indistinguishable from plain ScaledSReLU.
+        """
+
+        # Random data
+        x_ref, x_test = make_reference_and_test_tensors(
+            in_shape,
+            test_dtype=dtype,
+            test_device=device,
+            requires_grad=input_requires_grad,
+        )
+        scales_ref, scales_test = make_reference_and_test_tensors(
+            in_shape[:-1],
+            test_dtype=dtype,
+            test_device=device,
+            requires_grad=scales_requires_grad,
+        )
+        dy_ref, dy_test = make_reference_and_test_tensors(
+            in_shape,
+            test_dtype=dtype,
+            test_device=device,
+            requires_grad=False,
+        )
+
+        # Plain PyTorch implementation. Autograd supplies the reference gradients, so
+        # the op's hand-written backward is checked against a derivative it played no
+        # part in computing.
+        y = (
+            tanh_clamp_scale * torch.tanh(torch.nn.functional.relu(x_ref) / tanh_clamp_scale)
+        ).square()
+        y_ref = scales_ref.unsqueeze(-1) * y
+        if input_requires_grad or scales_requires_grad:
+            y_ref.backward(dy_ref)
+
+        # Implementation with fusible operation
+        op = te_ops.ScaledTanhSReLU(tanh_clamp_scale=tanh_clamp_scale)
+        y_test = op(x_test, scales_test)
+        if input_requires_grad or scales_requires_grad:
+            y_test.backward(dy_test)
+
+        # Check results
+        tols = dtype_tols(dtype)
+        y_test = y_test.to(dtype=torch.float64, device="cpu")
+        assert_close(y_test, y_ref, **tols)
+        if input_requires_grad:
+            assert_close_grads(x_test, x_ref, **tols)
+        if scales_requires_grad:
+            assert_close_grads(scales_test, scales_ref, **tols)
+
+    def test_scaled_tanh_srelu_saturates(self) -> None:
+        """Large inputs pin the output at tanh_clamp_scale**2, unlike plain SReLU."""
+        s = 2.0
+        x = torch.full((4, 8), 1.0e3, device="cuda", dtype=torch.float32)
+        scales = torch.ones((4,), device="cuda", dtype=torch.float32)
+
+        y = te_ops.ScaledTanhSReLU(tanh_clamp_scale=s)(x, scales)
+        torch.testing.assert_close(y, torch.full_like(y, s * s))
+
+        # Plain SReLU on the same input is ~250000x larger, so this cannot pass by
+        # accident if the clamp were silently dropped.
+        y_unclamped = te_ops.ScaledSReLU()(x, scales)
+        assert y_unclamped.min().item() > 1.0e5
+
+    @pytest.mark.parametrize("tanh_clamp_scale", (0.0, -1.0, float("inf"), float("nan")))
+    def test_scaled_tanh_srelu_rejects_bad_clamp_scale(self, tanh_clamp_scale) -> None:
+        """The clamp scale must be finite and positive."""
+        with pytest.raises(ValueError, match="tanh_clamp_scale"):
+            te_ops.ScaledTanhSReLU(tanh_clamp_scale=tanh_clamp_scale)
+
+    def test_scaled_tanh_srelu_activation_recompute_in_mlp_config(self) -> None:
+        """Tanh SReLU exposes the same activation recompute knob as ScaledSReLU."""
+        op = te_ops.ScaledTanhSReLU(tanh_clamp_scale=2.0)
+        assert op.activation_recompute_in_mlp is False
+        assert te_ops.ScaledTanhSReLU(
+            tanh_clamp_scale=2.0, activation_recompute_in_mlp=True
+        ).activation_recompute_in_mlp
+
     def test_interleaved_scaled_swiglu(self):
         """SwiGLU with post-scale and block interleaved input format"""
         self.test_scaled_swiglu(

tests/pytorch/test_grouped_mlp.py

@@ -57,6 +57,12 @@
 mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True)
 nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True)
 
+# Soft-clamp scale for ScaledTanhSReLU coverage. Deliberately small relative to the
+# FC1 outputs these tests produce, so tanh actually saturates -- a large scale would
+# make the activation numerically indistinguishable from plain ScaledSReLU and the
+# test would pass even if the clamp were dropped.
+_TANH_SRELU_CLAMP_SCALE: float = 2.0
+
 # Supported data types
 _dtypes: list[torch.dtype] = [torch.float32, torch.float16]
 if is_bf16_available():  # bf16 requires sm_80 or higher
@@ -1072,6 +1078,7 @@ class TestGroupedMLPFusedOp:
             "scaled_clamped_qgeglu",
             "scaled_clamped_qgeglu_custom",
             "scaled_srelu",
+            "scaled_tanh_srelu",
         ),
     )
     def test_grouped_mlp(
@@ -1140,12 +1147,25 @@ def test_grouped_mlp(
             pytest.skip("Unary activations do not use GLU interleaving")
         if quantization == "nvfp4_4over6":
             pytest.skip("NVFP4 4over6 grouped quantization is not supported")
-        if activation == "scaled_srelu" and quantization in ("nvfp4", "nvfp4_rht") and bias:
+        if (
+            activation in ("scaled_srelu", "scaled_tanh_srelu")
+            and quantization in ("nvfp4", "nvfp4_rht")
+            and bias
+        ):
             pytest.skip("NVFP4 SReLU grouped MLP coverage is limited to no-bias")
         if quantization == "nvfp4_rht":
             if activation == "scaled_swiglu" and (bias or glu_interleave_size != 32):
                 pytest.skip("NVFP4 RHT SwiGLU grouped MLP coverage is limited to no-bias")
-            if activation not in ("scaled_swiglu", "scaled_situglu", "scaled_srelu"):
+            # tanh-SReLU is included deliberately: NVFP4 is the only path that exercises
+            # the dsrelu-family fc2 alpha (full product rather than sqrt), so excluding
+            # it here would leave that branch untested. It runs without the hadamard
+            # sub-kernel, which the fused op handles via the generic quantize fallback.
+            if activation not in (
+                "scaled_swiglu",
+                "scaled_situglu",
+                "scaled_srelu",
+                "scaled_tanh_srelu",
+            ):
                 pytest.skip(
                     "NVFP4 RHT grouped MLP coverage is limited to SwiGLU, SiTU-GLU, and SReLU"
                 )
@@ -1271,6 +1291,11 @@ def _apply_activation(x: torch.Tensor) -> torch.Tensor:
                 return (x2c + geglu_offset) * (x1c * torch.sigmoid(geglu_alpha * x1c))
             if activation == "scaled_srelu":
                 return torch.nn.functional.relu(x).square()
+            if activation == "scaled_tanh_srelu":
+                s = _TANH_SRELU_CLAMP_SCALE
+                return (
+                    (s * torch.tanh(torch.nn.functional.relu(x.float()) / s)).square().to(x.dtype)
+                )
             raise ValueError(f"Unexpected grouped MLP activation ({activation})")
 
         # Reference implementation
@@ -1314,6 +1339,8 @@ def _make_scaled_act():
                 return te.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size)
             if activation == "scaled_srelu":
                 return te.ops.ScaledSReLU()
+            if activation == "scaled_tanh_srelu":
+                return te.ops.ScaledTanhSReLU(tanh_clamp_scale=_TANH_SRELU_CLAMP_SCALE)
             raise ValueError(f"Unexpected grouped MLP activation ({activation})")
 
         def _make_module():
@@ -1400,15 +1427,20 @@ def _make_module():
             fc2.backward_dw()
 
         # Check for expected fusions
-        cudnn_frontend_supports_grouped_mlp = (
-            grouped_mlp_module._cudnn_frontend_supports_grouped_gemm_situglu()
-            if activation == "scaled_situglu"
-            else (
+        if activation == "scaled_situglu":
+            cudnn_frontend_supports_grouped_mlp = (
+                grouped_mlp_module._cudnn_frontend_supports_grouped_gemm_situglu()
+            )
+        elif activation == "scaled_srelu":
+            cudnn_frontend_supports_grouped_mlp = _cudnn_frontend_supports_grouped_gemm_srelu()
+        elif activation == "scaled_tanh_srelu":
+            # Needs both the base srelu kernels and the tanh_clamp_scale parameter.
+            cudnn_frontend_supports_grouped_mlp = (
                 _cudnn_frontend_supports_grouped_gemm_srelu()
-                if activation == "scaled_srelu"
-                else _cudnn_frontend_version_supported()
+                and grouped_mlp_module._cudnn_frontend_supports_grouped_gemm_srelu_tanh()
             )
-        )
+        else:
+            cudnn_frontend_supports_grouped_mlp = _cudnn_frontend_version_supported()
         expected_grouped_mlp_fusion = cudnn_frontend_supports_grouped_mlp and (
             (
                 quantization == "mxfp8"
@@ -1908,6 +1940,7 @@ def _run(grad):
             "scaled_clamped_qgeglu",
             "scaled_clamped_qgeglu_custom",
             "scaled_srelu",
+            "scaled_tanh_srelu",
         ),
     )
     def test_grouped_mlp_fp16(

transformer_engine/pytorch/ops/basic/__init__.py

@@ -14,6 +14,7 @@
     ReGLU,
     SReLU,
     ScaledSReLU,
+    ScaledTanhSReLU,
     SReGLU,
     SiLU,
 )

transformer_engine/pytorch/ops/basic/activation.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 import abc
+import math
 from collections.abc import Iterable, Sequence
 from typing import Any, Optional
 
@@ -29,6 +30,7 @@
     "ReGLU",
     "SReLU",
     "ScaledSReLU",
+    "ScaledTanhSReLU",
     "SReGLU",
     "SiLU",
 ]
@@ -508,6 +510,88 @@ def _scaled_unary_backward(
         )
 
 
+class ScaledTanhSReLU(_ScaledUnary):
+    r"""Tanh soft-clamped squared ReLU with per-row post-scaling.
+
+    Identical to :class:`ScaledSReLU` except that the ReLU output is soft-clamped
+    by a tanh before squaring, which bounds the activation by
+    ``tanh_clamp_scale ** 2``:
+
+    .. math::
+        y = \left( s \cdot \tanh\left( \frac{\mathrm{relu}(x)}{s} \right) \right)^2 \cdot \mathrm{scales}
+
+    "Tanh" rather than "Clamped" is deliberate: in these ops ``Clamped`` already
+    means a hard min/max clamp (see :class:`ScaledClampedQGeGLU`), whereas this is
+    a smooth saturating clamp.
+
+    Parameters
+    ----------
+    tanh_clamp_scale : float
+        The soft-clamp scale ``s``; must be finite and positive. Required, since
+        an unclamped tanh-SReLU is just :class:`ScaledSReLU`.
+    activation_recompute_in_mlp : bool, default = ``False``
+        Enable fused grouped MLP kernels to recompute activation outputs
+        during backward when supported instead of saving them.
+    """
+
+    def __init__(
+        self,
+        *,
+        tanh_clamp_scale: float,
+        activation_recompute_in_mlp: bool = False,
+    ) -> None:
+        super().__init__(activation_recompute_in_mlp=activation_recompute_in_mlp)
+        self.tanh_clamp_scale: float = float(tanh_clamp_scale)
+        if not math.isfinite(self.tanh_clamp_scale) or self.tanh_clamp_scale <= 0.0:
+            raise ValueError(
+                f"tanh_clamp_scale must be finite and positive, got {tanh_clamp_scale}"
+            )
+
+    def _tanh_srelu_terms(self, input_: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """``b = s*tanh(relu(x)/s)`` and ``t``, in fp32.
+
+        Matches the fused kernels, which evaluate the whole epilogue in fp32 and
+        round once at the store rather than at each step.
+        """
+        s = self.tanh_clamp_scale
+        t = torch.tanh(torch.relu(input_.float()) / s)
+        return s * t, t
+
+    def _scaled_unary_forward(
+        self,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+    ) -> torch.Tensor:
+        # No fused CUDA kernel for this activation yet, so the unfused path is
+        # expressed in torch. It only serves non-SM100 / non-fused configurations;
+        # the fused grouped-MLP path goes straight to the cuDNN srelu_tanh epilogue.
+        b, _ = self._tanh_srelu_terms(input_)
+        out = b.square() * scales.float().unsqueeze(-1)
+        return out.to(input_.dtype)
+
+    def _scaled_unary_backward(
+        self,
+        grad_output: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+        *,
+        compute_scale_grad: bool,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        # d/dx (b^2) = 2*b*(1 - t^2); the scale multiplies it, exactly as the
+        # scaled-unary CUDA kernel does for the plain squared ReLU.
+        b, t = self._tanh_srelu_terms(input_)
+        grad_output_f32 = grad_output.float()
+        grad_input = grad_output_f32 * scales.float().unsqueeze(-1) * (2.0 * b * (1.0 - t * t))
+
+        grad_scales = None
+        if compute_scale_grad:
+            # Gradient wrt the per-row scale is sum over the last dim of the
+            # unscaled activation times the incoming gradient.
+            grad_scales = (b.square() * grad_output_f32).sum(dim=-1).to(scales.dtype)
+
+        return grad_input.to(input_.dtype), grad_scales
+
+
 class SReGLU(_ActivationOperation):
     r"""Squared Rectified Gated Linear Unit