[TRTLLM-9920][feat] Add support for arbitrary KVCache transfer

cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h

@@ -228,6 +228,12 @@ class BaseCacheTransceiver
 
     virtual bool cancelRequest(std::shared_ptr<LlmRequest> llmRequest) = 0;
 
+    /// Get the serialized DataTransceiverState (CacheState + CommState) for this transceiver.
+    [[nodiscard]] virtual std::vector<char> getSerializedDataTransceiverState() const
+    {
+        return {};
+    }
+
     [[nodiscard]] virtual bool hasPoisonedTransferBuffer() const
     {
         return false;
@@ -277,6 +283,8 @@ class CacheTransceiver : public BaseCacheTransceiver
 
     virtual bool cancelRequest(std::shared_ptr<LlmRequest> llmRequest) override;
 
+    [[nodiscard]] std::vector<char> getSerializedDataTransceiverState() const override;
+
     [[nodiscard]] bool hasPoisonedTransferBuffer() const override;
 
 private:

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h

@@ -1251,14 +1251,29 @@ class WindowBlockManager
         return mEnablePartialReuse;
     }
 
+    //! \brief Look up the block chain matching blockKey in the reuse tree.
     [[nodiscard]] std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKey(BlockKey const& blockKey);
 
+    //! \brief Same lookup, additionally pinning matched blocks; on a miss all pins are
+    //! rolled back and pinnedBlockIds is cleared.
+    [[nodiscard]] std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKey(
+        BlockKey const& blockKey, std::vector<KVCacheBlock::IdType>& pinnedBlockIds);
+
     [[nodiscard]] std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKeys(
         std::vector<BlockKey> const& blockKeys);
 
     //! \brief Unpin blocks by block ids directly
     void unpinBlocksById(std::vector<KVCacheBlock::IdType> const& blockIds);
 
+    //! \brief Pin a block: claim it from the eviction policy if free, then take a reference.
+    //! Safe to call from cache-transceiver threads: block bookkeeping is serialized by the
+    //! lookup-tree mutex, which every mutating entry point acquires.
+    void pinBlock(BlockPtr const& block);
+
+    //! \brief Inverse of pinBlock: drop one reference and release the block back to the
+    //! eviction policy once no references remain.
+    void unpinBlock(BlockPtr const& block);
+
     void truncateBlocks(LlmRequest::VecTokens const& targetTokens, SizeType32 numTokensToKeep);
 
     void resetReuseState()
@@ -1274,6 +1289,10 @@ class WindowBlockManager
     }
 
 private:
+    //! \brief Shared implementation of the findBlocksInReuseTreeByBlockKey overloads.
+    [[nodiscard]] std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKeyImpl(
+        BlockKey const& blockKey, bool pinBlocks, std::vector<KVCacheBlock::IdType>& pinnedBlockIds);
+
     //! \brief Walk the reuse tree with precomputed per-block keys (no lock; callers must hold mLookupTree->getMutex()).
     [[nodiscard]] std::shared_ptr<KVCacheBlock> searchReuseTree(std::vector<BlockKey> const& blockKeys);
 
@@ -1836,6 +1855,12 @@ class BlockManager
         return mWindowBlockManagers.at(windowSize).findBlocksInReuseTreeByBlockKey(blockKey);
     }
 
+    [[nodiscard]] std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKey(
+        BlockKey const& blockKey, SizeType32 windowSize, std::vector<KVCacheBlock::IdType>& pinnedBlockIds)
+    {
+        return mWindowBlockManagers.at(windowSize).findBlocksInReuseTreeByBlockKey(blockKey, pinnedBlockIds);
+    }
+
     [[nodiscard]] std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKeys(
         std::vector<BlockKey> const& blockKeys, SizeType32 windowSize)
     {
@@ -2210,6 +2235,11 @@ class BaseKVCacheManager
         BlockKey const& blockKey, SizeType32 windowSize)
         = 0;
 
+    //! \brief Pinning lookup: pins matched blocks and records their ids for unpinBlocksById.
+    [[nodiscard]] virtual std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKey(
+        BlockKey const& blockKey, SizeType32 windowSize, std::vector<KVCacheBlock::IdType>& pinnedBlockIds)
+        = 0;
+
     [[nodiscard]] virtual std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKeys(
         std::vector<BlockKey> const& blockKeys, SizeType32 windowSize)
         = 0;
@@ -2650,6 +2680,12 @@ class KVCacheManager : public BaseKVCacheManager
         return mBlockManager.findBlocksInReuseTreeByBlockKey(blockKey, windowSize);
     }
 
+    std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKey(
+        BlockKey const& blockKey, SizeType32 windowSize, std::vector<KVCacheBlock::IdType>& pinnedBlockIds) override
+    {
+        return mBlockManager.findBlocksInReuseTreeByBlockKey(blockKey, windowSize, pinnedBlockIds);
+    }
+
     std::shared_ptr<KVCacheBlock> findBlocksInReuseTreeByBlockKeys(
         std::vector<BlockKey> const& blockKeys, SizeType32 windowSize) override
     {

cpp/include/tensorrt_llm/executor/dataTransceiverState.h

@@ -619,9 +619,22 @@ class DataTransceiverState final
         return mCacheState.has_value() && mCacheState->hasRnnConfig();
     }
 
+    /// @brief Set only when exported via CacheTransceiver::getSerializedDataTransceiverState:
+    /// transfers driven by such a state have no LlmRequest on the sender.
+    [[nodiscard]] bool isArbitraryTransferState() const noexcept
+    {
+        return mIsArbitraryTransferState;
+    }
+
+    void setIsArbitraryTransferState(bool isArbitraryTransferState) noexcept
+    {
+        mIsArbitraryTransferState = isArbitraryTransferState;
+    }
+
     [[nodiscard]] bool operator==(DataTransceiverState const& other) const noexcept
     {
-        return mCacheState == other.mCacheState && mCommState == other.mCommState;
+        return mCacheState == other.mCacheState && mCommState == other.mCommState
+            && mIsArbitraryTransferState == other.mIsArbitraryTransferState;
     }
 
     [[nodiscard]] std::string toString() const
@@ -642,6 +655,7 @@ class DataTransceiverState final
     friend class Serialization;
     std::optional<kv_cache::CacheState> mCacheState;
     std::optional<kv_cache::CommState> mCommState;
+    bool mIsArbitraryTransferState{false};
 };
 
 } // namespace tensorrt_llm::executor

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -183,7 +183,7 @@ void sendAllBuffers(TransferSession& session, int deviceId,
 namespace tensorrt_llm::batch_manager::kv_cache_manager
 {
 
-BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest const& llmRequest,
+BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, std::optional<LlmRequest const*> llmRequest,
     BlockKey const& lastBlockKey, int32_t indexFromEnd, bool recvSideHasCP, SizeType32 ppSize)
 {
     auto poolNum = cacheManager->getBlockManager().getNumPools(
@@ -197,9 +197,10 @@ BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest
         || lastBlockKey.uniqueTokens.size() == 0 || recvSideHasCP || ppSize > 1)
     {
         // disable reuse path, and vwsa don't support reuse.
+        TLLM_CHECK_WITH_INFO(llmRequest.has_value(), "LlmRequest required for non-reuse-tree transfer path");
         bool needSendAllForWindow = common::getEnvKVCacheTransferAllBlocksForWindow();
 
-        auto blockRange = BlockRange::fromAllBlockIds(*cacheManager, llmRequest.mRequestId);
+        auto blockRange = BlockRange::fromAllBlockIds(*cacheManager, (*llmRequest)->mRequestId);
 
         auto const& windowsMetadata = cacheManager->getBlockManager().getWindowSizesMetadata();
 
@@ -235,16 +236,20 @@ BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest
 
     TLLM_CHECK_WITH_INFO(lastBlockKey.uniqueTokens.size() > 0, "lastBlockKey must be non-empty when reuse is enabled");
 
-    auto multimodalHashes = llmRequest.getMultimodalHashes();
-    bool isMultimodal = multimodalHashes.has_value() && *multimodalHashes && !(*multimodalHashes)->empty();
-    if (isMultimodal)
+    // No request on the reuse-tree path: fall through to the plain lastBlockKey lookup.
+    if (llmRequest.has_value())
     {
-        auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock();
-        auto const usableSize = static_cast<SizeType32>(lastBlockKey.uniqueTokens.size());
-        auto blockedUniqueTokens = chopVectorIntoBlocks<UniqueToken>(
-            lastBlockKey.uniqueTokens, usableSize, tokensPerBlock, /*allowPartial=*/true);
-        auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest);
-        return BlockRange::fromReuseTree(*cacheManager, blockKeys, indexFromEnd);
+        auto multimodalHashes = (*llmRequest)->getMultimodalHashes();
+        bool isMultimodal = multimodalHashes.has_value() && *multimodalHashes && !(*multimodalHashes)->empty();
+        if (isMultimodal)
+        {
+            auto tokensPerBlock = cacheManager->getBlockManager().getTokensPerBlock();
+            auto const usableSize = static_cast<SizeType32>(lastBlockKey.uniqueTokens.size());
+            auto blockedUniqueTokens = chopVectorIntoBlocks<UniqueToken>(
+                lastBlockKey.uniqueTokens, usableSize, tokensPerBlock, /*allowPartial=*/true);
+            auto blockKeys = buildBlockKeys(blockedUniqueTokens, **llmRequest);
+            return BlockRange::fromReuseTree(*cacheManager, blockKeys, indexFromEnd);
+        }
     }
 
     return BlockRange::fromReuseTree(*cacheManager, lastBlockKey, indexFromEnd);
@@ -363,11 +368,15 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
 {
     NVTX3_SCOPED_RANGE(CacheFormatter_format);
     session.setTime(TransferSession::kTimeFormatter);
-    auto const& llmRequest = session.getLlmRequest();
-    TLLM_LOG_DEBUG(
-        mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", llmRequest.mRequestId);
+    auto llmRequest = session.getLlmRequest();
+    if (llmRequest.has_value())
+    {
+        TLLM_LOG_DEBUG(
+            mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", (*llmRequest)->mRequestId);
+        TLLM_CHECK_WITH_INFO(
+            (*llmRequest)->mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported.");
+    }
 
-    TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported.");
     auto const& connections = session.getConnections();
     auto const& selfConfig = session.getSelfState().getCacheState().value();
     auto const& destConfig = session.getOtherState().getCacheState().value();
@@ -385,7 +394,6 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
     size_t targetNum = pickUpConnections.size();
     if (targetNum == 0)
     {
-        TLLM_LOG_DEBUG("No targets to send KV cache to for request ID: %ld", llmRequest.mRequestId);
         return;
     }
 
@@ -413,13 +421,15 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
 
     SizeType32 const numKvPools = static_cast<SizeType32>(kvWindowSizes.size());
 
-    TLLM_LOG_DEBUG("CacheFormatter::format: allWindowSizes=%zu, kvWindowSizes=%d, numPools=%d, requestId=%lu",
-        allWindowSizes.size(), numKvPools, numPools, llmRequest.mRequestId);
+    TLLM_LOG_DEBUG("CacheFormatter::format: allWindowSizes=%zu, kvWindowSizes=%d, numPools=%d, requestId=%s",
+        allWindowSizes.size(), numKvPools, numPools,
+        llmRequest.has_value() ? std::to_string((*llmRequest)->mRequestId).c_str() : "<request-free>");
 
     bool layerWise = common::getEnvDisaggLayerwise() && numKvPools == 1;
     if (layerWise)
     {
-        auto& progress = llmRequest.getContextProgress();
+        TLLM_CHECK_WITH_INFO(llmRequest.has_value(), "LlmRequest required for layer-wise transfer");
+        auto& progress = (*llmRequest)->getContextProgress();
         SizeType32 const numLayers = blockManager.getNumLayers();
         runtime::ITensor::Shape offset = runtime::ITensor::makeShape({0, 0});
         for (SizeType32 layerIdx = 0; layerIdx < numLayers; layerIdx++)
@@ -515,8 +525,11 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
                     }
                 }
             }
-            TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.",
-                llmRequest.mRequestId);
+            if (llmRequest.has_value())
+            {
+                TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID: %ld.",
+                    (*llmRequest)->mRequestId);
+            }
 
             return;
         }
@@ -635,15 +648,20 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio
         sendHolder.release();
         session.setTime(TransferSession::kTimePostprocess);
     }
-    TLLM_LOG_DEBUG(
-        mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID:%ld ", llmRequest.mRequestId);
+    if (llmRequest.has_value())
+    {
+        TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "End the sending of KV cache for the request ID:%ld ",
+            (*llmRequest)->mRequestId);
+    }
 }
 
 void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& session)
 {
     NVTX3_SCOPED_RANGE(CacheFormatter_unformat);
     session.setTime(TransferSession::kTimeFormatter);
-    auto const& llmRequest = session.getLlmRequest();
+    auto llmRequestOpt = session.getLlmRequest();
+    TLLM_CHECK_WITH_INFO(llmRequestOpt.has_value(), "LlmRequest required for receiving KV cache");
+    auto const& llmRequest = **llmRequestOpt;
     auto const ctxReqId = llmRequest.getContextPhaseParams().value().getReqId();
     TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(),
         "Start receiving KV cache for request ID: %ld, context request ID: %ld.", llmRequest.mRequestId, ctxReqId);

cpp/tensorrt_llm/batch_manager/cacheFormatter.h

@@ -176,7 +176,7 @@ inline std::pair<std::vector<size_t>, std::vector<size_t>> pickRecvConnections(s
 
 namespace tensorrt_llm::batch_manager::kv_cache_manager
 {
-BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest const& llmRequest,
+BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, std::optional<LlmRequest const*> llmRequest,
     BlockKey const& lastBlockKey, SizeType32 indexFromEnd, bool recvSideHasCP = false, SizeType32 ppSize = 1);
 
 using DataContext = tensorrt_llm::executor::kv_cache::DataContext;

cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp

@@ -50,6 +50,7 @@
 #include "tensorrt_llm/common/tllmDataType.h"
 #include "tensorrt_llm/executor/cache_transmission/mpi_utils/connection.h"
 #include "tensorrt_llm/executor/dataTransceiverState.h"
+#include "tensorrt_llm/executor/serialization.h"
 #include "tensorrt_llm/executor/serializeUtils.h"
 #include "tensorrt_llm/runtime/utils/mpiUtils.h"
 #include "tensorrt_llm/runtime/utils/pgUtils.h"
@@ -591,6 +592,17 @@ void CacheTransceiver::initializeCommState()
     mCommState = std::addressof(mCacheSender->getCommState());
 }
 
+std::vector<char> CacheTransceiver::getSerializedDataTransceiverState() const
+{
+    TLLM_CHECK(mCommState != nullptr && mCacheState != nullptr);
+    executor::DataTransceiverState state;
+    state.setCommState(*mCommState);
+    state.setCacheState(*mCacheState);
+    // Only this API marks the state; context responses leave it unset.
+    state.setIsArbitraryTransferState(true);
+    return executor::Serialization::serialize(state);
+}
+
 void CacheTransceiver::setContextState(LlmRequest* llmRequest)
 {
     TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest());