Avoid redundant CUDA graph RNG registration

tests/pytorch/test_cuda_graphs.py

@@ -5,6 +5,7 @@
 from typing import Callable, Dict, Iterable, List, Tuple, Union
 import pytest
 import copy
+import contextlib
 import gc
 import weakref
 
@@ -28,6 +29,7 @@
 )
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 import transformer_engine.pytorch.ops as te_ops
+import transformer_engine.pytorch.graph as graph_module
 from transformer_engine.common import recipe
 from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override
 
@@ -53,6 +55,55 @@ def nvfp4_vanilla():
     return nvfp4_recipe
 
 
+@pytest.mark.parametrize(
+    "advances, expected",
+    (
+        (False, True),
+        (True, False),
+    ),
+)
+def test_graph_rng_registration_capability(monkeypatch, advances, expected):
+    """Detect whether graph replay advances RNG without explicit registration."""
+
+    class FakeGraph:
+        def __init__(self):
+            self.replays = 0
+
+        def replay(self):
+            self.replays += 1
+
+        def reset(self):
+            pass
+
+    class FakeOutput:
+        def random_(self):
+            pass
+
+        def clone(self):
+            return 0
+
+    class FakeGenerator:
+        def get_state(self):
+            return None
+
+        def set_state(self, state):
+            pass
+
+    monkeypatch.setattr(torch.cuda, "CUDAGraph", FakeGraph)
+    monkeypatch.setattr(torch.cuda, "default_generators", [FakeGenerator()])
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(torch.cuda, "graph", lambda graph: contextlib.nullcontext())
+    monkeypatch.setattr(torch.cuda, "synchronize", lambda: None)
+    monkeypatch.setattr(torch, "empty", lambda *args, **kwargs: FakeOutput())
+    monkeypatch.setattr(torch, "equal", lambda first, second: not advances)
+    monkeypatch.setattr(graph_module, "graph_safe_rng_available", lambda: True)
+    graph_module._graph_rng_registration_required.cache_clear()
+    try:
+        assert graph_module._graph_rng_registration_required() is expected
+    finally:
+        graph_module._graph_rng_registration_required.cache_clear()
+
+
 def nvfp4_rht_and_2d_quantization():
     nvfp4_recipe = recipe.NVFP4BlockScaling()
     nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams(

transformer_engine/pytorch/graph.py

@@ -5,6 +5,7 @@
 """Functions for CUDA Graphs support in FP8"""
 from collections.abc import Iterable
 import contextlib
+import functools
 import gc
 import warnings
 from math import ceil
@@ -44,6 +45,32 @@
 )
 
 
+@functools.lru_cache(maxsize=1)
+def _graph_rng_registration_required() -> bool:
+    """Return whether CUDA graph RNG state must be registered explicitly."""
+    if not graph_safe_rng_available():
+        return False
+
+    generator = torch.cuda.default_generators[torch.cuda.current_device()]
+    original_state = generator.get_state()
+    graph = torch.cuda.CUDAGraph()
+    output = torch.empty(8, dtype=torch.int64, device="cuda")
+    try:
+        with torch.cuda.graph(graph):
+            output.random_()
+        graph.replay()
+        torch.cuda.synchronize()
+        first = output.clone()
+        graph.replay()
+        torch.cuda.synchronize()
+        return torch.equal(first, output)
+    finally:
+        try:
+            graph.reset()
+        finally:
+            generator.set_state(original_state)
+
+
 class _GraphedCallableHelpers(NamedTuple):
     """Lifecycle helpers owned by one graphed callable invocation."""
 
@@ -482,7 +509,7 @@ def _returned_param_grad_clone_slots(static_grad_inputs, module_params):
         )
 
     # For cases with multiple active RNG states, e.g. TP.
-    if graph_safe_rng_available():
+    if _graph_rng_registration_required():
         for _, state in get_all_rng_states().items():
             for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs):
                 fwd_graph.register_generator_state(state)