From e28c0cca12157a03f23fda434ed56262bbf544a4 Mon Sep 17 00:00:00 2001 From: Wanying Wang Date: Wed, 2 Sep 2026 12:22:10 -0700 Subject: [PATCH] [PyTorch] Add ScaledTanhSReLU and wire it into the fused grouped MLP 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 --- tests/pytorch/test_fusible_ops.py | 95 +++++++++++++++++++ tests/pytorch/test_grouped_mlp.py | 51 ++++++++-- .../pytorch/ops/basic/__init__.py | 1 + .../pytorch/ops/basic/activation.py | 84 ++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 84 ++++++++++++++-- 5 files changed, 299 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 0083fd9c0f..6cd1fc3065 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/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( diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 5049f58ce4..d48f7afae6 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/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( diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 4eb2796b2c..dae15330c1 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -14,6 +14,7 @@ ReGLU, SReLU, ScaledSReLU, + ScaledTanhSReLU, SReGLU, SiLU, ) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index a974d41ef9..5c33c08b44 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/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 diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index dca1119800..61f80b9d9f 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -47,6 +47,7 @@ ScaledSiTUGLU, ScaledSReLU, ScaledSwiGLU, + ScaledTanhSReLU, ) from ..fuser import register_forward_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext @@ -98,6 +99,34 @@ def _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() -> bool: return _cudnn_frontend_version_at_least("1.26.0") +@functools.lru_cache(maxsize=None) +def _cudnn_frontend_supports_grouped_gemm_srelu_tanh() -> bool: + """Feature-detect complete cuDNN frontend grouped tanh-SReLU support. + + Both directions are required: a frontend with only the forward clamp would + train against an unclamped backward. Detected by signature rather than + version so this can be developed against an editable cuDNN FE checkout; a + min-version constant can replace it once the feature is in a release. + """ + try: + from cudnn import ( # pylint: disable=import-outside-toplevel + grouped_gemm_dsrelu_wrapper_sm100, + grouped_gemm_srelu_wrapper_sm100, + ) + except ImportError: + return False + try: + wrappers = ( + grouped_gemm_srelu_wrapper_sm100, + grouped_gemm_dsrelu_wrapper_sm100, + ) + return all( + "tanh_clamp_scale" in inspect.signature(wrapper).parameters for wrapper in wrappers + ) + except (TypeError, ValueError): + return False + + @functools.lru_cache(maxsize=None) def _cudnn_frontend_supports_grouped_gemm_situglu() -> bool: """Feature-detect complete cuDNN frontend grouped SiTU-GLU support.""" @@ -132,9 +161,13 @@ def _cudnn_frontend_supports_single_group_runtime_offsets( activation_type: type[FusibleOperation], ) -> bool: """Check cuDNN FE support for single-group runtime offsets.""" - return not issubclass(activation_type, ScaledSReLU) and _cudnn_frontend_version_at_least( - "1.27.0" - ) + # The srelu/dsrelu wrappers take no use_single_group_runtime_offsets argument, + # so every activation in that family has to be excluded here, not just + # ScaledSReLU -- passing it through would raise TypeError on the single-group + # (shared expert) path. + return not issubclass( + activation_type, (ScaledSReLU, ScaledTanhSReLU) + ) and _cudnn_frontend_version_at_least("1.27.0") def _wrap_single_quantized_as_grouped( @@ -794,7 +827,7 @@ def validate_grouped_mlp_dims(fc1, activation_op, fc2) -> None: ) if is_glu_activation(activation_op): expected_fc1_out_features = 2 * fc2.in_features - elif isinstance(activation_op, ScaledSReLU): + elif isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)): expected_fc1_out_features = fc2.in_features else: raise TypeError(f"Unsupported grouped MLP activation ({activation_op.__class__.__name__}).") @@ -1008,6 +1041,21 @@ def __init__( self._cudnn_situ_beta1: float = activation.beta1 self._cudnn_situ_beta2: float = activation.beta2 + # Set unconditionally: the forward/backward paths read this attribute for + # every activation, so leaving it undefined would break plain ScaledSReLU. + self._pass_srelu_tanh_params: bool = isinstance(activation, ScaledTanhSReLU) + if self._pass_srelu_tanh_params: + # Fail at construction rather than silently running unclamped, which + # would train a different model than the config asks for. + if not _cudnn_frontend_supports_grouped_gemm_srelu_tanh(): + raise RuntimeError( + "ScaledTanhSReLU requires a cuDNN frontend whose " + "grouped_gemm_srelu_wrapper_sm100 and grouped_gemm_dsrelu_wrapper_sm100 " + "accept tanh_clamp_scale. The installed frontend does not, and running " + "without it would silently apply an unclamped squared ReLU." + ) + self._cudnn_tanh_clamp_scale: float = activation.tanh_clamp_scale + def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -1386,6 +1434,10 @@ def fuser_forward( and fc2_input_quantizer.with_rht and fc2_input_quantizer.with_post_rht_amax ) + # Deliberately ScaledSReLU only: the hadamard kernel this selects is the GLU + # one, which has no soft-clamp support, so ScaledTanhSReLU must not reach it. + # The cost is that NVFP4 RHT gives up hadamard fusion for tanh-SReLU -- a + # performance limitation, not a correctness one. activation_is_srelu = isinstance(activation_op, ScaledSReLU) activation_supports_hadamard = self._cudnn_act_func in ("swiglu", "situglu") or ( activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() @@ -1434,6 +1486,8 @@ def fuser_forward( situ_beta1=self._cudnn_situ_beta1, situ_beta2=self._cudnn_situ_beta2, ) + if self._pass_srelu_tanh_params: + fc1_activation_kwargs.update(tanh_clamp_scale=self._cudnn_tanh_clamp_scale) if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. @@ -1749,6 +1803,10 @@ def fuser_forward( mark_grouped_tensor(saved_fc1_x, activation_in, scales, grouped_fc2_x) activation_op = self.basic_ops[1] cpu_offloading = is_cpu_offload_enabled() + # Deliberately ScaledSReLU only for now: ScaledTanhSReLU falls back to + # saving fc2_x, which costs memory but stays correct. The cuDNN dsrelu + # d_srelu regeneration does honour the clamp, so enabling recompute here + # is a viable follow-up rather than a blocker. activation_is_srelu = isinstance(activation_op, ScaledSReLU) activation_recompute_in_mlp = bool( getattr(activation_op, "activation_recompute_in_mlp", False) @@ -1835,7 +1893,9 @@ def fuser_backward( # Get basic operations fc1_op, activation_op, fc2_op = self.basic_ops - activation_is_srelu = isinstance(activation_op, ScaledSReLU) + # Selects how the NVFP4 fc2 alpha is folded below: the whole dsrelu family + # applies alpha once, unlike the gated kernels which need sqrt(product). + activation_is_srelu = isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)) fc1_ctx, _activation_ctx, fc2_ctx = basic_op_ctxs # Tensor properties @@ -2098,6 +2158,8 @@ def fuser_backward( situ_beta1=self._cudnn_situ_beta1, situ_beta2=self._cudnn_situ_beta2, ) + if self._pass_srelu_tanh_params: + fc2_dactivation_kwargs.update(tanh_clamp_scale=self._cudnn_tanh_clamp_scale) fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 if is_distributed_weight(fc2_leader): @@ -2629,13 +2691,21 @@ def fuse_srelu_ops( recipe: Optional[Recipe] = None, **unused, # pylint: disable=unused-argument ) -> list[FusibleOperation]: - """Apply joint GroupedLinear + ScaledSReLU + GroupedLinear fusion.""" + """Apply joint GroupedLinear + scaled unary activation + GroupedLinear fusion.""" + + # ScaledTanhSReLU joins only when the installed cuDNN frontend can actually + # clamp. Listing it unconditionally would let the op fuse and then raise from + # _GroupedMLP_CuTeGEMMBase.__init__; leaving it out simply declines the fusion + # and runs the correct unfused activation instead. + activation_op_types: tuple[type[FusibleOperation], ...] = (ScaledSReLU,) + if _cudnn_frontend_supports_grouped_gemm_srelu_tanh(): + activation_op_types += (ScaledTanhSReLU,) return fuse_grouped_mlp_ops( ops, recipe=recipe, fused_op_cls=GroupedMLP_CuTeGEMMUnary, - activation_op_types=(ScaledSReLU,), + activation_op_types=activation_op_types, )