[None][feat] Add NIXL-LIBFABRIC

Allows users to use LIBFABRIC as NIXL backend for KV cache transmission.

Usage: TRTLLM_NIXL_KVCACHE_BACKEND=LIBFABRIC  # (or UCX)

Currently supports only UCX and LIBFABRIC as NIXL KV cahce backends.

Unsupported backend types will fallback to Default (UCX).

Signed-off-by: Yoray Zack <[email protected]>

Author: zackyoray

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

@@ -348,11 +348,11 @@ NixlTransferAgent::NixlTransferAgent(BaseAgentConfig const& config)
 
     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"};
+    static const std::set<std::string> kSUPPORTED_BACKENDS = {"UCX", "LIBFABRIC"};
 
     if (kSUPPORTED_BACKENDS.find(nixlBackend) == kSUPPORTED_BACKENDS.end())
     {
-        TLLM_LOG_ERROR("Unsupported NIXL backend: %s, fallback to UCX", nixlBackend.c_str());
+        TLLM_LOG_WARNING("Unsupported NIXL backend: %s, fallback to UCX", nixlBackend.c_str());
         nixlBackend = "UCX";
     }
 

cpp/tests/unit_tests/executor/transferAgentTest.cpp

@@ -133,6 +133,57 @@ TEST_F(TransferAgentTest, Basic2)
     nixlAgent0->invalidateRemoteAgent(agent1);
 }
 
+TEST_F(TransferAgentTest, LibfabricBackendBasic)
+{
+    // Save original environment variable
+    char const* originalBackend = std::getenv("TRTLLM_NIXL_KVCACHE_BACKEND");
+    std::string savedBackend = originalBackend ? originalBackend : "";
+
+    // Set LIBFABRIC backend
+    setenv("TRTLLM_NIXL_KVCACHE_BACKEND", "LIBFABRIC", 1);
+
+    std::string const agent0{"libfabric_agent0"}, agent1{"libfabric_agent1"};
+    BaseAgentConfig config0{agent0, true}, config1{agent1, true};
+    auto nixlAgent0 = makeTransferAgent(config0);
+    auto nixlAgent1 = makeTransferAgent(config1);
+
+    TLLM_CHECK(nixlAgent0);
+    TLLM_CHECK(nixlAgent1);
+
+    std::vector<char> memory0(100, 10);
+    std::vector<char> memory1(100, 1);
+
+    RegisteredHostMemory regMem0(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory0}}}, nixlAgent0.get());
+    RegisteredHostMemory regMem1(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{memory1}}}, nixlAgent1.get());
+
+    // nixlAgent0->loadRemoteAgent(agent1);
+    auto connectionInfo = nixlAgent1->getLocalConnectionInfo();
+    nixlAgent0->loadRemoteAgent(agent1, connectionInfo);
+    bool checked = false;
+    do
+    {
+        checked = nixlAgent0->checkRemoteDescs(agent1, regMem1.getDescs());
+        // wait for regMem is unpacked by nixlAgent0
+    } while (!checked);
+
+    TransferRequest writeReq{TransferOp::kWRITE, regMem0.getDescs(), regMem1.getDescs(), agent1};
+    auto status = nixlAgent0->submitTransferRequests(writeReq);
+    status->wait();
+
+    TLLM_CHECK(memory0 == memory1);
+
+    nixlAgent0->invalidateRemoteAgent(agent1);
+    // Restore original environment variable
+    if (savedBackend.empty())
+    {
+        unsetenv("TRTLLM_NIXL_KVCACHE_BACKEND");
+    }
+    else
+    {
+        setenv("TRTLLM_NIXL_KVCACHE_BACKEND", savedBackend.c_str(), 1);
+    }
+}
+
 TEST_F(TransferAgentTest, DeviceMemory)
 {
 

docker/common/install_nixl.sh

@@ -4,7 +4,8 @@ set -ex
 GITHUB_URL="https://github.com"
 UCX_INSTALL_PATH="/usr/local/ucx/"
 CUDA_PATH="/usr/local/cuda"
-NIXL_VERSION="0.5.0"
+LIBFABRIC_INSTALL_PATH="/opt/libfabric"
+NIXL_VERSION="0.7.1"
 NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
 OLD_LD_LIBRARY_PATH=$LD_LIBRARY_PATH
 
@@ -38,6 +39,7 @@ meson setup builddir \
     -Dcudapath_lib="$CUDA_PATH/lib64" \
     -Dcudapath_inc="$CUDA_PATH/include" \
     -Dgds_path="$GDS_PATH" \
+    -Dlibfabric_path="$LIBFABRIC_INSTALL_PATH" \
     -Dinstall_headers=true
 
 cd builddir && ninja install