From a00b7f2c7d146a8a71224d081bb21b8f39413861 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 19 Aug 2026 14:42:12 -0700 Subject: [PATCH 01/14] [PyTorch] Ask cuDNN for a deterministic dprob under NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 The cuDNN grouped-GEMM dactivation backward that the CuTe DSL fused grouped MLP calls accumulates the scale gradient (dprob) with cross-CTA atomic adds, so its floating-point summation order follows the tile scheduler and varies run to run. Until now there was no way to switch that off, and NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 did not reach it: the run trained fine and was silently not reproducible. cuDNN frontend 1.28.0 (NVIDIA/cudnn-frontend#521) added a `deterministic` argument to grouped_gemm_dsrelu_wrapper_sm100 that parks each N-subtile's partial result in its own slot and sums the slots in a canonical order, for dprob and for dbias. Pass it from the TE flag. Passed as True or not at all, never as False. The wrapper's own default is None, which follows torch.use_deterministic_algorithms; sending an explicit False would override that and take determinism away from a caller who asked torch for it without setting the TE variable. The capability is reported per subclass rather than per environment variable, because grouped_gemm_dglu_wrapper_sm100 has no equivalent argument -- a GLU activation stays non-deterministic however new the installed front-end is. That case, and an SReLU op on a front-end older than 1.28.0, warn instead, once per distinct reason since the remedies differ. The warning is raised from where dprob is actually produced: with a unit activation scale the epilogue never runs its atomic accumulation, so there is nothing to make deterministic and nothing to warn about. Tests: TestGroupedMLPDeterminism covers the env-var parse, that only the SReLU op reports the capability and that it tracks the front-end version (no GPU or cuDNN needed for either), that the warning fires once per reason, and an MXFP8 end-to-end run under determinism for both SwiGLU and SReLU that checks numerics and pins which of the two arms warns. Signed-off-by: Zhiyu Li --- docs/envvars.rst | 9 +- tests/pytorch/test_grouped_mlp.py | 113 ++++++++++++++++++ .../pytorch/ops/fused/grouped_mlp.py | 82 +++++++++++++ 3 files changed, 203 insertions(+), 1 deletion(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 97eaed5ddc..a644d56c8b 100644 --- a/docs/envvars.rst +++ b/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 diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index d195eb2f78..b51ec7d967 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/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) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 0113833647..f0ddb55ee0 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/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], From c770fc4b72f050565a69dbef44c3fcebbdee559f Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 19 Aug 2026 14:47:55 -0700 Subject: [PATCH 02/14] Honor torch.use_deterministic_algorithms too, not just the env variable _deterministic_algorithms_required() copied the narrow check from transformer_engine.pytorch.triton.grouped_dbias_dscales, which reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. DotProductAttention takes the union instead -- the variable OR torch.use_deterministic_algorithms -- and that is the right precedent here. The two knobs answer different questions. The variable is set once in a job launcher, applies uniformly across ranks, and is the only one TE's C++ layer can read. The torch flag is the framework standard, is togglable at runtime, and is what a user who wants reproducibility usually reaches for; most have never heard of the variable. Keying on the variable alone left the torch flag half-honored. The SReLU path happened to come out right, but by delegation rather than by decision: TE passed nothing and the wrapper's own default read torch.are_deterministic_algorithms_enabled(). The GLU path did not -- TE stayed silent about an atomic dprob it cannot fix, for a user who had asked torch for reproducibility. That silence is the exact failure mode the warning exists to prevent, so it was the one case that most needed to warn. Passing the argument only as True, never as False, now needs a different justification than the one the first commit gave: with the union in place the two are equivalent, since the wrapper's default reads the same torch flag TE just read. The reason that survives is narrower and firmer -- the argument does not exist on the dGLU wrapper or on a front-end older than 1.28.0, where passing it at all, even as False, is a TypeError. Tests: the env-var parametrization becomes the two-knob truth table, including the row that motivates the change (torch flag set, NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -- the variable's default is the absence of a request, not a request for non-determinism, so the torch flag still wins). A fixture restores the process-global torch flag. Signed-off-by: Zhiyu Li --- docs/envvars.rst | 9 ++++- tests/pytorch/test_grouped_mlp.py | 31 ++++++++++++---- .../pytorch/ops/fused/grouped_mlp.py | 37 ++++++++++++------- 3 files changed, 55 insertions(+), 22 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index a644d56c8b..3851788cf6 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -208,12 +208,17 @@ backend-selection overview. :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. 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:: + + As for attention, the fused grouped MLP treats determinism as requested when *either* + this variable is ``0`` or :func:`torch.use_deterministic_algorithms` is enabled. + .. 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. + either way. Transformer Engine warns once when a determinism request cannot be + honored for either reason. .. envvar:: NVTE_FUSED_RING_ATTENTION_USE_SCAN diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index b51ec7d967..3b352c62d7 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -1974,7 +1974,7 @@ def train_step( class TestGroupedMLPDeterminism: - """``NVTE_ALLOW_NONDETERMINISTIC_ALGO`` coverage for the CuTe DSL fused grouped MLP. + """Determinism 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 @@ -1995,26 +1995,43 @@ def _clean_caches(self): yield self._reset_caches() + @pytest.fixture + def _restore_torch_determinism(self): + """``use_deterministic_algorithms`` is process-global, so put it back.""" + previous = torch.are_deterministic_algorithms_enabled() + yield + torch.use_deterministic_algorithms(previous) + @pytest.mark.parametrize( - "allow_nondeterministic,expected", + "allow_nondeterministic,torch_flag,expected", ( - (None, False), # default: non-deterministic algorithms are allowed - ("1", False), - ("0", True), + (None, False, False), # default: non-deterministic algorithms are allowed + ("1", False, False), + ("0", False, True), # the TE variable alone + (None, True, True), # the torch flag alone, which TE must not ignore + ("1", True, True), # ... including when the TE variable says otherwise + ("0", True, True), ), ) - def test_env_var_selects_deterministic_mode( + def test_either_knob_requests_determinism( self, monkeypatch, + _restore_torch_determinism, *, allow_nondeterministic: Optional[str], + torch_flag: bool, expected: bool, ) -> None: - """Uncached, so flipping the variable mid-process takes effect.""" + """The same union DotProductAttention uses, and uncached so both stay live. + + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO=1`` is not a request for non-determinism, only + the absence of one, so the torch flag still wins that row. + """ if allow_nondeterministic is None: monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) else: monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", allow_nondeterministic) + torch.use_deterministic_algorithms(torch_flag) assert grouped_mlp_module._deterministic_algorithms_required() is expected def test_only_the_srelu_path_can_be_deterministic(self) -> None: diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index f0ddb55ee0..ecd386c1dd 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -121,12 +121,21 @@ def _cudnn_frontend_supports_deterministic_dprob() -> bool: def _deterministic_algorithms_required() -> bool: - """Whether the user has asked for bit-exact reproducibility. + """Whether the user has asked for bit-exact reproducibility, by either route. - Same check as ``transformer_engine.pytorch.triton.grouped_dbias_dscales``. - Deliberately uncached, so a test can flip the variable. + Same union as ``DotProductAttention``: the two knobs answer different questions. + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO`` is set once in a job launcher, applies uniformly + across ranks and is the only one TE's C++ layer can read; ``use_deterministic_algorithms`` + is the framework standard, is togglable at runtime, and is what a user who wants + reproducibility usually reaches for. Honoring only the first would leave the second + silently unenforced. + + Deliberately uncached: both inputs can change during the process. """ - return not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + return ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() + ) @functools.lru_cache(maxsize=None) @@ -136,11 +145,11 @@ def _warn_nondeterministic_cudnn_dprob(reason: str) -> None: 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}.", + "Deterministic execution was requested (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" + " torch.use_deterministic_algorithms), 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" + f" tile scheduler. That op is therefore not bit-exact here: {reason}.", UserWarning, ) @@ -2070,10 +2079,12 @@ def fuser_backward( } 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. + # Only ever set to True, and only once known to be accepted: the argument does + # not exist on the dGLU wrapper or on a front-end older than 1.28.0, where + # passing it at all -- even as False -- is a TypeError. Leaving it unset is also + # the same request: the wrapper's own default reads + # ``torch.are_deterministic_algorithms_enabled()``, which is half of what + # _deterministic_algorithms_required() just read. fc2_dactivation_kwargs["deterministic"] = True if _cudnn_frontend_supports_single_group_runtime_offsets(): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 From f88d6468a4c57299f2d04671b67453085443de60 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Wed, 19 Aug 2026 14:59:27 -0700 Subject: [PATCH 03/14] Test that dprob is actually bit-exact, not just within tolerance Review caught that nothing in the suite tested the property this change exists for. The end-to-end test runs the op once and checks numerics against a reference with rtol=0.125 / atol=0.25; reordering the same atomic adds moves dprob by about an ulp, so a run that is silently not reproducible passes it comfortably. The tolerance check proves the deterministic path is correct, which is worth keeping, but it cannot prove the path is deterministic. Add a second run. Same module, same inputs, grads cleared between passes, probs.grad compared with torch.equal. Three things the test has to get right to be worth having: * hidden_size 1024, not the 128 used elsewhere. dprob's reduction is over that extent and the tile is 256 wide, so 128 gives a single N-tile, one writer per token, and nothing to reorder -- the assertion would hold by construction and test nothing. * No bias. With an FC2 scale_bias the scale gradient is finished by the Triton grouped dbias/dscales kernel, which refuses to run under determinism, and probs.grad would stop being the dprob under test. * An assertion that the fusion happened, since dprob only comes from the cuDNN epilogue on the fused path. Skipped rather than xfailed on a front-end older than 1.28.0: there the kernel has no deterministic mode and is expected to vary, which is not a failure of this change. Weight gradients are deliberately left out of the comparison -- the CuTe DSL wgrad kernel has its own K-split atomics that this PR does not address. Signed-off-by: Zhiyu Li --- tests/pytorch/test_grouped_mlp.py | 95 +++++++++++++++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 3b352c62d7..aee16e12ea 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2101,6 +2101,101 @@ def test_deterministic_dactivation_is_numerically_correct( warned = any("dprob" in str(w.message) for w in caught) assert warned is expect_warning + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: + """The actual claim: identical inputs produce an identical ``dprob``. + + The tolerance-based check above cannot see this. Reordering the same atomic adds + moves the result by about an ulp, which every tolerance in this file accepts, so a + run that is silently not reproducible passes it. Only an exact comparison of two + runs can tell the difference. + + ``dbias`` rides the same slot mechanism, but reaching it needs an FC1 bias, and an + FC2 ``scale_bias`` then routes the scale gradient through a Triton kernel that + refuses to run under determinism at all. Left to the op-level tests. + """ + fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + if not fused_cls.grouped_gemm_dactivation_is_deterministic(): + pytest.skip( + "grouped_gemm_dsrelu_wrapper_sm100 has no deterministic mode before cuDNN" + " frontend 1.28.0, so dprob is expected to vary run to run here" + ) + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + self._reset_caches() + + device = torch.device("cuda") + dtype = torch.bfloat16 + group_size = 4 + # Wide enough that the dprob reduction spans several N-tiles. With a single tile + # there is one writer per token, nothing to reorder, and the test is vacuous. + hidden_size = 1024 + split_sizes = torch.tensor([256] * group_size, dtype=torch.int, device=device) + num_tokens = int(split_sizes.sum().item()) + + recipe = make_recipe("mxfp8") + _, x = make_reference_and_test_tensors( + (num_tokens, hidden_size), + min=-0.25, + max=0.25, + quantization="mxfp8", + test_dtype=dtype, + test_device=device, + ) + _, dy = make_reference_and_test_tensors( + (num_tokens, hidden_size), + min=-0.25, + max=0.25, + quantization="mxfp8", + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + _, probs = make_reference_and_test_tensors( + (num_tokens,), + test_dtype=dtype, + test_device=device, + ) + + # No bias on either linear: with bias the FC2 scale gradient goes through the + # Triton grouped-dbias kernel instead of arriving straight from the cuDNN epilogue, + # and probs.grad would no longer be the dprob under test. + with te.quantized_model_init(enabled=True, recipe=recipe): + module = te.ops.Sequential( + te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ), + te.ops.ScaledSReLU(), + te.ops.GroupedLinear( + group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype + ), + ) + + def _run() -> torch.Tensor: + x.grad = None + probs.grad = None + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + y.backward(dy) + return probs.grad.detach().clone() + + first = _run() + # The fusion has to have happened, or dprob never came from the cuDNN epilogue and + # the comparison below proves nothing. + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], fused_cls) + second = _run() + + # Exact, not assert_close. Weight gradients are deliberately not compared: the CuTe + # DSL wgrad kernel has its own K-split atomics that this change does not address. + assert torch.equal(first, second), ( + "dprob differed between two identical runs under determinism; max |delta| =" + f" {(first.float() - second.float()).abs().max().item()}" + ) + def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: if not mxfp8_available: From 8e09843c24385232b5565acd8c0e2c5da0527c15 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 20 Aug 2026 10:34:11 -0700 Subject: [PATCH 04/14] Probe the dsrelu wrapper's signature instead of the frontend version _cudnn_frontend_supports_deterministic_dprob() gated on _cudnn_frontend_version_at_least("1.28.0"). That check is too coarse to answer the question it is asked, and would have raised at runtime on a build TE is actually run against. #521 merged after v1.27.0 was tagged, so `deterministic` ships in 1.28.0. But cudnn-frontend's develop branch has called itself 1.28.0 since shortly after that tag -- eleven days before the merge. Any front-end built from develop in that window reports 1.28.0 and does not accept the argument, so the version check passes, TE adds `deterministic=True` to the call, and the backward dies with TypeError: grouped_gemm_dsrelu_wrapper_sm100() got an unexpected keyword argument 'deterministic' This is not hypothetical, and not new. The same coarseness already bit use_single_group_runtime_offsets: a cuDNN reporting 1.27.0 that did not implement 1.27.0's arguments failed the identical way, in fuser_forward, before any backward code ran. Version numbers describe a release; they do not describe whatever happens to be installed. Ask the function instead. `"deterministic" in inspect.signature(...).parameters` is exact, cannot drift, and needs no maintenance when the release lands. The import is wrapped the way _grouped_gemm_dsrelu_backward_supported() already wraps it, so a missing cuDNN answers False rather than raising. Cached, since the call site runs every backward. This also removes the version constant from the code path entirely -- 1.28.0 now appears only in user-facing text, where a release number is the useful thing to say. Tests: a smoke test that the probe returns a bool without raising, with or without cuDNN installed, since reading a signature has more ways to fail than comparing two version strings. It deliberately does not assert which answer -- that depends on the installed front-end, and pinning it would only restate the implementation. Signed-off-by: Zhiyu Li --- 3rdparty/nccl-extensions | 2 +- tests/pytorch/test_grouped_mlp.py | 15 +++++++-- .../pytorch/ops/fused/grouped_mlp.py | 31 +++++++++++++------ 3 files changed, 36 insertions(+), 12 deletions(-) diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions index 705ca8eb38..2c6135a721 160000 --- a/3rdparty/nccl-extensions +++ b/3rdparty/nccl-extensions @@ -1 +1 @@ -Subproject commit 705ca8eb38297f3a8c4af6adf4185176a1116cb9 +Subproject commit 2c6135a721824ff792af7b72900b0ab758fa1f98 diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index aee16e12ea..65bbf61512 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2037,8 +2037,8 @@ def test_either_knob_requests_determinism( 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. + Needs no GPU. Without cuDNN installed the capability probe returns False and both + assertions still hold, so this runs anywhere. """ assert ( grouped_mlp_module.GroupedMLP_CuTeGEMMGLU.grouped_gemm_dactivation_is_deterministic() @@ -2049,6 +2049,17 @@ def test_only_the_srelu_path_can_be_deterministic(self) -> None: is grouped_mlp_module._cudnn_frontend_supports_deterministic_dprob() ) + def test_capability_probe_answers_without_cudnn(self) -> None: + """The probe reads a signature, so it has more ways to fail than a version compare. + + A missing cuDNN must give ``False``, not an ImportError, and a present one must not + raise out of ``inspect.signature``. Deliberately does not assert *which* answer: + that depends on the installed front-end, and pinning it would just restate the + implementation. + """ + grouped_mlp_module._cudnn_frontend_supports_deterministic_dprob.cache_clear() + assert isinstance(grouped_mlp_module._cudnn_frontend_supports_deterministic_dprob(), bool) + 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"): diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index ecd386c1dd..7bc9e1e92d 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -8,6 +8,7 @@ from collections.abc import Callable, Iterable import functools +import inspect import os import warnings from importlib.metadata import PackageNotFoundError, version as get_pkg_version @@ -107,17 +108,29 @@ def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool: return _cudnn_frontend_version_at_least("1.27.0") +@functools.lru_cache(maxsize=None) 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. + """Whether the installed dSReLU backward wrapper takes a ``deterministic`` argument. + + That argument (NVIDIA/cudnn-frontend#521) parks each N-subtile's partial result in its + own slot and sums the slots 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, so the GLU activations cannot + take this path at all. + + Asked of the signature rather than of the package version, because the version cannot + answer it. #521 merged after ``v1.27.0`` was tagged, so it ships in 1.28.0 -- but + ``develop`` already called itself ``1.28.0`` for the eleven days before that merge, and + a build from that window passes a version check and then raises ``TypeError: ... + unexpected keyword argument 'deterministic'``. The same coarseness bit + ``use_single_group_runtime_offsets`` on a cuDNN that reported 1.27.0 without + implementing it. """ - return _cudnn_frontend_version_at_least("1.28.0") + try: + from cudnn import grouped_gemm_dsrelu_wrapper_sm100 # pylint: disable=no-name-in-module + except ImportError: + return False + return "deterministic" in inspect.signature(grouped_gemm_dsrelu_wrapper_sm100).parameters def _deterministic_algorithms_required() -> bool: From a659cf5069dc846efc57f8ded99e5ae8e231291b Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 20 Aug 2026 10:34:38 -0700 Subject: [PATCH 05/14] Revert the unrelated nccl-extensions submodule bump `git add -u` in the previous commit swept in a local 3rdparty/nccl-extensions pointer change that has nothing to do with this PR. Restore it to main's commit so the branch touches only the three files it means to. Signed-off-by: Zhiyu Li --- 3rdparty/nccl-extensions | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions index 2c6135a721..705ca8eb38 160000 --- a/3rdparty/nccl-extensions +++ b/3rdparty/nccl-extensions @@ -1 +1 @@ -Subproject commit 2c6135a721824ff792af7b72900b0ab758fa1f98 +Subproject commit 705ca8eb38297f3a8c4af6adf4185176a1116cb9 From 83877a779dcf246c38d715bf05d28aff266f0238 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 20 Aug 2026 17:56:54 -0700 Subject: [PATCH 06/14] Raise instead of warning, and cut the change down to what it needs Review asked for two things on the unsupported path: make it an error rather than a warning, and stop branching on self._cudnn_dact_func to pick a message. Both are right, and taking them removes most of the machinery this PR had accumulated. Raising matches what TE already does elsewhere: the Triton grouped dbias/dscales kernel refuses to run under determinism rather than running non-deterministically. It also matches what the variable documents -- "only deterministic algorithms are allowed" is not "prefer deterministic algorithms". A silently non-reproducible run is the failure this PR exists to prevent, so continuing past a request TE cannot honor was the wrong default. Checked that no existing determinism test hits this path: test_hybrid_quantization sets the variable for an attention recipe, and test_fusible_ops_with_userbuffers for linear ops. One message, no branch. The two cases did have different remedies, which is why the branch was there, but a single sentence states both facts -- "needs the scaled-SReLU activation and nvidia-cudnn-frontend 1.28.0 or later" -- without telling a SwiGLU user to go upgrade. What that let me delete: * _warn_nondeterministic_cudnn_dprob and its per-reason lru_cache, the two reason strings and the branch selecting them: 30 lines at the call site and above it, down to a single raise. * _cudnn_frontend_supports_deterministic_dprob as a standalone function. The probe now lives in GroupedMLP_CuTeGEMMUnary.grouped_gemm_dactivation_is_deterministic(), which reaches the wrapper through grouped_gemm_dactivation_kernel() -- the import and its ImportError handling already existed there, so folding it in dropped a duplicate import and an indirection. * The warn-once cache-clearing fixture in the tests, and the two tests that existed only to cover the warning. Tests: test_deterministic_dactivation_is_numerically_correct becomes test_determinism_either_runs_or_refuses -- it expects RuntimeError where the request cannot be honored and runs the full numerical check where it can, so both arms assert something either way. The bit-exactness and two-knob tests are unchanged in substance. Net: transformer_engine/pytorch/ops/fused/grouped_mlp.py goes from +106 to +68, all of it addition, no line of pre-existing code touched. Signed-off-by: Zhiyu Li --- docs/envvars.rst | 12 +- tests/pytorch/test_grouped_mlp.py | 145 +++++------------- .../pytorch/ops/fused/grouped_mlp.py | 94 ++++-------- 3 files changed, 74 insertions(+), 177 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 3851788cf6..2b1aafc834 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -211,14 +211,10 @@ backend-selection overview. .. note:: As for attention, the fused grouped MLP treats determinism as requested when *either* - this variable is ``0`` or :func:`torch.use_deterministic_algorithms` is enabled. - - .. 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 a determinism request cannot be - honored for either reason. + this variable is ``0`` or :func:`torch.use_deterministic_algorithms` is enabled. A + bit-exact ``dprob`` needs the scaled-SReLU activation and cuDNN frontend 1.28.0 or + later; the GLU activations have no deterministic counterpart. A request the fused + grouped MLP cannot honor raises ``RuntimeError`` rather than running anyway. .. envvar:: NVTE_FUSED_RING_ATTENTION_USE_SCAN diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 65bbf61512..c153c35229 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -8,7 +8,6 @@ import os import math import random -import warnings from typing import Optional import pytest @@ -1979,22 +1978,10 @@ class TestGroupedMLPDeterminism: 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. + replaces those atomics with per-subtile slots reduced in a canonical order. The dGLU + wrapper has no equivalent, so a determinism request it cannot honor is an error. """ - @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.fixture def _restore_torch_determinism(self): """``use_deterministic_algorithms`` is process-global, so put it back.""" @@ -2037,68 +2024,34 @@ def test_either_knob_requests_determinism( def test_only_the_srelu_path_can_be_deterministic(self) -> None: """The capability is a property of the wrapper, not of the environment. - Needs no GPU. Without cuDNN installed the capability probe returns False and both - assertions still hold, so this runs anywhere. + Needs no GPU, and no cuDNN either: the SReLU probe answers False when the import + fails, so both assertions hold anywhere. """ - 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_capability_probe_answers_without_cudnn(self) -> None: - """The probe reads a signature, so it has more ways to fail than a version compare. - - A missing cuDNN must give ``False``, not an ImportError, and a present one must not - raise out of ``inspect.signature``. Deliberately does not assert *which* answer: - that depends on the installed front-end, and pinning it would just restate the - implementation. - """ - grouped_mlp_module._cudnn_frontend_supports_deterministic_dprob.cache_clear() - assert isinstance(grouped_mlp_module._cudnn_frontend_supports_deterministic_dprob(), bool) - - 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") + glu = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU + unary = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + assert glu.grouped_gemm_dactivation_is_deterministic() is False + assert isinstance(unary.grouped_gemm_dactivation_is_deterministic(), bool) @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. + def test_determinism_either_runs_or_refuses(self, monkeypatch, *, activation: str) -> None: + """A request TE cannot honor must fail loudly, not train on silently. - ``group_size > 1`` here, so the unit-activation-scale shortcut is off and the - cuDNN epilogue really does produce ``dprob``. + ``group_size > 1`` here, so the unit-activation-scale shortcut is off and the cuDNN + epilogue really does produce ``dprob``. When the request *can* be honored, the + wrapped test checks numerics as usual. """ 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 = ( + fused_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") + def _run() -> None: TestGroupedMLPFusedOp().test_grouped_mlp( bias=False, hidden_size=128, @@ -2106,73 +2059,59 @@ def test_deterministic_dactivation_is_numerically_correct( 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 + + if fused_cls.grouped_gemm_dactivation_is_deterministic(): + _run() + else: + with pytest.raises(RuntimeError, match="dprob"): + _run() @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: """The actual claim: identical inputs produce an identical ``dprob``. - The tolerance-based check above cannot see this. Reordering the same atomic adds - moves the result by about an ulp, which every tolerance in this file accepts, so a - run that is silently not reproducible passes it. Only an exact comparison of two - runs can tell the difference. - - ``dbias`` rides the same slot mechanism, but reaching it needs an FC1 bias, and an - FC2 ``scale_bias`` then routes the scale gradient through a Triton kernel that - refuses to run under determinism at all. Left to the op-level tests. + The end-to-end test above cannot see this. Reordering the same atomic adds moves + the result by about an ulp, which every tolerance in this file accepts, so a run + that is silently not reproducible passes it. Only comparing two runs exactly can + tell the difference. """ fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary if not fused_cls.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") if not fused_cls.grouped_gemm_dactivation_is_deterministic(): - pytest.skip( - "grouped_gemm_dsrelu_wrapper_sm100 has no deterministic mode before cuDNN" - " frontend 1.28.0, so dprob is expected to vary run to run here" - ) + pytest.skip("dSReLU determinism needs cuDNN frontend 1.28.0 or later") monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") - self._reset_caches() device = torch.device("cuda") dtype = torch.bfloat16 group_size = 4 - # Wide enough that the dprob reduction spans several N-tiles. With a single tile - # there is one writer per token, nothing to reorder, and the test is vacuous. + # Wide enough that the dprob reduction spans several 256-wide N-tiles. With a + # single tile there is one writer per token, nothing to reorder, and the test is + # vacuous. hidden_size = 1024 split_sizes = torch.tensor([256] * group_size, dtype=torch.int, device=device) num_tokens = int(split_sizes.sum().item()) recipe = make_recipe("mxfp8") - _, x = make_reference_and_test_tensors( - (num_tokens, hidden_size), - min=-0.25, - max=0.25, - quantization="mxfp8", - test_dtype=dtype, - test_device=device, - ) + tensor_kwargs = { + "min": -0.25, + "max": 0.25, + "quantization": "mxfp8", + "test_dtype": dtype, + "test_device": device, + } + _, x = make_reference_and_test_tensors((num_tokens, hidden_size), **tensor_kwargs) _, dy = make_reference_and_test_tensors( - (num_tokens, hidden_size), - min=-0.25, - max=0.25, - quantization="mxfp8", - test_dtype=dtype, - test_device=device, - requires_grad=False, + (num_tokens, hidden_size), requires_grad=False, **tensor_kwargs ) _, probs = make_reference_and_test_tensors( - (num_tokens,), - test_dtype=dtype, - test_device=device, + (num_tokens,), test_dtype=dtype, test_device=device ) - # No bias on either linear: with bias the FC2 scale gradient goes through the - # Triton grouped-dbias kernel instead of arriving straight from the cuDNN epilogue, - # and probs.grad would no longer be the dprob under test. + # No bias: with an FC2 scale_bias the scale gradient is finished by the Triton + # grouped-dbias kernel instead of arriving straight from the cuDNN epilogue, and + # probs.grad would no longer be the dprob under test. with te.quantized_model_init(enabled=True, recipe=recipe): module = te.ops.Sequential( te.ops.GroupedLinear( diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 7bc9e1e92d..fb777421f1 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -10,7 +10,6 @@ import functools import inspect import os -import warnings from importlib.metadata import PackageNotFoundError, version as get_pkg_version from typing import Any, Optional @@ -108,31 +107,6 @@ def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool: return _cudnn_frontend_version_at_least("1.27.0") -@functools.lru_cache(maxsize=None) -def _cudnn_frontend_supports_deterministic_dprob() -> bool: - """Whether the installed dSReLU backward wrapper takes a ``deterministic`` argument. - - That argument (NVIDIA/cudnn-frontend#521) parks each N-subtile's partial result in its - own slot and sums the slots 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, so the GLU activations cannot - take this path at all. - - Asked of the signature rather than of the package version, because the version cannot - answer it. #521 merged after ``v1.27.0`` was tagged, so it ships in 1.28.0 -- but - ``develop`` already called itself ``1.28.0`` for the eleven days before that merge, and - a build from that window passes a version check and then raises ``TypeError: ... - unexpected keyword argument 'deterministic'``. The same coarseness bit - ``use_single_group_runtime_offsets`` on a cuDNN that reported 1.27.0 without - implementing it. - """ - try: - from cudnn import grouped_gemm_dsrelu_wrapper_sm100 # pylint: disable=no-name-in-module - except ImportError: - return False - return "deterministic" in inspect.signature(grouped_gemm_dsrelu_wrapper_sm100).parameters - - def _deterministic_algorithms_required() -> bool: """Whether the user has asked for bit-exact reproducibility, by either route. @@ -151,22 +125,6 @@ def _deterministic_algorithms_required() -> bool: ) -@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( - "Deterministic execution was requested (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" - " torch.use_deterministic_algorithms), 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" - f" tile scheduler. That op is therefore not bit-exact here: {reason}.", - UserWarning, - ) - - def _wrap_single_quantized_as_grouped( tensor: torch.Tensor, quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage, @@ -2023,27 +1981,21 @@ def fuser_backward( 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 + # Checked 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) + raise RuntimeError( + "Deterministic execution was requested" + " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" + " torch.use_deterministic_algorithms), but the cuDNN grouped-GEMM" + " dactivation backward accumulates the scale gradient (dprob) with" + " cross-CTA atomic adds whose summation order is set by the tile" + " scheduler. A bit-exact dprob needs the scaled-SReLU activation and" + " nvidia-cudnn-frontend 1.28.0 or later." + ) fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: @@ -2092,12 +2044,9 @@ def fuser_backward( } dactivation_kernel = self.grouped_gemm_dactivation_kernel() if deterministic_dactivation: - # Only ever set to True, and only once known to be accepted: the argument does - # not exist on the dGLU wrapper or on a front-end older than 1.28.0, where - # passing it at all -- even as False -- is a TypeError. Leaving it unset is also - # the same request: the wrapper's own default reads - # ``torch.are_deterministic_algorithms_enabled()``, which is half of what - # _deterministic_algorithms_required() just read. + # Only reachable when the installed wrapper takes the argument: the check above + # raises otherwise, and passing it to a wrapper that does not accept it -- the + # dGLU one, or a front-end older than 1.28.0 -- is a TypeError. fc2_dactivation_kwargs["deterministic"] = True if _cudnn_frontend_supports_single_group_runtime_offsets(): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 @@ -2627,9 +2576,22 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: return grouped_gemm_dsrelu_wrapper_sm100 @classmethod + @functools.lru_cache(maxsize=None) 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() + """Whether the installed dSReLU wrapper takes ``deterministic`` (cuDNN FE 1.28.0+). + + Asked of the signature rather than of the package version, because the version + cannot answer it: NVIDIA/cudnn-frontend#521 merged after ``v1.27.0`` was tagged, but + ``develop`` had already called itself ``1.28.0`` for the eleven days before that, so + a build from that window passes a version check and then raises ``TypeError: ... + unexpected keyword argument 'deterministic'``. The same coarseness bit + ``use_single_group_runtime_offsets`` on a cuDNN reporting 1.27.0 without it. + """ + try: + kernel = cls.grouped_gemm_dactivation_kernel() + except ImportError: + return False + return "deterministic" in inspect.signature(kernel).parameters def fuse_ops( From 6ba169e622392d8533943fd9e80a4c0d540e2501 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Thu, 20 Aug 2026 21:47:31 -0700 Subject: [PATCH 07/14] Apply suggestion from @vthumbe1503 Signed-off-by: vthumbe1503 --- docs/envvars.rst | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 2b1aafc834..97eaed5ddc 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -206,15 +206,7 @@ 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. 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:: - - As for attention, the fused grouped MLP treats determinism as requested when *either* - this variable is ``0`` or :func:`torch.use_deterministic_algorithms` is enabled. A - bit-exact ``dprob`` needs the scaled-SReLU activation and cuDNN frontend 1.28.0 or - later; the GLU activations have no deterministic counterpart. A request the fused - grouped MLP cannot honor raises ``RuntimeError`` rather than running anyway. + :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. .. envvar:: NVTE_FUSED_RING_ATTENTION_USE_SCAN From fe80b4d7b66a827209516a1321f746870de95fb3 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Thu, 20 Aug 2026 21:53:55 -0700 Subject: [PATCH 08/14] Match the feature-detection idiom main just landed The SiTU-GLU merge (#3402) brought _cudnn_frontend_supports_grouped_gemm_situglu() into this file, which asks inspect.signature(wrapper).parameters for the arguments it needs rather than comparing frontend versions -- the same conclusion this branch reached independently, now the house style. Two things to match. Guard the signature call with `except (TypeError, ValueError)`: a callable that is not introspectable answers "no" instead of raising out of a backward pass. I had left this out on the grounds that the wrapper is a plain undecorated function, which is true today but is not a property this code controls. And say "feature-detect" in the docstring summary, as the neighbor does. Also dropped the sentence about use_single_group_runtime_offsets from the docstring. The neighbor now demonstrates the pattern in the same file, so the cautionary tale is no longer what makes the choice legible. `import inspect` came in with the merge, so this branch no longer adds it. Signed-off-by: Zhiyu Li --- .../pytorch/ops/fused/grouped_mlp.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index a3714f4e04..4bac4f4cff 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -2659,20 +2659,22 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: @classmethod @functools.lru_cache(maxsize=None) def grouped_gemm_dactivation_is_deterministic(cls) -> bool: - """Whether the installed dSReLU wrapper takes ``deterministic`` (cuDNN FE 1.28.0+). - - Asked of the signature rather than of the package version, because the version - cannot answer it: NVIDIA/cudnn-frontend#521 merged after ``v1.27.0`` was tagged, but - ``develop`` had already called itself ``1.28.0`` for the eleven days before that, so - a build from that window passes a version check and then raises ``TypeError: ... - unexpected keyword argument 'deterministic'``. The same coarseness bit - ``use_single_group_runtime_offsets`` on a cuDNN reporting 1.27.0 without it. + """Feature-detect the dSReLU wrapper's ``deterministic`` argument (cuDNN FE 1.28.0+). + + Detected rather than version-checked, because the version cannot answer it: + NVIDIA/cudnn-frontend#521 merged after ``v1.27.0`` was tagged, but ``develop`` had + already called itself ``1.28.0`` for the eleven days before that, so a build from + that window passes a version check and then raises ``TypeError: ... unexpected + keyword argument 'deterministic'``. """ try: kernel = cls.grouped_gemm_dactivation_kernel() except ImportError: return False - return "deterministic" in inspect.signature(kernel).parameters + try: + return "deterministic" in inspect.signature(kernel).parameters + except (TypeError, ValueError): + return False def fuse_ops( From 14d2bb969647b0ff58ac19fcbbd37eafb6555fd9 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 01:05:31 -0700 Subject: [PATCH 09/14] Cut the comments down to the file's own register The new code carried multi-paragraph docstrings into a file whose 44 functions have a median docstring of one line. Measured before and after: grouped_mlp.py _deterministic_algorithms_required 10 -> 3 lines grouped_gemm_dactivation_is_deterministic (base) 5 -> 1 grouped_gemm_dactivation_is_deterministic (unary) 7 -> 1 test_grouped_mlp.py four new tests 4-6 -> 1-4 four inline comment blocks 2-3 -> 1 each Before this, the three new functions were the 2nd, 3rd and 5th longest docstrings in grouped_mlp.py; only fuse_grouped_mlp_ops, which has a full Parameters block, was longer. In the test file, 63 pre-existing tests have a median docstring of zero lines. Most of what came out was rationale, not explanation: why the union matches DotProductAttention, why feature detection beats a version compare, which cuDNN release window motivated it. That belongs in the commits that made those choices, where it already is, and it reads as noise next to _cudnn_frontend_supports_grouped_gemm_situglu -- the neighbor doing the very same feature detection in a one-line docstring with no rationale at all. What stayed is what the code cannot say itself: that the check sits inside the non-unit-scale branch because a unit scale produces no dprob; that hidden_size must exceed one N-tile or the bit-exactness test is vacuous; that bias would reroute probs.grad through Triton; that weight grads are excluded because wgrad has its own atomics. Each is now one line. No behavior change -- comments, docstrings and one local variable's reading order only. Signed-off-by: Zhiyu Li --- tests/pytorch/test_grouped_mlp.py | 48 +++++-------------- .../pytorch/ops/fused/grouped_mlp.py | 35 +++----------- 2 files changed, 18 insertions(+), 65 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index f035547e6b..7054908e85 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2775,11 +2775,8 @@ def train_step( class TestGroupedMLPDeterminism: """Determinism 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, so a determinism request it cannot honor is an error. + Only the dSReLU wrapper can make ``dprob`` bit-exact, and only from cuDNN FE 1.28.0 on. + Anything else must refuse a determinism request rather than run non-deterministically. """ @pytest.fixture @@ -2809,11 +2806,7 @@ def test_either_knob_requests_determinism( torch_flag: bool, expected: bool, ) -> None: - """The same union DotProductAttention uses, and uncached so both stay live. - - ``NVTE_ALLOW_NONDETERMINISTIC_ALGO=1`` is not a request for non-determinism, only - the absence of one, so the torch flag still wins that row. - """ + """``=1`` is the absence of a request, not a request for non-determinism.""" if allow_nondeterministic is None: monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) else: @@ -2822,11 +2815,7 @@ def test_either_knob_requests_determinism( 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 no GPU, and no cuDNN either: the SReLU probe answers False when the import - fails, so both assertions hold anywhere. - """ + """The capability belongs to the wrapper, not the environment. Needs no GPU.""" glu = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU unary = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary assert glu.grouped_gemm_dactivation_is_deterministic() is False @@ -2835,12 +2824,7 @@ def test_only_the_srelu_path_can_be_deterministic(self) -> None: @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_srelu")) def test_determinism_either_runs_or_refuses(self, monkeypatch, *, activation: str) -> None: - """A request TE cannot honor must fail loudly, not train on silently. - - ``group_size > 1`` here, so the unit-activation-scale shortcut is off and the cuDNN - epilogue really does produce ``dprob``. When the request *can* be honored, the - wrapped test checks numerics as usual. - """ + """A request TE cannot honor must fail loudly; one it can must still be correct.""" if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") @@ -2868,12 +2852,10 @@ def _run() -> None: @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: - """The actual claim: identical inputs produce an identical ``dprob``. + """Two identical runs must give a bit-identical ``dprob``. - The end-to-end test above cannot see this. Reordering the same atomic adds moves - the result by about an ulp, which every tolerance in this file accepts, so a run - that is silently not reproducible passes it. Only comparing two runs exactly can - tell the difference. + An ulp of reordering passes every tolerance in this file, so only an exact + comparison of two runs can see it. """ fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary if not fused_cls.is_supported(): @@ -2886,9 +2868,7 @@ def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: device = torch.device("cuda") dtype = torch.bfloat16 group_size = 4 - # Wide enough that the dprob reduction spans several 256-wide N-tiles. With a - # single tile there is one writer per token, nothing to reorder, and the test is - # vacuous. + # >256 so dprob spans several N-tiles; one tile means one writer and no reordering. hidden_size = 1024 split_sizes = torch.tensor([256] * group_size, dtype=torch.int, device=device) num_tokens = int(split_sizes.sum().item()) @@ -2909,9 +2889,7 @@ def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: (num_tokens,), test_dtype=dtype, test_device=device ) - # No bias: with an FC2 scale_bias the scale gradient is finished by the Triton - # grouped-dbias kernel instead of arriving straight from the cuDNN epilogue, and - # probs.grad would no longer be the dprob under test. + # No bias, or probs.grad comes from the Triton dbias kernel instead of cuDNN. with te.quantized_model_init(enabled=True, recipe=recipe): module = te.ops.Sequential( te.ops.GroupedLinear( @@ -2932,15 +2910,13 @@ def _run() -> torch.Tensor: return probs.grad.detach().clone() first = _run() - # The fusion has to have happened, or dprob never came from the cuDNN epilogue and - # the comparison below proves nothing. + # Without the fusion there is no cuDNN dprob and the comparison proves nothing. forward_ops = module._module_groups[0]._forward_ops assert len(forward_ops) == 1 assert isinstance(forward_ops[0][0], fused_cls) second = _run() - # Exact, not assert_close. Weight gradients are deliberately not compared: the CuTe - # DSL wgrad kernel has its own K-split atomics that this change does not address. + # Weight grads are excluded: the CuTe DSL wgrad kernel has its own K-split atomics. assert torch.equal(first, second), ( "dprob differed between two identical runs under determinism; max |delta| =" f" {(first.float() - second.float()).abs().max().item()}" diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 4bac4f4cff..168a92a626 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -138,16 +138,9 @@ def _cudnn_frontend_supports_single_group_runtime_offsets( def _deterministic_algorithms_required() -> bool: - """Whether the user has asked for bit-exact reproducibility, by either route. + """Whether bit-exact reproducibility was asked for. Same union as ``DotProductAttention``. - Same union as ``DotProductAttention``: the two knobs answer different questions. - ``NVTE_ALLOW_NONDETERMINISTIC_ALGO`` is set once in a job launcher, applies uniformly - across ranks and is the only one TE's C++ layer can read; ``use_deterministic_algorithms`` - is the framework standard, is togglable at runtime, and is what a user who wants - reproducibility usually reaches for. Honoring only the first would leave the second - silently unenforced. - - Deliberately uncached: both inputs can change during the process. + Uncached: both knobs can change during the process. """ return ( not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) @@ -937,12 +930,7 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: @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. - """ + """Whether this op's dactivation kernel can produce a bit-exact ``dprob``.""" return False @classmethod @@ -2061,9 +2049,7 @@ def fuser_backward( scales_f32 = scales.detach().to(dtype=torch.float32) scales_tensor = scales_f32.reshape(-1, 1, 1) dscales_tensor = torch.zeros_like(scales_tensor) - # Checked 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. + # Only inside this branch: a unit activation scale produces no dprob at all. if _deterministic_algorithms_required(): deterministic_dactivation = self.grouped_gemm_dactivation_is_deterministic() if not deterministic_dactivation: @@ -2124,9 +2110,7 @@ def fuser_backward( } dactivation_kernel = self.grouped_gemm_dactivation_kernel() if deterministic_dactivation: - # Only reachable when the installed wrapper takes the argument: the check above - # raises otherwise, and passing it to a wrapper that does not accept it -- the - # dGLU one, or a front-end older than 1.28.0 -- is a TypeError. + # Never passed to a wrapper that would reject it -- the check above raises first. fc2_dactivation_kwargs["deterministic"] = True if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)): fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 @@ -2659,14 +2643,7 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: @classmethod @functools.lru_cache(maxsize=None) def grouped_gemm_dactivation_is_deterministic(cls) -> bool: - """Feature-detect the dSReLU wrapper's ``deterministic`` argument (cuDNN FE 1.28.0+). - - Detected rather than version-checked, because the version cannot answer it: - NVIDIA/cudnn-frontend#521 merged after ``v1.27.0`` was tagged, but ``develop`` had - already called itself ``1.28.0`` for the eleven days before that, so a build from - that window passes a version check and then raises ``TypeError: ... unexpected - keyword argument 'deterministic'``. - """ + """Feature-detect the dSReLU wrapper's ``deterministic`` argument (cuDNN FE 1.28.0+).""" try: kernel = cls.grouped_gemm_dactivation_kernel() except ImportError: From 2e45ac4ed3bb5077da8f68f372ce5ffb606226d7 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 01:08:34 -0700 Subject: [PATCH 10/14] Flatten the determinism check and the runs-or-refuses test Structural cleanups from the review pass. grouped_mlp.py: the check was nested two deep inside `if not unit_activation_scale`, and assigned deterministic_dactivation only to immediately test its own assignment. Hoisted to two flat statements right after unit_activation_scale is computed. `not unit_activation_scale and _deterministic_algorithms_required()` now says in the expression what the comment had to say in prose, and the separate `= False` initializer is gone. The local itself stays -- the kwargs dict is built about sixty lines further down. Also shortened the error: the tile-scheduler detail was not actionable, and "this activation's cuDNN dactivation kernel" is more accurate than naming the grouped-GEMM backward, since which kernel it is depends on the activation. test_grouped_mlp.py: fused_cls was derived from `activation` by a five-line conditional inside the test; it is now the second half of the parametrize pair. That also fixes the skip guard, which asked GroupedMLP_CuTeGEMMGLU.is_supported() on both parametrizations including the SReLU one -- the sibling test three functions down already gets this right. The _run closure existed only so an if/else could call it twice; a contextlib.nullcontext / pytest.raises choice removes the closure and the branch. nullcontext is used in ten test files here, so it is the local idiom rather than a new one. Not taken: dropping the `isinstance(..., bool)` assertion. It looks vacuous but it is the only coverage of the ImportError branch in the capability probe, which is the branch that runs on every machine without cuDNN -- including CI. Signed-off-by: Zhiyu Li --- tests/pytorch/test_grouped_mlp.py | 32 ++++++++++--------- .../pytorch/ops/fused/grouped_mlp.py | 27 ++++++++-------- 2 files changed, 30 insertions(+), 29 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 7054908e85..c546db14ab 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -5,6 +5,7 @@ from __future__ import annotations from collections.abc import Iterable +import contextlib import functools import os import math @@ -2822,20 +2823,27 @@ def test_only_the_srelu_path_can_be_deterministic(self) -> None: assert isinstance(unary.grouped_gemm_dactivation_is_deterministic(), bool) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) - @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_srelu")) - def test_determinism_either_runs_or_refuses(self, monkeypatch, *, activation: str) -> None: + @pytest.mark.parametrize( + "activation,fused_cls", + ( + ("scaled_srelu", grouped_mlp_module.GroupedMLP_CuTeGEMMUnary), + ("scaled_swiglu", grouped_mlp_module.GroupedMLP_CuTeGEMMGLU), + ), + ) + def test_determinism_either_runs_or_refuses( + self, monkeypatch, *, activation, fused_cls + ) -> None: """A request TE cannot honor must fail loudly; one it can must still be correct.""" - if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported(): + if not fused_cls.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") - fused_cls = ( - grouped_mlp_module.GroupedMLP_CuTeGEMMUnary - if activation == "scaled_srelu" - else grouped_mlp_module.GroupedMLP_CuTeGEMMGLU + expectation = ( + contextlib.nullcontext() + if fused_cls.grouped_gemm_dactivation_is_deterministic() + else pytest.raises(RuntimeError, match="dprob") ) - - def _run() -> None: + with expectation: TestGroupedMLPFusedOp().test_grouped_mlp( bias=False, hidden_size=128, @@ -2844,12 +2852,6 @@ def _run() -> None: activation=activation, ) - if fused_cls.grouped_gemm_dactivation_is_deterministic(): - _run() - else: - with pytest.raises(RuntimeError, match="dprob"): - _run() - @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: """Two identical runs must give a bit-identical ``dprob``. diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 168a92a626..77d16129e5 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -2041,27 +2041,26 @@ def fuser_backward( current_stream = torch.cuda.current_stream().cuda_stream unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) + # A unit activation scale produces no dprob, so there is nothing to make deterministic. + deterministic_dactivation = ( + not unit_activation_scale and _deterministic_algorithms_required() + ) + if deterministic_dactivation and not self.grouped_gemm_dactivation_is_deterministic(): + raise RuntimeError( + "Deterministic execution was requested" + " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" + " torch.use_deterministic_algorithms), but this activation's cuDNN dactivation" + " kernel accumulates the scale gradient (dprob) with nondeterministic atomics." + " A bit-exact dprob requires the scaled-SReLU activation and" + " nvidia-cudnn-frontend 1.28.0 or later." + ) 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) - # Only inside this branch: a unit activation scale produces no dprob at all. - if _deterministic_algorithms_required(): - deterministic_dactivation = self.grouped_gemm_dactivation_is_deterministic() - if not deterministic_dactivation: - raise RuntimeError( - "Deterministic execution was requested" - " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" - " torch.use_deterministic_algorithms), but the cuDNN grouped-GEMM" - " dactivation backward accumulates the scale gradient (dprob) with" - " cross-CTA atomic adds whose summation order is set by the tile" - " scheduler. A bit-exact dprob needs the scaled-SReLU activation and" - " nvidia-cudnn-frontend 1.28.0 or later." - ) fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn if use_nvfp4: From 34d8de57da0b0c2f2d502c0608facc83c8574584 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 01:10:50 -0700 Subject: [PATCH 11/14] Close the second dprob producer, and stop discarding fp64 test tensors Two findings from the review pass. dprob has two producers in this backward, and the check only covered one. The cuDNN epilogue produces grad_scales at fuser_backward, and when scale_bias is set compute_grouped_dbias_dscales accumulates into it further down -- the Triton kernel that grouped_dbias_dscales.py documents as nondeterministic atomic adds. That kernel's own guard reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. So the hole opened exactly where this branch widened the trigger. With torch.use_deterministic_algorithms(True) and the variable unset -- the case the union exists to start honoring -- SReLU on a 1.28.0 front-end with scale_bias passed the new check, set deterministic=True, raised nothing, and then routed dprob through the nondeterministic path anyway. Env-var users were never exposed: the Triton guard fires for them. It was reachable only via the torch flag, which is to say only through what this branch added. The test picked bias=False and so never crossed it. scale_bias is computed ~130 lines earlier in the same scope, so the fix is to require both producers rather than one. Still one condition and one message, per review -- the message now lists all three requirements instead of two. Separately, the bit-exactness test built its tensors with make_reference_and_test_tensors and discarded the reference every time. That helper allocates an fp64 CPU companion, quantizes and dequantizes for MXFP8 representability, then copies back D2H with an implicit sync -- about 16 MB of host allocation across the two (1024, 1024) calls, for a test that compares run 1 against run 2 and never against a reference. Twelve of the file's other fifteen uses keep the reference; this one had no use for it. Plain uniform_ tensors instead. Also dropped a .item() sync for a token count already known in Python. Not taken, with reasons: * Hoisting _deterministic_algorithms_required into pytorch/utils.py so the Triton guard reads the same union. That is the deeper fix and it is correct, but broadening that guard changes behavior for callers this PR does not touch (ops/basic/grouped_linear.py, module/grouped_linear.py) -- users who set only the torch flag would start seeing RuntimeError where they now get silent nondeterminism. Worth doing deliberately, not as a side effect of this branch. * Extracting the signature-probe shared with _cudnn_frontend_supports_grouped_gemm_situglu. The overlap is about four lines and the two are not interchangeable; refactoring working code outside the diff to save them is not this PR's job. Signed-off-by: Zhiyu Li --- tests/pytorch/test_grouped_mlp.py | 30 +++++++++---------- .../pytorch/ops/fused/grouped_mlp.py | 15 ++++++---- 2 files changed, 24 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index c546db14ab..e50c799293 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2872,24 +2872,22 @@ def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: group_size = 4 # >256 so dprob spans several N-tiles; one tile means one writer and no reordering. hidden_size = 1024 - split_sizes = torch.tensor([256] * group_size, dtype=torch.int, device=device) - num_tokens = int(split_sizes.sum().item()) + tokens_per_group = 256 + split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device) + num_tokens = tokens_per_group * group_size recipe = make_recipe("mxfp8") - tensor_kwargs = { - "min": -0.25, - "max": 0.25, - "quantization": "mxfp8", - "test_dtype": dtype, - "test_device": device, - } - _, x = make_reference_and_test_tensors((num_tokens, hidden_size), **tensor_kwargs) - _, dy = make_reference_and_test_tensors( - (num_tokens, hidden_size), requires_grad=False, **tensor_kwargs - ) - _, probs = make_reference_and_test_tensors( - (num_tokens,), test_dtype=dtype, test_device=device - ) + + # Plain random tensors, not make_reference_and_test_tensors: this test compares two + # runs against each other, never against a reference, so the fp64 companion and the + # MXFP8 representability round-trip would both be allocated and thrown away. + def _rand(*shape, requires_grad=True) -> torch.Tensor: + out = torch.empty(shape, dtype=dtype, device=device).uniform_(-0.25, 0.25) + return out.requires_grad_() if requires_grad else out + + x = _rand(num_tokens, hidden_size) + dy = _rand(num_tokens, hidden_size, requires_grad=False) + probs = _rand(num_tokens) # No bias, or probs.grad comes from the Triton dbias kernel instead of cuDNN. with te.quantized_model_init(enabled=True, recipe=recipe): diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index 77d16129e5..cc6f7a6bb1 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -2045,14 +2045,19 @@ def fuser_backward( deterministic_dactivation = ( not unit_activation_scale and _deterministic_algorithms_required() ) - if deterministic_dactivation and not self.grouped_gemm_dactivation_is_deterministic(): + # dprob has two producers here: the cuDNN epilogue below, and -- when scale_bias is + # set -- the Triton kernel that accumulates into it further down. Both must be + # deterministic, and the Triton one never is. + if deterministic_dactivation and not ( + self.grouped_gemm_dactivation_is_deterministic() and not scale_bias + ): raise RuntimeError( "Deterministic execution was requested" " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" - " torch.use_deterministic_algorithms), but this activation's cuDNN dactivation" - " kernel accumulates the scale gradient (dprob) with nondeterministic atomics." - " A bit-exact dprob requires the scaled-SReLU activation and" - " nvidia-cudnn-frontend 1.28.0 or later." + " torch.use_deterministic_algorithms), but the scale gradient (dprob) is" + " accumulated with nondeterministic atomics on this configuration." + " A bit-exact dprob requires the scaled-SReLU activation," + " nvidia-cudnn-frontend 1.28.0 or later, and an FC2 without scale_bias." ) scales_f32 = None scales_tensor = None From e14b08e21aa5a3e2cd71a7a8de64631a3aa0f6ce Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 14:03:52 -0700 Subject: [PATCH 12/14] Add the regression test for the scale_bias hole The previous commit fixed a real bug and shipped it with no test. Every test in the class used bias=False and every end-to-end one set the env var, so neither half of the bug was reachable: not scale_bias, and not the torch-flag-only trigger. Both halves are load-bearing. With the env var the Triton kernel raises on its own, so an env-var test would have passed before the fix as well as after and pinned nothing. Only torch.use_deterministic_algorithms with the variable unset reaches the state where this op's check said yes and the Triton reduction then ran nondeterministically. warn_only=True so torch's own enforcement cannot raise first and be mistaken for TE's refusal. are_deterministic_algorithms_enabled() still reports True in that mode -- the separate is_deterministic_algorithms_warn_only_enabled() getter exists precisely because the two are independent -- so the predicate under test sees what it should. Not executed: no GPU or torch on the machine this was written on. Formatting and syntax only, like the rest of the branch. Signed-off-by: Zhiyu Li --- tests/pytorch/test_grouped_mlp.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index e50c799293..f601ac06c7 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2852,6 +2852,31 @@ def test_determinism_either_runs_or_refuses( activation=activation, ) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_scale_bias_refuses_under_the_torch_flag( + self, monkeypatch, _restore_torch_determinism + ) -> None: + """``scale_bias`` finishes ``dprob`` in a Triton kernel that reads only the env var. + + So the torch flag alone is the combination that used to pass this op's own check and + then reduce nondeterministically anyway, on a front-end new enough to say yes. + """ + fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) + # warn_only so torch's own enforcement cannot raise first and mask what TE does. + torch.use_deterministic_algorithms(True, warn_only=True) + with pytest.raises(RuntimeError, match="dprob"): + TestGroupedMLPFusedOp().test_grouped_mlp( + bias=True, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation="scaled_srelu", + ) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: """Two identical runs must give a bit-identical ``dprob``. From aff3b05f6012e896cb7a254f62e033c2f1e28ed2 Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 14:22:41 -0700 Subject: [PATCH 13/14] Make the bit-exactness test capable of failing Followed cudnn-frontend#521's own test work and found this test had the flaw its commit 88c7fab was written to fix, at the same config. That commit measured 16 launches per shape and found that at l=4 / [256]*4 / n=512 the NONDETERMINISTIC dprob is already bit-stable: the assertion cannot fail there, so a pass certifies nothing. It varies 15/15 at l=8 / [1024]*8 / n=2048. This test used l=4 / [256]*4 / n=1024 -- the vacuous shape, one power of two along n. Moved to the shape that actually varies. n > 256 was necessary but not sufficient, which is what the old comment got wrong. Spanning several N-tiles exercises the within-CTA subtile ordering; making the cross-CTA reduction unstable needs the larger token count and expert count too. Also took the rest of #521's discipline for these comparisons: * Repeat rather than compare a pair. The order determinism removes is set by the tile scheduler, so two runs can match by luck. Four by default, NVTE_TEST_DETERMINISM_REPEATS to raise it, matching that file's DETERMINISM_REPEATS. * Compare bytes, not values. torch.equal treats +0.0 and -0.0 as equal, and a change in reduction order produces exactly that; upstream's bitwise_bits views as uint8 for the same reason. * Assert the output is finite first, so a NaN run cannot be read as a determinism result. Not copied: asserting that the nondeterministic path *does* vary. It is the thing that makes the config meaningful, but as an assertion it is timing-dependent and would flake. Upstream settled this by measuring once and pinning the config; the comment now cites that measurement so the next person does not shrink the shape back. Not run yet: job 535935 is building the previous revision of this test. Signed-off-by: Zhiyu Li --- tests/pytorch/test_grouped_mlp.py | 38 ++++++++++++++++++++----------- 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index f601ac06c7..835ec752be 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2879,10 +2879,10 @@ def test_scale_bias_refuses_under_the_torch_flag( @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: - """Two identical runs must give a bit-identical ``dprob``. + """Repeated identical runs must give a bit-identical ``dprob``. An ulp of reordering passes every tolerance in this file, so only an exact - comparison of two runs can see it. + comparison across runs can see it. """ fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary if not fused_cls.is_supported(): @@ -2894,10 +2894,13 @@ def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: device = torch.device("cuda") dtype = torch.bfloat16 - group_size = 4 - # >256 so dprob spans several N-tiles; one tile means one writer and no reordering. - hidden_size = 1024 - tokens_per_group = 256 + # l=8 / [1024]*8 / n=2048, the shape cudnn-frontend#521 measured as varying 15/15 + # without the fix. Its own tests originally used l=4 / [256]*4 / n=512, where the + # nondeterministic dprob is already bit-stable -- there the assertion below cannot + # fail, and a pass means nothing. + group_size = 8 + hidden_size = 2048 + tokens_per_group = 1024 split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device) num_tokens = tokens_per_group * group_size @@ -2934,18 +2937,27 @@ def _run() -> torch.Tensor: y.backward(dy) return probs.grad.detach().clone() - first = _run() + runs = [_run()] # Without the fusion there is no cuDNN dprob and the comparison proves nothing. forward_ops = module._module_groups[0]._forward_ops assert len(forward_ops) == 1 assert isinstance(forward_ops[0][0], fused_cls) - second = _run() + # More than two, as cudnn-frontend#521 does: the cross-CTA order that determinism + # removes is set by the scheduler, so two runs can agree by luck. + runs += [_run() for _ in range(int(os.getenv("NVTE_TEST_DETERMINISM_REPEATS", "4")) - 1)] + torch.cuda.synchronize() - # Weight grads are excluded: the CuTe DSL wgrad kernel has its own K-split atomics. - assert torch.equal(first, second), ( - "dprob differed between two identical runs under determinism; max |delta| =" - f" {(first.float() - second.float()).abs().max().item()}" - ) + assert torch.isfinite(runs[0]).all(), "dprob is not finite; the comparison would be moot" + # Bytes, not values: torch.equal calls +0.0 and -0.0 equal, and a change in reduction + # order can produce exactly that. Weight grads are excluded from the comparison -- + # the CuTe DSL wgrad kernel has its own K-split atomics, which this change leaves. + for index, later in enumerate(runs[1:], start=1): + assert torch.equal( + runs[0].contiguous().view(torch.uint8), later.contiguous().view(torch.uint8) + ), ( + f"dprob differs between run 0 and run {index} under determinism; max |delta| =" + f" {(runs[0].float() - later.float()).abs().max().item()}" + ) def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: From bedd363bc62747df0a0712ec4a3a80a15e80dc3b Mon Sep 17 00:00:00 2001 From: Zhiyu Li Date: Fri, 21 Aug 2026 23:12:09 -0700 Subject: [PATCH 14/14] Pick the bit-exactness config by measuring it, not by borrowing one Measured on GB300 across five shapes, determinism off, 8 launches each (job 538058), counting how many runs differ from run 0: l=8 tok/grp=1024 n=2048 2/7 max|d| 5.96e-08 l=8 tok/grp=1024 n=4096 2/7 max|d| 9.54e-07 l=16 tok/grp=1024 n=2048 7/7 max|d| 1.19e-07 l=8 tok/grp=2048 n=2048 5/7 max|d| 7.63e-06 l=4 tok/grp=512 n=8192 6/7 max|d| 1.91e-06 Moved to l=16, the only shape where every run differs, so the assertion cannot pass by luck. The previous choice, l=8, varies 2/7 -- an eight-run sample calls it stable often enough to be a poor detector, and an earlier control run (537313) did exactly that and reported 0/7 at this shape. I took that single sample as proof the config was vacuous and said so; it was a sampling artifact, and the shape does vary, just weakly. That earlier shape came from cudnn-frontend#521's own measurement, which was taken on its direct wrapper test. It does not transfer to TE's path -- different scheduler settings, different quantization -- so borrowing the number was the mistake underneath both errors. This config is measured through the fused grouped MLP itself. Two things the same job settled that are worth recording: * Without #521 the values genuinely move: 6e-08 to 8e-06 absolute across these shapes. Small, but nonzero every time, and the reason the refusal exists rather than a warning. * The refusal cannot be exercised against the stock 1.27.0 frontend on this image at all. TE's forward passes prob_tensor=None because _cudnn_frontend_version_at_least("1.27.0") reports optional-prob support that a stock 1.27.0 does not implement, so the op dies in fuser_forward with "prob_tensor is required" before any determinism code runs. Same version-gate-too-coarse failure this PR avoids for its own argument, on a gate it does not own. Signed-off-by: Zhiyu Li --- tests/pytorch/test_grouped_mlp.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 835ec752be..4623bcb98a 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -2894,11 +2894,12 @@ def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: device = torch.device("cuda") dtype = torch.bfloat16 - # l=8 / [1024]*8 / n=2048, the shape cudnn-frontend#521 measured as varying 15/15 - # without the fix. Its own tests originally used l=4 / [256]*4 / n=512, where the - # nondeterministic dprob is already bit-stable -- there the assertion below cannot - # fail, and a pass means nothing. - group_size = 8 + # Measured on GB300, determinism off, 8 launches per shape (job 538058): this shape + # gives 7/7 runs differing from run 0, so the assertion below can actually fail. + # Shapes matter more than they look -- l=8 with the same n and tokens/group varies + # only 2/7, which an 8-run sample reports as stable often enough to be useless, and + # cudnn-frontend#521 measured its own l=4 / [256]*4 / n=512 as never varying. + group_size = 16 hidden_size = 2048 tokens_per_group = 1024 split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device)