Use cuDNN's deterministic dprob in the fused grouped MLP

docs/envvars.rst

@@ -206,7 +206,14 @@ backend-selection overview.
 
    :Type: ``int`` (0 or 1)
    :Default: ``1``
-   :Description: Allow non-deterministic algorithms for Transformer Engine execution. When set to ``0``, only deterministic algorithms are allowed. This is relevant for both PyTorch and JAX attention implementations.
+   :Description: Allow non-deterministic algorithms for Transformer Engine execution. When set to ``0``, only deterministic algorithms are allowed. This is relevant for both PyTorch and JAX attention implementations. In PyTorch it also asks the cuDNN grouped-GEMM dSReLU backward used by the CuTe DSL fused grouped MLP for a bit-exact scale gradient (``dprob``) and bias gradient, in place of the default cross-CTA atomic accumulation whose summation order follows the tile scheduler.
+
+   .. note::
+
+      The deterministic dSReLU backward needs cuDNN frontend 1.28.0 or newer, and has no
+      counterpart for the GLU activations -- their ``dprob`` is accumulated with atomics
+      either way. Transformer Engine warns once when ``NVTE_ALLOW_NONDETERMINISTIC_ALGO=0``
+      cannot be honored for this reason.
 
 .. envvar:: NVTE_FUSED_RING_ATTENTION_USE_SCAN
 

tests/pytorch/test_grouped_mlp.py

@@ -8,6 +8,7 @@
 import os
 import math
 import random
+import warnings
 from typing import Optional
 
 import pytest
@@ -1972,6 +1973,118 @@ def train_step(
                 assert_close(graph_grad, param.grad, **tols)
 
 
+class TestGroupedMLPDeterminism:
+    """``NVTE_ALLOW_NONDETERMINISTIC_ALGO`` coverage for the CuTe DSL fused grouped MLP.
+
+    cuDNN's grouped-GEMM dactivation epilogue accumulates the scale gradient (``dprob``)
+    with cross-CTA atomic adds, so its summation order follows the tile scheduler. The
+    dSReLU wrapper takes a ``deterministic`` argument from cuDNN frontend 1.28.0 on that
+    replaces those atomics with per-subtile slots reduced in a canonical order; the dGLU
+    wrapper has no equivalent, and neither does an older front-end, so those two cases
+    warn instead.
+    """
+
+    @staticmethod
+    def _reset_caches() -> None:
+        """Drop the warn-once state, which is an lru_cache keyed on the reason string."""
+        grouped_mlp_module._warn_nondeterministic_cudnn_dprob.cache_clear()
+
+    @pytest.fixture(autouse=True)
+    def _clean_caches(self):
+        self._reset_caches()
+        yield
+        self._reset_caches()
+
+    @pytest.mark.parametrize(
+        "allow_nondeterministic,expected",
+        (
+            (None, False),  # default: non-deterministic algorithms are allowed
+            ("1", False),
+            ("0", True),
+        ),
+    )
+    def test_env_var_selects_deterministic_mode(
+        self,
+        monkeypatch,
+        *,
+        allow_nondeterministic: Optional[str],
+        expected: bool,
+    ) -> None:
+        """Uncached, so flipping the variable mid-process takes effect."""
+        if allow_nondeterministic is None:
+            monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False)
+        else:
+            monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", allow_nondeterministic)
+        assert grouped_mlp_module._deterministic_algorithms_required() is expected
+
+    def test_only_the_srelu_path_can_be_deterministic(self) -> None:
+        """The capability is a property of the wrapper, not of the environment.
+
+        Needs neither a GPU nor cuDNN: both answers come from the installed
+        ``nvidia-cudnn-frontend`` version.
+        """
+        assert (
+            grouped_mlp_module.GroupedMLP_CuTeGEMMGLU.grouped_gemm_dactivation_is_deterministic()
+            is False
+        )
+        assert (
+            grouped_mlp_module.GroupedMLP_CuTeGEMMUnary.grouped_gemm_dactivation_is_deterministic()
+            is grouped_mlp_module._cudnn_frontend_supports_deterministic_dprob()
+        )
+
+    def test_dprob_warning_is_emitted_once_per_reason(self) -> None:
+        """TE flags a determinism request it cannot honor -- but not on every backward."""
+        with pytest.warns(UserWarning, match="dprob"):
+            grouped_mlp_module._warn_nondeterministic_cudnn_dprob("first reason")
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            grouped_mlp_module._warn_nondeterministic_cudnn_dprob("first reason")
+        assert not caught, "the dprob warning must not repeat every backward"
+        # A different reason is a different remedy, so it is worth saying once too.
+        with pytest.warns(UserWarning, match="second reason"):
+            grouped_mlp_module._warn_nondeterministic_cudnn_dprob("second reason")
+
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_srelu"))
+    def test_deterministic_dactivation_is_numerically_correct(
+        self,
+        monkeypatch,
+        *,
+        activation: str,
+    ) -> None:
+        """End-to-end under determinism, and pins which of the two arms warns.
+
+        ``group_size > 1`` here, so the unit-activation-scale shortcut is off and the
+        cuDNN epilogue really does produce ``dprob``.
+        """
+        if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+
+        monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
+        self._reset_caches()
+        fused_op_cls = (
+            grouped_mlp_module.GroupedMLP_CuTeGEMMUnary
+            if activation == "scaled_srelu"
+            else grouped_mlp_module.GroupedMLP_CuTeGEMMGLU
+        )
+        expect_warning = not fused_op_cls.grouped_gemm_dactivation_is_deterministic()
+
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            TestGroupedMLPFusedOp().test_grouped_mlp(
+                bias=False,
+                hidden_size=128,
+                quantization="mxfp8",
+                single_grouped_weight=False,
+                activation=activation,
+            )
+        # Asserted after the call rather than inside pytest.warns: on exit pytest.warns
+        # raises its own "DID NOT WARN" and chains it over whatever the body raised, so a
+        # real failure inside the op would report as a warning assertion and bury the cause.
+        warned = any("dprob" in str(w.message) for w in caught)
+        assert warned is expect_warning
+
+
 def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None:
     if not mxfp8_available:
         pytest.skip(reason_for_no_mxfp8)

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -9,6 +9,7 @@
 from collections.abc import Callable, Iterable
 import functools
 import os
+import warnings
 from importlib.metadata import PackageNotFoundError, version as get_pkg_version
 from typing import Any, Optional
 
@@ -106,6 +107,44 @@ def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool:
     return _cudnn_frontend_version_at_least("1.27.0")
 
 
+def _cudnn_frontend_supports_deterministic_dprob() -> bool:
+    """Check cuDNN FE min version for deterministic ``dprob`` in the dSReLU backward.
+
+    cuDNN frontend 1.28.0 (NVIDIA/cudnn-frontend#521) gave
+    ``grouped_gemm_dsrelu_wrapper_sm100`` a ``deterministic`` argument: each N-subtile's
+    partial result parks in its own slot and the slots are summed in a canonical order,
+    replacing the cross-CTA atomic accumulation whose summation order follows the tile
+    scheduler. ``grouped_gemm_dglu_wrapper_sm100`` has no equivalent argument, so the GLU
+    activations cannot take this path.
+    """
+    return _cudnn_frontend_version_at_least("1.28.0")
+
+
+def _deterministic_algorithms_required() -> bool:
+    """Whether the user has asked for bit-exact reproducibility.
+
+    Same check as ``transformer_engine.pytorch.triton.grouped_dbias_dscales``.
+    Deliberately uncached, so a test can flip the variable.
+    """
+    return not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")))
+
+
+@functools.lru_cache(maxsize=None)
+def _warn_nondeterministic_cudnn_dprob(reason: str) -> None:
+    """Warn once per reason that determinism did not reach cuDNN's ``dprob`` accumulation.
+
+    Cached because the call site runs on every backward pass.
+    """
+    warnings.warn(
+        "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 selects deterministic kernels inside"
+        " Transformer Engine, but the cuDNN grouped-GEMM dactivation backward that the"
+        " CuTe DSL fused grouped MLP calls still accumulates the scale gradient (dprob)"
+        " with cross-CTA atomic adds, whose summation order is set by the tile scheduler."
+        f" That op is therefore not bit-exact here: {reason}.",
+        UserWarning,
+    )
+
+
 def _wrap_single_quantized_as_grouped(
     tensor: torch.Tensor,
     quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage,
@@ -883,6 +922,16 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable:
         """Fused kernel for grouped GEMM, activation backward, and scale grad."""
         raise NotImplementedError
 
+    @classmethod
+    def grouped_gemm_dactivation_is_deterministic(cls) -> bool:
+        """Whether this op's dactivation kernel can be asked for a bit-exact ``dprob``.
+
+        Reported per subclass rather than per environment variable: the argument exists on
+        the dSReLU wrapper only, so a GLU activation stays non-deterministic no matter how
+        new the installed cuDNN front-end is.
+        """
+        return False
+
     @classmethod
     @functools.lru_cache(maxsize=None)
     def grouped_gemm_quant_kernel(cls) -> Callable:
@@ -1947,10 +1996,32 @@ def fuser_backward(
         scales_f32 = None
         scales_tensor = None
         dscales_tensor = None
+        deterministic_dactivation = False
         if not unit_activation_scale:
             scales_f32 = scales.detach().to(dtype=torch.float32)
             scales_tensor = scales_f32.reshape(-1, 1, 1)
             dscales_tensor = torch.zeros_like(scales_tensor)
+            # Decided here rather than up front because this is where dprob is actually
+            # produced: with a unit activation scale the cuDNN epilogue never runs its
+            # atomic dprob accumulation, so there is nothing to make deterministic.
+            if _deterministic_algorithms_required():
+                deterministic_dactivation = self.grouped_gemm_dactivation_is_deterministic()
+                if not deterministic_dactivation:
+                    # Two ways to get here, with different remedies: no deterministic mode
+                    # exists for this activation, or one exists but the installed front-end
+                    # predates it.
+                    if self._cudnn_dact_func is not None:
+                        reason = (
+                            "grouped_gemm_dglu_wrapper_sm100 has no deterministic mode, so"
+                            " only the scaled-SReLU activation can be made bit-exact"
+                        )
+                    else:
+                        reason = (
+                            "grouped_gemm_dsrelu_wrapper_sm100 takes a deterministic argument"
+                            " only from cuDNN frontend 1.28.0 on; upgrade"
+                            " nvidia-cudnn-frontend to get a bit-exact dprob"
+                        )
+                    _warn_nondeterministic_cudnn_dprob(reason)
 
         fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn
         if use_nvfp4:
@@ -1998,6 +2069,12 @@ def fuser_backward(
             "use_dynamic_sched": True,
         }
         dactivation_kernel = self.grouped_gemm_dactivation_kernel()
+        if deterministic_dactivation:
+            # Only ever set to True. Left unset otherwise so the wrapper keeps its own
+            # default, which follows ``torch.use_deterministic_algorithms``; passing False
+            # would override that and silently take determinism away from a caller who
+            # asked torch for it without setting NVTE_ALLOW_NONDETERMINISTIC_ALGO.
+            fc2_dactivation_kwargs["deterministic"] = True
         if _cudnn_frontend_supports_single_group_runtime_offsets():
             fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
         if self._cudnn_dact_func is not None:
@@ -2525,6 +2602,11 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable:
 
         return grouped_gemm_dsrelu_wrapper_sm100
 
+    @classmethod
+    def grouped_gemm_dactivation_is_deterministic(cls) -> bool:
+        """``grouped_gemm_dsrelu_wrapper_sm100`` takes ``deterministic`` from 1.28.0 on."""
+        return _cudnn_frontend_supports_deterministic_dprob()
+
 
 def fuse_ops(
     ops: list[FusibleOperation],