[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

@@ -1236,6 +1236,7 @@ def tearDownClass(cls) -> None:
     def _run(self, model: nn.Module, inputs: list[torch.Tensor]) -> None:
         """Compile with torch_tensorrt and verify output matches PyTorch."""
         import torch_tensorrt
+        from torch_tensorrt.dynamo.utils import ATOL, RTOL
 
         model = model.cuda().eval()
         inputs_cuda = [t.cuda() for t in inputs]
@@ -1254,7 +1255,12 @@ def _run(self, model: nn.Module, inputs: list[torch.Tensor]) -> None:
             )
             out = trt_model(*inputs_cuda)
 
-        torch.testing.assert_close(ref, out, atol=1e-4, rtol=1e-4)
+        # Project-standard tolerance (torch_tensorrt.dynamo.utils), the same one
+        # every converter test uses via tests/py/dynamo/conversion/harness.py.
+        # TensorRT selects tensor-core kernels for batched matmuls, giving ~2^-11
+        # (~4.9e-4) rounding versus PyTorch's FP32 reference -- well inside 5e-3
+        # but over the 1e-4 this file previously hard-coded.
+        torch.testing.assert_close(ref, out, atol=ATOL, rtol=RTOL)
 
     def test_all_reduce_single_rank(self) -> None:
         """all_reduce compiles and produces correct output on a single rank."""
@@ -1315,6 +1321,7 @@ def _run_dynamic(
     ) -> None:
         """Mark a dim dynamic (min/max), compile at the opt shape, verify at other shapes."""
         import torch_tensorrt  # noqa: F401
+        from torch_tensorrt.dynamo.utils import ATOL, RTOL
 
         model = model.cuda().eval()
         opt_cuda = [t.cuda() for t in opt_inputs]
@@ -1338,7 +1345,10 @@ def _run_dynamic(
                 check_cuda = [t.cuda() for t in check]
                 ref = model(*check_cuda)
                 out = trt_model(*check_cuda)
-                torch.testing.assert_close(ref, out, atol=1e-4, rtol=1e-4)
+                # See _run: project-standard tolerance. These checks run at
+                # batch > 1, where TensorRT picks a tensor-core kernel, so the
+                # error is ~4.9e-4 rather than the ~1e-6 seen at batch 1.
+                torch.testing.assert_close(ref, out, atol=ATOL, rtol=RTOL)
 
     def test_all_reduce_single_rank_dynamic(self) -> None:
         """all_reduce compiles with a dynamic seq dim and is correct at other shapes."""