fix: add CUDA 13 support for TRT-LLM plugins

py/torch_tensorrt/_utils.py

@@ -15,7 +15,11 @@
 logger = logging.getLogger(__name__)
 
 _WHL_CPYTHON_VERSION = "cp310"
-_TENSORRT_LLM_VERSION_ = "0.17.0.post1"
+# Auto-downloaded TensorRT-LLM wheels must match PyTorch's CUDA major version.
+_TENSORRT_LLM_VERSION_BY_CUDA_MAJOR = {
+    12: "0.17.0.post1",
+    13: "1.2.0",
+}
 
 
 def sanitized_torch_version() -> Any:
@@ -115,7 +119,7 @@ def is_platform_supported_for_trtllm() -> bool:
         - Windows platforms
         - Jetson/Orin/Xavier (aarch64 architecture + 'tegra' in platform release)
         - Thor devices
-        - CUDA 13 not supported
+        - PyTorch builds without CUDA support
     """
     system = platform.system().lower()
     machine = platform.machine().lower()
@@ -127,24 +131,22 @@ def is_platform_supported_for_trtllm() -> bool:
         )
         return False
 
-    if machine == "aarch64" and "tegra" in release or is_thor():
+    if machine == "aarch64" and "tegra" in release:
         logger.info(
-            "TensorRT-LLM plugins for NCCL backend are not supported on Jetson/Orin/Xavier (Tegra) or Thor devices."
+            "TensorRT-LLM plugins for NCCL backend are not supported on Jetson/Orin/Xavier (Tegra) devices."
         )
         return False
 
     try:
-        cuda_version = torch.version.cuda  # e.g., "12.4" or "13.0"
-        if cuda_version is None:
+        if torch.version.cuda is None:
             logger.error(
                 "This pytorch build does not support CUDA, please reinstall pytorch with CUDA support"
             )
             return False
 
-        major, minor = map(int, cuda_version.split("."))
-        if major != 12:
-            logger.error(
-                "CUDA 13 is not currently supported for TRT-LLM plugins. Please install pytorch with CUDA 12.x support"
+        if is_thor():
+            logger.info(
+                "TensorRT-LLM plugins for NCCL backend are not supported on Thor devices."
             )
             return False
 
@@ -154,47 +156,78 @@ def is_platform_supported_for_trtllm() -> bool:
         logger.info(f"Failed to detect CUDA version: {e}")
         return False
 
-    return True
+
+def _get_trtllm_version_for_cuda(cuda_version: Optional[str]) -> Optional[str]:
+    if cuda_version is None:
+        logger.error(
+            "This pytorch build does not support CUDA, please reinstall pytorch with CUDA support"
+        )
+        return None
+
+    try:
+        cuda_major = int(cuda_version.split(".", maxsplit=1)[0])
+    except (AttributeError, ValueError):
+        logger.error(f"Failed to parse CUDA version: {cuda_version}")
+        return None
+
+    trtllm_version = _TENSORRT_LLM_VERSION_BY_CUDA_MAJOR.get(cuda_major)
+    if trtllm_version is None:
+        supported_cuda_versions = ", ".join(
+            str(version) for version in _TENSORRT_LLM_VERSION_BY_CUDA_MAJOR
+        )
+        logger.error(
+            f"TensorRT-LLM plugin auto-download is not configured for CUDA {cuda_version}. "
+            f"Supported CUDA major versions are {supported_cuda_versions}. To use a compatible "
+            "plugin supplied by another source, set TRTLLM_PLUGINS_PATH."
+        )
+        return None
+
+    return trtllm_version
 
 
 def _cache_root() -> Path:
     username = getpass.getuser()
     return Path(tempfile.gettempdir()) / f"torch_tensorrt_{username}"
 
 
-def _extracted_dir_trtllm(platform_system: str, platform_machine: str) -> Path:
+def _extracted_dir_trtllm(
+    trtllm_version: str, platform_system: str, platform_machine: str
+) -> Path:
     return (
         _cache_root()
         / "trtllm"
-        / f"{_TENSORRT_LLM_VERSION_}_{platform_system}_{platform_machine}"
+        / f"{trtllm_version}_{platform_system}_{platform_machine}"
     )
 
 
 def download_and_get_plugin_lib_path() -> Optional[str]:
     """
     Returns the path to the TensorRT‑LLM shared library, downloading and extracting if necessary.
 
-    Args:
-        platform (str): Platform identifier (e.g., 'linux_x86_64')
-
     Returns:
         Optional[str]: Path to shared library or None if operation fails.
     """
+    trtllm_version = _get_trtllm_version_for_cuda(torch.version.cuda)
+    if trtllm_version is None:
+        return None
+
     platform_system = platform.system().lower()
     platform_machine = platform.machine().lower()
     wheel_filename = (
-        f"tensorrt_llm-{_TENSORRT_LLM_VERSION_}-{_WHL_CPYTHON_VERSION}-"
+        f"tensorrt_llm-{trtllm_version}-{_WHL_CPYTHON_VERSION}-"
         f"{_WHL_CPYTHON_VERSION}-{platform_system}_{platform_machine}.whl"
     )
     wheel_path = _cache_root() / wheel_filename
-    extract_dir = _extracted_dir_trtllm(platform_system, platform_machine)
+    extract_dir = _extracted_dir_trtllm(
+        trtllm_version, platform_system, platform_machine
+    )
     # else will never be met though
     lib_filename = (
         "libnvinfer_plugin_tensorrt_llm.so"
         if "linux" in platform_system
         else "libnvinfer_plugin_tensorrt_llm.dll"
     )
-    # eg: /tmp/torch_tensorrt_<username>/trtllm/0.17.0.post1_linux_x86_64/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so
+    # eg: /tmp/torch_tensorrt_<username>/trtllm/<version>_linux_x86_64/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so
     plugin_lib_path = extract_dir / "tensorrt_llm" / "libs" / lib_filename
 
     if plugin_lib_path.exists():
@@ -351,8 +384,10 @@ def check_native_trt_collectives(
 def load_tensorrt_llm_for_nccl() -> bool:
     """
     Attempts to load the TensorRT-LLM plugin and initialize it.
-    Either the env variable TRTLLM_PLUGINS_PATH can specify the path
-    Or the user can specify USE_TRTLLM_PLUGINS as either of (1, true, yes, on) to download the TRT-LLM distribution and load it
+
+    TRTLLM_PLUGINS_PATH can specify a user-provided plugin. Alternatively,
+    USE_TRTLLM_PLUGINS can be set to one of (1, true, yes, on) to download
+    the compatible TensorRT-LLM distribution and load its plugin.
 
     Returns:
         bool: True if the plugin was successfully loaded and initialized, False otherwise.
@@ -373,10 +408,13 @@ def load_tensorrt_llm_for_nccl() -> bool:
         )
         if not use_trtllm_plugin:
             logger.info(
-                "Neither TRTLLM_PLUGIN_PATH is set nor is it directed to download the shared library. Please set either of the two to use TRT-LLM libraries in torchTRT"
+                "Neither TRTLLM_PLUGINS_PATH nor USE_TRTLLM_PLUGINS is set. "
+                "Please set one of them to use TensorRT-LLM libraries in Torch-TensorRT."
             )
             return False
 
         plugin_lib_path = download_and_get_plugin_lib_path()
-        return load_and_initialize_trtllm_plugin(plugin_lib_path)  # type: ignore[arg-type]
+        if plugin_lib_path is None:
+            return False
+        return load_and_initialize_trtllm_plugin(plugin_lib_path)
     return False

tests/py/core/test_trtllm_utils.py

@@ -0,0 +1,96 @@
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+import torch
+from torch_tensorrt import _utils
+
+
+class TestTensorRTLLMPlatformSupport(unittest.TestCase):
+    @mock.patch.object(_utils.platform, "system", return_value="Linux")
+    @mock.patch.object(_utils.platform, "machine", return_value="x86_64")
+    @mock.patch.object(_utils.platform, "release", return_value="generic")
+    @mock.patch.object(_utils, "is_thor", return_value=False)
+    def test_cuda_13_is_supported(self, *unused_mocks):
+        with mock.patch.object(torch.version, "cuda", "13.0"):
+            self.assertTrue(_utils.is_platform_supported_for_trtllm())
+
+    @mock.patch.object(_utils.platform, "system", return_value="Linux")
+    @mock.patch.object(_utils.platform, "machine", return_value="x86_64")
+    @mock.patch.object(_utils.platform, "release", return_value="generic")
+    @mock.patch.object(_utils, "is_thor")
+    def test_cpu_only_pytorch_is_not_supported(self, mock_is_thor, *unused_mocks):
+        with mock.patch.object(torch.version, "cuda", None):
+            self.assertFalse(_utils.is_platform_supported_for_trtllm())
+        mock_is_thor.assert_not_called()
+
+
+class TestTensorRTLLMVersionSelection(unittest.TestCase):
+    def test_cuda_12_uses_existing_plugin_version(self):
+        self.assertEqual(_utils._get_trtllm_version_for_cuda("12.8"), "0.17.0.post1")
+
+    def test_cuda_13_uses_cuda_13_plugin_version(self):
+        self.assertEqual(_utils._get_trtllm_version_for_cuda("13.0"), "1.2.0")
+
+    @mock.patch.object(_utils.platform, "system", return_value="Linux")
+    @mock.patch.object(_utils.platform, "machine", return_value="x86_64")
+    def test_cuda_13_auto_download_uses_cuda_13_artifact(self, *unused_mocks):
+        with tempfile.TemporaryDirectory() as cache_root:
+            plugin_path = (
+                Path(cache_root)
+                / "trtllm"
+                / "1.2.0_linux_x86_64"
+                / "tensorrt_llm"
+                / "libs"
+                / "libnvinfer_plugin_tensorrt_llm.so"
+            )
+            plugin_path.parent.mkdir(parents=True)
+            plugin_path.touch()
+
+            with mock.patch.object(
+                _utils, "_cache_root", return_value=Path(cache_root)
+            ):
+                with mock.patch.object(torch.version, "cuda", "13.0"):
+                    self.assertEqual(
+                        _utils.download_and_get_plugin_lib_path(), str(plugin_path)
+                    )
+
+    def test_unknown_cuda_version_requires_user_supplied_plugin(self):
+        with self.assertLogs(_utils.logger, level="ERROR") as captured_logs:
+            self.assertIsNone(_utils._get_trtllm_version_for_cuda("14.0"))
+        self.assertIn("set TRTLLM_PLUGINS_PATH", captured_logs.output[0])
+
+
+class TestTensorRTLLMLoading(unittest.TestCase):
+    @mock.patch.object(_utils, "is_platform_supported_for_trtllm", return_value=True)
+    @mock.patch.object(_utils, "load_and_initialize_trtllm_plugin", return_value=True)
+    @mock.patch.object(_utils, "download_and_get_plugin_lib_path")
+    def test_explicit_plugin_path_does_not_use_auto_download(
+        self, mock_download, mock_load, *unused_mocks
+    ):
+        plugin_path = "/opt/trtllm/libnvinfer_plugin_tensorrt_llm.so"
+        with mock.patch.dict(
+            os.environ, {"TRTLLM_PLUGINS_PATH": plugin_path}, clear=True
+        ):
+            self.assertTrue(_utils.load_tensorrt_llm_for_nccl())
+
+        mock_load.assert_called_once_with(plugin_path)
+        mock_download.assert_not_called()
+
+    @mock.patch.object(_utils, "is_platform_supported_for_trtllm", return_value=True)
+    @mock.patch.object(_utils, "load_and_initialize_trtllm_plugin")
+    @mock.patch.object(_utils, "download_and_get_plugin_lib_path", return_value=None)
+    def test_auto_download_failure_does_not_try_to_load_none(
+        self, mock_download, mock_load, *unused_mocks
+    ):
+        with mock.patch.dict(os.environ, {"USE_TRTLLM_PLUGINS": "1"}, clear=True):
+            self.assertFalse(_utils.load_tensorrt_llm_for_nccl())
+
+        mock_download.assert_called_once_with()
+        mock_load.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()