Addressing review comments- include in enabled feature and error logging. Pending- check support on Thor and sbsa

Author: apbose

py/torch_tensorrt/_features.py

@@ -9,6 +9,7 @@
     check_cross_compile_trt_win_lib,
     sanitized_torch_version,
 )
+from torch_tensorrt.dynamo.utils import load_tensorrt_llm_for_nccl
 
 from packaging import version
 
@@ -23,6 +24,7 @@
         "qdp_plugin",
         "windows_cross_compile",
         "tensorrt_rtx",
+        "trtllm_for_nccl",
     ],
 )
 
@@ -48,6 +50,7 @@
 _FX_FE_AVAIL = False if _TENSORRT_RTX else True
 _REFIT_AVAIL = True
 _WINDOWS_CROSS_COMPILE = check_cross_compile_trt_win_lib()
+_TRTLLM_AVAIL = load_tensorrt_llm_for_nccl()
 
 if importlib.util.find_spec("tensorrt.plugin"):
     _QDP_PLUGIN_AVAIL = True
@@ -63,6 +66,7 @@
     _QDP_PLUGIN_AVAIL,
     _WINDOWS_CROSS_COMPILE,
     _TENSORRT_RTX,
+    _TRTLLM_AVAIL,
 )
 
 T = TypeVar("T")
@@ -158,6 +162,18 @@ def not_implemented(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
     return wrapper
 
 
+def needs_trtllm_for_nccl(f: Callable[..., Any]) -> Callable[..., Any]:
+    def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+        if ENABLED_FEATURES.trtllm_for_nccl:
+            return f(*args, **kwargs)
+        else:
+            raise NotImplementedError(
+                "TensorRT-LLM plugins for NCCL backend could not be loaded"
+            )
+
+    return wrapper
+
+
 def for_all_methods(
     decorator: Callable[..., Any], exclude: Optional[List[str]] = None
 ) -> Callable[..., Any]:

py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py

@@ -5,6 +5,7 @@
 
 import tensorrt as trt
 from torch.fx.node import Argument, Target
+from torch_tensorrt._features import needs_trtllm_for_nccl
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion import impl
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
@@ -15,45 +16,41 @@
     tensorrt_fused_nccl_all_gather_op,
     tensorrt_fused_nccl_reduce_scatter_op,
 )
-from torch_tensorrt.dynamo.utils import load_tensorrt_llm_for_nccl
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
-if load_tensorrt_llm_for_nccl():
 
-    @dynamo_tensorrt_converter(tensorrt_fused_nccl_all_gather_op)
-    def fused_nccl_gather(
-        ctx: ConversionContext,
-        target: Target,
-        args: Tuple[Argument, ...],
-        kwargs: Dict[str, Argument],
-        name: str,
-    ) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
-        return impl.nccl_ops.nccl_gather(
-            ctx,
-            target,
-            SourceIR.ATEN,
-            name,
-            [args[0]],
-        )
+@needs_trtllm_for_nccl
+@dynamo_tensorrt_converter(tensorrt_fused_nccl_all_gather_op)
+def fused_nccl_gather(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
+    return impl.nccl_ops.nccl_gather(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        [args[0]],
+    )
 
-    @dynamo_tensorrt_converter(tensorrt_fused_nccl_reduce_scatter_op)
-    def fused_nccl_reduce_scatter(
-        ctx: ConversionContext,
-        target: Target,
-        args: Tuple[Argument, ...],
-        kwargs: Dict[str, Argument],
-        name: str,
-    ) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
-        return impl.nccl_ops.nccl_reduce_scatter(
-            ctx,
-            target,
-            SourceIR.ATEN,
-            name,
-            [args[0]],
-        )
 
-else:
-    _LOGGER.debug(
-        "Did not load torch.distributed converters since TensorRT-LLM is not available"
+@needs_trtllm_for_nccl
+@dynamo_tensorrt_converter(tensorrt_fused_nccl_reduce_scatter_op)
+def fused_nccl_reduce_scatter(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
+    return impl.nccl_ops.nccl_reduce_scatter(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        [args[0]],
     )

py/torch_tensorrt/dynamo/utils.py

@@ -907,12 +907,16 @@ def is_platform_supported_for_trtllm() -> bool:
     try:
         cuda_version = torch.version.cuda  # e.g., "12.4" or "13.0"
         if cuda_version is None:
-            logger.warning("No CUDA runtime detected — TRT-LLM plugins unavailable.")
+            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.warning("CUDA 13 is not supported for TRT-LLM plugins.")
+            logger.error(
+                "CUDA 13 is not supported for TRT-LLM plugins. Please install pytorch with CUDA 12.x support"
+            )
             return False
 
         return True