[None][feat] Add TRTLLM_NIXL_KVCACHE_BACKEND environment variable for NIXL backend selection

Allows users to specify which NIXL backend to use for KV cache transmission.

Usage: TRTLLM_NIXL_KVCACHE_BACKEND=<backend>

Currently supports only UCX backend. Unsupported backend types will fallback to Default (UCX).

Signed-off-by: Yoray Zack <62789610+zackyoray@users.noreply.github.com>

Author: zackyoray

cpp/tensorrt_llm/common/envUtils.cpp

@@ -318,6 +318,28 @@ std::string getEnvNixlInterface()
     return nixlInterface;
 }
 
+std::string getEnvNixlBackend()
+{
+    static std::once_flag flag;
+    static std::string nixlBackend;
+
+    std::call_once(flag,
+        [&]()
+        {
+            char const* nixl_backend = std::getenv("TRTLLM_NIXL_KVCACHE_BACKEND");
+            if (nixl_backend)
+            {
+                nixlBackend = nixl_backend;
+            }
+            else
+            {
+                // Default to UCX if not specified
+                nixlBackend = "UCX";
+            }
+        });
+    return nixlBackend;
+}
+
 bool getEnvDisaggLayerwise()
 {
     static bool const disaggLayerwise = getBoolEnv("TRTLLM_DISAGG_LAYERWISE");

cpp/tensorrt_llm/common/envUtils.h

@@ -88,6 +88,8 @@ std::string getEnvUCXInterface();
 
 std::string getEnvNixlInterface();
 
+std::string getEnvNixlBackend();
+
 bool getEnvDisaggLayerwise();
 
 bool getEnvParallelCacheSend();

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp

@@ -28,6 +28,7 @@
 #include <netdb.h>
 #include <netinet/in.h>
 #include <nixl_types.h>
+#include <set>
 #include <sys/file.h>
 #include <sys/stat.h>
 #include <unistd.h>
@@ -345,15 +346,27 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config)
         mRawAgent = std::make_unique<nixlAgent>(config.mName, std::move(nixlConfig));
     }
 
+    std::string nixlBackend = common::getEnvNixlBackend();
+    // List of supported backends - extend this list as new backends are added
+    static const std::set<std::string> kSUPPORTED_BACKENDS = {"UCX"};
+
+    if (kSUPPORTED_BACKENDS.find(nixlBackend) == kSUPPORTED_BACKENDS.end())
+    {
+        TLLM_LOG_ERROR("Unsupported NIXL backend: %s, fallback to UCX", nixlBackend.c_str());
+        nixlBackend = "UCX";
+    }
+
+    TLLM_LOG_INFO("NixlTransferAgent::NixlTransferAgent using NIXL backend: %s", nixlBackend.c_str());
+
     nixl_b_params_t init1;
     nixl_mem_list_t mems1;
-    status = mRawAgent->getPluginParams("UCX", mems1, init1);
+    status = mRawAgent->getPluginParams(nixlBackend.c_str(), mems1, init1);
     TLLM_CHECK(status == NIXL_SUCCESS);
 
-    status = mRawAgent->createBackend("UCX", init1, mRawBackend);
+    status = mRawAgent->createBackend(nixlBackend.c_str(), init1, mRawBackend);
     if (status != NIXL_SUCCESS || !mRawBackend)
     {
-        TLLM_THROW("Failed to create NIXL backend");
+        TLLM_THROW("Failed to create NIXL backend: %s", nixlBackend.c_str());
     }
     mExtraParams.backends.push_back(mRawBackend);
     TLLM_LOG_INFO("NixlTransferAgent::NixlTransferAgent mAddress: %s", mAddress.c_str());