add flag for EPLB to force using GDRCopy

Signed-off-by: Dongxu Yang <[email protected]>

Author: dongxuy04

cpp/tensorrt_llm/common/envUtils.cpp

@@ -450,4 +450,9 @@ bool getEnvDisableChunkedAttentionInGenPhase()
     return getBoolEnv("TRTLLM_DISABLE_CHUNKED_ATTENTION_IN_GEN_PHASE");
 }
 
+bool getEnvEplbForceGdrcopy()
+{
+    return getBoolEnv("TRTLLM_EPLB_FORCE_GDRCOPY");
+}
+
 } // namespace tensorrt_llm::common

cpp/tensorrt_llm/common/envUtils.h

@@ -136,4 +136,6 @@ bool getEnvDisaggBenchmarkGenOnly();
 // Whether to disable the chunked-attention in the generation phase.
 bool getEnvDisableChunkedAttentionInGenPhase();
 
+bool getEnvEplbForceGdrcopy();
+
 } // namespace tensorrt_llm::common

cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.cpp

@@ -27,6 +27,7 @@
 #include "topologyDetector.h"
 
 #include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/envUtils.h"
 #include "tensorrt_llm/common/logger.h"
 
 namespace tensorrt_llm::runtime
@@ -169,7 +170,8 @@ bool HostAccessibleDeviceAllocator::mAllowManagedFallback = false;
 
 bool HostAccessibleDeviceAllocator::isSupported()
 {
-    if (TopologyDetector::getInstance().getCurrentGpuMemoryNumaId() >= 0)
+    if (!tensorrt_llm::common::getEnvEplbForceGdrcopy()
+        && TopologyDetector::getInstance().getCurrentGpuMemoryNumaId() >= 0)
     {
         // we are on systems that GPU memory is also a NUMA node.
         return true;
@@ -195,7 +197,16 @@ void HostAccessibleDeviceAllocator::init()
     }
 
     TLLM_CUDA_CHECK(cudaGetDevice(&mDevId));
-    mGpuMemNumaId = TopologyDetector::getInstance().getCurrentGpuMemoryNumaId();
+    if (tensorrt_llm::common::getEnvEplbForceGdrcopy())
+    {
+        mGpuMemNumaId = -1;
+        TLLM_LOG_INFO("Force using GDRCopy for EPLB, ignore NUMA node for GPU memory.");
+    }
+    else
+    {
+        mGpuMemNumaId = TopologyDetector::getInstance().getCurrentGpuMemoryNumaId();
+    }
+
     if (mGpuMemNumaId < 0)
     {
         // We only use GDRCopy when there is no NUMA node for GPU memory.