Guard NCCL EP lazy init during graph capture

Fail fast if the lazy NCCL EP context or dispatch handle would be initialized while CUDA graph capture is active, and cover those guard paths in the communication factory tests.

Signed-off-by: Ludwig Schneider <lschneider@nvidia.com>

Author: nv-lschneider

tensorrt_llm/_torch/modules/fused_moe/communication/nccl_ep.py

@@ -112,6 +112,11 @@ def supports_post_quant_dispatch(self) -> bool:
 
     def _get_context(self):
         if self._ctx is None:
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError(
+                    "NcclEP context must be initialized before CUDA graph capture. "
+                    "Run an eager warmup forward before enabling or capturing CUDA graphs."
+                )
             from nccl.ep import Layout
 
             from tensorrt_llm._torch.modules.fused_moe.nccl_ep_utils import get_nccl_ep_context
@@ -130,6 +135,11 @@ def _get_context(self):
     def _setup_handle(self, ctx, topk_nd, stream):
         """Ensure self._handle exists; rebind topk via handle.update on subsequent calls."""
         if self._handle is None:
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError(
+                    "NcclEP dispatch handle must be initialized before CUDA graph capture. "
+                    "Run an eager warmup forward before enabling or capturing CUDA graphs."
+                )
             self._handle = ctx.ep_group.create_handle(
                 ctx.layout,
                 topk_nd,

tests/unittest/_torch/modules/moe/test_communication_factory.py

@@ -24,6 +24,7 @@
 from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import (
     AllGatherReduceScatter,
 )
+from tensorrt_llm._torch.modules.fused_moe.communication.nccl_ep import NcclEP
 
 
 def _make_model_config(
@@ -180,3 +181,35 @@ def fail_if_called(*args, **kwargs):
     )
 
     assert isinstance(strategy, AllGatherReduceScatter)
+
+
+def test_nccl_ep_context_init_rejects_cuda_graph_capture(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    strategy = object.__new__(NcclEP)
+    strategy._ctx = None
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+
+    with pytest.raises(RuntimeError, match="context must be initialized before CUDA graph capture"):
+        strategy._get_context()
+
+
+def test_nccl_ep_handle_init_rejects_cuda_graph_capture(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    strategy = object.__new__(NcclEP)
+    strategy._handle = None
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+
+    def fail_create_handle(*args, **kwargs):
+        raise AssertionError("create_handle should not run during CUDA graph capture")
+
+    ctx = SimpleNamespace(
+        ep_group=SimpleNamespace(create_handle=fail_create_handle),
+        layout=object(),
+    )
+
+    with pytest.raises(
+        RuntimeError, match="dispatch handle must be initialized before CUDA graph capture"
+    ):
+        strategy._setup_handle(ctx, object(), 0)