From ebe69c4dc5b26ff7e1a33554d24e77380fdc0305 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Fri, 7 Aug 2026 19:03:29 +0000 Subject: [PATCH] fix cudnn version gurad Signed-off-by: Varun Thumbe --- tests/pytorch/test_grouped_mlp.py | 6 ++++-- .../pytorch/ops/fused/grouped_mlp.py | 16 ++++++++++------ 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d09c92ad49..499c22509c 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1559,13 +1559,15 @@ def test_grouped_mlp_single_group_mxfp8( """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8.""" if ( runtime_offsets_supported - and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets() + and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets( + te.ops.ScaledSwiGLU + ) ): pytest.skip("Requires cuDNN frontend >= 1.27.0") monkeypatch.setattr( grouped_mlp_module, "_cudnn_frontend_supports_single_group_runtime_offsets", - lambda: runtime_offsets_supported, + lambda _activation_type: runtime_offsets_supported, ) self.test_grouped_mlp( group_size=1, diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 76d51673f0..909e5a8a9b 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -101,9 +101,13 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() -def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool: - """Check cuDNN FE min version for single-group runtime offsets.""" - return _cudnn_frontend_version_at_least("1.27.0") +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" + ) def _wrap_single_quantized_as_grouped( @@ -1073,7 +1077,7 @@ def fuser_forward( activation_kernel = self.grouped_gemm_activation_kernel() supports_single_group_runtime_offsets = ( - _cudnn_frontend_supports_single_group_runtime_offsets() + _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) ) # Shared experts have one dense group and all optimized kernels derive M @@ -2032,7 +2036,7 @@ def fuser_backward( "use_dynamic_sched": True, } dactivation_kernel = self.grouped_gemm_dactivation_kernel() - if _cudnn_frontend_supports_single_group_runtime_offsets(): + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if self._cudnn_dact_func is not None: fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor @@ -2377,7 +2381,7 @@ def fuser_backward( "use_dynamic_sched": True, } fc1_dgrad_kernel = self.grouped_gemm_quant_kernel() - if _cudnn_frontend_supports_single_group_runtime_offsets(): + if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 if fc1_op.single_grouped_weight: