GDS_MT backend support for LoopbackAgent

Signed-off-by: Tomer Shmilovich <[email protected]>

Author: tshmilnvidia

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -874,7 +874,7 @@ class BlockManager
         SizeType32 sinkBubbleLength, bool onboardBlocks, CacheType cacheType = CacheType::kSELF,
         std::optional<executor::RetentionPriority> secondaryOffloadMinPriority = std::nullopt,
         std::shared_ptr<KVCacheEventManager> eventManager = nullptr, bool enableHashKey = false,
-        bool enablePartialReuse = true, bool copyOnPartialReuse = true);
+        bool enablePartialReuse = true, bool copyOnPartialReuse = true, bool multiThreadReuse = false);
 
     BlockManager(BlockManager const&) = delete;
     BlockManager& operator=(BlockManager const&) = delete;

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -270,6 +270,7 @@ struct BaseAgentConfig
 {
     std::string mName;
     bool useProgThread;
+    bool multiThread;
 };
 
 class BaseTransferAgent

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -407,15 +407,15 @@ BlockManager::BlockManager(std::vector<SizeType32> const& numKvHeadsPerLayer, Si
     SizeType32 sinkBubbleLength, bool onboardBlocks, CacheType cacheType,
     std::optional<executor::RetentionPriority> secondaryOffloadMinPriority,
     std::shared_ptr<KVCacheEventManager> eventManager, bool enableHashKey, bool enablePartialReuse,
-    bool copyOnPartialReuse)
+    bool copyOnPartialReuse, bool multiThreadReuse)
     : mNumLayers{static_cast<SizeType32>(numKvHeadsPerLayer.size())}
     , mTokensPerBlock{tokensPerBlock}
     , mEventManager{std::move(eventManager)}
     , mStream{stream}
     , mCacheType{cacheType}
 {
     mAgentName = std::string("GDSAgent");
-    BaseAgentConfig config{mAgentName, true};
+    BaseAgentConfig config{mAgentName, true, multiThreadReuse};
     mLoopbackAgent = makeLoopbackAgent("nixl", &config);
 
     auto const uniqueWindowSizeToLayers

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

@@ -501,11 +501,19 @@ NixlLoopbackAgent::NixlLoopbackAgent(BaseAgentConfig const& config)
     init["batch_limit"] = std::to_string(128);
     init["max_request_size"] = std::to_string(16 * 1024 * 1024);
 
-    status = mRawAgent->createBackend("GDS", init, backend);
-    if (status != NIXL_SUCCESS || !backend)
+    if (config.multiThread)
     {
-        TLLM_THROW("Failed to create NIXL backend, status = %d", status);
+        status = mRawAgent->createBackend("GDS_MT", init, backend);
+        if (status != NIXL_SUCCESS || !backend)
+            TLLM_THROW("Failed to create NIXL GDS_MT backend, status = %d", status);
     }
+    else
+    {
+        status = mRawAgent->createBackend("GDS", init, backend);
+        if (status != NIXL_SUCCESS || !backend)
+            TLLM_THROW("Failed to create NIXL GDS backend, status = %d", status);
+    }
+
     TLLM_LOG_INFO("NixlLoopbackAgent::NixlLoopbackAgent mAddress: %s", mAddress.c_str());
 }
 

docker/common/install_nixl.sh

@@ -49,7 +49,7 @@ fi
 pip3 install --no-cache-dir meson ninja pybind11
 git clone --depth 1 -b ${NIXL_VERSION} ${NIXL_REPO}
 cd nixl
-meson setup builddir -Ducx_path=${UCX_INSTALL_PATH} -Dstatic_plugins=UCX,GDS ${EXTRA_NIXL_ARGS}
+meson setup builddir -Ducx_path=${UCX_INSTALL_PATH} -Dstatic_plugins=UCX,GDS,GDS_MT ${EXTRA_NIXL_ARGS}
 cd builddir && ninja install
 cd ../..
 rm -rf nixl*  # Remove NIXL source tree to save space