Skip to content

[Bug] use_single_group_runtime_offsets arg is not supported for ScaledSrelu #3328

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vthumbe1503 merged 2 commits into from
Aug 10, 2026
+14 −8
Merged

[Bug] use_single_group_runtime_offsets arg is not supported for ScaledSrelu #3328

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ebe69c4
fix cudnn version gurad
vthumbe1503 Aug 7, 2026
e6a6816
Merge branch 'main' into update_arg_condition
ksivaman Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions tests/pytorch/test_grouped_mlp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
16 changes: 10 additions & 6 deletions transformer_engine/pytorch/ops/fused/grouped_mlp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
Loading