adding SBSA to the supported platform of TRT-LLM libs and installing MPI libs for the distributed tests

Author: apbose

.github/scripts/install-mpi-linux-x86.sh

@@ -12,12 +12,6 @@ if [[ $(uname -m) == "aarch64" ]]; then
     install_cuda_aarch64
 fi
 
-if [[ "$(uname -s)" == "Linux" && "$(uname -m)" == "x86_64" ]]; then
-    # install MPI for Linux x86_64
-    source .github/scripts/install-mpi-linux-x86.sh
-    install_mpi_linux_x86
-fi
-
 # Install all the dependencies required for Torch-TensorRT
 pip install --pre -r ${PWD}/tests/py/requirements.txt
 # dependencies in the tests/py/requirements.txt might install a different version of torch or torchvision

.github/scripts/install-torch-tensorrt.sh

@@ -363,6 +363,7 @@ jobs:
         export USE_HOST_DEPS=1
         export CI_BUILD=1
         export USE_TRTLLM_PLUGINS=1
+        dnf install -y mpich mpich-devel openmpi openmpi-devel
         pushd .
         cd tests/py
         cd dynamo

.github/workflows/build-test-linux-x86_64.yml

@@ -878,18 +878,21 @@ def is_platform_supported_for_trtllm(platform: str) -> bool:
     Note:
         TensorRT-LLM plugins for NCCL backend are not supported on:
         - Windows platforms
-        - Jetson devices (aarch64 architecture)
+        - Orin, Xavier, or Tegra devices (aarch64 architecture)
+
     """
     if "windows" in platform:
         logger.info(
             "TensorRT-LLM plugins for NCCL backend are not supported on Windows"
         )
         return False
-    if "aarch64" in platform:
+    if torch.cuda.is_available():
+        device_name = torch.cuda.get_device_name().lower()
+        if any(keyword in device_name for keyword in ["orin", "xavier", "tegra"]):
+            return False
         logger.info(
-            "TensorRT-LLM plugins for NCCL backend are not supported on Jetson devices (aarch64)"
+            "TensorRT-LLM plugins for NCCL backend are not supported on Jetson devices"
         )
-        return False
     return True
 
 

py/torch_tensorrt/dynamo/utils.py

@@ -9,6 +9,7 @@
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt._enums import Platform
+from torch_tensorrt.dynamo.utils import is_platform_supported_for_trtllm
 
 
 class DistributedGatherModel(nn.Module):
@@ -44,10 +45,10 @@ def forward(self, x):
 platform_str = str(Platform.current_platform()).lower()
 
 
-class TestGatherNcclOpsConverter(DispatchTestCase):
+class TestNcclOpsConverter(DispatchTestCase):
     @unittest.skipIf(
-        "win" in platform_str or "aarch64" in platform_str,
-        "Skipped on Windows and Jetson: NCCL backend is not supported.",
+        not is_platform_supported_for_trtllm(platform_str),
+        "Skipped on Windows, Jetson: NCCL backend is not supported.",
     )
     @classmethod
     def setUpClass(cls):