GDS_MT backend support for LoopbackAgent

Author: tshmilnvidia

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -872,7 +872,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);
 
     //! \brief Calculate the proportional share each window size receives of the total memory pool
     //! \details Example:       (uniqueWindowSizeToLayers={1024: [1], 4096: [0, 4, 5], 8192: [2, 3]})

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());
 }