rework how nvlink is checked

Signed-off-by: Ludwig Schneider <[email protected]>

Author: nv-lschneider

cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp

@@ -712,15 +712,19 @@ void AllreducePlugin::setGroupTopology() noexcept
     if (isMultiNode)
     {
         TLLM_LOG_INFO("Found inter-node TP group for rank %d", rank);
-        // For MNNVL (Multi-Node NVLink), we need to check P2P/NVLINK even in multi-node
-        // For other strategies, multi-node means no P2P/NVLINK support
-        if (mStrategy != AllReduceStrategyType::MNNVL)
+        // Strategies that don't support multi-node P2P/NVLINK
+        // MIN_LATENCY, ONESHOT, TWOSHOT, LOWPRECISION are designed for intra-node only
+        if (mStrategy == AllReduceStrategyType::MIN_LATENCY || mStrategy == AllReduceStrategyType::ONESHOT
+            || mStrategy == AllReduceStrategyType::TWOSHOT || mStrategy == AllReduceStrategyType::LOWPRECISION)
         {
             mIsP2PSupported = false;
             mIsNVLINKSupported = false;
+            TLLM_LOG_INFO("Strategy %s does not support multi-node, setting P2P/NVLINK to false for rank %d",
+                tensorrt_llm::kernels::toString(mStrategy).c_str(), rank);
             return;
         }
-        TLLM_LOG_INFO("MNNVL strategy detected - checking multi-node P2P/NVLINK for rank %d", rank);
+        // Other strategies (like MNNVL) will continue to check P2P/NVLINK in multi-node
+        TLLM_LOG_INFO("Multi-node strategy detected - checking P2P/NVLINK across nodes for rank %d", rank);
     }
     else
     {
@@ -732,10 +736,9 @@ void AllreducePlugin::setGroupTopology() noexcept
     mIsP2PSupported = true;
     mIsNVLINKSupported = true;
 
-    // For MNNVL in multi-node, check all devices in the full group
+    // For multi-node with supported strategies, check all devices in the full group
     // For intra-node, check only localGroup
-    std::set<int> const& devicesToCheck
-        = (isMultiNode && mStrategy == AllReduceStrategyType::MNNVL) ? mGroup : localGroup;
+    std::set<int> const& devicesToCheck = isMultiNode ? mGroup : localGroup;
 
     // Use cudaDeviceCanAccessPeer to determine whether p2p is supported,
     // and use nvml to determine whether there are nvlink links between ranks.

cpp/tensorrt_llm/thop/allreduceOp.cpp

@@ -1009,15 +1009,19 @@ class AllreduceOp
         if (isMultiNode)
         {
             TLLM_LOG_INFO("Found inter-node TP group for rank %d", rank);
-            // For MNNVL (Multi-Node NVLink), we need to check P2P/NVLINK even in multi-node
-            // For other strategies, multi-node means no P2P/NVLINK support
-            if (mStrategy != AllReduceStrategyType::MNNVL)
+            // Strategies that don't support multi-node P2P/NVLINK
+            // MIN_LATENCY, ONESHOT, TWOSHOT, LOWPRECISION are designed for intra-node only
+            if (mStrategy == AllReduceStrategyType::MIN_LATENCY || mStrategy == AllReduceStrategyType::ONESHOT
+                || mStrategy == AllReduceStrategyType::TWOSHOT || mStrategy == AllReduceStrategyType::LOWPRECISION)
             {
                 mIsP2PSupported = false;
                 mIsNVLINKSupported = false;
+                TLLM_LOG_INFO("Strategy %s does not support multi-node, setting P2P/NVLINK to false for rank %d",
+                    tensorrt_llm::kernels::toString(mStrategy).c_str(), rank);
                 return;
             }
-            TLLM_LOG_INFO("MNNVL strategy detected - checking multi-node P2P/NVLINK for rank %d", rank);
+            // Other strategies (like MNNVL) will continue to check P2P/NVLINK in multi-node
+            TLLM_LOG_INFO("Multi-node strategy detected - checking P2P/NVLINK across nodes for rank %d", rank);
         }
         else
         {
@@ -1028,10 +1032,9 @@ class AllreduceOp
         mIsP2PSupported = true;
         mIsNVLINKSupported = true;
 
-        // For MNNVL in multi-node, check all devices in the full group
+        // For multi-node with supported strategies, check all devices in the full group
         // For intra-node, check only local_group
-        std::set<int> const& devices_to_check
-            = (isMultiNode && mStrategy == AllReduceStrategyType::MNNVL) ? mGroup : local_group;
+        std::set<int> const& devices_to_check = isMultiNode ? mGroup : local_group;
 
         // TODO(ytong): Should we provide group topology info instead of querying it here?
         // Use cudaDeviceCanAccessPeer to determine whether p2p is supported,