[None][feat] NIXL support for hybrid model cache transfer (#11608)

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

Author: NVShreyas

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -288,7 +288,7 @@ class CacheTransceiver : public BaseCacheTransceiver
     std::unique_ptr<executor::kv_cache::ConnectionManager> mManager;
     std::optional<executor::CacheTransceiverConfig> mCacheTransceiverConfig;
     std::vector<std::unique_ptr<kv_cache_manager::CacheTransBufferManager>> mCacheTransBufferManagers;
-    std::vector<kv_cache_manager::CacheTransBufferManager*> mCacheTransBufferManagerPtrs;
+    std::vector<BaseTransBufferManager*> mCacheTransBufferManagerPtrs;
 
     rnn_state_manager::RnnStateManager* mRnnStateManager{nullptr};
     // TODO(shreyasm): update this to use same container as kv by using base trans buffers instead

cpp/include/tensorrt_llm/executor/cacheCommunicator.h

@@ -18,6 +18,8 @@
 
 #include "tensorrt_llm/executor/serialization.h"
 #include <atomic>
+#include <cstdint>
+#include <optional>
 #include <vector>
 
 namespace tensorrt_llm::executor::kv_cache
@@ -63,6 +65,13 @@ class Connection
     {
         return false;
     }
+
+    virtual void activateBuffer(uint8_t /*kind*/) const {}
+
+    [[nodiscard]] virtual std::optional<size_t> getPreAssignedBufferId(uint8_t /*kind*/) const
+    {
+        return std::nullopt;
+    }
 };
 
 class ConnectionManager

cpp/tensorrt_llm/batch_manager/baseTransBuffer.h

@@ -23,6 +23,7 @@
 #include <atomic>
 #include <condition_variable>
 #include <cstddef>
+#include <cstdint>
 #include <memory>
 #include <mutex>
 #include <optional>
@@ -38,6 +39,13 @@ class FabricMemory;
 namespace tensorrt_llm::batch_manager
 {
 
+enum class BufferKind : uint8_t
+{
+    kKV = 0,
+    kKV_INDEXER = 1,
+    kRNN = 2
+};
+
 /// @brief Base class for cache transfer buffer management.
 /// Handles buffer pool allocation, index assignment, and slicing.
 /// Derived classes provide cache-specific size calculations.
@@ -46,6 +54,8 @@ class BaseTransBufferManager
 public:
     virtual ~BaseTransBufferManager() = default;
 
+    [[nodiscard]] virtual BufferKind getBufferKind() const = 0;
+
     /// @brief Assign a buffer index for sending.
     /// @return Assigned buffer index, or nullopt if using dynamic buffers.
     std::optional<int> assignBufferIndexForSend();

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -539,9 +539,9 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
             "bufferCoverTargetNum:%d pickUpConnections.size():%ld",
             bufferTargetNum, targetNum, peerDuplicateHeadFactor, targetInfo.mDupHeadFactor, bufferCoverTargetNum,
             pickUpConnections.size());
-        auto* agentConnnecion
+        auto const* agentConnection
             = dynamic_cast<executor::kv_cache::AgentConnection const*>(connections[pickUpConnections[0]]);
-        if (agentConnnecion != nullptr)
+        if (agentConnection != nullptr)
         {
             TLLM_CHECK_WITH_INFO(bufferCoverTargetNum == bufferTargetNum, "Agent need all buffer pre-allocated");
             TLLM_CHECK(onlyUseDynamicBuffer == false);
@@ -792,12 +792,11 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess
 
                 TLLM_CHECK(blockNum > 0);
 
-                auto* agentConnnecion
-                    = dynamic_cast<executor::kv_cache::AgentConnection const*>(connections[pickUpConnections[0]]);
-                if (agentConnnecion != nullptr)
+                auto preAssignedKvId
+                    = connections[pickUpConnections[0]]->getPreAssignedBufferId(static_cast<uint8_t>(BufferKind::kKV));
+                if (preAssignedKvId.has_value())
                 {
-                    cacheBufferId = agentConnnecion->getCacheBufferId();
-                    TLLM_CHECK(cacheBufferId.has_value());
+                    cacheBufferId = static_cast<int>(*preAssignedKvId);
                 }
                 else
                 {
@@ -811,7 +810,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess
                 bufferCoverTargetNum = bufferCoverTargetNumtmp;
                 remainNoCoverTargetNum = targetNum > bufferCoverTargetNum ? targetNum - bufferCoverTargetNum : 0;
 
-                if (agentConnnecion != nullptr)
+                if (preAssignedKvId.has_value())
                 {
                     TLLM_CHECK_WITH_INFO(bufferCoverTargetNum == targetNum, "Agent need buffer pre-allocated");
                     TLLM_CHECK(onlyUseDynamicBuffer == false);

cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp

@@ -249,6 +249,7 @@ CacheTransBufferManager::CacheTransBufferManager(
                               : cacheManager->getPrimaryPool(0)->getDataType(),
         maxNumTokens)
     , mCacheManager{cacheManager}
+    , mTransferIndexerKCache{transferIndexerKCache}
 {
     // TODO: FP4 dataSize
     TLLM_CHECK(mCacheManager);

cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h

@@ -74,12 +74,18 @@ class CacheTransBufferManager : public BaseTransBufferManager
         return mCacheManager;
     }
 
+    [[nodiscard]] BufferKind getBufferKind() const override
+    {
+        return mTransferIndexerKCache ? BufferKind::kKV_INDEXER : BufferKind::kKV;
+    }
+
 private:
     /// @brief Compute transfer buffer size from KV cache configuration.
     static size_t computeTransferBufferSize(KVCacheManager::BaseKVCacheManager* cacheManager,
         std::optional<size_t> maxNumTokens, bool transferIndexerKCache);
 
     KVCacheManager::BaseKVCacheManager* mCacheManager;
+    bool mTransferIndexerKCache;
 };
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -185,26 +185,13 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
         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());
-    }
 
     // RNN specific setup
     if (mRnnStateManager != nullptr)
     {
         TLLM_LOG_DEBUG("Setting up RNN cache transfer components.");
         TLLM_CHECK(!rnnLayerNumPerPP.empty());
 
-        if (backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL
-            || backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE)
-        {
-            TLLM_LOG_ERROR("RNN cache transfer is not supported for NIXL and MOONCAKE yet");
-            return;
-        }
-
         mRnnCacheTransBufferManager
             = std::make_unique<rnn_state_manager::RnnCacheTransBufferManager>(mRnnStateManager, maxNumTokens);
 
@@ -218,6 +205,17 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
         TLLM_LOG_INFO("RNN cache transfer components initialized.");
     }
 
+    mCacheTransBufferManagerPtrs.clear();
+    mCacheTransBufferManagerPtrs.reserve(mCacheTransBufferManagers.size() + (mRnnCacheTransBufferManager ? 1 : 0));
+    for (auto& manager : mCacheTransBufferManagers)
+    {
+        mCacheTransBufferManagerPtrs.push_back(manager.get());
+    }
+    if (mRnnCacheTransBufferManager)
+    {
+        mCacheTransBufferManagerPtrs.push_back(mRnnCacheTransBufferManager.get());
+    }
+
     if (backendType.value() == executor::CacheTransceiverConfig::BackendType::UCX)
     {
         std::lock_guard<std::mutex> lock(mDllMutex);
@@ -239,14 +237,18 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
     }
     else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::NIXL)
     {
+        auto rnnState
+            = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt;
         mManager = std::make_unique<tensorrt_llm::executor::kv_cache::AgentConnectionManager>(
-            mCacheTransBufferManagerPtrs, *mCacheState, "nixl");
+            mCacheTransBufferManagerPtrs, *mCacheState, "nixl", rnnState);
         TLLM_LOG_INFO("NIXL Connection Manager created");
     }
     else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MOONCAKE)
     {
+        auto rnnState
+            = mCacheState->hasRnnConfig() ? std::make_optional(mCacheState->getRnnCacheState()) : std::nullopt;
         mManager = std::make_unique<tensorrt_llm::executor::kv_cache::AgentConnectionManager>(
-            mCacheTransBufferManagerPtrs, *mCacheState, "mooncake");
+            mCacheTransBufferManagerPtrs, *mCacheState, "mooncake", rnnState);
         TLLM_LOG_INFO("MOONCAKE Connection Manager created");
     }
     else if (backendType.value() == executor::CacheTransceiverConfig::BackendType::MPI)
@@ -261,7 +263,15 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
     }
 
     auto makeFormatter = [cacheManager, isMLA, this]()
-    { return createCacheFormatter(cacheManager, mCacheTransBufferManagerPtrs, isMLA); };
+    {
+        std::vector<kv_cache_manager::CacheTransBufferManager*> kvBufferPtrs;
+        kvBufferPtrs.reserve(mCacheTransBufferManagers.size());
+        for (auto& mgr : mCacheTransBufferManagers)
+        {
+            kvBufferPtrs.push_back(mgr.get());
+        }
+        return createCacheFormatter(cacheManager, kvBufferPtrs, isMLA);
+    };
 
     auto makeRnnFormatter = [this]() -> std::unique_ptr<RnnCacheFormatter>
     {

cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp

@@ -21,6 +21,7 @@
 #include "tensorrt_llm/batch_manager/rnnCacheFormatter.h"
 #include "tensorrt_llm/common/assert.h"
 #include "tensorrt_llm/common/logger.h"
+#include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h"
 #include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h"
 
 #include <algorithm>
@@ -95,6 +96,13 @@ void CacheTransferLayer::format(TransferSession& session) const
     mKvFormatter->format(session);
     if (mRnnFormatter)
     {
+        for (auto const* conn : session.getConnections())
+        {
+            if (conn != nullptr)
+            {
+                conn->activateBuffer(static_cast<uint8_t>(BufferKind::kRNN));
+            }
+        }
         mRnnFormatter->format(session);
     }
 }

cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp

@@ -26,6 +26,7 @@
 #include "tensorrt_llm/common/tllmException.h"
 #include "tensorrt_llm/common/utils.h"
 #include "tensorrt_llm/executor/cache_transmission/agent_utils/connection.h"
+#include "tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h"
 #include "tensorrt_llm/runtime/common.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 #include <chrono>
@@ -384,9 +385,10 @@ class CacheSender::Impl
         auto allCounterparts = mCacheTransferLayer.computeCounterparts(
             mSelfState.getCommState().value().getSelfIdx(), info.getTransState());
 
-        auto peerSelfIdx = info.getTransState().getCommState()->getSelfIdx(); // Index of self in peer's comm state
+        auto peerSelfIdx = info.getTransState().getCommState()->getSelfIdx();
         int peerIdx = std::distance(
             allCounterparts.begin(), std::find(allCounterparts.begin(), allCounterparts.end(), peerSelfIdx));
+
         TLLM_CHECK_WITH_INFO(peerIdx < static_cast<int>(allCounterparts.size()),
             "Peer rank %d not found in expected counterparts", peerSelfIdx);
         {
@@ -861,6 +863,19 @@ class CacheReceiver::Impl
         auto allCounterparts
             = mCacheTransferLayer.computeCounterparts(mSelfState.getCommState().value().getSelfIdx(), contextState);
 
+        auto kvCounterParts = mCacheTransferLayer.getKvFormatter()->getCounterparts(
+            mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx(), destCacheState);
+
+        bool hasRnn = mCacheTransferLayer.getCacheState().hasRnnConfig() && destCacheState.hasRnnConfig();
+
+        std::vector<SizeType32> rnnCounterParts;
+        if (hasRnn)
+        {
+            rnnCounterParts = executor::kv_cache::targetIRanksForRnn(
+                destCacheState, mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx())
+                                  .mIRanks;
+        }
+
         auto connections = mManager->getConnections(commState);
         std::vector<executor::kv_cache::Connection const*> allConnections;
         for (auto index : allCounterparts)
@@ -869,24 +884,59 @@ class CacheReceiver::Impl
             allConnections.emplace_back(connection);
         }
 
-        for (size_t i = 0; i < allConnections.size(); i++)
+        for (size_t ci = 0; ci < allCounterparts.size(); ci++)
         {
-            auto const* connection = allConnections[i];
-            // if Manager is agentConnectionManager, then send request info to agent
-            auto* agentConnectionManager = dynamic_cast<executor::kv_cache::AgentConnectionManager*>(mManager);
+            auto rank = allCounterparts[ci];
+            auto const* connection = connections.at(rank);
+
+            bool isKvCounterpart
+                = std::find(kvCounterParts.begin(), kvCounterParts.end(), rank) != kvCounterParts.end();
+            bool isRnnCounterpart
+                = hasRnn && std::find(rnnCounterParts.begin(), rnnCounterParts.end(), rank) != rnnCounterParts.end();
+
             if (agentConnectionManager)
             {
-                // TODO: index -> validConnectionIdx conversion
-                // TODO(shreyasm): this will not work for RNN. Will error out in the constructor if used with RNN.
-                auto [pickUpIdx, localRankIdx] = mCacheTransferLayer.getKvFormatter()->pickRecvConnections(
-                    allCounterparts.size(), mSelfState.getCacheState().value(),
-                    mSelfState.getCommState().value().getSelfIdx(), destCacheState, allCounterparts);
-                auto validConnectionIdx = std::find(localRankIdx.begin(), localRankIdx.end(), i) - localRankIdx.begin();
+                auto idsForRank = cacheBufferIds;
+                auto const& managers = agentConnectionManager->getCacheTransBufferManagers();
+                for (size_t i = 0; i < idsForRank.size(); i++)
+                {
+                    auto kind = managers[i]->getBufferKind();
+                    bool include = (kind != BufferKind::kRNN) ? isKvCounterpart : isRnnCounterpart;
+                    if (!include)
+                    {
+                        idsForRank[i] = std::nullopt;
+                    }
+                }
+
+                int validConnectionIdx = 0;
+                if (isKvCounterpart)
+                {
+                    auto kvCpIdx
+                        = std::find(kvCounterParts.begin(), kvCounterParts.end(), rank) - kvCounterParts.begin();
+                    auto [pickUpIdx, localRankIdx] = mCacheTransferLayer.getKvFormatter()->pickRecvConnections(
+                        allCounterparts.size(), mSelfState.getCacheState().value(),
+                        mSelfState.getCommState().value().getSelfIdx(), destCacheState, allCounterparts);
+                    validConnectionIdx
+                        = std::find(localRankIdx.begin(), localRankIdx.end(), kvCpIdx) - localRankIdx.begin();
+                }
+                else if (isRnnCounterpart)
+                {
+                    auto rnnTargetInfo = executor::kv_cache::targetIRanksForRnn(destCacheState,
+                        mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx());
+                    auto rnnCpIdx
+                        = std::find(rnnCounterParts.begin(), rnnCounterParts.end(), rank) - rnnCounterParts.begin();
+                    auto [pickUpIdx, localRankIdx] = cache_formatter_utils::pickRecvConnections(rnnCounterParts.size(),
+                        mCacheTransferLayer.getCacheState(), mSelfState.getCommState().value().getSelfIdx(),
+                        destCacheState, rnnCounterParts, rnnTargetInfo);
+                    validConnectionIdx
+                        = std::find(localRankIdx.begin(), localRankIdx.end(), rnnCpIdx) - localRankIdx.begin();
+                }
+
                 auto* agentConnection = dynamic_cast<executor::kv_cache::AgentConnection const*>(connection);
                 TLLM_CHECK(agentConnection != nullptr);
-                TLLM_CHECK(!cacheBufferIds.empty());
+
                 const_cast<executor::kv_cache::AgentConnection*>(agentConnection)
-                    ->sendRequestAndBufferInfo(requestInfo, cacheBufferIds, validConnectionIdx);
+                    ->sendRequestAndBufferInfo(requestInfo, idsForRank, validConnectionIdx);
             }
             else
             {