GDS_MT backend support for LoopbackAgent

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

Author: tshmilnvidia

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -873,7 +873,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 enablePartialReuse = true,
-        bool copyOnPartialReuse = 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

@@ -505,15 +505,16 @@ BlockManager::BlockManager(std::vector<SizeType32> const& numKvHeadsPerLayer, Si
     std::optional<TempAttentionWindowInputs> const& tempAttentionWindowInputs, nvinfer1::DataType dtype,
     SizeType32 sinkBubbleLength, bool onboardBlocks, CacheType cacheType,
     std::optional<executor::RetentionPriority> secondaryOffloadMinPriority,
-    std::shared_ptr<KVCacheEventManager> eventManager, bool enablePartialReuse, bool copyOnPartialReuse)
+    std::shared_ptr<KVCacheEventManager> eventManager, bool enablePartialReuse, 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

@@ -490,10 +490,17 @@ 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);
     }
 }
 

docker/common/install_nixl.sh

@@ -32,7 +32,7 @@ meson setup builddir \
     -Dcudapath_inc="$CUDA_PATH/include" \
     -Dgds_path="$GDS_PATH" \
     -Dinstall_headers=true \
-    -Dstatic_plugins=UCX,GDS
+    -Dstatic_plugins=UCX,GDS,GDS_MT
 
 cd builddir && ninja install
 cd ../..