Address review comments

Signed-off-by: Iman Tabrizian <[email protected]>

Author: Tabrizian

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -318,12 +318,14 @@ def __init__(
         self._torch_compile_backend = None
 
         try:
-            if pytorch_backend_config.allreduce_strategy == "NCCL_SYMMETRIC":
-                self._init_userbuffers(self.model.config.hidden_size)
+            use_ub_for_nccl = (
+                pytorch_backend_config.allreduce_strategy == "NCCL_SYMMETRIC"
+                and self._init_userbuffers(self.model.config.hidden_size))
             if pytorch_backend_config.torch_compile_enabled:
                 set_torch_compiling(True)
-                use_ub = pytorch_backend_config.torch_compile_enable_userbuffers and self._init_userbuffers(
-                    self.model.config.hidden_size)
+                use_ub = use_ub_for_nccl or (
+                    pytorch_backend_config.torch_compile_enable_userbuffers
+                    and self._init_userbuffers(self.model.config.hidden_size))
                 self._torch_compile_backend = Backend(
                     pytorch_backend_config.torch_compile_inductor_enabled,
                     enable_userbuffers=use_ub,
@@ -2230,10 +2232,11 @@ def _init_userbuffers(self, hidden_size):
         # Disable UB for unsupported platforms
         if not ub.ub_supported():
             return False
+        use_nccl_symmetric = self.pytorch_backend_config.allreduce_strategy == "NCCL_SYMMETRIC"
         ub.initialize_userbuffers_manager(
             self.mapping.tp_size, self.mapping.pp_size, self.mapping.cp_size,
             self.mapping.rank, self.mapping.gpus_per_node,
-            hidden_size * self.max_num_tokens * 2, True)
+            hidden_size * self.max_num_tokens * 2, use_nccl_symmetric)
 
         return True