diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 9d28fa26c4ee..62a7cbdf8edf 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -228,6 +228,12 @@ class BaseCacheTransceiver virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; + /// Get the serialized DataTransceiverState (CacheState + CommState) for this transceiver. + [[nodiscard]] virtual std::vector 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) override; + [[nodiscard]] std::vector getSerializedDataTransceiverState() const override; + [[nodiscard]] bool hasPoisonedTransferBuffer() const override; private: diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 01adf276f878..04b6c230e75a 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/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 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 findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, std::vector& pinnedBlockIds); + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys); //! \brief Unpin blocks by block ids directly void unpinBlocksById(std::vector 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 findBlocksInReuseTreeByBlockKeyImpl( + BlockKey const& blockKey, bool pinBlocks, std::vector& pinnedBlockIds); + //! \brief Walk the reuse tree with precomputed per-block keys (no lock; callers must hold mLookupTree->getMutex()). [[nodiscard]] std::shared_ptr searchReuseTree(std::vector const& blockKeys); @@ -1836,6 +1855,12 @@ class BlockManager return mWindowBlockManagers.at(windowSize).findBlocksInReuseTreeByBlockKey(blockKey); } + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, SizeType32 windowSize, std::vector& pinnedBlockIds) + { + return mWindowBlockManagers.at(windowSize).findBlocksInReuseTreeByBlockKey(blockKey, pinnedBlockIds); + } + [[nodiscard]] std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector 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 findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, SizeType32 windowSize, std::vector& pinnedBlockIds) + = 0; + [[nodiscard]] virtual std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys, SizeType32 windowSize) = 0; @@ -2650,6 +2680,12 @@ class KVCacheManager : public BaseKVCacheManager return mBlockManager.findBlocksInReuseTreeByBlockKey(blockKey, windowSize); } + std::shared_ptr findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, SizeType32 windowSize, std::vector& pinnedBlockIds) override + { + return mBlockManager.findBlocksInReuseTreeByBlockKey(blockKey, windowSize, pinnedBlockIds); + } + std::shared_ptr findBlocksInReuseTreeByBlockKeys( std::vector const& blockKeys, SizeType32 windowSize) override { diff --git a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h index b00d44d129e7..578e53b81dbf 100644 --- a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h +++ b/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 mCacheState; std::optional mCommState; + bool mIsArbitraryTransferState{false}; }; } // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 9dc44531d409..6165ec7386b4 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/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, 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(lastBlockKey.uniqueTokens.size()); - auto blockedUniqueTokens = chopVectorIntoBlocks( - 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(lastBlockKey.uniqueTokens.size()); + auto blockedUniqueTokens = chopVectorIntoBlocks( + 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(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() : ""); 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); diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h index 92d1b9d58a27..21356f3dd6c0 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @@ -176,7 +176,7 @@ inline std::pair, std::vector> pickRecvConnections(s namespace tensorrt_llm::batch_manager::kv_cache_manager { -BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, LlmRequest const& llmRequest, +BlockRange getBlockRangeForSending(BaseKVCacheManager* cacheManager, std::optional llmRequest, BlockKey const& lastBlockKey, SizeType32 indexFromEnd, bool recvSideHasCP = false, SizeType32 ppSize = 1); using DataContext = tensorrt_llm::executor::kv_cache::DataContext; diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index f730fb2aaf1e..02c101a43d5a 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/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 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()); diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 851f0abb04ef..e406d6620468 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -37,6 +37,7 @@ #include #include #include +#include namespace tensorrt_llm::batch_manager { @@ -81,6 +82,11 @@ void TransferSession::send(size_t idx, void const* data, size_t size) } catch (std::exception const& e) { + // Request-free (llmRequest-agnostic) transfer: there is no valid ID to attach. + if (mRequest == nullptr) + { + TLLM_THROW("%s", e.what()); + } throw common::RequestSpecificException( __FILE__, __LINE__, e.what(), mRequest->mRequestId, common::RequestErrorCode::kNETWORK_ERROR); } @@ -94,15 +100,23 @@ void TransferSession::recv(size_t idx, void* data, size_t size) } catch (std::exception const& e) { + // Request-free (llmRequest-agnostic) transfer: there is no valid ID to attach. + if (mRequest == nullptr) + { + TLLM_THROW("%s", e.what()); + } throw common::RequestSpecificException( __FILE__, __LINE__, e.what(), mRequest->mRequestId, common::RequestErrorCode::kNETWORK_ERROR); } } -LlmRequest const& TransferSession::getLlmRequest() const +std::optional TransferSession::getLlmRequest() const { - TLLM_CHECK(mRequest != nullptr); - return *mRequest; + if (mRequest == nullptr) + { + return std::nullopt; + } + return mRequest; } void TransferSession::setLlmRequest(LlmRequest const& llmRequest) @@ -171,7 +185,8 @@ void TransferSession::poisonReservedRecvBuffers() noexcept void TransferSession::exportMeasure(std::ofstream& outFile, bool isContext) const { - if (!mTimes || mTimes->measures.empty()) + // Request-free transfers are excluded: the exported row is keyed by the LlmRequest. + if (!mTimes || mTimes->measures.empty() || mRequest == nullptr) { return; } @@ -275,7 +290,7 @@ RequestInfo::RequestInfo(LlmRequest::RequestIdType requestId, executor::DataTran bool RequestInfo::operator==(RequestInfo const& rhs) const { return mRequestId == rhs.mRequestId && mIndexFromEnd == rhs.mIndexFromEnd && mLastBlockKey == rhs.mLastBlockKey - && mTransState == rhs.mTransState; + && mIsArbitraryTransfer == rhs.mIsArbitraryTransfer && mTransState == rhs.mTransState; } LlmRequest::RequestIdType RequestInfo::getRequestId() const noexcept @@ -294,6 +309,7 @@ void RequestInfo::serialize(RequestInfo const& requestInfo, std::ostream& os) su::serialize(requestInfo.mRequestId, os); su::serialize(requestInfo.mIndexFromEnd, os); su::serialize(requestInfo.mLastBlockKey, os); + su::serialize(requestInfo.mIsArbitraryTransfer, os); su::serialize(requestInfo.mTransState, os); } @@ -303,8 +319,11 @@ RequestInfo RequestInfo::deserialize(std::istream& is) auto requestId = su::deserialize(is); auto indexFromEnd = su::deserialize(is); auto lastBlockKey = su::deserialize(is); + auto isArbitraryTransfer = su::deserialize(is); auto transState = su::deserialize(is); - return RequestInfo{requestId, std::move(transState), indexFromEnd, lastBlockKey}; + auto requestInfo = RequestInfo{requestId, std::move(transState), indexFromEnd, lastBlockKey}; + requestInfo.setIsArbitraryTransfer(isArbitraryTransfer); + return requestInfo; } std::size_t RequestInfo::serializedSize(RequestInfo const& requestInfo) @@ -314,6 +333,7 @@ std::size_t RequestInfo::serializedSize(RequestInfo const& requestInfo) totalSize += su::serializedSize(requestInfo.mRequestId); totalSize += su::serializedSize(requestInfo.mIndexFromEnd); totalSize += su::serializedSize(requestInfo.mLastBlockKey); + totalSize += su::serializedSize(requestInfo.mIsArbitraryTransfer); totalSize += su::serializedSize(requestInfo.mTransState); return totalSize; } @@ -467,7 +487,6 @@ class CacheSender::Impl : mManager->recvConnect(DataContext{TransceiverTag::kID_TAG, mTerminate}, &id, sizeof(id)); if (connection == nullptr) { - TLLM_LOG_WARNING("recvRequestInfo connection is nullptr, maybe the server is terminating"); return std::nullopt; } @@ -625,10 +644,23 @@ class CacheSender::Impl private: struct Response { - // shared_ptr so this struct co-owns the request until the promise resolves; - // protects worker-side dereferences and the promise itself from premature destruction. - std::shared_ptr mRequest; + // An LlmRequest (co-owned until the promise resolves) for normal transfers, or + // just the request id for llmRequest-agnostic reuse-tree transfers. + std::variant, RequestIdType> mRequestOrId; std::promise mPromise; + std::vector mPinnedBlockIds; + + [[nodiscard]] LlmRequest* getRequest() const + { + auto const* request = std::get_if>(&mRequestOrId); + return request != nullptr ? request->get() : nullptr; + } + + [[nodiscard]] RequestIdType getRequestId() const + { + auto const* request = getRequest(); + return request != nullptr ? request->mRequestId : std::get(mRequestOrId); + } }; struct AsyncSendResource @@ -661,21 +693,44 @@ class CacheSender::Impl resp = std::move(resource.mSendQueue.front()); resource.mSendQueue.pop_front(); } - // Sequence the read before the move: argument initializations - // are indeterminately sequenced, so inlining resp.mRequest->... - // alongside std::move(resp) is UB once mRequest is a shared_ptr. - TLLM_CHECK(resp.mRequest != nullptr); - auto const reqId = resp.mRequest->mRequestId; - sendAndRemoveResponse(reqId, std::move(resp)); + // Read before std::move(resp): argument evaluations are indeterminately sequenced. + auto const requestId = resp.getRequestId(); + sendAndRemoveResponse(requestId, std::move(resp)); } } + //! Must not throw: called from noexcept send/failure paths. + void releasePinnedBlocks(Response& response) noexcept + { + if (response.mPinnedBlockIds.empty()) + { + return; + } + try + { + mCacheTransferLayer.getCacheManager()->unpinBlocksById(response.mPinnedBlockIds); + } + catch (std::exception const& err) + { + TLLM_LOG_ERROR("Failed to unpin reuse-tree blocks: %s", err.what()); + } + response.mPinnedBlockIds.clear(); + } + void sendAndRemoveResponse(RequestIdType id, Response resp) noexcept { try { TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - sendSync(*resp.mRequest); + if (auto* llmRequest = resp.getRequest(); llmRequest != nullptr) + { + sendSync(*llmRequest); + } + else + { + // Reuse tree path — no LlmRequest + sendSyncFromReuseTree(id); + } release(id); resp.mPromise.set_value(); } @@ -700,6 +755,7 @@ class CacheSender::Impl discardTransferState(id); failResponse(resp, exception); } + releasePinnedBlocks(resp); } void asyncSendAndRemoveResponse(RequestIdType id, Response resp) noexcept @@ -806,6 +862,37 @@ class CacheSender::Impl } } + void sendSyncFromReuseTree(RequestIdType requestId) + { + TransferSession* session = nullptr; + { + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + TLLM_CHECK(it != mRequestToSession.end()); + session = std::addressof(it->second); + } + // READY was already sent by response(); the receiver consumes exactly one per transfer. + mCacheTransferLayer.format(*session); + } + + // Pin the requested chain in the reuse tree; an empty result means no full match. + // The caller must unpin once the transfer settles. + std::vector pinReuseTreeBlocks(RequestIdType requestId) + { + std::unique_lock lk(mMtxForMap); + auto it = mRequestToSession.find(requestId); + auto const& lastBlockKey = it->second.getLastBlockKey(); + auto* cacheManager = mCacheTransferLayer.getCacheManager(); + auto windowSize = cacheManager->getBlockManager().getWindowSizesMetadata().begin()->first; + std::vector pinnedIds; + auto lastBlock = cacheManager->findBlocksInReuseTreeByBlockKey(lastBlockKey, windowSize, pinnedIds); + if (lastBlock == nullptr) + { + return {}; + } + return pinnedIds; + } + void response() noexcept { std::exception_ptr responseException; @@ -815,16 +902,13 @@ class CacheSender::Impl TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); while (true) { + if (mTerminate) { - std::unique_lock lock(mSenderMutex); - mSenderCv.wait(lock, - [this]() { return mTerminate || !mReadyResponses.empty() || !mCancelledRequests.empty(); }); - if (mTerminate) - { - break; - } + break; } + // Arbitrary transfers arrive without a pre-registered response; do not gate on + // mReadyResponses. auto requestInfo = recvRequestInfo(); if (!requestInfo.has_value() || mTerminate || !mManager->isRunning()) { @@ -837,22 +921,70 @@ class CacheSender::Impl mRemainSendCount[reqId] = getCounterpartsCount(reqId); } + if (requestInfo->isArbitraryTransfer()) { - std::unique_lock lock(mSenderMutex); - mCurrentRequest = reqId; - mSenderCv.wait(lock, - [this, reqId]() + // No LlmRequest will ever be registered; serve from the reuse tree off-thread. + { + std::scoped_lock lock(mSenderMutex); + mCurrentRequest = reqId; + } + auto countIt = mRemainSendCount.find(reqId); + auto const count = --countIt->second; + TLLM_CHECK(count >= 0); + if (count == 0) + { + mRemainSendCount.erase(countIt); + auto pinnedIds = pinReuseTreeBlocks(reqId); + if (pinnedIds.empty()) + { + TLLM_LOG_ERROR( + "Requested blocks do not exist in the source's reuse tree (request id: %lu). Notifying " + "receiver.", + reqId); + sendReadySignal(reqId, false); + discardTransferState(reqId); + } + else { - return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end() - || mCancelledRequests.find(reqId) != mCancelledRequests.end(); - }); - if (mTerminate) + sendReadySignal(reqId, true); + std::promise promise; + // Id-only response: the reuse-tree path has no LlmRequest. + Response resp{reqId, std::move(promise), std::move(pinnedIds)}; + if (dynamic_cast(mManager) != nullptr) + { + sendAndRemoveResponse(reqId, std::move(resp)); + } + else + { + asyncSendAndRemoveResponse(reqId, std::move(resp)); + } + } + } { + std::scoped_lock lock(mSenderMutex); mCurrentRequest = std::nullopt; - break; } } - sendResponse(reqId); + else + { + // The RequestInfo may race ahead of sendAsync; wait for the specific response. + { + std::unique_lock lock(mSenderMutex); + mCurrentRequest = reqId; + mSenderCv.wait(lock, + [this, reqId]() + { + return mTerminate || mReadyResponses.find(reqId) != mReadyResponses.end() + || mCancelledRequests.find(reqId) != mCancelledRequests.end(); + }); + if (mTerminate) + { + mCurrentRequest = std::nullopt; + break; + } + } + sendResponse(reqId); + } } } catch (std::exception const& err) @@ -930,6 +1062,7 @@ class CacheSender::Impl { TLLM_LOG_ERROR("Failed to set CacheSender response exception: %s", err.what()); } + releasePinnedBlocks(response); } void failPendingResponses(std::exception_ptr const& exception) noexcept @@ -1115,6 +1248,9 @@ class CacheReceiver::Impl requestInfo = RequestInfo(requestId, mSelfState, indexFromEnd, lastBlockKey); } } + // The state's provenance marks llmRequest-agnostic transfers: only + // getSerializedDataTransceiverState sets it; context responses leave it unset. + requestInfo.setIsArbitraryTransfer(contextState.isArbitraryTransferState()); auto* agentConnectionManager = dynamic_cast(mManager); std::vector recvHolders; diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index c1da646916dc..fffe55dceab5 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include @@ -121,7 +122,7 @@ class TransferSession void recv(size_t idx, void* data, size_t size); - [[nodiscard]] LlmRequest const& getLlmRequest() const; + [[nodiscard]] std::optional getLlmRequest() const; // in CacheSender, the LlmRequest is not available until the sendSync is called void setLlmRequest(LlmRequest const& llmRequest); @@ -234,6 +235,17 @@ class RequestInfo return mLastBlockKey; } + /// @brief Arbitrary (llmRequest-agnostic) transfer served from the sender's reuse tree. + [[nodiscard]] bool isArbitraryTransfer() const noexcept + { + return mIsArbitraryTransfer; + } + + void setIsArbitraryTransfer(bool isArbitraryTransfer) noexcept + { + mIsArbitraryTransfer = isArbitraryTransfer; + } + /// @brief Serialization. /// @param requestInfo Request information to be serialized. /// @param os The output stream to which the serialization result points. @@ -257,6 +269,9 @@ class RequestInfo // Last block key, used to derive other block keys on receiver BlockKey mLastBlockKey{}; + // True for arbitrary (llmRequest-agnostic) transfers served from the sender's reuse tree. + bool mIsArbitraryTransfer{false}; + // The state of the data transceiver. executor::DataTransceiverState mTransState; }; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 266cdafc8d06..8c1ffb70e372 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -937,6 +937,7 @@ bool BlockManager::verifyQueueIntegrity(SizeType32 windowSize) const bool WindowBlockManager::verifyQueueIntegrity() const { + std::lock_guard lock(mLookupTree->getMutex()); return mEvictionPolicy->verifyQueueIntegrity(); } @@ -1219,6 +1220,7 @@ void BlockManager::startScheduling() void WindowBlockManager::startScheduling() { + std::lock_guard lock(mLookupTree->getMutex()); mSchedulingNumFreeBlocks = mEvictionPolicy->getNumFreeBlocks(kPrimaryLevel); for (auto& [requestId, slotAllocatedBlocks] : mAllocatedBlocksPerSeq) { @@ -1238,6 +1240,7 @@ void WindowBlockManager::freeLeafBlock(BlockPtr const& block) void WindowBlockManager::releaseSubtree(BlockPtr const& block) { + std::lock_guard lock(mLookupTree->getMutex()); // Iterative pre-order DFS over `block` and all its descendants. Collect // first, then detach in reverse order: cascade-prune in freeLeafBlock() // removes empty parent nodes from the trie, so leaves must be detached @@ -1275,6 +1278,7 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: std::optional durationMs, executor::KvCacheTransferMode mode, std::string const& directory, bool wantPlaceholder) { + std::lock_guard lock(mLookupTree->getMutex()); // eviction policy get free primary block auto [block, canOffload] = mEvictionPolicy->getFreeBlock(kPrimaryLevel, wantPlaceholder); if (block->getUniqueTokens().empty()) @@ -1398,6 +1402,7 @@ void BlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& off void WindowBlockManager::onboardBlock(GenerationRequest& sequence, BlockPtr const& offloadBlock, executor::KvCacheTransferMode mode, std::string const& directory) { + std::lock_guard lock(mLookupTree->getMutex()); if (!offloadBlock->isPlaceholder() && !offloadBlock->isPrimary()) { auto block = getFreeBlock( @@ -1426,6 +1431,7 @@ void BlockManager::offloadBlock( void WindowBlockManager::offloadBlock( BlockPtr const& block, executor::KvCacheTransferMode mode, std::string const& directory) { + std::lock_guard lock(mLookupTree->getMutex()); // The current default behavior is to offload the out-of-window block // to secondary block pool to allow more free primary blocks for reuse. // However, such behavior does not take account whether the offloaded @@ -2131,6 +2137,19 @@ bool WindowBlockManager::blockInRadixTree(BlockPtr const& block) } std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKey(BlockKey const& blockKey) +{ + std::vector unusedPinnedBlockIds; + return findBlocksInReuseTreeByBlockKeyImpl(blockKey, /*pinBlocks=*/false, unusedPinnedBlockIds); +} + +std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKey( + BlockKey const& blockKey, std::vector& pinnedBlockIds) +{ + return findBlocksInReuseTreeByBlockKeyImpl(blockKey, /*pinBlocks=*/true, pinnedBlockIds); +} + +std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKeyImpl( + BlockKey const& blockKey, bool pinBlocks, std::vector& pinnedBlockIds) { std::lock_guard lock(mLookupTree->getMutex()); auto blockedUniqueTokens @@ -2143,7 +2162,42 @@ std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKe blockKeys.emplace_back(blockKey.usesExtraIds, blockKey.loraTaskId, blockedUniqueTokensList, blockKey.extraKeys, blockKey.cacheSalt); } - return searchReuseTree(blockKeys); + auto searchRoot = mCachedBlocksRoot; + std::vector pinnedInScope; + for (auto const& blockKey : blockKeys) + { + auto [partialMatch, numMatched, matchingBlock] = searchRoot != nullptr + ? searchRoot->findMatchingBlock(blockKey, true, true) + : std::make_tuple(false, 0, nullptr); + + // A prefix-only match lacks KV for the requested tail, and pinning (transfer) + // lookups read primary-pool buffers directly, so offloaded and placeholder blocks + // are misses too (isPlaceholder first: isPrimary asserts on placeholders). + bool const fullyMatched + = matchingBlock != nullptr && numMatched == static_cast(blockKey.uniqueTokens.size()); + bool const transferable + = fullyMatched && (!pinBlocks || (!matchingBlock->isPlaceholder() && matchingBlock->isPrimary())); + if (!transferable) + { + // Roll back pins taken during the partial walk. + for (auto const& block : pinnedInScope) + { + unpinBlock(block); + } + pinnedBlockIds.clear(); + return nullptr; + } + + if (pinBlocks) + { + pinBlock(matchingBlock); + pinnedInScope.push_back(matchingBlock); + pinnedBlockIds.push_back(matchingBlock->getBlockId()); + } + + searchRoot = std::move(matchingBlock); + } + return searchRoot; } std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKeys( @@ -2192,6 +2246,7 @@ void BlockManager::refreshBlocks() void WindowBlockManager::refreshBlocks() { + std::lock_guard lock(mLookupTree->getMutex()); mEvictionPolicy->refresh(); mTransferManager->syncTransfers(); } @@ -2322,6 +2377,7 @@ bool BlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequ bool WindowBlockManager::tryAllocatePlaceholderForLinearAttention(GenerationRequest& sequence, bool shareAmongBeams) { + std::lock_guard lock(mLookupTree->getMutex()); auto const beamWidth = sequence.getBeamWidth(); auto const newBlockIdx = sequence.getCacheBlockIds(mWindowSize).at(0).size(); // The first block is not a placeholder. @@ -2411,6 +2467,7 @@ bool WindowBlockManager::tryAllocatePlaceholderForLinearAttention(GenerationRequ void WindowBlockManager::allocateBlock(GenerationRequest& sequence, bool shareAmongBeams) { + std::lock_guard lock(mLookupTree->getMutex()); auto const beamWidth = sequence.getBeamWidth(); auto const requiredBlocks = shareAmongBeams ? 1 : beamWidth; @@ -2692,16 +2749,7 @@ std::pair> WindowBlockManager::sto if (pinBlocks) { - // If the block has no refs it sits in the eviction policy's free - // queue. Claim it first so that the later unpinBlocksById / - // releaseBlock cycle does not create a duplicate queue entry. - // Pass the block's existing priority and duration so that - // claimBlock does not clear its retention/expiry metadata. - if (!prevBlock->hasRefs()) - { - mEvictionPolicy->claimBlock(prevBlock, prevBlock->getPriority(), prevBlock->getDurationMs()); - } - prevBlock->incRefCount(); + pinBlock(prevBlock); pinnedBlockIds.push_back(prevBlock->getBlockId()); } } @@ -2736,6 +2784,7 @@ void BlockManager::replaceSharedBlock(GenerationRequest& sequence, SizeType32 wi void WindowBlockManager::replaceSharedBlock(GenerationRequest& sequence, SizeType32 blockIdx) { + std::lock_guard lock(mLookupTree->getMutex()); auto const requestId = sequence.getRequestId(); auto const beamWidth = sequence.getBeamWidth(); auto& allocatedBlocks = mAllocatedBlocksPerSeq.at(requestId); @@ -2787,6 +2836,7 @@ void BlockManager::releaseLastBlock(GenerationRequest& sequence, SizeType32 wind void WindowBlockManager::releaseLastBlock(GenerationRequest& sequence) { + std::lock_guard lock(mLookupTree->getMutex()); if (isRecurrentState()) { // In recurrent state, the last block always contains the current state and should not be released. @@ -2835,6 +2885,7 @@ void WindowBlockManager::releaseLastBlock(GenerationRequest& sequence) [[nodiscard]] SizeType32 WindowBlockManager::getNumFreeBlocks() const { + std::lock_guard lock(mLookupTree->getMutex()); auto const numFree = mEvictionPolicy->getNumFreeBlocks(kPrimaryLevel); TLLM_CHECK_WITH_INFO(numFree <= getMaxNumBlocks(), "%s::getNumFreeBlocks - primary free block count (%d) exceeds total block count (%d). " @@ -2846,6 +2897,7 @@ void WindowBlockManager::releaseLastBlock(GenerationRequest& sequence) [[nodiscard]] SizeType32 WindowBlockManager::getNumFreeSecondaryBlocks() const noexcept { + std::lock_guard lock(mLookupTree->getMutex()); return mEvictionPolicy->getNumFreeBlocks(kSecondaryLevel); } @@ -2976,18 +3028,41 @@ void BlockManager::unpinBlocksById(std::vector const& bloc firstManager.unpinBlocksById(blockIds); } +void WindowBlockManager::pinBlock(BlockPtr const& block) +{ + std::lock_guard lock(mLookupTree->getMutex()); + // Claim free blocks out of the eviction queue first, keeping their retention metadata. + if (!block->hasRefs()) + { + mEvictionPolicy->claimBlock(block, block->getPriority(), block->getDurationMs()); + } + block->incRefCount(); +} + +void WindowBlockManager::unpinBlock(BlockPtr const& block) +{ + std::lock_guard lock(mLookupTree->getMutex()); + block->decRefCount(); + if (!block->hasRefs()) + { + mEvictionPolicy->releaseBlock(block); + } +} + void WindowBlockManager::pinBlocks(GenerationRequest& sequence) { + std::lock_guard lock(mLookupTree->getMutex()); auto const requestId = sequence.getRequestId(); auto& allocatedBlocks = mAllocatedBlocksPerSeq.at(requestId); for (auto& block : allocatedBlocks) { - block->incRefCount(); + pinBlock(block); } } void WindowBlockManager::unpinBlocksById(std::vector const& blockIds) { + std::lock_guard lock(mLookupTree->getMutex()); if (blockIds.empty()) { return; @@ -3000,11 +3075,7 @@ void WindowBlockManager::unpinBlocksById(std::vector const auto block = mAllBlocksById[blockId]; if (block && block->getBlockId() != KVCacheBlock::kCachedBlocksRootId) { - block->decRefCount(); - if (!block->hasRefs()) - { - mEvictionPolicy->releaseBlock(block); - } + unpinBlock(block); } } } @@ -3020,6 +3091,7 @@ void BlockManager::storeNewBlock(GenerationRequest& sequence, OptionalRef llmRequest) { + std::lock_guard lock(mLookupTree->getMutex()); auto constexpr beamIdx = 0; auto const& uniqueTokens = llmRequest->getUniqueTokens(beamIdx); @@ -3126,6 +3198,7 @@ std::vector WindowBlockManager::storeBlocksForReuse( std::optional WindowBlockManager::releaseBlocks( GenerationRequest& sequence, OptionalRef llmRequest) { + std::lock_guard lock(mLookupTree->getMutex()); auto const requestId = sequence.getRequestId(); TLLM_LOG_DEBUG("%s::releaseBlocks - requestId=%lu, llmRequest.id=%s", mLogPrefix.c_str(), requestId, llmRequest.has_value() ? std::to_string(llmRequest->mRequestId).c_str() : "null"); @@ -3712,6 +3785,7 @@ bool KVCacheManager::copyLinearAttentionBlockBatch(std::vector lock(mLookupTree->getMutex()); // streamLLM is not supported at the moment. The out of window block will // always be the 0th block. TLLM_CHECK_WITH_INFO( diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index 24b7182fe628..2631df9aa9a2 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -124,9 +124,13 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses { NVTX3_SCOPED_RANGE(MLACacheFormatter_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 supports beam width 1."); + } auto const& selfConfig = session.getSelfState().getCacheState().value(); auto const& destConfig = session.getOtherState().getCacheState().value(); auto const selfIdx = session.getSelfState().getCommState().value().getSelfIdx(); @@ -134,7 +138,6 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses auto const& lastBlockKey = session.getLastBlockKey(); auto const& connections = session.getConnections(); auto& bufferManager = session.getBufferManager(); - TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently only supports beam width 1."); TLLM_CHECK(!connections.empty()); if (!needSendCache(selfConfig, destConfig, selfIdx)) { @@ -145,7 +148,6 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses auto targetNum = pickUpConnections.size(); if (targetNum == 0) { - TLLM_LOG_DEBUG("No targets to send KV cache to for request ID: %ld", llmRequest.mRequestId); return; } @@ -219,8 +221,11 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses } } - 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; } @@ -412,15 +417,20 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses session.setTime(TransferSession::kTimeTransmissions); 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 MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& session) { NVTX3_SCOPED_RANGE(MLACacheFormatter_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; TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently only supports beam width 1."); auto const ctxReqId = llmRequest.getContextPhaseParams().value().getReqId(); TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp index c58733f25885..b7ceb727efa3 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cpp @@ -47,7 +47,9 @@ void RnnCacheFormatter::format(TransferSession& session) NVTX3_SCOPED_RANGE(RnnCacheFormatter_formatUnifiedPool); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); + auto llmRequestOpt = session.getLlmRequest(); + TLLM_CHECK_WITH_INFO(llmRequestOpt.has_value(), "LlmRequest required for RNN state transfer"); + auto const& llmRequest = **llmRequestOpt; TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start sending unified pool RNN state for request ID: %ld.", llmRequest.mRequestId); TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); @@ -79,7 +81,7 @@ void RnnCacheFormatter::format(TransferSession& session) bool const recvSideHasCP = destConfig.getParallelConfig().mContextParallelism > 1; auto const indexFromEnd = session.getIndexFromEnd(); auto blockRange = kv_cache_manager::getBlockRangeForSending( - mKvCacheManager, llmRequest, lastBlockKey, indexFromEnd, recvSideHasCP, ppSize); + mKvCacheManager, llmRequestOpt, lastBlockKey, indexFromEnd, recvSideHasCP, ppSize); auto const& blockIdsPerWindow = blockRange.getBlockIdsPerWindow(); auto const allWindowSizes = blockRange.getWindowSizes(); @@ -251,7 +253,9 @@ void RnnCacheFormatter::unformat(TransferSession& session) NVTX3_SCOPED_RANGE(RnnCacheFormatter_unformatUnifiedPool); session.setTime(TransferSession::kTimeFormatter); - auto const& llmRequest = session.getLlmRequest(); + auto llmRequestOpt = session.getLlmRequest(); + TLLM_CHECK_WITH_INFO(llmRequestOpt.has_value(), "LlmRequest required for RNN state transfer"); + auto const& llmRequest = **llmRequestOpt; TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "Start receiving unified pool RNN state for request ID: %ld.", llmRequest.mRequestId); TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); diff --git a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp index 4ad1e7bffc86..3cd388625640 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp @@ -611,23 +611,41 @@ UcxConnection::ConnectionIdType UcxConnectionManager::getNewConnectionId(std::sh Connection const* UcxConnectionManager::recvConnect(DataContext const& ctx, void* data, size_t size) { - std::vector buffer(size + sizeof(UcxConnection::ConnectionIdType)); - std::promise promise; - std::future future = promise.get_future(); - auto completionCallback = [&](ucs_status_t, ucxx::RequestCallbackUserData) -> void { promise.set_value(); }; + // Co-owned by the completion callback: the terminate path returns while the + // cancelled request may still write the buffer and fire the callback. + struct RecvState + { + std::vector mBuffer; + std::promise mPromise; + }; - std::shared_ptr req = mWorkersPool.front()->tagRecv( - buffer.data(), buffer.size(), ucxx::Tag(ctx.getTag()), ucxx::TagMask(0xFFFFFFFF), false, completionCallback); + auto state = std::make_shared(); + state->mBuffer.resize(size + sizeof(UcxConnection::ConnectionIdType)); + std::future future = state->mPromise.get_future(); + auto completionCallback + = [state](ucs_status_t, ucxx::RequestCallbackUserData) -> void { state->mPromise.set_value(); }; + + std::shared_ptr req = mWorkersPool.front()->tagRecv(state->mBuffer.data(), state->mBuffer.size(), + ucxx::Tag(ctx.getTag()), ucxx::TagMask(0xFFFFFFFF), false, completionCallback); if (!req->isCompleted()) { - future.get(); + // Poll with timeout to allow checking the terminate flag + auto const& terminate = ctx.getTransferTerminate(); + while (future.wait_for(std::chrono::milliseconds(100)) != std::future_status::ready) + { + if (terminate.load()) + { + req->cancel(); + return nullptr; + } + } } TLLM_CHECK_WITH_INFO(req->isCompleted(), "recv SendConnectionId should be completed"); req->checkError(); - memcpy(data, buffer.data(), size); + memcpy(data, state->mBuffer.data(), size); UcxConnection::ConnectionIdType connectionId - = *reinterpret_cast(buffer.data() + size); + = *reinterpret_cast(state->mBuffer.data() + size); std::scoped_lock lock(mConnectionsMutex, mConnectionFuturesMutex); TLLM_CHECK_WITH_INFO(mConnectionFutures.find(connectionId) != mConnectionFutures.end(), "connectionFuture not found In recvConnect connectionId : %lu , worldRank: %d", connectionId, mRank); @@ -642,7 +660,7 @@ Connection const* UcxConnectionManager::recvConnect(DataContext const& ctx, void TLLM_CHECK(!mConnections[connectionId]->isFromRequester()); TLLM_LOG_DEBUG(mRank, "recvConnect connectionId: %lu , sendIDData:%lu", connectionId, - *reinterpret_cast(buffer.data())); + *reinterpret_cast(state->mBuffer.data())); return mConnections[connectionId].get(); } diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index 4c2998c25d26..02077b6857ad 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -765,6 +765,8 @@ DataTransceiverState Serialization::deserializeDataTransceiverState(std::istream { state.setCacheState(std::move(cacheState).value()); } + auto isArbitraryTransferState = su::deserialize(is); + state.setIsArbitraryTransferState(isArbitraryTransferState); return state; } @@ -772,6 +774,7 @@ void Serialization::serialize(DataTransceiverState const& state, std::ostream& o { su::serialize(state.mCommState, os); su::serialize(state.mCacheState, os); + su::serialize(state.mIsArbitraryTransferState, os); } std::vector Serialization::serialize(DataTransceiverState const& state) @@ -790,6 +793,7 @@ size_t Serialization::serializedSize(DataTransceiverState const& state) size_t totalSize = 0; totalSize += su::serializedSize(state.mCommState); totalSize += su::serializedSize(state.mCacheState); + totalSize += su::serializedSize(state.mIsArbitraryTransferState); return totalSize; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index d11838f68af5..2268e6f8ed03 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -92,6 +92,12 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) .def("request_and_receive_sync", &BaseCacheTransceiver::requestAndReceiveSync, nb::call_guard()) .def("request_and_receive_async", &BaseCacheTransceiver::requestAndReceiveAsync) + .def("get_serialized_data_transceiver_state", + [](tb::BaseCacheTransceiver& self) + { + auto serialized = self.getSerializedDataTransceiverState(); + return nb::bytes(serialized.data(), serialized.size()); + }) .def( "check_context_transfer_status", [](tb::BaseCacheTransceiver& self, std::optional const& atLeastRequestNum, bool markComplete = false) diff --git a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp index e294cdff9d49..37283750d114 100644 --- a/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp +++ b/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp @@ -194,7 +194,7 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- return mWorldSize; } - void setUpCacheManager() + void setUpCacheManager(bool enableBlockReuse = false) { auto constexpr numLayers = 4; auto constexpr numHeads = 2; @@ -216,7 +216,6 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- auto totalNumBlocks = mMaxNumSequences * numBlocksPerSeq; auto constexpr blocksInSecondaryPool = 0; - auto constexpr enableBlockReuse = false; auto constexpr dataType = tensorrt_llm::DataType::kFLOAT; using BlocksPerWindow = std::map>; @@ -330,6 +329,21 @@ class SymmetricalCacheTest : public ::testing::Test // NOLINT(cppcoreguidelines- return std::make_unique(mRequestId++, std::move(request)); } + // Generation-only request whose DataTransceiverState carries the arbitrary-transfer + // provenance marker, as getSerializedDataTransceiverState would produce. + auto makeArbitraryLlmRequest(SizeType32 length, LlmRequest::RequestIdType arbitraryId) + { + constexpr SizeType32 maxNewTokens{1}; + texec::Request request{VecTokens(length, length), maxNewTokens}; + auto state = std::make_unique(); + state->setCommState(*mContextCommState); + state->setCacheState(*mCacheState); + state->setIsArbitraryTransferState(true); + auto stats = texec::ContextPhaseParams({}, arbitraryId, state.release(), std::nullopt); + request.setContextPhaseParams(std::move(stats)); + return std::make_unique(arbitraryId, std::move(request)); + } + void addRequestAndTransportCache(std::shared_ptr const& llmRequest) { auto constexpr beamIdx{0}; @@ -427,6 +441,79 @@ TEST_F(SymmetricalCacheTest, SimpleTest) } } +TEST_F(SymmetricalCacheTest, ArbitraryTransferTest) +{ + auto worldSize = setUpCommunicator(); + if (worldSize != 2) + { + GTEST_SKIP() << "mpirun 2 processes is required to run this test."; + } + setUpCacheManager(/*enableBlockReuse=*/true); + setUpCacheTransceiver(); + + // 3 full blocks (tokensPerBlock = 8) so every requested chunk fully matches a stored block. + constexpr SizeType32 promptLen = 24; + constexpr SizeType32 missPromptLen = 40; + constexpr LlmRequest::RequestIdType arbitraryId = 4242; + auto constexpr beamIdx{0}; + auto constexpr beamWidth{1}; + + if (isSender) + { + // Store a patterned request in the reuse tree; no LlmRequest exists on the + // sender for the transfers below. + auto request = makeLlmRequest(promptLen); + mManager->addSequenceBatch( + {{{request->mRequestId, request->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*request)}); + auto blockRange = BlockRange::fromAllBlockIds(*mManager, request->mRequestId); + for (auto const& windowSize : blockRange.getWindowSizes()) + { + auto blockRangeForWindow = blockRange.getBlockRangeForWindow(windowSize); + for (auto it = blockRangeForWindow.begin(); it != blockRangeForWindow.end(); ++it) + { + TLLM_CUDA_CHECK(cudaMemset(it->data(), request->getPromptLen(), it->getSizeInBytes())); + } + } + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*request); + // A completed context request carries one generated token beyond the prompt; + // without it, storeBlocksForReuse drops the trailing prompt token and the + // final block never enters the reuse tree. + request->addNewToken(0, beamIdx); + mManager->removeSequence(request->mRequestId, request); + } + else + { + // Hit: the requested tokens are stored in the sender's reuse tree. + std::shared_ptr request = makeArbitraryLlmRequest(promptLen, arbitraryId); + mManager->addSequenceBatch( + {{{request->mRequestId, request->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*request)}); + auto future = mRequester->receiveAsync(request); + future.get(); + TLLM_CUDA_CHECK(cudaDeviceSynchronize()); + auto blockRange = BlockRange::fromAllBlockIds(*mManager, request->mRequestId); + for (auto const& windowSize : blockRange.getWindowSizes()) + { + auto blockRangeForWindow = blockRange.getBlockRangeForWindow(windowSize); + for (auto it = blockRangeForWindow.begin(); it != blockRangeForWindow.end(); ++it) + { + std::vector bytes(it->getSizeInBytes()); + TLLM_CUDA_CHECK(cudaMemcpy(bytes.data(), it->data(), it->getSizeInBytes(), cudaMemcpyDeviceToHost)); + EXPECT_TRUE(std::all_of(bytes.begin(), bytes.end(), [](uint8_t i) { return i == (promptLen & 0xff); })); + } + } + + // Miss: tokens never stored on the sender are rejected; the rejection + // surfaces as an exception on the receive future. + std::shared_ptr missRequest = makeArbitraryLlmRequest(missPromptLen, arbitraryId + 1); + mManager->addSequenceBatch( + {{{missRequest->mRequestId, missRequest->getNumTokens(beamIdx), beamWidth}}}, {std::ref(*missRequest)}); + auto missFuture = mRequester->receiveAsync(missRequest); + EXPECT_THROW(missFuture.get(), tensorrt_llm::common::TllmException); + } + // The sender must not tear down while the receiver's transfers are in flight. + tensorrt_llm::mpi::MpiComm::world().barrier(); +} + #if ENABLE_MULTI_DEVICE using AsymmetricTestParam = std::tuple Dict[str, Any]: def commit_blocks_for_reuse(self, req: LlmRequest) -> None: """Commit received KV blocks to the radix tree for prefix reuse. No-op by default.""" + def get_data_transceiver_state(self) -> bytes: + """Get the serialized DataTransceiverState (CacheState + CommState).""" + return b"" + def shutdown(self): """Shut down the transceiver and release registered resources.""" @@ -335,6 +339,9 @@ def prepare_context_requests(self, requests: List[LlmRequest]): # not implemented, an empty placeholder to allow being invoked unconditionally ... + def get_data_transceiver_state(self) -> bytes: + return self.impl.get_serialized_data_transceiver_state() + def get_disaggregated_params(self): # Cpp kv cache transceiver will set the disaggregated params to context response # Only new py cache transceiver will support gen-first disagg diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index d83c0251d6de..3d11d235bdde 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -52,6 +52,7 @@ is_llm_response) if TYPE_CHECKING: + from .._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver from ..disaggregated_params import DisaggregatedParams __all__ = [ @@ -960,6 +961,16 @@ def get_disaggregated_params(self) -> dict: return {} return self.engine.kv_cache_transceiver.get_disaggregated_params() + def get_cache_transceiver(self) -> Optional["KvCacheTransceiver"]: + if self.engine is None: + return None + return self.engine.kv_cache_transceiver + + def get_data_transceiver_state(self) -> bytes: + if self.engine is None or self.engine.kv_cache_transceiver is None: + return b"" + return self.engine.kv_cache_transceiver.get_data_transceiver_state() + @staticmethod def _stats_serializer(stats) -> str: # Per-rank path: stats is ("per_rank_dict", {..., "rank": N}). diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index ba5b843ff347..8863862b3b1b 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -41,6 +41,7 @@ from .utils import IntraProcessQueue, ProcessPoolExecutorSession, RequestError if TYPE_CHECKING: + from .._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver from .proxy import GenerationExecutorProxy from .worker import GenerationExecutorWorker @@ -462,6 +463,12 @@ def aget_kv_events(self, timeout=None) -> IterationResult: def get_disaggregated_params(self) -> dict: return {} + def get_cache_transceiver(self) -> Optional["KvCacheTransceiver"]: + return None + + def get_data_transceiver_state(self) -> bytes: + return b"" + @staticmethod def _create_ray_executor( worker_kwargs: Dict, diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 07d0ea0414fa..f1ca4a49e02c 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -848,6 +848,16 @@ def get_disaggregated_params(self) -> dict: logger.warning(f"Error fetching disaggregated params via RPC: {e}") return {} + def get_data_transceiver_state(self) -> bytes: + """Get serialized DataTransceiverState from worker runtime via RPC.""" + if self.rpc_client is None: + return b"" + try: + return self.rpc_client.get_data_transceiver_state().remote() + except RPCError as e: + logger.error(f"Error fetching data transceiver state via RPC: {e}") + raise + def aget_stats(self, timeout: float) -> IterationResult: """Get iteration statistics from the runtime via RPC (async). diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index 680b8d5c84e4..0763e09ad032 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/tensorrt_llm/executor/rpc_proxy.py @@ -254,6 +254,10 @@ def collective_rpc( return [remote_call.remote_future()] return [remote_call.remote()] + def get_data_transceiver_state(self) -> bytes: + """Get serialized DataTransceiverState from worker runtime via RPC.""" + return self.rpc_client.get_data_transceiver_state().remote() + def setup_engine_remote(self): return self.rpc_client.setup_engine().remote(need_response=True) diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 3d606e152267..f590acb717bd 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -324,6 +324,17 @@ def llm_id(self) -> str: return self._llm_id + @set_api_status("prototype") + def get_data_transceiver_state(self) -> bytes: + """Get the serialized DataTransceiverState for arbitrary KV cache transfer. + + Returns: + bytes: Serialized DataTransceiverState, or empty bytes if no transceiver is configured. + """ + if self._executor is None: + return b"" + return self._executor.get_data_transceiver_state() + @property @set_api_status("beta") def disaggregated_params(self) -> dict: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 3e1aa0d644db..85764b05c42f 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -791,6 +791,9 @@ def register_routes(self): methods=["GET"]) self.app.add_api_route("/version", self.version, methods=["GET"]) self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) + self.app.add_api_route("/v1/data_transceiver_state", + self.data_transceiver_state, + methods=["GET"]) # TODO: the metrics endpoint only reports iteration stats, not the runtime stats for now self.app.add_api_route("/metrics", self.get_iteration_stats, @@ -1097,6 +1100,19 @@ def register_visual_gen_routes(self): self.delete_video, methods=["DELETE"]) + async def data_transceiver_state(self) -> JSONResponse: + """Return the serialized DataTransceiverState as base64-encoded JSON.""" + state = self.generator.get_data_transceiver_state() + if not state: + return JSONResponse( + status_code=404, + content={"error": "No transceiver state available"}) + import base64 + return JSONResponse(content={ + "data_transceiver_state": + base64.b64encode(state).decode("utf-8") + }) + async def health(self) -> Response: if self._check_health(): return Response(status_code=200) diff --git a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py index 2d5b5000005a..97634b686ea5 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py +++ b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py @@ -138,6 +138,7 @@ async def run_worker(kv_cache_config, print(f"Sending ready signal to main process") intercomm.send(intercomm.Get_rank(), dest=0, tag=MPI_READY) + last_cached_tokens = 0 print(f"Waiting for requests") while True: try: @@ -146,6 +147,15 @@ async def run_worker(kv_cache_config, if requests is None: break + # Handle special commands + if requests == "GET_STATE": + state = llm.get_data_transceiver_state() + intercomm.send(state, dest=0, tag=MPI_RESULT) + continue + if requests == "GET_LAST_CACHED_TOKENS": + intercomm.send(last_cached_tokens, dest=0, tag=MPI_RESULT) + continue + request_metas = [] futures = [] for i, request in enumerate(requests): @@ -214,8 +224,12 @@ async def run_worker(kv_cache_config, f"Worker {rank}: awaiting future {i}/{len(futures)}", flush=True) result = await future - print(f"Worker {rank}: got result {i}, sending", - flush=True) + last_cached_tokens = getattr( + result, 'cached_tokens', 0) + print( + f"Worker {rank}: got result {i}, " + f"cached_tokens={last_cached_tokens}, sending", + flush=True) intercomm.send(result.outputs, dest=0, tag=MPI_RESULT) @@ -612,6 +626,7 @@ def test_disaggregated_spec_dec_batch_slot_limit(model, spec_dec_model_path, @pytest.mark.parametrize("generation_overlap", [False, True]) def test_disaggregated_logprobs(model, generation_overlap): """Verify that logprobs propagate correctly from prefill to decode. + Ensures first_gen_log_probs is carried in DisaggregatedParams so the generation_only worker receives one logprob per token. """ @@ -956,5 +971,282 @@ def test_disaggregated_logits(model, generation_overlap): print("All workers terminated.") +@pytest.mark.parametrize("model", ["TinyLlama-1.1B-Chat-v1.0"]) +@pytest.mark.parametrize("generation_overlap", [False]) +def test_arbitrary_kv_cache_transfer(model, generation_overlap): + """Test KV cache transfer from the reuse tree. + + Flow: + 1. Worker 0 runs a normal generate to fill the KV cache reuse tree. + 2. Retrieve the data_transceiver_state from worker 0. + 3. Worker 1 sends a generation_only request using the state from step 2. + The sender (worker 0) serves blocks directly from its reuse tree. + """ + worker_pytorch_configs = [] + + # Worker 0 (sender) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig())) + + # Worker 1 (receiver) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=not generation_overlap, + cuda_graph_config=CudaGraphConfig())) + + kv_cache_configs = [ + KvCacheConfig(max_tokens=2048 * 8, enable_block_reuse=True) + for _ in range(2) + ] + cache_transceiver_configs = [ + CacheTransceiverConfig(backend="DEFAULT") for _ in range(2) + ] + model_names = [model_path(model) for _ in range(2)] + ranks = [0, 1] + worker_args = list( + zip(kv_cache_configs, cache_transceiver_configs, worker_pytorch_configs, + model_names, ranks)) + + port_name = mpi_publish_name() + + prompt = "What is the capital of Germany?" + + with MPIPoolExecutor(max_workers=2, + env={ + "UCX_TLS": "^ib,gdr_copy", + "UCX_MM_ERROR_HANDLING": "y" + }) as executor: + futures = [] + try: + for worker_arg in worker_args: + future = executor.submit(worker_entry_point, *worker_arg) + futures.append(future) + except Exception as e: + print(f"Error in worker {worker_arg}: {e}") + raise e + + intercomm = None + try: + print("Launched all the workers.", flush=True) + intercomm = mpi_initialize_intercomm(port_name) + + for _ in range(2): + intercomm.recv(tag=MPI_READY) + print("Received ready signal.") + + # Normal generate on worker 0 to fill KV cache reuse tree + print("Filling KV cache reuse tree on worker 0", flush=True) + requests = [(prompt, SamplingParams(max_tokens=10, + ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 0, intercomm) + assert len(responses) == 1 + print(f"Worker 0 output: {responses[0][0].text}", flush=True) + + # Get data_transceiver_state from worker 0 + print("Getting data_transceiver_state from worker 0", flush=True) + intercomm.send("GET_STATE", dest=0, tag=MPI_REQUEST) + state = intercomm.recv(source=0, tag=MPI_RESULT) + assert isinstance(state, bytes) + assert len(state) > 0 + print(f"Got data_transceiver_state: {len(state)} bytes", flush=True) + + # generation_only on worker 1 using state from worker 0. + # max_tokens=1 ensures only one decode step runs, so the test + # relies on the KV cache transfer delivering all prompt tokens. + print("Sending generation_only to worker 1", flush=True) + DISAGG_REQ_ID = 42 + disagg_params = DisaggregatedParams( + request_type="generation_only", + opaque_state=state, + first_gen_tokens=[0], + disagg_request_id=DISAGG_REQ_ID, + ) + requests = [(prompt, SamplingParams(max_tokens=1, ignore_eos=True), + disagg_params)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + output = responses[0][0] + print(f"Worker 1 output: {output.text}", flush=True) + print(f"Worker 1 token_ids: {output.token_ids}", flush=True) + assert len(output.token_ids) > 0, \ + "generation_only request should produce output tokens" + + # Send a normal (non-disagg) request to worker 1 to verify + # it still works properly after the KV cache transfer. + print("Normal request on worker 1 after transfer", flush=True) + requests = [(prompt, SamplingParams(max_tokens=10, + ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + output2 = responses[0][0] + print(f"Worker 1 normal output: {output2.text}", flush=True) + assert len(output2.token_ids) > 0, \ + "Normal request should produce output tokens after transfer" + + # Query cached_tokens from worker 1's normal request + intercomm.send("GET_LAST_CACHED_TOKENS", dest=1, tag=MPI_REQUEST) + cached1 = intercomm.recv(source=1, tag=MPI_RESULT) + print(f"Worker 1 normal cached_tokens: {cached1}", flush=True) + + # Send a completely different prompt to worker 1 to confirm + # cached_tokens is 0 when there is no reuse match. + different_prompt = "Bonjour le monde, comment allez-vous aujourd'hui" + print("Different prompt on worker 1 (no match expected)", + flush=True) + requests = [(different_prompt, + SamplingParams(max_tokens=5, ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + intercomm.send("GET_LAST_CACHED_TOKENS", dest=1, tag=MPI_REQUEST) + cached2 = intercomm.recv(source=1, tag=MPI_RESULT) + print(f"Worker 1 different prompt cached_tokens: {cached2}", + flush=True) + # BOS token may be shared, so at most 1 cached token is expected + assert cached2 <= 1, \ + f"Expected at most 1 cached token for unrelated prompt, got {cached2}" + assert cached2 < cached1, \ + "Unrelated prompt should have fewer cached tokens than transferred prompt" + + except Exception as e: + print(f"Exception encountered: {e}", flush=True) + raise + finally: + print("Sending termination request", flush=True) + mpi_send_termination_request(intercomm) + + print("Waiting for all workers to terminate. ", flush=True) + for future in futures: + future.result() + print("All workers terminated.") + + +@pytest.mark.parametrize("model", ["TinyLlama-1.1B-Chat-v1.0"]) +@pytest.mark.parametrize("generation_overlap", [False]) +def test_arbitrary_kv_cache_transfer_missing_blocks(model, generation_overlap): + """Test that missing-block transfers fail. + + When the receiver asks for blocks that don't exist on the sender, + the sender must notify the receiver so the request surfaces a + clear error instead of waiting for the hang detector. + + Flow: + 1. Worker 0 (sender) keeps an EMPTY KV cache reuse tree (no prior + generate call). Its reuse tree therefore has no blocks matching + any prompt the receiver might request. + 2. Retrieve the data_transceiver_state from worker 0. + 3. Worker 1 (receiver) sends a generation_only request using that + state. The sender attempts to look up the receiver's prompt + blocks in its empty reuse tree and fails with + "Couldn't find the requested block in the reuse tree". + 4. The receiver surfaces the failure as an error string. + """ + worker_pytorch_configs = [] + + # Worker 0 (sender) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig())) + + # Worker 1 (receiver) + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=not generation_overlap, + cuda_graph_config=CudaGraphConfig())) + + kv_cache_configs = [ + KvCacheConfig(max_tokens=2048 * 8, enable_block_reuse=True) + for _ in range(2) + ] + cache_transceiver_configs = [ + CacheTransceiverConfig(backend="DEFAULT") for _ in range(2) + ] + model_names = [model_path(model) for _ in range(2)] + ranks = [0, 1] + worker_args = list( + zip(kv_cache_configs, cache_transceiver_configs, worker_pytorch_configs, + model_names, ranks)) + + port_name = mpi_publish_name() + + prompt = "What is the capital of Germany?" + + with MPIPoolExecutor(max_workers=2, + env={ + "UCX_TLS": "^ib,gdr_copy", + "UCX_MM_ERROR_HANDLING": "y" + }) as executor: + futures = [] + try: + for worker_arg in worker_args: + future = executor.submit(worker_entry_point, *worker_arg) + futures.append(future) + except Exception as e: + print(f"Error in worker {worker_arg}: {e}") + raise e + + intercomm = None + try: + print("Launched all the workers.", flush=True) + intercomm = mpi_initialize_intercomm(port_name) + + for _ in range(2): + intercomm.recv(tag=MPI_READY) + print("Received ready signal.") + + # Get state from worker 0 + print("Getting data_transceiver_state from empty worker 0", + flush=True) + intercomm.send("GET_STATE", dest=0, tag=MPI_REQUEST) + state = intercomm.recv(source=0, tag=MPI_RESULT) + assert isinstance(state, bytes) + assert len(state) > 0 + print(f"Got data_transceiver_state: {len(state)} bytes", flush=True) + + # generation_only on worker 1 — sender has no blocks to serve. + print("Sending generation_only to worker 1 (sender has no blocks)", + flush=True) + DISAGG_REQ_ID = 43 + disagg_params = DisaggregatedParams( + request_type="generation_only", + opaque_state=state, + first_gen_tokens=[0], + disagg_request_id=DISAGG_REQ_ID, + ) + requests = [(prompt, SamplingParams(max_tokens=10, ignore_eos=True), + disagg_params)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + response = responses[0] + print(f"Worker 1 response: {response}", flush=True) + assert isinstance(response, str), \ + f"Expected error string when sender has no blocks, got: {response!r}" + assert "cache transfer" in response.lower(), \ + f"Expected cache-transfer error, got: {response!r}" + + # Verify worker 1 still works for normal (non-disagg) requests + # after the failed transfer attempt. + print("Normal request on worker 1 after failed transfer", + flush=True) + requests = [(prompt, SamplingParams(max_tokens=10, + ignore_eos=True), None)] + responses = send_requests_to_worker(requests, 1, intercomm) + assert len(responses) == 1 + output = responses[0][0] + print(f"Worker 1 normal output: {output.text}", flush=True) + assert len(output.token_ids) > 0, \ + "Worker 1 should still serve normal requests after failed transfer" + + except Exception as e: + print(f"Exception encountered: {e}", flush=True) + raise + finally: + print("Sending termination request", flush=True) + mpi_send_termination_request(intercomm) + + print("Waiting for all workers to terminate. ", flush=True) + for future in futures: + future.result() + print("All workers terminated.") + + if __name__ == "__main__": pytest.main() diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 02c43b89f8dd..f0855d268f9c 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -880,6 +880,8 @@ disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_perf_metrics[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_single_gpu[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer[False-TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer_missing_blocks[False-TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_cancel_gen_requests[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[False-TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index bf8c6d131d50..57027ff7feb5 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -104,6 +104,8 @@ l0_a10: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_llama[False-True-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_llama[True-False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_llama[True-True-TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer[False-TinyLlama-1.1B-Chat-v1.0] + - disaggregated/test_disaggregated_single_gpu.py::test_arbitrary_kv_cache_transfer_missing_blocks[False-TinyLlama-1.1B-Chat-v1.0] - test_e2e.py::test_get_ci_container_port - test_e2e.py::test_openai_chat_multimodal_example ISOLATION - test_e2e.py::test_openai_mmencoder_example diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index 2a82d11fe3ee..c9b9a42388a7 100644 --- a/tests/unittest/api_stability/api_stability_core.py +++ b/tests/unittest/api_stability/api_stability_core.py @@ -157,9 +157,18 @@ def from_inspect(cls, method: MethodType): return_annotation = eval(return_annotation) return cls(parameters, return_annotation) + @classmethod + def _strip_api_status_tag(cls, docstring: str) -> str: + """Strip the :tag:`...` prefix added by @set_api_status decorator.""" + if docstring and docstring.startswith(":tag:"): + import re + docstring = re.sub(r'^:tag:`[^`]*`\s*', '', docstring) + return docstring + @classmethod def from_docstring(cls, method: MethodType): - doc = docstring_parser.parse(method.__doc__) + docstring = cls._strip_api_status_tag(method.__doc__) + doc = docstring_parser.parse(docstring) parameters = {} for param in doc.params: if param.args[0] == 'param': diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 7ff208c41ce3..4f9ca95bf520 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -437,6 +437,10 @@ methods: default: 2 return_annotation: tensorrt_llm.executor.result.IterationResult status: beta + get_data_transceiver_state: + parameters: {} + return_annotation: bytes + status: prototype shutdown: parameters: {} return_annotation: None