[DRAFT] feat: enable native multi-device TensorRT on TensorRT-RTX builds

.github/workflows/_test-linux.yml

@@ -142,6 +142,8 @@ jobs:
       build-matrix: ${{ needs.filter-matrix.outputs.matrix }}
       pre-script: packaging/pre_build_script.sh
       use-rtx: ${{ inputs.use-rtx }}
+      # Per-suite runner from the manifest; "" falls back to matrix.validation_runner.
+      runner: ${{ matrix.runner }}
       fail-on-empty: true
       script: |
         set -euo pipefail

core/runtime/TRTEngine.h

@@ -22,14 +22,29 @@
 #include "core/runtime/TensorRTBindingNames.h"
 #include "core/util/prelude.h"
 
-// TensorRT 10.16+ has native NCCL collective support via IExecutionContext::setCommunicator()
-#if NV_TENSORRT_MAJOR > 10 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 16)
+// Native NCCL collective support is exposed via IExecutionContext::setCommunicator()
+// together with IDistCollectiveLayer. Two independent release lines ship that API:
+//
+//   * TensorRT-RTX 1.5+ -- NvInferVersion.h defines TRT_MAJOR_RTX/TRT_MINOR_RTX and
+//     then aliases NV_TENSORRT_MAJOR/MINOR to them, so NV_TENSORRT_MAJOR is 1 on RTX
+//     and the mainline ">= 10.16" comparison below can never match. Detect the RTX
+//     package first and version-check against its own numbering.
+//   * TensorRT 10.16+   -- mainline.
+//
+// Do not collapse these into a single NV_TENSORRT_MAJOR/MINOR test: the two lines use
+// incompatible numbering schemes. See is_tensorrt_version_supported() in
+// py/torch_tensorrt/_utils.py for the Python-side equivalent of the same problem.
+#if defined(TRT_MAJOR_RTX)
+#if TRT_MAJOR_RTX > 1 || (TRT_MAJOR_RTX == 1 && TRT_MINOR_RTX >= 5)
+#define TRT_HAS_NATIVE_NCCL 1
+#endif
+#elif NV_TENSORRT_MAJOR > 10 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 16)
 #define TRT_HAS_NATIVE_NCCL 1
 #endif
 
 // Full TRT NCCL collectives support requires both:
 // 1. PyTorch built with NCCL (USE_C10D_NCCL defined via Bazel)
-// 2. TensorRT 10.16+ (TRT_HAS_NATIVE_NCCL defined above)
+// 2. A TensorRT exposing the native collectives API (TRT_HAS_NATIVE_NCCL above)
 #if defined(USE_C10D_NCCL) && defined(TRT_HAS_NATIVE_NCCL)
 #define ENABLE_TRT_NCCL_COLLECTIVES 1
 #endif

tests/ci/runner.py

@@ -318,6 +318,10 @@ def matrix(**filters: str | None) -> list[dict[str, str]]:
             "variant": var,
             "tier": s.tier,
             "cwd": s.for_variant(var)["cwd"],
+            # "" means "no override" -- linux-test.yml falls back to
+            # matrix.validation_runner. Set on suites that need specific
+            # hardware (e.g. multi-GPU for distributed).
+            "runner": s.for_variant(var)["runner"] or "",
         }
         for s, var in select(**filters)
     ]

tests/ci/suites.py

@@ -81,6 +81,7 @@ class Suite:
     setup: tuple[str, ...] = ()  # named pre-steps: hub|executorch|cuda-core|mpi
     follow: tuple[tuple[str, ...], ...] = ()  # extra argv to run AFTER pytest
     env: dict[str, str] = field(default_factory=dict)
+    runner: str | None = None  # GHA runner label; None = matrix.validation_runner
     overrides: dict[str, dict[str, Any]] = field(default_factory=dict)  # per-variant
 
     def for_variant(self, variant: Variant) -> dict[str, Any]:
@@ -101,6 +102,7 @@ def for_variant(self, variant: Variant) -> dict[str, Any]:
                 "setup",
                 "follow",
                 "env",
+                "runner",
             )
         }
         base.update(self.overrides.get(variant, {}))
@@ -310,10 +312,30 @@ def for_variant(self, variant: Variant) -> dict[str, Any]:
         jobs="auto",
         verbose=True,
         reruns=False,
-        variants=("standard",),
+        variants=("standard", "rtx"),
         platforms=("linux-x86_64",),
         setup=("mpi",),
         env={"USE_HOST_DEPS": "1", "CI_BUILD": "1", "USE_TRTLLM_PLUGINS": "1"},
+        # The --multirank follow-ups need 2 GPUs, so this suite cannot run on
+        # the default single-GPU validation_runner.
+        runner="linux.g4dn.12xlarge.nvidia.gpu",
+        # TensorRT-RTX has no TensorRT-LLM plugin path, so multi-device runs
+        # entirely on the native TRT DistCollective API. Drop test_nccl_ops.py
+        # (every test in it is gated on ENABLED_FEATURES.trtllm_for_nccl and
+        # would no-op) and USE_TRTLLM_PLUGINS along with it.
+        overrides={
+            "rtx": {
+                "paths": (
+                    "distributed/test_native_nccl.py",
+                    "distributed/test_export_save_load.py",
+                ),
+                "env": {"USE_HOST_DEPS": "1", "CI_BUILD": "1"},
+                # Multi-GPU box: the --multirank follow-ups need 2 devices.
+                # g5 is A10G (SM 8.6) rather than g4dn's T4 (SM 7.5), since the
+                # TensorRT docs describe DistCollective as needing Ampere+.
+                "runner": "linux.g5.12xlarge.nvidia.gpu",
+            }
+        },
         follow=(
             (
                 "-m",

tests/py/dynamo/distributed/test_native_nccl.py

@@ -92,6 +92,16 @@ def has_nccl_collectives() -> bool:
         return False
 
 
+def is_tensorrt_rtx() -> bool:
+    """Check if this is a TensorRT-RTX build."""
+    try:
+        from torch_tensorrt._features import ENABLED_FEATURES
+
+        return bool(ENABLED_FEATURES.tensorrt_rtx)
+    except Exception:
+        return False
+
+
 def _cpp_runtime_available() -> bool:
     """Return True when the C++ Torch-TensorRT runtime extension is loaded.
 
@@ -2188,6 +2198,125 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     print(f"[Rank {rank}] PASS _multirank_pg_migration", flush=True)
 
 
+def _find_py_engine(mod: nn.Module) -> Any:
+    """Return the Python-runtime ``TRTEngine`` inside a compiled module, or None."""
+    from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule
+    from torch_tensorrt.dynamo.runtime._TRTEngine import TRTEngine
+
+    if isinstance(mod, TorchTensorRTModule) and isinstance(mod.engine, TRTEngine):
+        return mod.engine
+    for child in mod.children() if isinstance(mod, nn.Module) else []:
+        found = _find_py_engine(child)
+        if found is not None:
+            return found
+    return None
+
+
+def _multirank_rtx_cudagraph_keeps_nccl_comm(
+    rank: int, world_size: int, device: torch.device
+) -> None:
+    """TRT-RTX: enabling cudagraphs must not cost an MD engine its communicator.
+
+    On TensorRT-RTX, manual ``torch.cuda.CUDAGraph`` capture is unsafe (lazy
+    kernel specialization), so ``_TRTEngine._execute_standard`` switches the
+    engine to RTX-native CUDA graphs on the first execute when cudagraphs are
+    enabled.  That switch runs ``update_runtime_settings()`` ->
+    ``invalidate_context()``.
+
+    The NCCL communicator is bound to the *IExecutionContext* by
+    ``setup_nccl_comm()`` (``set_communicator()``).  Invalidation drops that
+    context but leaves ``_nccl_comm`` set, and the re-bind guard in
+    ``execute()`` is ``if self._nccl_comm is None`` -- so the replacement
+    context never receives the communicator and cannot be repaired.
+
+    Unlike standard TensorRT, where a context invalidation requires the user to
+    deliberately change runtime settings, on RTX this fires with no user action
+    beyond turning cudagraphs on.
+
+    The ``_nccl_comm`` assertion is deliberate: it fails fast rather than
+    hanging inside a collective, which matters in CI.
+    """
+    import torch_tensorrt
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.tensor.parallel import (
+        ColwiseParallel,
+        RowwiseParallel,
+        parallelize_module,
+    )
+    from torch_tensorrt._features import ENABLED_FEATURES
+    from torch_tensorrt.distributed._distributed import distributed_context
+    from torch_tensorrt.distributed._nccl_utils import setup_nccl_for_torch_tensorrt
+
+    if not ENABLED_FEATURES.tensorrt_rtx:
+        print(
+            f"[Rank {rank}] SKIP _multirank_rtx_cudagraph_keeps_nccl_comm "
+            "(not a TensorRT-RTX build)",
+            flush=True,
+        )
+        return
+
+    setup_nccl_for_torch_tensorrt()
+    device_mesh = init_device_mesh("cuda", (world_size,))
+
+    class TinyMLP(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.fc1 = nn.Linear(16, 32, bias=False)
+            self.relu = nn.ReLU()
+            self.fc2 = nn.Linear(32, 16, bias=False)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.fc2(self.relu(self.fc1(x)))
+
+    torch.manual_seed(42)
+    model = TinyMLP().to(device)
+    parallelize_module(
+        model, device_mesh, {"fc1": ColwiseParallel(), "fc2": RowwiseParallel()}
+    )
+
+    torch.manual_seed(0)
+    inp = torch.randn(4, 16, device=device)
+
+    # Python runtime: set_communicator() lives in _TRTEngine, not the C++ engine.
+    with distributed_context(dist.group.WORLD):
+        trt_model = torch.compile(
+            model,
+            backend="torch_tensorrt",
+            dynamic=False,
+            options={
+                "min_block_size": 1,
+                "use_distributed_mode_trace": True,
+                "use_python_runtime": True,
+            },
+        )
+        with torch.no_grad():
+            expected = trt_model(inp)
+
+    engine = _find_py_engine(trt_model)
+    if engine is None:
+        raise AssertionError("Could not locate a Python TRTEngine in compiled model")
+    assert engine._has_nccl_ops, "Engine is not a multi-device engine"
+    assert engine._nccl_comm is not None, "Communicator was never bound"
+
+    # Turning cudagraphs on is enough: the first execute auto-switches to
+    # RTX-native graphs and invalidates the context.
+    with torch_tensorrt.runtime.enable_cudagraphs(trt_model) as cg_model:
+        with torch.no_grad():
+            out = cg_model(inp)
+
+        assert engine._nccl_comm is not None, (
+            "the RTX cudagraph switch invalidated the IExecutionContext and the "
+            "communicator was not re-bound to the replacement context"
+        )
+
+    _check_close(out, expected, f"output after RTX cudagraph switch rank={rank}")
+
+    print(
+        f"[Rank {rank}] PASS _multirank_rtx_cudagraph_keeps_nccl_comm",
+        flush=True,
+    )
+
+
 # ============================================================================
 # Section 8 — Multi-rank pytest tests (MultiProcessTestCase, requires 2 GPUs)
 # ============================================================================
@@ -2319,6 +2448,15 @@ def test_pg_migration(self) -> None:
         device = self._init_dist()
         _multirank_pg_migration(self.rank, self.world_size, device)
 
+    @unittest.skipIf(not is_tensorrt_rtx(), "TensorRT-RTX only")
+    @unittest.skipIf(not has_nccl_collectives(), "No NCCL collective support available")
+    @requires_nccl()
+    @skip_if_lt_x_gpu(2)
+    def test_rtx_cudagraph_keeps_nccl_comm(self) -> None:
+        """RTX cudagraph auto-switch invalidates the context; the comm must survive."""
+        device = self._init_dist()
+        _multirank_rtx_cudagraph_keeps_nccl_comm(self.rank, self.world_size, device)
+
 
 # ============================================================================
 # Section 9 — torchrun / mpirun entry point (legacy multi-rank runner)
@@ -2346,6 +2484,8 @@ def run_multirank_tests() -> None:
         _multirank_cpp_runtime_bind_nccl,
         _multirank_distributed_mode_context_switch,
         _multirank_pg_migration,
+        # Self-skips on non-RTX builds.
+        _multirank_rtx_cudagraph_keeps_nccl_comm,
     ]
 
     failed = []