diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 452b90462448..ea061af74db7 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -35,6 +35,7 @@ #include #include #include +#include #include using SizeType32 = tensorrt_llm::runtime::SizeType32; @@ -204,13 +205,15 @@ class BaseCacheTransceiver { public: virtual ~BaseCacheTransceiver() = default; - virtual void respondAndSendAsync(LlmRequest* llmRequest) = 0; + // Async entry points take shared_ptr so the async worker holds a strong + // reference for the transfer lifetime; see CacheTransceiver::mSenderFutures. + virtual void respondAndSendAsync(std::shared_ptr llmRequest) = 0; virtual void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) = 0; - virtual void requestAndReceiveSync(LlmRequest* llmRequest) = 0; - virtual void requestAndReceiveAsync(LlmRequest* llmRequest) = 0; + virtual void requestAndReceiveSync(std::shared_ptr llmRequest) = 0; + virtual void requestAndReceiveAsync(std::shared_ptr llmRequest) = 0; /// Check all requests transferring context, and return the requests that have completed or encountered an error. virtual RequestStatuses checkContextTransferStatus( @@ -221,7 +224,7 @@ class BaseCacheTransceiver [[nodiscard]] virtual bool checkGenTransferComplete() const = 0; - virtual bool cancelRequest(LlmRequest* llmRequest) = 0; + virtual bool cancelRequest(std::shared_ptr llmRequest) = 0; }; class CacheTransceiver : public BaseCacheTransceiver @@ -252,13 +255,13 @@ class CacheTransceiver : public BaseCacheTransceiver virtual ~CacheTransceiver(); - void respondAndSendAsync(LlmRequest* llmRequest) override; + void respondAndSendAsync(std::shared_ptr llmRequest) override; void respondAndSendLayerWise( RequestVector const& requests, std::shared_ptr const& progress) override; - void requestAndReceiveSync(LlmRequest* llmRequest) override; - void requestAndReceiveAsync(LlmRequest* llmRequest) override; + void requestAndReceiveSync(std::shared_ptr llmRequest) override; + void requestAndReceiveAsync(std::shared_ptr llmRequest) override; RequestStatuses checkContextTransferStatus( std::optional const& atLeastRequestNum = std::nullopt, bool markComplete = false) override; @@ -267,7 +270,7 @@ class CacheTransceiver : public BaseCacheTransceiver [[nodiscard]] bool checkGenTransferComplete() const override; - virtual bool cancelRequest(LlmRequest* llmRequest) override; + virtual bool cancelRequest(std::shared_ptr llmRequest) override; private: void initializeCommState(); @@ -276,8 +279,14 @@ class CacheTransceiver : public BaseCacheTransceiver std::unique_ptr mCacheSender; std::unique_ptr mCacheReceiver; - std::vector>> mSenderFutures; - std::vector>> mRequesterFutures; + // shared_ptr (not raw LlmRequest*) so the futures hold a strong reference for + // the transfer lifetime; otherwise Python's _terminate_request can drop the + // request while a C++ status check still dereferences it. + std::vector, std::future>> mSenderFutures; + std::vector, std::future>> mRequesterFutures; + // Dedup sets so observe-only timeout WARN logs fire at most once per stuck request. + std::unordered_set mTimedOutSenderIds; + std::unordered_set mTimedOutRequesterIds; mpi::MpiComm const* mMpiWorldComm{nullptr}; std::shared_ptr mGroupComm; diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 58092897ebbe..3be8ee22bc1f 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -26,6 +26,30 @@ namespace tensorrt_llm::batch_manager { +void BufferIndexHolder::release() noexcept +{ + if (!mHeld || mMgr == nullptr) + { + return; + } + try + { + if (mIsRecv) + { + mMgr->freeBufferIndexForRecv(mIndex); + } + else + { + mMgr->freeBufferIndexForSend(mIndex); + } + } + catch (...) + { + // noexcept: swallow so the destructor can never throw. + } + mHeld = false; +} + BaseTransBufferManager::BaseTransBufferManager( size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens) : mDataType{dataType} diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 1efeb89ccc04..2cbf9f514bd7 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -46,6 +46,87 @@ enum class BufferKind : uint8_t kRNN = 2 }; +class BaseTransBufferManager; + +/// @brief RAII holder for an index from BaseTransBufferManager::assignBufferIndexFor{Send,Recv}. +/// Releases on destruction (incl. exception unwind). Move-only; call release() on +/// the happy path or detach() when ownership is handed off downstream. +class BufferIndexHolder +{ +public: + BufferIndexHolder() = default; + + BufferIndexHolder(BaseTransBufferManager& mgr, std::optional index, bool isRecv) noexcept + : mMgr(&mgr) + , mIndex(index) + , mHeld(index.has_value()) + , mIsRecv(isRecv) + { + } + + // Defined out-of-line in baseTransBuffer.cpp: release() calls + // BaseTransBufferManager methods, whose full definition appears later + // in this header. + ~BufferIndexHolder() + { + release(); + } + + BufferIndexHolder(BufferIndexHolder const&) = delete; + BufferIndexHolder& operator=(BufferIndexHolder const&) = delete; + + BufferIndexHolder(BufferIndexHolder&& other) noexcept + : mMgr(other.mMgr) + , mIndex(other.mIndex) + , mHeld(other.mHeld) + , mIsRecv(other.mIsRecv) + { + other.mHeld = false; + } + + BufferIndexHolder& operator=(BufferIndexHolder&& other) noexcept + { + if (this != &other) + { + release(); + mMgr = other.mMgr; + mIndex = other.mIndex; + mHeld = other.mHeld; + mIsRecv = other.mIsRecv; + other.mHeld = false; + } + return *this; + } + + [[nodiscard]] std::optional index() const noexcept + { + return mIndex; + } + + [[nodiscard]] bool held() const noexcept + { + return mHeld; + } + + /// @brief Relinquish ownership without releasing. Use when a downstream + /// owner (e.g. the formatter inside receiveSync) takes over the + /// release responsibility on the happy path. + std::optional detach() noexcept + { + mHeld = false; + return mIndex; + } + + /// @brief Release the slot now and disarm the destructor. Safe to call multiple times. + void release() noexcept; + +private: + BaseTransBufferManager* mMgr{nullptr}; + std::optional mIndex{}; + bool mHeld{false}; + bool mIsRecv{true}; +}; + /// @brief Base class for cache transfer buffer management. /// Handles buffer pool allocation, index assignment, and slicing. /// Derived classes provide cache-specific size calculations. diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp index 17dd557be1a3..68a398f1b52d 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp @@ -523,6 +523,7 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio // 5. send the buffer to the corresponding target. Ideally, we send only once (one buffer) for each target. auto cacheBufferId = mCacheTransBufferManager->assignBufferIndexForSend(); + BufferIndexHolder sendHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/false); int peerDuplicateHeadFactor = targetInfo.mPeerDupHeadFactor; auto bufferTargetNum = targetNum / peerDuplicateHeadFactor; auto ppRank = selfIdx @@ -606,7 +607,7 @@ void CacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& sessio session.setTime(TransferSession::kTimeTransmissions); - mCacheTransBufferManager->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); session.setTime(TransferSession::kTimePostprocess); } TLLM_LOG_DEBUG( @@ -849,6 +850,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess size_t remainNoCoverTargetNum = 0; size_t bufferCoverTargetNum = 0; std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; { NVTX3_SCOPED_RANGE(formatInputAllocBuffer); @@ -864,6 +866,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess { cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv(); } + recvHolder = BufferIndexHolder(*mCacheTransBufferManager, cacheBufferId, /*isRecv=*/true); auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer] = mCacheTransBufferManager->getOrAllocateRecvBuffers( cacheBufferId, static_cast(targetNum), bufferEleSizes, bufferManager); @@ -997,10 +1000,7 @@ void CacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& sess recvSplitCaches, outputBuffersPerWindow, destConfig, selfConfig, selfIdx, bufferManager); bufferManager.getStream().synchronize(); - if (cacheBufferId.has_value()) - { - mCacheTransBufferManager->freeBufferIndexForRecv(cacheBufferId); - } + recvHolder.release(); } session.setTime(TransferSession::kTimePostprocess); } diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 42fa3d9eac7c..de146a7652de 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -379,7 +379,7 @@ void CacheTransceiver::setContextState(LlmRequest* llmRequest) } } -void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) +void CacheTransceiver::respondAndSendAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isContextOnlyRequest()); llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_TRANS_IN_PROGRESS); @@ -393,9 +393,9 @@ void CacheTransceiver::respondAndSendAsync(LlmRequest* llmRequest) } return; } - setContextState(llmRequest); + setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest, std::move(future)); + mSenderFutures.emplace_back(std::move(llmRequest), std::move(future)); } void CacheTransceiver::respondAndSendLayerWise( @@ -411,11 +411,11 @@ void CacheTransceiver::respondAndSendLayerWise( llmRequest->setState(LlmRequestState::kDISAGG_CONTEXT_INIT_AND_TRANS); setContextState(llmRequest.get()); auto future = mCacheSender->sendAsync(*llmRequest); - mSenderFutures.emplace_back(llmRequest.get(), std::move(future)); + mSenderFutures.emplace_back(llmRequest, std::move(future)); } } -void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveSync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); { @@ -425,21 +425,23 @@ void CacheTransceiver::requestAndReceiveSync(LlmRequest* llmRequest) llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); } -void CacheTransceiver::requestAndReceiveAsync(LlmRequest* llmRequest) +void CacheTransceiver::requestAndReceiveAsync(std::shared_ptr llmRequest) { TLLM_CHECK(llmRequest && llmRequest->isGenerationOnlyRequest()); + auto const requestId = llmRequest->mRequestId; if (std::find_if(mRequesterFutures.begin(), mRequesterFutures.end(), - [llmRequest](auto const& pair) { return pair.first->mRequestId == llmRequest->mRequestId; }) + [requestId](auto const& pair) { return pair.first->mRequestId == requestId; }) != mRequesterFutures.end()) { - TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", llmRequest->mRequestId); + TLLM_LOG_WARNING("Request ID %zu is already in mRequestFutures.", requestId); return; } auto future = mCacheReceiver->receiveAsync(*llmRequest); - mRequesterFutures.emplace_back(llmRequest, std::move(future)); - llmRequest->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); + auto* requestPtr = llmRequest.get(); + mRequesterFutures.emplace_back(std::move(llmRequest), std::move(future)); + requestPtr->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_IN_PROGRESS); } std::vector gatherRequestIds( @@ -546,6 +548,13 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { senderFutureTimeoutMs = mCacheTransceiverConfig->getKvTransferSenderFutureTimeoutMs(); } + // Observe-only: WARN per-request when the wall-clock transfer time exceeds + // kvTransferTimeoutMs. No cancellation, eviction, or state transition. + std::optional kvTransferTimeoutMs = std::nullopt; + if (mCacheTransceiverConfig.has_value()) + { + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + } auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupTPInDPComm : mGroupTensorParaComm; std::vector contextCompleteRequestIds; @@ -602,6 +611,19 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( for (auto it = mSenderFutures.begin(); it != mSenderFutures.end();) { auto& [request, future] = *it; + if (kvTransferTimeoutMs.has_value()) + { + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutSenderIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Context KV cache transfer for request %ld exceeded configured timeout: " + "elapsed %ld ms > limit %d ms (observe-only).", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); + } + } if (blockAll || (toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end())) { try @@ -616,6 +638,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( { request->setState(LlmRequestState::kDISAGG_CONTEXT_COMPLETE); } + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } else if (status == std::future_status::timeout) @@ -631,6 +654,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } @@ -640,6 +664,7 @@ RequestStatuses CacheTransceiver::checkContextTransferStatus( "Error occurred during context transfer for request %ld: %s", request->mRequestId, e.what()); request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); requestsStatus.errorRequestIds.insert(request->mRequestId); + mTimedOutSenderIds.erase(request->mRequestId); it = mSenderFutures.erase(it); } } @@ -765,40 +790,61 @@ void CacheTransceiver::checkGenTransferStatus(std::optional const& atLeastR " checkGenTransferStatus toCompleteIdSet size: %zu, atLeastRequestNum: %d ", toCompleteIdSet.size(), atLeastRequestNum.value_or(0)); } + // Observe-only: gen-side mirror of the context-side timeout WARN. + std::optional kvTransferTimeoutMs = std::nullopt; + if (mCacheTransceiverConfig.has_value()) + { + kvTransferTimeoutMs = mCacheTransceiverConfig->getKvTransferTimeoutMs(); + } for (auto it = mRequesterFutures.begin(); it != mRequesterFutures.end();) { - if (blockAll || toCompleteIdSet.find(it->first->mRequestId) != toCompleteIdSet.end()) + auto& request = it->first; + if (kvTransferTimeoutMs.has_value()) + { + auto elapsed = std::chrono::duration_cast( + LlmRequest::getSteadyClockNow() - request->getKvCacheTransferStart()); + auto elapsedMs = static_cast(elapsed.count()); + if (elapsedMs > kvTransferTimeoutMs.value() && mTimedOutRequesterIds.insert(request->mRequestId).second) + { + TLLM_LOG_WARNING( + "Generation KV cache transfer for request %ld exceeded configured timeout: " + "elapsed %ld ms > limit %d ms (observe-only).", + request->mRequestId, elapsedMs, kvTransferTimeoutMs.value()); + } + } + if (blockAll || toCompleteIdSet.find(request->mRequestId) != toCompleteIdSet.end()) { try { it->second.get(); - it->first->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); + request->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); // Gather the kv cache transfer time from all workers and update to leader rank if (!common::getEnvKVCacheTimeOutputPath().empty()) { - auto syncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; - updateKVCacheTransferBW(syncComm, it->first); + auto bwSyncComm = mCacheState->getParallelConfig().mEnableAttentionDP ? mGroupDataComm : mGroupComm; + updateKVCacheTransferBW(bwSyncComm, request.get()); } } catch (std::exception const& e) { TLLM_LOG_ERROR( - "Error occurred during generation transfer for request %ld: %s", it->first->mRequestId, e.what()); - it->first->setState(LlmRequestState::kDISAGG_TRANS_ERROR); + "Error occurred during generation transfer for request %ld: %s", request->mRequestId, e.what()); + request->setState(LlmRequestState::kDISAGG_TRANS_ERROR); } if (useMPI()) { TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + request->mRequestId, request->getContextPhaseParams().value().getReqId()); } else { TLLM_LOG_DEBUG(tensorrt_llm::pg_utils::get_world_pg()->getRank(), "**** it->first->mRequestId: %ld, context request ID: %ld ******** get feature ***", - it->first->mRequestId, it->first->getContextPhaseParams().value().getReqId()); + request->mRequestId, request->getContextPhaseParams().value().getReqId()); } + mTimedOutRequesterIds.erase(request->mRequestId); it = mRequesterFutures.erase(it); } else @@ -813,7 +859,7 @@ bool CacheTransceiver::checkGenTransferComplete() const return mRequesterFutures.empty(); } -bool CacheTransceiver::cancelRequest(LlmRequest* llmRequest) +bool CacheTransceiver::cancelRequest(std::shared_ptr llmRequest) { if (llmRequest->isContextOnlyRequest()) { diff --git a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp index c72090867f29..812323e92b81 100644 --- a/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp +++ b/cpp/tensorrt_llm/batch_manager/mlaCacheFormatter.cpp @@ -254,6 +254,8 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses }; auto bufferEleSizes = getBufferSizeForTarget(); auto cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForSend(); + BufferIndexHolder sendHolder( + *mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/false); auto result = mCacheTransBufferManagers[transferIndexerKCache]->getOrAllocateSendBuffers( cacheBufferId, static_cast(pPDomainSize * cPDomainSize), bufferEleSizes, bufferManager); auto& outputSplitCaches = std::get<0>(result); @@ -380,7 +382,7 @@ void MLACacheFormatter::format(tensorrt_llm::batch_manager::TransferSession& ses { sendBufferFun(deviceId, pickUpConnections[0]); } - mCacheTransBufferManagers[transferIndexerKCache]->freeBufferIndexForSend(cacheBufferId); + sendHolder.release(); } session.setTime(TransferSession::kTimeTransmissions); session.setTime(TransferSession::kTimePostprocess); @@ -459,6 +461,7 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s int deviceId = bufferManager.getStream().getDevice(); std::optional cacheBufferId = std::nullopt; + BufferIndexHolder recvHolder; if (common::getEnvTryZCopyForKVCacheTransfer() && destConfig.getParallelConfig().mPipelineParallelism @@ -495,6 +498,8 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s { cacheBufferId = mCacheTransBufferManagers[transferIndexerKCache]->assignBufferIndexForRecv(); } + recvHolder + = BufferIndexHolder(*mCacheTransBufferManagers[transferIndexerKCache], cacheBufferId, /*isRecv=*/true); auto targetNum = pickUpConnections.size(); @@ -642,10 +647,7 @@ void MLACacheFormatter::unformat(tensorrt_llm::batch_manager::TransferSession& s bufferManager.getStream().synchronize(); } - if (cacheBufferId.has_value()) - { - mCacheTransBufferManagers[transferIndexerKCache]->freeBufferIndexForRecv(cacheBufferId); - } + recvHolder.release(); } session.setTime(TransferSession::kTimePostprocess); diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp index 100e96f03119..05ed827a9511 100644 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp @@ -923,7 +923,7 @@ void TrtGptModelInflightBatching::forwardSync() TLLM_CHECK_WITH_INFO(mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration of " "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq.get()); + mCacheTransceiver->respondAndSendAsync(llmReq); } mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); } @@ -1604,11 +1604,11 @@ void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); if (common::getEnvDisableKVCacheTransferOverlap()) { - mCacheTransceiver->requestAndReceiveSync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveSync(newGenReq); } else { - mCacheTransceiver->requestAndReceiveAsync(newGenReq.get()); + mCacheTransceiver->requestAndReceiveAsync(newGenReq); } } if (!common::getEnvDisableKVCacheTransferOverlap()) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp index 30ef42560c6e..220868893f08 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp @@ -230,7 +230,7 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) "submit_transfer_requests", [](kvc::BaseTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, - nb::arg("request"), nb::rv_policy::take_ownership) + nb::arg("request"), nb::rv_policy::take_ownership, nb::keep_alive<0, 1>()) .def( "notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) .def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages) @@ -263,7 +263,8 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m) "submit_transfer_requests", [](kvc::NixlTransferAgent& self, kvc::TransferRequest const& request) { return self.submitTransferRequests(request).release(); }, - nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard()) + nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard(), + nb::keep_alive<0, 1>()) .def( "notify_sync_message", &kvc::NixlTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message")) .def("get_notified_sync_messages", &kvc::NixlTransferAgent::getNotifiedSyncMessages) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 8adddf1c86c0..5257c7adb1c9 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -46,17 +46,17 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver // using BaseCacheTransceiver::BaseCacheTransceiver; // Inherit constructors NB_TRAMPOLINE(tb::BaseCacheTransceiver, 6); - void respondAndSendAsync(tb::LlmRequest* llmRequest) override + void respondAndSendAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(respondAndSendAsync, llmRequest); } - void requestAndReceiveSync(tb::LlmRequest* llmRequest) override + void requestAndReceiveSync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveSync, llmRequest); } - void requestAndReceiveAsync(tb::LlmRequest* llmRequest) override + void requestAndReceiveAsync(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(requestAndReceiveAsync, llmRequest); } @@ -77,7 +77,7 @@ class PyCacheTransceiver : public tb::BaseCacheTransceiver NB_OVERRIDE_PURE(checkGenTransferComplete); } - bool cancelRequest(tb::LlmRequest* llmRequest) override + bool cancelRequest(std::shared_ptr llmRequest) override { NB_OVERRIDE_PURE(cancelRequest, llmRequest); } diff --git a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt index 02e293b8d102..2ee994f76226 100644 --- a/cpp/tests/unit_tests/batch_manager/CMakeLists.txt +++ b/cpp/tests/unit_tests/batch_manager/CMakeLists.txt @@ -17,6 +17,7 @@ add_gtest(radixTreeTest radixTreeTest.cpp) add_gtest(blockKeyTest blockKeyTest.cpp) add_gtest(radixBlockTreeTest radixBlockTreeTest.cpp) add_gtest(cacheTransBufferTest cacheTransBufferTest.cpp) +add_gtest(bufferIndexHolderTest bufferIndexHolderTest.cpp) add_gtest(capacitySchedulerTest capacitySchedulerTest.cpp) add_gtest(contextProgressTest contextProgressTest.cu) add_gtest(evictionPolicyTest evictionPolicyTest.cpp) diff --git a/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp new file mode 100644 index 000000000000..d4d561f55d07 --- /dev/null +++ b/cpp/tests/unit_tests/batch_manager/bufferIndexHolderTest.cpp @@ -0,0 +1,291 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/batch_manager/baseTransBuffer.h" +#include "tensorrt_llm/batch_manager/cacheTransBuffer.h" +#include "tensorrt_llm/batch_manager/kvCacheManager.h" +#include +#include +#include +#include +#include +#include + +using namespace tensorrt_llm::batch_manager; +using namespace tensorrt_llm::batch_manager::kv_cache_manager; +using namespace tensorrt_llm::runtime; + +namespace +{ + +// Subclass that exposes concurrence counters so tests can assert how many +// slots are currently held without relying on indirect observable effects. +class ObservableTransBufferManager : public CacheTransBufferManager +{ +public: + using CacheTransBufferManager::CacheTransBufferManager; + + [[nodiscard]] int sendInUse() const + { + return mConcurrenceSendResource.mConcurrence.load(); + } + + [[nodiscard]] int recvInUse() const + { + return mConcurrenceRecvResource.mConcurrence.load(); + } +}; + +} // namespace + +enum class Side +{ + Send, + Recv +}; + +struct HolderCase +{ + std::string name; + Side side; +}; + +class BufferIndexHolderLifecycleTest : public ::testing::TestWithParam +{ +protected: + void SetUp() override + { + setenv("TRTLLM_USE_UCX_KVCACHE", "1", 1); + + int constexpr numLayers = 2; + int constexpr numHeads = 2; + int constexpr sizePerHead = 8; + int constexpr tokensPerBlock = 4; + SizeType32 constexpr maxBlocksPerSeq = 4; + SizeType32 constexpr maxBeamWidth = 1; + SizeType32 constexpr maxNumSequences = 4; + SizeType32 constexpr sinkTokenLength = 0; + auto stream = std::make_shared(); + auto const kvMaxNumTokens = tokensPerBlock * maxBlocksPerSeq; + auto const totalNumBlocks = maxNumSequences * maxBlocksPerSeq; + using BlocksPerWindow = std::map>; + auto const blocksPerWindow = BlocksPerWindow{{kvMaxNumTokens, {totalNumBlocks, 0}}}; + + mKv = std::make_unique(numLayers, numHeads, sizePerHead, tokensPerBlock, blocksPerWindow, + maxNumSequences, maxBeamWidth, std::vector{kvMaxNumTokens}, + nvinfer1::DataType::kFLOAT, sinkTokenLength, stream, kvMaxNumTokens, kvMaxNumTokens, + /*enableBlockReuse=*/true, CacheType::kSELF, std::nullopt, nullptr, true); + mKv->allocatePools(false); + mTrans = std::make_unique(mKv.get(), std::optional{kvMaxNumTokens}); + } + + void TearDown() override + { + mTrans.reset(); + mKv.reset(); + } + + [[nodiscard]] std::optional acquire() const + { + return GetParam().side == Side::Send ? mTrans->assignBufferIndexForSend() : mTrans->assignBufferIndexForRecv(); + } + + [[nodiscard]] int inUse() const + { + return GetParam().side == Side::Send ? mTrans->sendInUse() : mTrans->recvInUse(); + } + + [[nodiscard]] bool isRecv() const + { + return GetParam().side == Side::Recv; + } + + [[nodiscard]] ObservableTransBufferManager& mgr() const + { + return *mTrans; + } + + std::unique_ptr mKv; + std::unique_ptr mTrans; +}; + +// Default-constructed holder owns nothing; destruction is a no-op. +TEST_P(BufferIndexHolderLifecycleTest, DefaultConstructedHolderReleasesNothing) +{ + int const before = inUse(); + { + BufferIndexHolder holder; + EXPECT_FALSE(holder.held()); + EXPECT_FALSE(holder.index().has_value()); + } + EXPECT_EQ(inUse(), before); +} + +// Explicit release() on a default-constructed holder (no manager bound) +// must not dereference the null manager pointer. +TEST_P(BufferIndexHolderLifecycleTest, DefaultConstructedExplicitReleaseIsNoOp) +{ + int const before = inUse(); + BufferIndexHolder holder; + holder.release(); + EXPECT_FALSE(holder.held()); + EXPECT_EQ(inUse(), before); +} + +// Holder built from a nullopt index is disarmed (mHeld == false). +TEST_P(BufferIndexHolderLifecycleTest, NulloptIndexReleasesNothing) +{ + int const before = inUse(); + { + BufferIndexHolder holder{mgr(), std::nullopt, isRecv()}; + EXPECT_FALSE(holder.held()); + } + EXPECT_EQ(inUse(), before); +} + +// RAII: a held slot is released when the holder goes out of scope. +TEST_P(BufferIndexHolderLifecycleTest, ValidIndexReleasedOnDestruction) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + EXPECT_EQ(inUse(), before + 1); + { + BufferIndexHolder holder{mgr(), idx, isRecv()}; + EXPECT_TRUE(holder.held()); + EXPECT_EQ(holder.index(), idx); + } + EXPECT_EQ(inUse(), before); +} + +// release() frees the slot immediately and disarms the destructor. +TEST_P(BufferIndexHolderLifecycleTest, ExplicitReleaseFreesSlotEagerly) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + BufferIndexHolder holder{mgr(), idx, isRecv()}; + EXPECT_EQ(inUse(), before + 1); + holder.release(); + EXPECT_EQ(inUse(), before); + EXPECT_FALSE(holder.held()); +} + +// release() is idempotent: a second call is a no-op. +TEST_P(BufferIndexHolderLifecycleTest, DoubleReleaseIsSafe) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + BufferIndexHolder holder{mgr(), idx, isRecv()}; + holder.release(); + EXPECT_EQ(inUse(), before); + holder.release(); + EXPECT_EQ(inUse(), before); +} + +// detach() returns the index and disarms the destructor. The caller is +// responsible for freeing the slot — we free it manually to keep the pool +// balanced for subsequent tests. +TEST_P(BufferIndexHolderLifecycleTest, DetachDisarmsDestructor) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + { + BufferIndexHolder holder{mgr(), idx, isRecv()}; + auto detached = holder.detach(); + EXPECT_EQ(detached, idx); + EXPECT_FALSE(holder.held()); + } + // Destructor was disarmed: slot still in use. + EXPECT_EQ(inUse(), before + 1); + + // Manually free so the harness is balanced. + if (isRecv()) + { + mgr().freeBufferIndexForRecv(idx); + } + else + { + mgr().freeBufferIndexForSend(idx); + } + EXPECT_EQ(inUse(), before); +} + +// Move construction transfers ownership; the moved-from holder is disarmed. +TEST_P(BufferIndexHolderLifecycleTest, MoveConstructTransfersOwnership) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + { + BufferIndexHolder source{mgr(), idx, isRecv()}; + BufferIndexHolder sink{std::move(source)}; + EXPECT_FALSE(source.held()); + EXPECT_TRUE(sink.held()); + EXPECT_EQ(sink.index(), idx); + EXPECT_EQ(inUse(), before + 1); + } + // Sink released on scope exit; source was disarmed. + EXPECT_EQ(inUse(), before); +} + +// Move assignment releases the prior holding before taking new ownership. +TEST_P(BufferIndexHolderLifecycleTest, MoveAssignReleasesPriorThenTransfers) +{ + int const before = inUse(); + auto firstIdx = acquire(); + auto secondIdx = acquire(); + ASSERT_TRUE(firstIdx.has_value()); + ASSERT_TRUE(secondIdx.has_value()); + EXPECT_EQ(inUse(), before + 2); + { + BufferIndexHolder sink{mgr(), firstIdx, isRecv()}; + BufferIndexHolder source{mgr(), secondIdx, isRecv()}; + sink = std::move(source); + // Sink now owns secondIdx; firstIdx was released by the move-assign. + EXPECT_EQ(inUse(), before + 1); + EXPECT_EQ(sink.index(), secondIdx); + EXPECT_FALSE(source.held()); + } + EXPECT_EQ(inUse(), before); +} + +// Exception unwind through a scope containing a held holder still releases +// the slot. +TEST_P(BufferIndexHolderLifecycleTest, ExceptionUnwindStillReleases) +{ + int const before = inUse(); + auto idx = acquire(); + ASSERT_TRUE(idx.has_value()); + try + { + BufferIndexHolder holder{mgr(), idx, isRecv()}; + EXPECT_EQ(inUse(), before + 1); + throw std::runtime_error("forced unwind"); + } + catch (std::runtime_error const&) + { + // Holder destructor ran during unwind. + } + EXPECT_EQ(inUse(), before); +} + +INSTANTIATE_TEST_SUITE_P(SideVariants, BufferIndexHolderLifecycleTest, + ::testing::Values(HolderCase{"send", Side::Send}, HolderCase{"recv", Side::Recv}), + [](::testing::TestParamInfo const& info) { return info.param.name; }); diff --git a/tests/integration/test_lists/test-db/l0_sanity_check.yml b/tests/integration/test_lists/test-db/l0_sanity_check.yml index 792387bc6db3..bc947305f276 100644 --- a/tests/integration/test_lists/test-db/l0_sanity_check.yml +++ b/tests/integration/test_lists/test-db/l0_sanity_check.yml @@ -34,6 +34,9 @@ l0_sanity_check: - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[UCX-mha-ctx_fp16_gen_fp16] - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mha] - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mla] + - unittest/others/test_kv_cache_transceiver.py::test_async_transfer_keeps_llm_request_alive + - unittest/others/test_kv_cache_transceiver.py::test_kv_transfer_timeout_warns_once_per_request + - unittest/others/test_kv_cache_transceiver.py::test_kv_transfer_timeout_silent_when_unset - unittest/_torch/pyexecutor/test_model_loader_mx.py - condition: ranges: diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 1a38fd8b8d51..d593a572363b 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1,5 +1,8 @@ +import gc +import sys import time import uuid +import weakref import pytest import torch @@ -229,12 +232,208 @@ def test_cancel_request_in_transmission(attention_type): assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR +@pytest.mark.timeout(120) +def test_async_transfer_keeps_llm_request_alive(): + """Async entry points must hold a strong shared_ptr to the LlmRequest. + + A regression to raw-pointer storage in mSenderFutures / + mRequesterFutures lets Python GC the wrapper while a C++ status check + still dereferences it. + """ + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + kv_cache_manager_gen = create_kv_cache_manager(mapping, DataType.HALF) + + cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT", + max_tokens_in_buffer=512) + transceiver_ctx = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + transceiver_gen = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_gen, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + sampling_params = SamplingParams() + ctx_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + + # Snapshot ctx refcount *before* submission. add_sequence_batch takes + # std::reference_wrapper (no Python ref retained), so the + # baseline here is clean and the +1 below isolates exactly the + # shared_ptr captured by respond_and_send_async. + ctx_ref = weakref.ref(ctx_request) + baseline_ctx_refcount = sys.getrefcount(ctx_request) + + # respond_and_send_async also populates ctx_request.context_phase_params + # as a side effect — gen_request below requires this to be non-empty. + transceiver_ctx.respond_and_send_async(ctx_request) + assert sys.getrefcount(ctx_request) == baseline_ctx_refcount + 1, ( + f"respond_and_send_async did not capture a shared_ptr; " + f"refcount={sys.getrefcount(ctx_request)} " + f"(expected baseline {baseline_ctx_refcount} + 1)") + + gen_request = LlmRequest( + request_id=0, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_request.context_phase_params) + kv_cache_manager_gen.impl.add_sequence_batch( + [(gen_request.py_request_id, gen_request.prompt_len, 1)], [gen_request]) + + gen_ref = weakref.ref(gen_request) + baseline_gen_refcount = sys.getrefcount(gen_request) + + transceiver_gen.request_and_receive_async(gen_request) + assert sys.getrefcount(gen_request) == baseline_gen_refcount + 1, ( + f"request_and_receive_async did not capture a shared_ptr; " + f"refcount={sys.getrefcount(gen_request)} " + f"(expected baseline {baseline_gen_refcount} + 1)") + + # Drop the external Python refs. After this, the only owners are the + # nanobind-managed shared_ptr entries inside mSenderFutures / + # mRequesterFutures. + del ctx_request + del gen_request + gc.collect() + + assert ctx_ref() is not None, ( + "Sender-side LlmRequest was destroyed prematurely; " + "respond_and_send_async must hold a strong shared_ptr while the " + "sender future is in flight") + assert gen_ref() is not None, ( + "Receiver-side LlmRequest was destroyed prematurely; " + "request_and_receive_async must hold a strong shared_ptr while " + "the requester future is in flight") + + # Drive the transfer to completion so the harness tears down cleanly. + transceiver_ctx.check_context_transfer_status(1) + transceiver_gen.check_gen_transfer_status(1) + + +def _build_ctx_request_for_timeout_test(request_id: int) -> LlmRequest: + sampling_params = SamplingParams() + return LlmRequest( + request_id=request_id, + max_new_tokens=1, + input_tokens=list(range(256)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + + +@pytest.mark.timeout(60) +def test_kv_transfer_timeout_warns_once_per_request(capfd): + """Observe-only timeout WARN must fire exactly once per stuck request. + + checkContextTransferStatus emits a TLLM_LOG_WARNING when elapsed time + exceeds kv_transfer_timeout_ms; mTimedOutSenderIds dedup suppresses + repeat emissions across subsequent polls of the same in-flight future. + """ + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + + cache_transceiver_config = CacheTransceiverConfig( + backend="DEFAULT", max_tokens_in_buffer=512, kv_transfer_timeout_ms=100) + transceiver_ctx = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + ctx_request = _build_ctx_request_for_timeout_test(request_id=42) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + + capfd.readouterr() # drain prior noise + + transceiver_ctx.respond_and_send_async(ctx_request) + time.sleep(0.3) # > kv_transfer_timeout_ms + + transceiver_ctx.check_context_transfer_status(0) + first = capfd.readouterr() + marker = "Context KV cache transfer for request 42 exceeded configured timeout" + # TLLM_LOG_WARNING writes to stdout; check first.out (not first.err). + assert marker in first.out, ( + f"Expected observe-only WARN after timeout; stdout:\n{first.out}") + assert first.out.count(marker) == 1, ( + f"WARN should fire exactly once per poll for a single stuck request; " + f"got {first.out.count(marker)} emissions:\n{first.out}") + + transceiver_ctx.check_context_transfer_status(0) + second = capfd.readouterr() + assert marker not in second.out, ( + f"WARN re-emitted on a subsequent poll; dedup broken. " + f"stdout:\n{second.out}") + + transceiver_ctx.cancel_request(ctx_request) + + +@pytest.mark.timeout(60) +def test_kv_transfer_timeout_silent_when_unset(capfd): + """Without kv_transfer_timeout_ms the observe-only WARN must stay silent. + + The elapsed-time check is gated by the optional config field; absence + of the field must short-circuit the WARN path even on a long-running + transfer. + """ + mapping = Mapping(world_size=1, rank=0) + dist = Distributed.get(mapping) + kv_cache_manager_ctx = create_kv_cache_manager(mapping, DataType.HALF) + + cache_transceiver_config = CacheTransceiverConfig(backend="DEFAULT", + max_tokens_in_buffer=512) + transceiver_ctx = create_kv_cache_transceiver(mapping, dist, + kv_cache_manager_ctx, + AttentionTypeCpp.DEFAULT, + cache_transceiver_config) + + fill_kv_cache_buffer(kv_cache_manager_ctx) + + ctx_request = _build_ctx_request_for_timeout_test(request_id=99) + kv_cache_manager_ctx.impl.add_sequence_batch( + [(ctx_request.py_request_id, ctx_request.prompt_len, 1)], [ctx_request]) + + capfd.readouterr() + + transceiver_ctx.respond_and_send_async(ctx_request) + time.sleep(0.3) + transceiver_ctx.check_context_transfer_status(0) + + out = capfd.readouterr() + # TLLM_LOG_WARNING writes to stdout; check out.out (not out.err). + assert "exceeded configured timeout" not in out.out, ( + f"Observe-only WARN must not fire when kv_transfer_timeout_ms is " + f"unset; stdout:\n{out.out}") + + transceiver_ctx.cancel_request(ctx_request) + + def create_hybrid_cache_manager(mapping, dtype, mamba_conv_dtype=torch.float16, mamba_ssm_dtype=torch.float16): - """ - Create a MambaHybridCacheManager for testing hybrid models. + """Create a MambaHybridCacheManager for testing hybrid models. + This manager handles both KV cache (attention layers) and Mamba cache (RNN layers). Args: