diff --git a/py/torch_tensorrt/_utils.py b/py/torch_tensorrt/_utils.py index 5d43db0e57..0b1bd65185 100644 --- a/py/torch_tensorrt/_utils.py +++ b/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,7 +156,33 @@ 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: @@ -162,11 +190,13 @@ def _cache_root() -> Path: 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}" ) @@ -174,27 +204,30 @@ 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_/trtllm/0.17.0.post1_linux_x86_64/tensorrt_llm/libs/libnvinfer_plugin_tensorrt_llm.so + # eg: /tmp/torch_tensorrt_/trtllm/_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 diff --git a/tests/py/core/test_trtllm_utils.py b/tests/py/core/test_trtllm_utils.py new file mode 100644 index 0000000000..32b512fb72 --- /dev/null +++ b/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()