Fix nixl support for indexer cache transfer

Signed-off-by: Iman Tabrizian <[email protected]>

Author: Tabrizian

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -269,7 +269,8 @@ class CacheTransceiver : public BaseCacheTransceiver
     std::unique_ptr<executor::kv_cache::CacheState> mCacheState;
     std::unique_ptr<executor::kv_cache::ConnectionManager> mManager;
     std::optional<executor::CacheTransceiverConfig> mCacheTransceiverConfig;
-    std::unique_ptr<kv_cache_manager::CacheTransBufferManager> mCacheTransBufferManager;
+    std::vector<std::unique_ptr<kv_cache_manager::CacheTransBufferManager>> mCacheTransBufferManagers;
+    std::vector<kv_cache_manager::CacheTransBufferManager*> mCacheTransBufferManagerPtrs;
     // library handle to the communicator related features,
     // this is used to defer dependency resolution until needed.
     static std::mutex mDllMutex;

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -972,19 +972,14 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess
 }
 
 std::unique_ptr<BaseCacheFormatter> createCacheFormatter(
-    BaseKVCacheManager* cacheManager, CacheTransBufferManager* cacheTransBufferManager, bool isMLA)
+    BaseKVCacheManager* cacheManager, std::vector<CacheTransBufferManager*> const& cacheTransBufferManagers, bool isMLA)
 {
+    TLLM_CHECK(!cacheTransBufferManagers.empty());
     if (isMLA)
     {
-        std::vector<CacheTransBufferManager*> cacheTransBufferManagers = {cacheTransBufferManager};
-        auto maxNumTokens = cacheTransBufferManager->getMaxNumTokens();
-        if (cacheManager->isEnableIndexerKCache())
-        {
-            cacheTransBufferManagers.push_back(new CacheTransBufferManager(cacheManager, maxNumTokens, true));
-        }
         return std::make_unique<MLACacheFormatter>(cacheManager, cacheTransBufferManagers);
     }
-    return std::make_unique<CacheFormatter>(cacheManager, cacheTransBufferManager);
+    return std::make_unique<CacheFormatter>(cacheManager, cacheTransBufferManagers[0]);
 }
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager

cpp/tensorrt_llm/batch_manager/cacheFormatter.h

@@ -133,7 +133,7 @@ class CacheFormatter final : public BaseCacheFormatter
     CacheTransBufferManager* mCacheTransBufferManager;
 };
 
-std::unique_ptr<BaseCacheFormatter> createCacheFormatter(
-    BaseKVCacheManager* cacheManager, CacheTransBufferManager* cacheTransBufferManager, bool isMLA = false);
+std::unique_ptr<BaseCacheFormatter> createCacheFormatter(BaseKVCacheManager* cacheManager,
+    std::vector<CacheTransBufferManager*> const& cacheTransBufferManagers, bool isMLA = false);
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -168,7 +168,19 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
 
     std::optional<size_t> maxNumTokens = mCacheTransceiverConfig.value().getMaxTokensInBuffer();
 
-    mCacheTransBufferManager = std::make_unique<kv_cache_manager::CacheTransBufferManager>(cacheManager, maxNumTokens);
+    mCacheTransBufferManagers.push_back(
+        std::make_unique<kv_cache_manager::CacheTransBufferManager>(cacheManager, maxNumTokens));
+    if (isMLA && cacheManager->isEnableIndexerKCache())
+    {
+        mCacheTransBufferManagers.push_back(
+            std::make_unique<kv_cache_manager::CacheTransBufferManager>(cacheManager, maxNumTokens, true));
+    }
+    mCacheTransBufferManagerPtrs.clear();
+    mCacheTransBufferManagerPtrs.reserve(mCacheTransBufferManagers.size());
+    for (auto& manager : mCacheTransBufferManagers)
+    {
+        mCacheTransBufferManagerPtrs.push_back(manager.get());
+    }
     if (backendType.value() == executor::CacheTransceiverConfig::BackendType::UCX)
     {
         std::lock_guard<std::mutex> lock(mDllMutex);
@@ -191,7 +203,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
     else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL)
     {
         mManager = std::make_unique<tensorrt_llm::executor::kv_cache::AgentConnectionManager>(
-            mCacheTransBufferManager.get(), *mCacheState);
+            mCacheTransBufferManagerPtrs, *mCacheState);
         TLLM_LOG_INFO("NIXL Connection Manager created");
     }
     else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI)
@@ -206,7 +218,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
     }
 
     auto makeFormatter = [cacheManager, isMLA, this]()
-    { return createCacheFormatter(cacheManager, mCacheTransBufferManager.get(), isMLA); };
+    { return createCacheFormatter(cacheManager, mCacheTransBufferManagerPtrs, isMLA); };
 
     mCacheSender = std::make_unique<CacheSender>(mManager.get(), *mCacheState, worldConfig.getRank(), makeFormatter());
     mCacheReceiver

cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp

@@ -834,12 +834,14 @@ class CacheReceiver::Impl
         }
 
         auto* agentConnectionManager = dynamic_cast<executor::kv_cache::AgentConnectionManager*>(mManager);
-        std::optional<size_t> cacheBufferId = std::nullopt;
+        std::vector<std::optional<size_t>> cacheBufferIds;
         if (agentConnectionManager)
         {
-            cacheBufferId = agentConnectionManager->getCacheTransBufferManager()->assignBufferIndexForRecv();
-            TLLM_CHECK(cacheBufferId.has_value());
-            // memory Desp , validSegmentIdx send
+            for (auto& cacheTransBufferManager : agentConnectionManager->getCacheTransBufferManagers())
+            {
+                cacheBufferIds.push_back(cacheTransBufferManager->assignBufferIndexForRecv());
+            }
+            TLLM_CHECK(!cacheBufferIds.empty());
         }
         auto counterParts = mFormatter->getCounterparts(
             mSelfState.getCacheState().value(), mSelfState.getCommState().value().getSelfIdx(), destCacheState);
@@ -864,9 +866,9 @@ class CacheReceiver::Impl
                 auto validConnectionIdx = std::find(pickUpIdx.begin(), pickUpIdx.end(), i) - pickUpIdx.begin();
                 auto* agentConnection = dynamic_cast<executor::kv_cache::AgentConnection const*>(connection);
                 TLLM_CHECK(agentConnection != nullptr);
-                TLLM_CHECK(cacheBufferId.has_value());
+                TLLM_CHECK(!cacheBufferIds.empty());
                 const_cast<executor::kv_cache::AgentConnection*>(agentConnection)
-                    ->sendRequestAndBufferInfo(requestInfo, cacheBufferId, validConnectionIdx);
+                    ->sendRequestAndBufferInfo(requestInfo, cacheBufferIds, validConnectionIdx);
             }
             else
             {

cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp

@@ -131,6 +131,16 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses
 
     for (auto transferIndexerKCache : transferringIndexerKCache)
     {
+        auto activeBufferIdx = transferIndexerKCache ? 1UL : 0UL;
+        for (auto const* connection : connections)
+        {
+            if (auto const* agentConnection = dynamic_cast<executor::kv_cache::AgentConnection const*>(connection))
+            {
+                TLLM_CHECK(agentConnection->getSenderBufferCount() > activeBufferIdx);
+                const_cast<executor::kv_cache::AgentConnection*>(agentConnection)
+                    ->setActiveSenderBufferIdx(activeBufferIdx);
+            }
+        }
         int blockNum = 0;
         std::vector<runtime::ITensor::SharedPtr> inputKvCacheBlocks;
         if (!transferIndexerKCache)
@@ -417,9 +427,10 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s
         else
         {
             auto* agentConnnecion = dynamic_cast<executor::kv_cache::AgentConnection const*>(connections[0]);
+            size_t activeBufferIdx = transferIndexerKCache ? 1 : 0;
             if (agentConnnecion != nullptr)
             {
-                cacheBufferId = agentConnnecion->getCacheBufferId();
+                cacheBufferId = agentConnnecion->getCacheBufferId(activeBufferIdx);
                 TLLM_CHECK(cacheBufferId.has_value());
             }
             else