From 28db8c55cb51845d8c719cfbb1c77edb6f41c087 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Wed, 15 Apr 2026 20:20:21 -0700 Subject: [PATCH 1/3] [None][pref] Consolidate prefix reuse queries into single analyzePrefixReuse radix tree walk Signed-off-by: Simeng Liu --- .../batch_manager/kvCacheManager.h | 94 ++++---- .../batch_manager/capacityScheduler.cpp | 152 ++++++++----- .../batch_manager/kvCacheManager.cpp | 109 ++++----- .../batch_manager/scheduledBlocksManager.h | 41 ++-- .../nanobind/batch_manager/kvCacheManager.cpp | 28 ++- .../batch_manager/kvCacheManagerTest.cpp | 89 +++++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 214 +++++++++++------- .../_torch/pyexecutor/resource_manager.py | 26 ++- .../_torch/pyexecutor/scheduler/scheduler.py | 24 +- tensorrt_llm/llmapi/llm_args.py | 3 +- .../executor/test_kvcache_aware_router.py | 8 +- .../_torch/executor/test_py_scheduler.py | 38 +++- 12 files changed, 499 insertions(+), 327 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 64ff4c0d3fc8..f37005d391e0 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -559,6 +559,26 @@ class KVCacheBlockPool } }; +/// \brief Result of a single radix tree walk that collects all prefix reuse information. +/// \details Replaces separate findNewContextBlock() + countReusableBlocks(true) + +/// countReusableBlocks(false) calls with a single walk, reducing 2-3 walks to 1. +struct PrefixReuseSummary +{ + using SizeType32 = tensorrt_llm::runtime::SizeType32; + + /// Number of prefix blocks already allocated (refCount > 0). + /// Used by the block budget to avoid double-counting with the eviction free count. + SizeType32 reusableBlocksAllocated{0}; + + /// Total number of prefix blocks cached (allocated or free-cached). + /// Used by the token budget (NoEvict) since all cached tokens avoid recompute. + SizeType32 reusableBlocksAll{0}; + + /// First block key NOT found in the radix tree (nullopt = full prefix hit). + /// Used by the capacity scheduler's skip-check logic to decide whether to defer a request. + std::optional firstNewBlock{std::nullopt}; +}; + // The WindowBlockManager manages the metadata of KVCacheBlocks. // It manages multiple arrays of cache blocks called pools. // Layers with the same number of kv heads are grouped under the same pool. @@ -820,23 +840,10 @@ class WindowBlockManager void offloadBlock(BlockPtr const& block, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); - //! \brief Find first new block that must be allocated for context phase and return it's concatenated token vectors. - //! \details Only full blocks are considered. - [[nodiscard]] std::optional findNewContextBlock( + //! \brief Combined prefix reuse analysis — single radix tree walk. + [[nodiscard]] PrefixReuseSummary analyzePrefixReuse( VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const; - //! \brief Count the number of full blocks that can be reused from the KV cache for a given request. - //! \details Traverses the radix tree to count how many consecutive blocks from the beginning - //! of the request's context are already cached. - //! \param uniqueTokens The unique tokens representing the request's context. - //! \param llmRequest The request to check for reusable blocks. - //! \param onlyAllocated If true, only count blocks that have active references (already allocated - //! to another sequence). Free cached blocks are excluded because they are already counted - //! in the eviction policy's free count. - //! \return The number of full blocks that can be reused. - [[nodiscard]] SizeType32 countReusableBlocks( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest, bool onlyAllocated = false) const; - [[nodiscard]] runtime::BufferManager const& getBufferManager() const { return mBufferManager; @@ -1140,14 +1147,10 @@ class BlockManager void setOffsets(kernels::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId, SizeType32 windowSize) const; - // WILL NOT WORK FOR VARIABLE WINDOW ATTENTION - [[nodiscard]] std::optional findNewContextBlock( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const; - - //! \brief Count the number of full blocks that can be reused from the KV cache for a given request. + //! \brief Combined prefix reuse analysis — single radix tree walk. //! \details WILL NOT WORK FOR VARIABLE WINDOW ATTENTION. - [[nodiscard]] SizeType32 countReusableBlocks( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest, bool onlyAllocated = false) const; + [[nodiscard]] PrefixReuseSummary analyzePrefixReuse( + VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const; //! \brief Bring block from primary to secondary memory for window size. //! \details Does nothing if block is already in primary memory. @@ -1545,18 +1548,21 @@ class BaseKVCacheManager [[nodiscard]] virtual BlockManager const& getBlockManager() const = 0; /// @brief Function that computes the number of KV cache blocks needed to advance a request by one or two - /// iterations + /// iterations. /// @param req The request for which we need to calculate the number of needed KV cache blocks + /// @param cachedSummary Optional pre-computed PrefixReuseSummary to avoid redundant radix tree walks. /// @return The number of blocks - [[nodiscard]] virtual SizeType32 getNeededBlocksOneStep( - LlmRequest const& req, bool twoStepsLookAhead, SizeType32 windowSize) const + [[nodiscard]] virtual SizeType32 getNeededBlocksOneStep(LlmRequest const& req, bool twoStepsLookAhead, + SizeType32 windowSize, std::optional const& cachedSummary = std::nullopt) const = 0; /// @brief Function that computes the number of KV cache blocks needed to advance a request to completion (i.e. for - /// maxNewTokens) + /// maxNewTokens). /// @param req The request for which we need to calculate the number of needed KV cache blocks + /// @param cachedSummary Optional pre-computed PrefixReuseSummary to avoid redundant radix tree walks. /// @return The number of blocks - [[nodiscard]] virtual SizeType32 getRemainingBlocksToCompletion(LlmRequest const& req, SizeType32 windowSize) const + [[nodiscard]] virtual SizeType32 getRemainingBlocksToCompletion(LlmRequest const& req, SizeType32 windowSize, + std::optional const& cachedSummary = std::nullopt) const = 0; /// @brief Pin blocks associated with a request to prevent eviction. @@ -1613,23 +1619,12 @@ class BaseKVCacheManager [[nodiscard]] virtual bool isCrossKv() const = 0; - //! \brief Find first new block that must be allocated for context phase and return it's concatenated token vector. - //! \details Only full blocks are considered. - [[nodiscard]] virtual std::optional findNewContextBlock( + //! \brief Combined prefix reuse analysis — single radix tree walk. + //! \details Collects firstNewBlock + reusableBlocksAllocated + reusableBlocksAll in one pass. + [[nodiscard]] virtual PrefixReuseSummary analyzePrefixReuse( VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const = 0; - //! \brief Count the number of full blocks that can be reused from the KV cache for a given request. - //! \details Traverses the radix tree to count how many consecutive blocks from the beginning - //! of the request's context are already cached. - //! \param uniqueTokens The unique tokens representing the request's context. - //! \param llmRequest The request to check for reusable blocks. - //! \param onlyAllocated If true, only count blocks that have active references. - //! \return The number of full blocks that can be reused. - [[nodiscard]] virtual SizeType32 countReusableBlocks( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest, bool onlyAllocated = false) const - = 0; - //! \brief Store full context blocks contributed by llmRequest. //! \details These blocks become reusable from next step. virtual void storeContextBlocks(LlmRequest const& llmRequest) = 0; @@ -1904,15 +1899,15 @@ class KVCacheManager : public BaseKVCacheManager /// iterations /// @param req The request for which we need to calculate the number of needed KV cache blocks /// @return The number of blocks - [[nodiscard]] SizeType32 getNeededBlocksOneStep( - LlmRequest const& req, bool twoStepsLookAhead, SizeType32 windowSize) const override; + [[nodiscard]] SizeType32 getNeededBlocksOneStep(LlmRequest const& req, bool twoStepsLookAhead, + SizeType32 windowSize, std::optional const& cachedSummary = std::nullopt) const override; /// @brief Function that computes the number of KV cache blocks remaining to advance a request to completion (i.e. /// for maxNewTokens); the allocated blocks are excluded /// @param req The request for which we need to calculate the number of needed KV cache blocks /// @return The number of blocks - [[nodiscard]] SizeType32 getRemainingBlocksToCompletion( - LlmRequest const& req, SizeType32 windowSize) const override; + [[nodiscard]] SizeType32 getRemainingBlocksToCompletion(LlmRequest const& req, SizeType32 windowSize, + std::optional const& cachedSummary = std::nullopt) const override; /// @brief Increase size for request with requestId. Allocate new KV cache block(s) if needed. void addToken(LlmRequest::RequestIdType requestId) override; @@ -1995,15 +1990,10 @@ class KVCacheManager : public BaseKVCacheManager return mBlockManager.getCacheType(); } - //! \brief Find first new block that must be allocated for context phase and return it's concatenated token vector. - //! \details Only full blocks are considered. - [[nodiscard]] std::optional findNewContextBlock( + //! \brief Combined prefix reuse analysis — single radix tree walk. + [[nodiscard]] PrefixReuseSummary analyzePrefixReuse( VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const override; - //! \brief Count the number of full blocks that can be reused from the KV cache for a given request. - [[nodiscard]] SizeType32 countReusableBlocks( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest, bool onlyAllocated = false) const override; - //! \brief Store full context blocks contributed by llmRequest. //! \details These blocks become reusable from next step. void storeContextBlocks(LlmRequest const& llmRequest) override; diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index d765bcf31730..478483be72c3 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -47,19 +47,19 @@ prefillWithChunkedContextsAlreadyExecuting(RequestList const& activeRequests, if (kvCacheManager.isEnableBlockReuse()) { auto uniqueTokens = req->getUniqueTokens(0); - auto newContextBlockOpt = kvCacheManager.findNewContextBlock(uniqueTokens, *req); - if (newContextBlockOpt.has_value()) + auto summary = kvCacheManager.analyzePrefixReuse(uniqueTokens, *req); + if (summary.firstNewBlock.has_value()) { - newlyContributedContextBlocks.insert(newContextBlockOpt.value()); + newlyContributedContextBlocks.insert(summary.firstNewBlock.value()); } } if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse()) { auto uniqueTokens = *(req->getEncoderUniqueTokens().value()); - auto newContextBlockOpt = crossKvCacheManager->findNewContextBlock(uniqueTokens, *req); - if (newContextBlockOpt.has_value()) + auto summary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req); + if (summary.firstNewBlock.has_value()) { - newlyContributedCrossContextBlocks.insert(newContextBlockOpt.value()); + newlyContributedCrossContextBlocks.insert(summary.firstNewBlock.value()); } } } @@ -67,60 +67,39 @@ prefillWithChunkedContextsAlreadyExecuting(RequestList const& activeRequests, return {std::move(newlyContributedContextBlocks), std::move(newlyContributedCrossContextBlocks)}; } -bool oneManagerBeneficialToSkip(tensorrt_llm::batch_manager::kv_cache_manager::BaseKVCacheManager const& kvCacheManager, - VecUniqueTokens const& uniqueTokens, std::shared_ptr const& llmRequest, +/// @brief Check if a single manager's summary indicates we should skip this request. +/// @details Returns true if the request's first new context block was already contributed +/// by an earlier scheduled request (so waiting would let us reuse it). Registers the +/// block as contributed when not skipping. +bool oneManagerBeneficialToSkip(std::optional const& summary, std::unordered_set& newlyContributedContextBlocks) { - // Find first context block that isn't already in KV cache - auto newContextBlockOpt = kvCacheManager.findNewContextBlock(uniqueTokens, *llmRequest); - if (newContextBlockOpt.has_value()) + if (summary.has_value() && summary->firstNewBlock.has_value()) { - auto const& newContextBlock = newContextBlockOpt.value(); - if (newlyContributedContextBlocks.count(newContextBlock) > 0) + if (newlyContributedContextBlocks.count(summary->firstNewBlock.value()) > 0) { - // newContextBlock was contributed by earlier scheduled request. - // Better to skip this step so we can reuse. return true; } - - // This request is contributing newContextBlock. - newlyContributedContextBlocks.insert(newContextBlock); + newlyContributedContextBlocks.insert(summary->firstNewBlock.value()); } - // Either all context blocks are already in KV cache, - // or no previously scheduled request has contributed newContextBlock. return false; } -//! \brief Check if it is beneficial to skip this request rather than schedule it. -//! \details One condition that makes it beneficial is if this request can reuse kv cache block(s) contributed by -//! already scheduled context requests. -bool beneficialToSkip(std::shared_ptr const& req, - kv_cache_manager::BaseKVCacheManager const& kvCacheManager, - OptionalRef crossKvCacheManager, +/// @brief Check if it is beneficial to skip this request rather than schedule it. +/// @details Returns true if this request can reuse KV cache block(s) that will be contributed +/// by already-scheduled context requests. Uses pre-computed PrefixReuseSummary values. +bool beneficialToSkip(std::optional const& summary, + std::optional const& crossSummary, std::unordered_set& newlyContributedContextBlocks, std::unordered_set& newlyContributedCrossContextBlocks) { - if (req->isContextInitState() && req->isFirstContextChunk()) + if (oneManagerBeneficialToSkip(summary, newlyContributedContextBlocks)) { - if (kvCacheManager.isEnableBlockReuse()) - { - auto uniqueTokens = req->getUniqueTokens(0); - if (oneManagerBeneficialToSkip(kvCacheManager, uniqueTokens, req, newlyContributedContextBlocks)) - { - return true; - } - } - if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse()) - { - auto uniqueTokens = *(req->getEncoderUniqueTokens().value()); - if (oneManagerBeneficialToSkip(*crossKvCacheManager, uniqueTokens, req, newlyContributedCrossContextBlocks)) - { - return true; - } - } + return true; } - return false; + return oneManagerBeneficialToSkip(crossSummary, newlyContributedCrossContextBlocks); } + } // namespace MaxRequestsScheduler::MaxRequestsScheduler( @@ -251,9 +230,13 @@ std::tuple GuaranteedNoEvictScheduler::impl( if (req->isGenerationInProgressState()) { scheduledRequests.emplace_back(req); - reservedBlocks.decrementReservedBlocks(*req); + reservedBlocks.enoughAvailableBlocks(*req); + reservedBlocks.commitBlocks(); if (reservedCrossBlocks) - reservedCrossBlocks->decrementReservedBlocks(*req); + { + reservedCrossBlocks->enoughAvailableBlocks(*req); + reservedCrossBlocks->commitBlocks(); + } bool const reqHasLora = req->getLoraTaskId().has_value(); bool const isNewTask = reqHasLora && !uniqTaskIds.count(req->getLoraTaskId().value()); if (isNewTask) @@ -285,10 +268,32 @@ std::tuple GuaranteedNoEvictScheduler::impl( { for (auto const& req : requests) { - // if context request can reuse blocks contributed by another context request, skip - if (!StaticBatchScheduling && skippingIsRelevant && !req->isDisaggGenerationInitState() - && beneficialToSkip(req, kvCacheManager, crossKvCacheManager, newlyContributedContextBlocks, - newlyContributedCrossContextBlocks)) + // For first-chunk context requests with block reuse, compute the prefix reuse + // summary once. This single radix tree walk serves both the beneficial-to-skip + // check and the block budget estimation in getRemainingBlocksToCompletion, + // eliminating 2 redundant walks per request. + bool const isFirstChunkContext + = req->isContextInitState() && req->isFirstContextChunk() && !req->isDisaggGenerationInitState(); + std::optional summary; + std::optional crossSummary; + if (isFirstChunkContext) + { + if (kvCacheManager.isEnableBlockReuse()) + { + auto uniqueTokens = req->getUniqueTokens(0); + summary = kvCacheManager.analyzePrefixReuse(uniqueTokens, *req); + } + if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse()) + { + auto uniqueTokens = *(req->getEncoderUniqueTokens().value()); + crossSummary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req); + } + } + + // Beneficial-to-skip check using the cached summary + if (!StaticBatchScheduling && skippingIsRelevant && isFirstChunkContext + && beneficialToSkip( + summary, crossSummary, newlyContributedContextBlocks, newlyContributedCrossContextBlocks)) { continue; } @@ -300,9 +305,14 @@ std::tuple GuaranteedNoEvictScheduler::impl( if (req->isContextInitState() || req->isDisaggGenerationInitState()) { - bool enoughBlocks = reservedBlocks.enoughAvailableBlocks(*req); - bool enoughCrossBlocks - = reservedCrossBlocks ? reservedCrossBlocks->enoughAvailableBlocks(*req) : true; + // Check block availability using the cached summary when available. + // enoughAvailableBlocks is check-only (no decrement) — safe if cross check fails. + bool enoughBlocks = reservedBlocks.enoughAvailableBlocks(*req, summary); + bool enoughCrossBlocks = true; + if (reservedCrossBlocks) + { + enoughCrossBlocks = reservedCrossBlocks->enoughAvailableBlocks(*req, crossSummary); + } bool reqHasLora = req->getLoraTaskId().has_value(); bool isNewTask = reqHasLora && !uniqTaskIds.count(req->getLoraTaskId().value()); auto neededPeftPages = isNewTask && peftCacheManager ? peftCacheManager->determineNumPages(req) : 0; @@ -310,9 +320,12 @@ std::tuple GuaranteedNoEvictScheduler::impl( if (enoughBlocks && enoughCrossBlocks && neededPeftPages <= availablePeftPages) { scheduledRequests.emplace_back(req); - reservedBlocks.decrementReservedBlocks(*req); + // Decrement using the cached values computed by enoughAvailableBlocks. + reservedBlocks.commitBlocks(); if (reservedCrossBlocks) - reservedCrossBlocks->decrementReservedBlocks(*req); + { + reservedCrossBlocks->commitBlocks(); + } availablePeftPages -= neededPeftPages; if (isNewTask) { @@ -336,7 +349,8 @@ std::tuple GuaranteedNoEvictScheduler::impl( bool trySchedulingRequestMaxUtilization(std::shared_ptr const& req, SizeType32 maxNumRequests, RequestVector& scheduledRequests, kv_cache_manager::MaxUtilizationScheduledBlocksManager& blocksManager, OptionalRef peftCacheManager, SizeType32& numScheduledPeftPages, - std::unordered_set& seenTaskIds); + std::unordered_set& seenTaskIds, + std::optional const& cachedSummary); std::tuple MaxUtilizationScheduler::operator()( kv_cache_manager::BaseKVCacheManager& kvCacheManager, OptionalRef peftCacheManager, @@ -383,17 +397,29 @@ std::tuple MaxUtilizationScheduler::operator()( continue; } - // if context request can reuse blocks contributed by another context request, skip - if (skippingIsRelevant + // For first-chunk context requests with block reuse, compute the prefix reuse + // summary once. This single radix tree walk serves both the beneficial-to-skip + // check and the block budget estimation in getNeededBlocksOneStep. + bool const isFirstChunkContext + = req->isContextInitState() && req->isFirstContextChunk() && !req->isDisaggGenerationInitState(); + std::optional summary; + if (isFirstChunkContext && kvCacheManager.isEnableBlockReuse()) + { + auto uniqueTokens = req->getUniqueTokens(0); + summary = kvCacheManager.analyzePrefixReuse(uniqueTokens, *req); + } + + // Beneficial-to-skip check using the cached summary (no cross KV cache for MaxUtil) + if (skippingIsRelevant && isFirstChunkContext && beneficialToSkip( - req, kvCacheManager, std::nullopt, newlyContributedContextBlocks, newlyContributedCrossContextBlocks)) + summary, std::nullopt, newlyContributedContextBlocks, newlyContributedCrossContextBlocks)) { reqIt++; continue; } bool const wasScheduled = trySchedulingRequestMaxUtilization(req, mMaxNumRequests, scheduledRequests, - scheduledBlocksManager, peftCacheManager, numScheduledPeftPages, seenTaskIds); + scheduledBlocksManager, peftCacheManager, numScheduledPeftPages, seenTaskIds, summary); if (wasScheduled) { TLLM_LOG_DEBUG("MaxUtilizationScheduler: request ID %lu -> start", req->mRequestId); @@ -427,7 +453,7 @@ std::tuple MaxUtilizationScheduler::operator()( bool trySchedulingRequestMaxUtilization(std::shared_ptr const& req, SizeType32 maxNumRequests, RequestVector& scheduledRequests, kv_cache_manager::MaxUtilizationScheduledBlocksManager& blocksManager, OptionalRef peftCacheManager, SizeType32& numScheduledPeftPages, - std::unordered_set& seenTaskIds) + std::unordered_set& seenTaskIds, std::optional const& cachedSummary) { if (scheduledRequests.size() < static_cast(maxNumRequests)) { @@ -437,7 +463,9 @@ bool trySchedulingRequestMaxUtilization(std::shared_ptr const& req, = (isNewTask && peftCacheManager) ? peftCacheManager->determineNumPages(req) : 0; TLLM_LOG_DEBUG( "MaxUtilizationScheduler: request ID %lu required peft pages: %i", req->mRequestId, numRequiredPeftPages); - auto const scheduledBlocksIfFitsKvCache = blocksManager.prepareNewNumberOfBlocksIfWeEndUpScheduling(*req); + // Use the cached summary when available to avoid a redundant tree walk + auto const scheduledBlocksIfFitsKvCache + = blocksManager.prepareNewNumberOfBlocksIfWeEndUpScheduling(*req, cachedSummary); bool fitsPeft = (peftCacheManager ? numRequiredPeftPages + numScheduledPeftPages <= peftCacheManager->getMaxDevicePages() : true); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 107c703b0b79..b8d6096ad870 100755 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1136,86 +1136,55 @@ void WindowBlockManager::offloadBlock( } } -[[nodiscard]] std::optional BlockManager::findNewContextBlock( +PrefixReuseSummary BlockManager::analyzePrefixReuse( VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const { - TLLM_CHECK_WITH_INFO( - !isVariableWindow(), "The optimization of delaying requests won't work for variable window attention"); + TLLM_CHECK_WITH_INFO(!isVariableWindow(), "analyzePrefixReuse does not work for variable window attention"); auto const& onlyManager = mWindowBlockManagers.cbegin()->second; - return onlyManager.findNewContextBlock(uniqueTokens, llmRequest); + return onlyManager.analyzePrefixReuse(uniqueTokens, llmRequest); } -SizeType32 BlockManager::countReusableBlocks( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest, bool onlyAllocated) const -{ - TLLM_CHECK_WITH_INFO(!isVariableWindow(), "countReusableBlocks does not work for variable window attention"); - auto const& onlyManager = mWindowBlockManagers.cbegin()->second; - return onlyManager.countReusableBlocks(uniqueTokens, llmRequest, onlyAllocated); -} - -std::optional WindowBlockManager::findNewContextBlock( +PrefixReuseSummary WindowBlockManager::analyzePrefixReuse( VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const { auto blockedUniqueTokens = chopVectorIntoBlocks(uniqueTokens, uniqueTokens.size(), mTokensPerBlock, false); auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); - BlockKey ret; - ret.loraTaskId = llmRequest.getLoraTaskId(); - std::lock_guard lock(mCachedBlocksRootMutex); - auto searchRoot = mCachedBlocksRoot; - for (auto const& blockKey : blockKeys) - { - ret.uniqueTokens.insert(ret.uniqueTokens.end(), blockKey.uniqueTokens.begin(), blockKey.uniqueTokens.end()); - auto [partialMatch, numMatched, matchingBlock] = searchRoot != nullptr - ? searchRoot->findMatchingBlock(blockKey, false, false) - : std::make_tuple(false, 0, nullptr); - if (matchingBlock == nullptr) - { - return ret; - } - searchRoot = std::move(matchingBlock); - } - return std::nullopt; -} -SizeType32 WindowBlockManager::countReusableBlocks( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest, bool onlyAllocated) const -{ - // Chop tokens into full blocks only (allowPartial=false) - auto blockedUniqueTokens - = chopVectorIntoBlocks(uniqueTokens, uniqueTokens.size(), mTokensPerBlock, false); - auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); + PrefixReuseSummary summary; + BlockKey accumKey; + accumKey.loraTaskId = llmRequest.getLoraTaskId(); - SizeType32 reusableBlocks = 0; std::lock_guard lock(mCachedBlocksRootMutex); auto searchRoot = mCachedBlocksRoot; for (auto const& blockKey : blockKeys) { + accumKey.uniqueTokens.insert( + accumKey.uniqueTokens.end(), blockKey.uniqueTokens.begin(), blockKey.uniqueTokens.end()); + auto [partialMatch, numMatched, matchingBlock] = searchRoot != nullptr ? searchRoot->findMatchingBlock(blockKey, false, false) : std::make_tuple(false, 0, nullptr); if (matchingBlock == nullptr) { - // No more matching blocks found + summary.firstNewBlock = accumKey; break; } - // When onlyAllocated is true, only count blocks that are already allocated - // to an active sequence (have refs). Sharing these blocks doesn't consume - // from the free pool. Free cached blocks (no refs) are already counted in - // the eviction policy's free count and must not be double-counted. - if (!onlyAllocated || matchingBlock->hasRefs()) + ++summary.reusableBlocksAll; + if (matchingBlock->hasRefs()) { - ++reusableBlocks; + ++summary.reusableBlocksAllocated; } + searchRoot = std::move(matchingBlock); } - TLLM_LOG_DEBUG("%s::countReusableBlocks - Found %d reusable blocks (onlyAllocated=%d)", mLogPrefix.c_str(), - reusableBlocks, onlyAllocated); - return reusableBlocks; + TLLM_LOG_DEBUG("%s::analyzePrefixReuse - reusableAllocated=%d, reusableAll=%d, hasNewBlock=%d", mLogPrefix.c_str(), + summary.reusableBlocksAllocated, summary.reusableBlocksAll, summary.firstNewBlock.has_value()); + return summary; } bool WindowBlockManager::blockInRadixTree(BlockPtr const& block) @@ -2231,8 +2200,8 @@ void KVCacheManager::startScheduling() mBlockManager.startScheduling(); } -SizeType32 KVCacheManager::getNeededBlocksOneStep( - LlmRequest const& req, bool twoStepsLookAhead, SizeType32 windowSize) const +SizeType32 KVCacheManager::getNeededBlocksOneStep(LlmRequest const& req, bool twoStepsLookAhead, SizeType32 windowSize, + std::optional const& cachedSummary) const { // Default to zero; overwritten below when block reuse is active for a first-chunk context request. req.setEstimatedReusableTokens(0); @@ -2255,8 +2224,16 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep( if (mEnableBlockReuse && !mBlockManager.isVariableWindow() && !isCrossKv() && !req.isDisaggGenerationInitState()) { - auto const uniqueTokens = req.getUniqueTokens(0); - auto const numReusableBlocks = countReusableBlocks(uniqueTokens, req, /*onlyAllocated=*/true); + // Use the cached summary if provided; otherwise perform a fresh tree walk. + auto const summary + = cachedSummary.has_value() ? cachedSummary.value() : analyzePrefixReuse(req.getUniqueTokens(0), req); + auto const numReusableBlocks = summary.reusableBlocksAllocated; + auto const promptInputLen = std::min(req.mPromptLen, windowSize + chunkSize); + // `addSequence()` ignores the last prompt token because its KV cannot be recovered. + // When the prompt lands exactly on a block boundary, counting reusable full blocks from + // all unique tokens can over-credit one extra shared block. + TLLM_CHECK_WITH_INFO(promptInputLen > 0, "Unexpected: promptInputLen == 0"); + auto const maxRecoverableSharedBlocks = (promptInputLen - 1) / getTokensPerBlock(); // Only subtract from shared blocks (reusable blocks are always shared) auto const reusableSharedBlocks = std::min(numReusableBlocks, numSharedBlocks); numRequiredBlocks -= reusableSharedBlocks; @@ -2294,7 +2271,8 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep( return 0; } -SizeType32 KVCacheManager::getRemainingBlocksToCompletion(LlmRequest const& req, SizeType32 windowSize) const +SizeType32 KVCacheManager::getRemainingBlocksToCompletion( + LlmRequest const& req, SizeType32 windowSize, std::optional const& cachedSummary) const { // Default to zero; overwritten below when block reuse is active for a first-chunk context request. req.setEstimatedReusableTokens(0); @@ -2338,21 +2316,21 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion(LlmRequest const& req, if (mEnableBlockReuse && !mBlockManager.isVariableWindow() && req.isContextInitState() && req.isFirstContextChunk() && numAllocBlocksPerBeam == 0) { - auto const uniqueTokens = req.getUniqueTokens(0); + // Use the cached summary if provided; otherwise perform a fresh tree walk. + auto const summary + = cachedSummary.has_value() ? cachedSummary.value() : analyzePrefixReuse(req.getUniqueTokens(0), req); // Block budget: only subtract blocks that are already allocated (have active refs). // Free cached blocks are already counted in the eviction policy's free pool and // must not be double-counted against the capacity estimate. - auto const numReusableBlocksAllocated = countReusableBlocks(uniqueTokens, req, /*onlyAllocated=*/true); - numReusableContextBlocks = std::min(numReusableBlocksAllocated, numContextBlocks); + numReusableContextBlocks = std::min(summary.reusableBlocksAllocated, numContextBlocks); // Token budget: count all reusable blocks (free or allocated). Cached tokens need // not be recomputed regardless of whether their blocks currently have active refs. - auto const numReusableBlocksAll = countReusableBlocks(uniqueTokens, req, /*onlyAllocated=*/false); - req.setEstimatedReusableTokens(std::min(numReusableBlocksAll, numContextBlocks) * getTokensPerBlock()); + req.setEstimatedReusableTokens(std::min(summary.reusableBlocksAll, numContextBlocks) * getTokensPerBlock()); TLLM_LOG_DEBUG( "getRemainingBlocksToCompletion: request ID %lu, numContextBlocks=%d, " "numReusableBlocksAllocated=%d, numReusableBlocksAll=%d, " "numReusableContextBlocks=%d, numGenBlocksPerBeam=%d", - req.mRequestId, numContextBlocks, numReusableBlocksAllocated, numReusableBlocksAll, + req.mRequestId, numContextBlocks, summary.reusableBlocksAllocated, summary.reusableBlocksAll, numReusableContextBlocks, numGenBlocksPerBeam); } @@ -2481,17 +2459,10 @@ void WindowBlockManager::detachFrontBlock(GenerationRequest& sequence) sequence.removeFrontBlock(mWindowSize); } -std::optional KVCacheManager::findNewContextBlock( +PrefixReuseSummary KVCacheManager::analyzePrefixReuse( VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const { - auto newContextBlockOpt = mBlockManager.findNewContextBlock(uniqueTokens, llmRequest); - return newContextBlockOpt; -} - -SizeType32 KVCacheManager::countReusableBlocks( - VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest, bool onlyAllocated) const -{ - return mBlockManager.countReusableBlocks(uniqueTokens, llmRequest, onlyAllocated); + return mBlockManager.analyzePrefixReuse(uniqueTokens, llmRequest); } void KVCacheManager::addSequence( diff --git a/cpp/tensorrt_llm/batch_manager/scheduledBlocksManager.h b/cpp/tensorrt_llm/batch_manager/scheduledBlocksManager.h index c7852ef0a688..212543889d8d 100644 --- a/cpp/tensorrt_llm/batch_manager/scheduledBlocksManager.h +++ b/cpp/tensorrt_llm/batch_manager/scheduledBlocksManager.h @@ -22,6 +22,7 @@ #include #include #include +#include namespace tensorrt_llm::batch_manager::kv_cache_manager { @@ -37,28 +38,40 @@ class NoEvictScheduledBlocksManager { } - void decrementReservedBlocks(LlmRequest const& req) + /// @brief Check whether enough blocks are available for the request. + /// Caches the per-window block counts for a subsequent commitBlocks() call. + bool enoughAvailableBlocks( + LlmRequest const& req, std::optional const& cachedSummary = std::nullopt) { - for (auto& [windowSize, availableBlocks] : mAvailableBlocks) + mCachedNeededBlocks.clear(); + bool enough = true; + for (auto const& [windowSize, availableBlocks] : mAvailableBlocks) { - availableBlocks -= mKvCacheManager.getRemainingBlocksToCompletion(req, windowSize); + auto const needed = mKvCacheManager.getRemainingBlocksToCompletion(req, windowSize, cachedSummary); + mCachedNeededBlocks.emplace_back(windowSize, needed); + if (needed > availableBlocks) + { + enough = false; + } } + return enough; } - bool enoughAvailableBlocks(LlmRequest const& req) + /// @brief Subtract the block counts cached by the last enoughAvailableBlocks() call. + void commitBlocks() { - return std::all_of(mAvailableBlocks.cbegin(), mAvailableBlocks.cend(), - [this, &req](auto const& pair) - { - auto const& [windowSize, availableBlocks] = pair; - auto const neededBlocks = mKvCacheManager.getRemainingBlocksToCompletion(req, windowSize); - return neededBlocks <= availableBlocks; - }); + for (auto const& [windowSize, needed] : mCachedNeededBlocks) + { + mAvailableBlocks[windowSize] -= needed; + } + mCachedNeededBlocks.clear(); } private: BaseKVCacheManager const& mKvCacheManager; std::map mAvailableBlocks; + // Cache from last enoughAvailableBlocks call, consumed by commitBlocks + std::vector> mCachedNeededBlocks; }; class MaxUtilizationScheduledBlocksManager @@ -77,12 +90,14 @@ class MaxUtilizationScheduledBlocksManager } } - std::optional> prepareNewNumberOfBlocksIfWeEndUpScheduling(LlmRequest const& req) + std::optional> prepareNewNumberOfBlocksIfWeEndUpScheduling( + LlmRequest const& req, std::optional const& cachedSummary = std::nullopt) { std::map blocksIfScheduled; for (auto const& [windowSize, numScheduled] : mNumScheduledBlocks) { - auto const required = mKvCacheManager.getNeededBlocksOneStep(req, mTwoStepsLookAhead, windowSize); + auto const required + = mKvCacheManager.getNeededBlocksOneStep(req, mTwoStepsLookAhead, windowSize, cachedSummary); TLLM_LOG_DEBUG("MaxUtilizationScheduler: request ID %lu required blocks %i for %i window size", req.mRequestId, required, windowSize); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 9ff754cf1e99..1fdb58ee5325 100755 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -167,6 +167,12 @@ class PyKvCacheManager : public tbk::BaseKVCacheManager NB_OVERRIDE_PURE(isEnableBlockReuse); } + tbk::PrefixReuseSummary analyzePrefixReuse( + tbk::VecUniqueTokens const& uniqueTokens, tb::LlmRequest const& llmRequest) const override + { + NB_OVERRIDE_PURE(analyzePrefixReuse, uniqueTokens, llmRequest); + } + bool isEnablePartialReuse() const override { NB_OVERRIDE_PURE(isEnablePartialReuse); @@ -182,12 +188,6 @@ class PyKvCacheManager : public tbk::BaseKVCacheManager NB_OVERRIDE_PURE(isCrossKv); } - std::optional findNewContextBlock( - VecUniqueTokens const& uniqueTokens, tb::LlmRequest const& llmRequest) const override - { - NB_OVERRIDE_PURE(findNewContextBlock, uniqueTokens, llmRequest); - } - void storeContextBlocks(tb::LlmRequest const& llmRequest) override { NB_OVERRIDE_PURE(storeContextBlocks, llmRequest); @@ -255,12 +255,6 @@ class PyKvCacheManager : public tbk::BaseKVCacheManager { NB_OVERRIDE_PURE(flushIterationEvents); } - - SizeType32 countReusableBlocks(VecUniqueTokens const& uniqueTokens, tb::LlmRequest const& llmRequest, - bool onlyAllocated = false) const override - { - NB_OVERRIDE_PURE(countReusableBlocks, uniqueTokens, llmRequest, onlyAllocated); - } }; // TODO: Deduplicate executor bindings KvCacheStats @@ -316,6 +310,12 @@ class PyBasePeftCacheManager : public tb::BasePeftCacheManager void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) { + nb::class_(m, "PrefixReuseSummary") + .def(nb::init<>()) + .def_ro("reusable_blocks_allocated", &tbk::PrefixReuseSummary::reusableBlocksAllocated) + .def_ro("reusable_blocks_all", &tbk::PrefixReuseSummary::reusableBlocksAll) + .def_ro("first_new_block", &tbk::PrefixReuseSummary::firstNewBlock); + nb::class_(m, "KvCacheStats") .def(nb::init<>()) .def_rw("max_num_blocks", &tbk::KvCacheStats::maxNumBlocks) @@ -497,10 +497,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) .def("store_context_blocks", &BaseKVCacheManager::storeContextBlocks, nb::call_guard()) .def("store_blocks_for_reuse", &BaseKVCacheManager::storeBlocksForReuse, nb::call_guard()) - .def("find_new_context_block", &BaseKVCacheManager::findNewContextBlock, nb::arg("unique_tokens"), + .def("analyze_prefix_reuse", &BaseKVCacheManager::analyzePrefixReuse, nb::arg("unique_tokens"), nb::arg("llm_request"), nb::call_guard()) - .def("count_reusable_blocks", &BaseKVCacheManager::countReusableBlocks, nb::arg("unique_tokens"), - nb::arg("llm_request"), nb::arg("only_allocated") = false, nb::call_guard()) .def("get_cache_block_ids", &BaseKVCacheManager::getCacheBlockIds, nb::call_guard()) .def("get_batch_cache_block_ids", &BaseKVCacheManager::getBatchCacheBlockIds, nb::call_guard()) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 995453fd669d..833456470b5c 100755 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -5804,9 +5804,9 @@ TEST(KVCacheManagerReuseAccountingTest, ReuseAwareBlockEstimatesStayConsistentAf // Note: storeContextBlocks only stores (length - 1) tokens worth of blocks // For 64 tokens with 16 tokens/block, only 63/16 = 3 full blocks are stored - auto const reusableBlocks = kvCacheManager->countReusableBlocks(req1.getUniqueTokens(0), req1); + auto const summary = kvCacheManager->analyzePrefixReuse(req1.getUniqueTokens(0), req1); auto const expectedReusableBlocks = (promptLength - 1) / tokensPerBlock; // 3 blocks - EXPECT_EQ(reusableBlocks, expectedReusableBlocks); + EXPECT_EQ(summary.reusableBlocksAll, expectedReusableBlocks); // After removeSequence, reusable blocks have no active refs and are free in the eviction queue. // The scheduling functions must NOT subtract free reusable blocks to avoid double-counting @@ -5837,6 +5837,77 @@ TEST(KVCacheManagerReuseAccountingTest, ReuseAwareBlockEstimatesStayConsistentAf EXPECT_EQ(remainingAfterContextAlloc, maxNewTokens / tokensPerBlock); } +TEST(KVCacheManagerReuseAccountingTest, NeededBlocksOneStepCapsAllocatedReuseAtExactBlockBoundary) +{ + auto const stream = std::make_shared(); + auto constexpr tokensPerBlock = 16; + auto constexpr promptLength = 48; // 3 full context blocks + auto constexpr reusablePrefixLength = promptLength + 1; + auto constexpr maxNewTokens = 32; + auto constexpr maxBeamWidth = 1; + auto constexpr maxAttentionWindow = 512; + auto constexpr maxNumTokens = 1024; + + auto kvCacheManager = createKvCacheManager( + KvCacheManagerInstantiationParameters{ + /* numLayers */ 1, + /* numHeads */ 1, + /* sizePerHead */ 1, + /* tokensPerBlock */ tokensPerBlock, + /* blocksPerWindow */ blocksAndWindow(/* numPrimaryBlocks */ 256, /* windowSize */ maxAttentionWindow), + /* sinkTokenLength */ 0, + /* maxAttentionWindow */ maxAttentionWindow, + /* maxBeamWidth */ maxBeamWidth, + /* maxNumTokens */ maxNumTokens, + /* kvCacheBlockReuse */ true, + }, + stream); + kvCacheManager->allocatePools(/*useUvm=*/false); + auto const onlyWindowSize = theOnlyWindowSize(*kvCacheManager); + auto const samplingConfig = tensorrt_llm::runtime::SamplingConfig{maxBeamWidth}; + auto constexpr isStreaming = true; + auto makeRequest = [&](LlmRequest::RequestIdType requestId, std::vector const& tokens) + { + return LlmRequest( + requestId, maxNewTokens, std::make_shared>(tokens), samplingConfig, isStreaming); + }; + + auto reusableTokens = std::vector(static_cast(reusablePrefixLength)); + std::iota(reusableTokens.begin(), reusableTokens.end(), 0); + + auto seedReq = makeRequest(0, reusableTokens); + kvCacheManager->addSequence(seedReq.mRequestId, seedReq.getPromptLen(), maxBeamWidth, seedReq); + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(seedReq); + kvCacheManager->removeSequence(seedReq.mRequestId, seedReq); + + // Keep a request with 49 prompt tokens active so all 3 full prefix blocks remain + // both reusable and allocated. + auto holderReq = makeRequest(1, reusableTokens); + kvCacheManager->addSequence(holderReq.mRequestId, holderReq.getPromptLen(), maxBeamWidth, holderReq); + EXPECT_EQ(holderReq.getContextCurrentPosition(), promptLength); + + auto promptTokens = std::vector(static_cast(promptLength)); + std::iota(promptTokens.begin(), promptTokens.end(), 0); + + auto req1 = makeRequest(2, promptTokens); + + // Simulate a recompute-style context request: prompt length stays at the exact block + // boundary, but one generated token already exists in the token history. + req1.addNewToken(promptLength, 0); + + auto const summaryAlloc = kvCacheManager->analyzePrefixReuse(req1.getUniqueTokens(0), req1); + EXPECT_EQ(summaryAlloc.reusableBlocksAllocated, promptLength / tokensPerBlock); + + auto const neededOneStep + = kvCacheManager->getNeededBlocksOneStep(req1, /*twoStepsLookAhead=*/false, onlyWindowSize); + EXPECT_EQ(neededOneStep, 1); + + auto const numAllocBlocksBeforeAdd = kvCacheManager->getNumAllocTotalBlocks(); + kvCacheManager->addSequence(req1.mRequestId, req1.getPromptLen(), maxBeamWidth, req1); + auto const numAllocBlocksAfterAdd = kvCacheManager->getNumAllocTotalBlocks(); + EXPECT_EQ(numAllocBlocksAfterAdd - numAllocBlocksBeforeAdd, neededOneStep); +} + TEST(KVCacheManagerReuseAccountingTest, CountReusableBlocksNoMatchReturnsZero) { auto const stream = std::make_shared(); @@ -5874,8 +5945,8 @@ TEST(KVCacheManagerReuseAccountingTest, CountReusableBlocksNoMatchReturnsZero) }; // No blocks should be reusable since nothing has been cached - auto const reusableBlocks = kvCacheManager->countReusableBlocks(req.getUniqueTokens(0), req); - EXPECT_EQ(reusableBlocks, 0); + auto const summaryEmpty = kvCacheManager->analyzePrefixReuse(req.getUniqueTokens(0), req); + EXPECT_EQ(summaryEmpty.reusableBlocksAll, 0); } TEST(KVCacheManagerReuseAccountingTest, CountReusableBlocksPartialMatch) @@ -5937,8 +6008,8 @@ TEST(KVCacheManagerReuseAccountingTest, CountReusableBlocksPartialMatch) }; // Should find 2 reusable blocks (first 32 tokens match) - auto const reusableBlocks = kvCacheManager->countReusableBlocks(req1.getUniqueTokens(0), req1); - EXPECT_EQ(reusableBlocks, 2); + auto const summaryShared = kvCacheManager->analyzePrefixReuse(req1.getUniqueTokens(0), req1); + EXPECT_EQ(summaryShared.reusableBlocksAll, 2); // After removeSequence, reusable blocks are free (no active refs). // getNeededBlocksOneStep must NOT subtract free reusable blocks to avoid double-counting. @@ -6203,9 +6274,9 @@ TEST(KVCacheManagerReuseAccountingTest, MultipleRequestsWithSharedPrefix) true, }; - // Should reuse 2 blocks (shared prefix) — public API counts all reusable regardless of ref state - auto const reusableBlocks = kvCacheManager->countReusableBlocks(req1.getUniqueTokens(0), req1); - EXPECT_EQ(reusableBlocks, sharedPrefixLength / tokensPerBlock); + // Should reuse 2 blocks (shared prefix) — analyzePrefixReuse counts all reusable regardless of ref state + auto const summaryPrefix = kvCacheManager->analyzePrefixReuse(req1.getUniqueTokens(0), req1); + EXPECT_EQ(summaryPrefix.reusableBlocksAll, sharedPrefixLength / tokensPerBlock); // After removeSequence, reusable blocks are free (no active refs). // getNeededBlocksOneStep must NOT subtract free reusable blocks. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 8f1a983db9d2..2103e55abb78 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1332,6 +1332,17 @@ def _executor_loop_pp(self): self._add_inflight_ids(scheduled_batch) + # Snapshot the request IDs we just added so we can roll + # back any that prepare_resources drops (e.g., via the + # reuse-budget skip). _add_inflight_ids only inserts + # last-chunk context and generation request IDs, so those + # are the only ones we need to track. + added_inflight_ids = { + req.request_id + for req in scheduled_batch.context_requests_last_chunk + + scheduled_batch.generation_requests + } + if self.kv_cache_transceiver: # For generation requests which have completed KV cache transfer self._prepare_disagg_gen_transmission_complete( @@ -1341,88 +1352,122 @@ def _executor_loop_pp(self): self.resource_manager.prepare_resources(scheduled_batch) - # The generation requests that do not have batch_idx - # need to be in front of the batch due to the assumptions - # made in model_engine.py::_forward_step. This is only important - # for disaggregated serving. For non-disaggregated serving, - # the generation requests always have batch_idx. - scheduled_batch.generation_requests = sorted( # stable sort - scheduled_batch.generation_requests, - key=lambda req: int(req.py_batch_idx is not None), - ) - - if self.kv_cache_transceiver: - # Return the first token to the client - self._handle_first_token_response(scheduled_batch) - - # Stage 1.1: Async forward (all ranks) and decoding pass (last rank only) - if not self.dist.is_last_pp_rank: - with torch.cuda.nvtx.range( - f"_forward_step_inter_pp pp_rank {self.dist.pp_rank}" - ): - sample_state = self._forward_step_inter_pp( - scheduled_batch) + # Erase inflight IDs for any request dropped by + # prepare_resources. Without this, a skipped request + # stays in inflight_req_ids forever and the scheduler + # (which skips inflight IDs) never reschedules it. + remaining_inflight_ids = { + req.request_id + for req in scheduled_batch.context_requests_last_chunk + + scheduled_batch.generation_requests + } + for dropped_id in (added_inflight_ids - + remaining_inflight_ids): + self.inflight_req_ids.erase(dropped_id) + + # prepare_resources may have skipped all requests (e.g., + # prefix-aware reuse budget skip removes over-admitted + # context requests), leaving an empty batch. Re-evaluate + # can_queue to avoid forwarding on an empty microbatch. + can_queue, _ = self._can_queue(scheduled_batch) + if not can_queue: + logger.debug( + f"microbatch {microbatch_id} emptied by prepare_resources, skipping forward" + ) + # Any remaining inflight IDs belong to requests still + # in the batch; erase those too since we're not + # forwarding. + self._remove_inflight_ids(scheduled_batch) + self.micro_batches[microbatch_id] = None else: - with torch.cuda.nvtx.range( - f"_forward_step_last_pp pp_rank {self.dist.pp_rank}" - ): - # init_disagg_gen_requests must be before engine forward, where the prev_seq_slot is updated. - if self.guided_decoder is not None and self.kv_cache_transceiver: - self.guided_decoder.add_batch(scheduled_batch) - self.guided_decoder.init_disagg_gen_requests() - - batch_outputs = self._forward_step(scheduled_batch) - - guided_decoder_failed_requests = None - if self.guided_decoder is not None: - self.guided_decoder.add_batch(scheduled_batch) - guided_decoder_failed_requests = self.guided_decoder.execute( - batch_outputs['logits']) - - if self.pp_multi_stream_sample: - # Wait for the previous sample to finish. - self.finish_sample_event.wait() - # Copy the batch outputs as sampler inputs - # to avoid next forward step overwriting them. - batch_outputs_copy = { - name: tensor.clone() - for name, tensor in batch_outputs.items() - } - self.sample_stream.wait_stream( - torch.cuda.current_stream()) - with torch.cuda.stream(self.sample_stream): - sample_state = self._sample_async( - scheduled_batch, batch_outputs_copy) - self.finish_sample_event.record() - else: - sample_state = self._sample_async( - scheduled_batch, batch_outputs) - - assert sample_state is not None, "Sampling failed" - - # Handle guided decoder errors after _sample_async to avoid state conflicts. - # If called before, failed requests would be marked as GENERATION_COMPLETE, - # causing _sample_async to fail when accessing context_chunk_size property. - self._handle_guided_decoder_errors( - scheduled_batch, guided_decoder_failed_requests) + # The generation requests that do not have batch_idx + # need to be in front of the batch due to the assumptions + # made in model_engine.py::_forward_step. This is only important + # for disaggregated serving. For non-disaggregated serving, + # the generation requests always have batch_idx. + scheduled_batch.generation_requests = sorted( # stable sort + scheduled_batch.generation_requests, + key=lambda req: int(req.py_batch_idx is not None), + ) - self._update_request_states(scheduled_batch) - if not self.disable_overlap_scheduler: - self._update_generation_requests_that_will_complete_next_iteration( - scheduled_batch.generation_requests) + if self.kv_cache_transceiver: + # Return the first token to the client + self._handle_first_token_response(scheduled_batch) + + # Stage 1.1: Async forward (all ranks) and decoding pass (last rank only) + if not self.dist.is_last_pp_rank: + with torch.cuda.nvtx.range( + f"_forward_step_inter_pp pp_rank {self.dist.pp_rank}" + ): + sample_state = self._forward_step_inter_pp( + scheduled_batch) + else: + with torch.cuda.nvtx.range( + f"_forward_step_last_pp pp_rank {self.dist.pp_rank}" + ): + # init_disagg_gen_requests must be before engine forward, where the prev_seq_slot is updated. + if self.guided_decoder is not None and self.kv_cache_transceiver: + self.guided_decoder.add_batch( + scheduled_batch) + self.guided_decoder.init_disagg_gen_requests( + ) + + batch_outputs = self._forward_step( + scheduled_batch) - if self.enable_iter_perf_stats: - iter_stats.inflight_batching_stats.num_ctx_tokens = self.model_engine.iter_states[ - 'num_ctx_tokens'] - batch_state = BatchStatePP( - scheduled_requests=scheduled_batch, - sample_state=sample_state, - iter_start_time=iter_start_time, - iter_stats=iter_stats, - microbatch_id=microbatch_id, - ) + guided_decoder_failed_requests = None + if self.guided_decoder is not None: + self.guided_decoder.add_batch( + scheduled_batch) + guided_decoder_failed_requests = self.guided_decoder.execute( + batch_outputs['logits']) + + if self.pp_multi_stream_sample: + # Wait for the previous sample to finish. + self.finish_sample_event.wait() + # Copy the batch outputs as sampler inputs + # to avoid next forward step overwriting them. + batch_outputs_copy = { + name: tensor.clone() + for name, tensor in + batch_outputs.items() + } + self.sample_stream.wait_stream( + torch.cuda.current_stream()) + with torch.cuda.stream(self.sample_stream): + sample_state = self._sample_async( + scheduled_batch, batch_outputs_copy) + self.finish_sample_event.record() + else: + sample_state = self._sample_async( + scheduled_batch, batch_outputs) + + assert sample_state is not None, "Sampling failed" + + # Handle guided decoder errors after _sample_async to avoid state conflicts. + # If called before, failed requests would be marked as GENERATION_COMPLETE, + # causing _sample_async to fail when accessing context_chunk_size property. + self._handle_guided_decoder_errors( + scheduled_batch, + guided_decoder_failed_requests) + + self._update_request_states(scheduled_batch) + if not self.disable_overlap_scheduler: + self._update_generation_requests_that_will_complete_next_iteration( + scheduled_batch.generation_requests) + + if self.enable_iter_perf_stats: + iter_stats.inflight_batching_stats.num_ctx_tokens = self.model_engine.iter_states[ + 'num_ctx_tokens'] + batch_state = BatchStatePP( + scheduled_requests=scheduled_batch, + sample_state=sample_state, + iter_start_time=iter_start_time, + iter_stats=iter_stats, + microbatch_id=microbatch_id, + ) - self.micro_batches[microbatch_id] = batch_state + self.micro_batches[microbatch_id] = batch_state # Stage 1.2: Sync sampler for previous microbatch to start new sample state comm chain. # For last PP rank, we must synchronize the previous batch @@ -1899,6 +1944,12 @@ def _executor_loop(self): self.resource_manager.prepare_resources(scheduled_batch) + # prepare_resources may have skipped all requests (e.g., + # prefix-aware reuse budget skip removes over-admitted + # context requests), leaving an empty batch. Re-evaluate + # can_queue to avoid forwarding on an empty batch. + can_queue, _ = self._can_queue(scheduled_batch) + if self.kv_connector_manager: self.kv_connector_manager.handle_metadata() @@ -2171,6 +2222,13 @@ def _executor_loop_overlap(self): self.resource_manager.prepare_resources(scheduled_batch) + # prepare_resources may have skipped all requests (e.g., + # prefix-aware reuse budget skip removes over-admitted + # context requests), leaving an empty batch. Re-evaluate + # can_queue to avoid forwarding on an empty batch. + can_queue, can_queue_this_rank = self._can_queue( + scheduled_batch) + if self.kv_connector_manager: self.kv_connector_manager.handle_metadata() diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index ad606440ce0b..cf9ec32fe63d 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -586,9 +586,8 @@ def probe_prefix_match_length(self, input_tokens, lora_task_id=None): sampling_config=SamplingConfig(), is_streaming=False, lora_task_id=lora_task_id) - num_blocks = self.impl.count_reusable_blocks(unique_tokens, dummy_req, - False) - return num_blocks * self.tokens_per_block + summary = self.impl.analyze_prefix_reuse(unique_tokens, dummy_req) + return summary.reusable_blocks_all * self.tokens_per_block def shutdown(self): self.impl.release_pools() @@ -657,9 +656,24 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): if req.is_first_context_chunk and self._kv_connector_should_add_sequence( req): if remaining_budget is not None: - unique_tokens = req.get_unique_tokens(0) - reusable_blocks = self.impl.count_reusable_blocks( - unique_tokens, req, False) + # analyze_prefix_reuse is not supported on + # variable-window KV cache managers (the C++ + # side asserts !isVariableWindow). Fall back + # to a conservative estimate of zero reusable + # blocks, matching probe_prefix_match_length. + # + # Use reusable_blocks_all (not _allocated) for the + # token-budget decision: free-cached blocks still + # avoid recompute, so they count against the + # recompute budget. reusable_blocks_allocated is + # a narrower block-budget quantity and would + # under-count reuse, causing false skips. + if getattr(self.impl, 'is_variable_window', False): + reusable_blocks = 0 + else: + unique_tokens = req.get_unique_tokens(0) + reusable_blocks = self.impl.analyze_prefix_reuse( + unique_tokens, req).reusable_blocks_all actual_reuse = (reusable_blocks * self.tokens_per_block) req_compute = self._estimate_post_reuse_compute( diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 82555e17e28f..2fe1f61928f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -60,7 +60,7 @@ def is_generation_only(self) -> bool: @property def can_run_cuda_graph(self) -> bool: - return self.num_context_requests == 0 + return self.num_context_requests == 0 and len(self.generation_requests) > 0 @property def batch_size(self) -> int: @@ -1263,18 +1263,18 @@ def _prefill_contributed_blocks(self, active_requests: RequestList) -> tuple[set # Chunked context request already executing if enable_block_reuse: unique_tokens = req.get_unique_tokens(0) - block_key = self.kv_cache_manager.find_new_context_block(unique_tokens, req) - if block_key is not None: - newly_contributed_context_blocks.add(block_key) + summary = self.kv_cache_manager.analyze_prefix_reuse(unique_tokens, req) + if summary.first_new_block is not None: + newly_contributed_context_blocks.add(summary.first_new_block) if cross_enable_reuse: encoder_unique_tokens = req.get_encoder_unique_tokens() if encoder_unique_tokens is not None: - block_key = self.cross_kv_cache_manager.find_new_context_block( + summary = self.cross_kv_cache_manager.analyze_prefix_reuse( encoder_unique_tokens, req ) - if block_key is not None: - newly_contributed_cross_context_blocks.add(block_key) + if summary.first_new_block is not None: + newly_contributed_cross_context_blocks.add(summary.first_new_block) return newly_contributed_context_blocks, newly_contributed_cross_context_blocks @@ -1283,13 +1283,13 @@ def _one_manager_beneficial_to_skip( ) -> bool: """ Check if skipping is beneficial for one KV cache manager. - C++ reference: capacityScheduler.cpp:70-92 (oneManagerBeneficialToSkip) + C++ reference: inlined skip-check logic in capacityScheduler.cpp """ - new_context_block = kv_cache_manager.find_new_context_block(unique_tokens, req) - if new_context_block is not None: - if new_context_block in newly_contributed_blocks: + summary = kv_cache_manager.analyze_prefix_reuse(unique_tokens, req) + if summary.first_new_block is not None: + if summary.first_new_block in newly_contributed_blocks: return True - newly_contributed_blocks.add(new_context_block) + newly_contributed_blocks.add(summary.first_new_block) return False def _beneficial_to_skip( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9b6b525d00e1..15b45b44c15a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2194,7 +2194,7 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): ) def _to_pybind(self): - return _KvCacheConfig( + config = _KvCacheConfig( enable_block_reuse=self.enable_block_reuse, max_tokens=self.max_tokens, max_attention_window=self.max_attention_window, @@ -2211,6 +2211,7 @@ def _to_pybind(self): attention_dp_events_gather_period_ms=self. attention_dp_events_gather_period_ms, max_gpu_total_bytes=self.max_gpu_total_bytes) + return config @field_validator('free_gpu_memory_fraction') @classmethod diff --git a/tests/unittest/_torch/executor/test_kvcache_aware_router.py b/tests/unittest/_torch/executor/test_kvcache_aware_router.py index 34518faa7549..62dcebaa1c7a 100644 --- a/tests/unittest/_torch/executor/test_kvcache_aware_router.py +++ b/tests/unittest/_torch/executor/test_kvcache_aware_router.py @@ -410,8 +410,8 @@ def test_variable_window_returns_zero(self): result = KVCacheManager.probe_prefix_match_length(mgr, input_tokens=[1, 2, 3]) assert result == 0 - # count_reusable_blocks should NOT be called (would crash) - mgr.impl.count_reusable_blocks.assert_not_called() + # analyze_prefix_reuse should NOT be called (would crash) + mgr.impl.analyze_prefix_reuse.assert_not_called() def test_empty_tokens_returns_zero(self): from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager @@ -432,7 +432,9 @@ def test_block_to_token_conversion(self): mgr.enable_block_reuse = True mgr.impl = Mock() mgr.impl.is_variable_window = False - mgr.impl.count_reusable_blocks = Mock(return_value=3) + mock_summary = Mock() + mock_summary.reusable_blocks_all = 3 + mgr.impl.analyze_prefix_reuse = Mock(return_value=mock_summary) mgr.tokens_per_block = 64 result = KVCacheManager.probe_prefix_match_length(mgr, input_tokens=list(range(200))) diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index 01b5c90edd30..a0dee5aa8d16 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -37,6 +37,15 @@ from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy +@dataclass +class MockPrefixReuseSummary: + """Mock for C++ PrefixReuseSummary returned by analyze_prefix_reuse.""" + + reusable_blocks_allocated: int = 0 + reusable_blocks_all: int = 0 + first_new_block: Optional[object] = None + + def _make_request( request_id: int, prompt_len: int = 10, @@ -165,8 +174,14 @@ def start_scheduling(self): def scheduling_remove_sequence(self, req_id: int): pass - def find_new_context_block(self, unique_tokens, req): - return None + def analyze_prefix_reuse(self, unique_tokens, req): + if self.enable_block_reuse and unique_tokens: + # Derive a deterministic block key from the tokens so that + # duplicate requests produce the same first_new_block and + # trigger the beneficial-to-skip logic. + block_key = tuple(t if isinstance(t, int) else t for t in unique_tokens) + return MockPrefixReuseSummary(first_new_block=block_key) + return MockPrefixReuseSummary() def get_max_resource_count(self) -> int: return self._num_free_blocks @@ -2180,8 +2195,10 @@ def test_delay_duplicate_request(self): r1 = _make_request(1, prompt_len=21, input_tokens=tokens) r2 = _make_request(2, prompt_len=21, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2]) - # With reuse enabled, beneficial_to_skip may delay r1 and r2 - assert len(fitting) >= 1 + # With reuse enabled, r1 and r2 share the same prefix as r0 and are delayed + assert len(fitting) == 1 + assert fitting[0].request_id == 0 + assert len(paused) == 2 def test_delay_duplicate_request_chunked(self): """C++ ref: DelayDuplicateRequestChunked""" @@ -2197,7 +2214,10 @@ def test_delay_duplicate_request_chunked(self): r1 = _make_request(1, prompt_len=50, input_tokens=tokens) r1.context_chunk_size = 20 fitting, disagg, paused = scheduler.schedule_request([r0, r1]) - assert len(fitting) >= 1 + # r1 shares the same prefix as r0 and is delayed + assert len(fitting) == 1 + assert fitting[0].request_id == 0 + assert len(paused) == 1 def test_delay_five_requests_complicated(self): """C++ ref: DelayFiveRequestsComplicated""" @@ -2227,7 +2247,10 @@ def test_reuse_aware_allows_more_requests(self): r0 = _make_request(0, prompt_len=20, input_tokens=tokens) r1 = _make_request(1, prompt_len=20, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) - assert len(fitting) >= 1 + # r1 shares the same prefix as r0 and is delayed + assert len(fitting) == 1 + assert fitting[0].request_id == 0 + assert len(paused) == 1 def test_reuse_aware_partial_prefix_match(self): """C++ ref: ReuseAwareSchedulingWithPartialPrefixMatch""" @@ -2242,7 +2265,8 @@ def test_reuse_aware_partial_prefix_match(self): r0 = _make_request(0, prompt_len=30, input_tokens=tokens0) r1 = _make_request(1, prompt_len=30, input_tokens=tokens1) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) - assert len(fitting) >= 1 + # Different token sequences produce different block keys — no skipping + assert len(fitting) == 2 def test_no_reuse_with_different_prompts(self): """C++ ref: NoReuseWithDifferentPrompts""" From 5508dcf89cf0e71c992f723fc932766e6abc2621 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Thu, 16 Apr 2026 13:05:50 -0700 Subject: [PATCH 2/3] Address coderabbit's comments. Signed-off-by: Simeng Liu --- .../batch_manager/capacityScheduler.cpp | 35 ++++++++----- .../batch_manager/kvCacheManager.cpp | 15 +++--- .../nanobind/batch_manager/kvCacheManager.cpp | 4 +- .../_torch/pyexecutor/scheduler/scheduler.py | 52 +++++++++---------- .../_torch/executor/test_py_scheduler.py | 26 ++++++---- 5 files changed, 74 insertions(+), 58 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 478483be72c3..0639787d0a16 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -69,25 +69,20 @@ prefillWithChunkedContextsAlreadyExecuting(RequestList const& activeRequests, /// @brief Check if a single manager's summary indicates we should skip this request. /// @details Returns true if the request's first new context block was already contributed -/// by an earlier scheduled request (so waiting would let us reuse it). Registers the -/// block as contributed when not skipping. +/// by an earlier scheduled request (so waiting would let us reuse it). Does NOT mutate +/// the set — registration is deferred to beneficialToSkip after both KV checks pass. bool oneManagerBeneficialToSkip(std::optional const& summary, - std::unordered_set& newlyContributedContextBlocks) + std::unordered_set const& newlyContributedContextBlocks) { - if (summary.has_value() && summary->firstNewBlock.has_value()) - { - if (newlyContributedContextBlocks.count(summary->firstNewBlock.value()) > 0) - { - return true; - } - newlyContributedContextBlocks.insert(summary->firstNewBlock.value()); - } - return false; + return summary.has_value() && summary->firstNewBlock.has_value() + && newlyContributedContextBlocks.count(summary->firstNewBlock.value()) > 0; } /// @brief Check if it is beneficial to skip this request rather than schedule it. /// @details Returns true if this request can reuse KV cache block(s) that will be contributed /// by already-scheduled context requests. Uses pre-computed PrefixReuseSummary values. +/// When the request is NOT skipped, registers its firstNewBlock contributions so that +/// subsequent duplicate requests can be deferred. bool beneficialToSkip(std::optional const& summary, std::optional const& crossSummary, std::unordered_set& newlyContributedContextBlocks, @@ -97,7 +92,21 @@ bool beneficialToSkip(std::optional const& { return true; } - return oneManagerBeneficialToSkip(crossSummary, newlyContributedCrossContextBlocks); + if (oneManagerBeneficialToSkip(crossSummary, newlyContributedCrossContextBlocks)) + { + return true; + } + // Request is NOT skipped — register its contributions so subsequent duplicate + // requests can be deferred correctly. + if (summary.has_value() && summary->firstNewBlock.has_value()) + { + newlyContributedContextBlocks.insert(summary->firstNewBlock.value()); + } + if (crossSummary.has_value() && crossSummary->firstNewBlock.has_value()) + { + newlyContributedCrossContextBlocks.insert(crossSummary->firstNewBlock.value()); + } + return false; } } // namespace diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index b8d6096ad870..fb1de83e7b4f 100755 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -1152,24 +1152,19 @@ PrefixReuseSummary WindowBlockManager::analyzePrefixReuse( auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); PrefixReuseSummary summary; - BlockKey accumKey; - accumKey.loraTaskId = llmRequest.getLoraTaskId(); std::lock_guard lock(mCachedBlocksRootMutex); auto searchRoot = mCachedBlocksRoot; for (auto const& blockKey : blockKeys) { - accumKey.uniqueTokens.insert( - accumKey.uniqueTokens.end(), blockKey.uniqueTokens.begin(), blockKey.uniqueTokens.end()); - auto [partialMatch, numMatched, matchingBlock] = searchRoot != nullptr ? searchRoot->findMatchingBlock(blockKey, false, false) : std::make_tuple(false, 0, nullptr); if (matchingBlock == nullptr) { - summary.firstNewBlock = accumKey; + summary.firstNewBlock = blockKey; break; } @@ -2322,10 +2317,14 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion( // Block budget: only subtract blocks that are already allocated (have active refs). // Free cached blocks are already counted in the eviction policy's free pool and // must not be double-counted against the capacity estimate. - numReusableContextBlocks = std::min(summary.reusableBlocksAllocated, numContextBlocks); + // Cap at (promptLen-1)/tpb to avoid over-counting when the prompt is exactly + // block-aligned (last full block has not been committed to the tree yet). + SizeType32 const maxRecoverableBlocks = (req.mPromptLen - 1) / getTokensPerBlock(); + numReusableContextBlocks = std::min({summary.reusableBlocksAllocated, numContextBlocks, maxRecoverableBlocks}); // Token budget: count all reusable blocks (free or allocated). Cached tokens need // not be recomputed regardless of whether their blocks currently have active refs. - req.setEstimatedReusableTokens(std::min(summary.reusableBlocksAll, numContextBlocks) * getTokensPerBlock()); + req.setEstimatedReusableTokens( + std::min({summary.reusableBlocksAll, numContextBlocks, maxRecoverableBlocks}) * getTokensPerBlock()); TLLM_LOG_DEBUG( "getRemainingBlocksToCompletion: request ID %lu, numContextBlocks=%d, " "numReusableBlocksAllocated=%d, numReusableBlocksAll=%d, " diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 1fdb58ee5325..b3a4db3f4d64 100755 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -344,7 +344,9 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::arg("lora_task_id"), nb::arg("unique_tokens")) .def_ro("uses_extra_ids", &tbk::BlockKey::usesExtraIds) .def_ro("lora_task_id", &tbk::BlockKey::loraTaskId) - .def_ro("unique_tokens", &tbk::BlockKey::uniqueTokens); + .def_ro("unique_tokens", &tbk::BlockKey::uniqueTokens) + .def("__eq__", &tbk::BlockKey::operator==) + .def("__hash__", [](tbk::BlockKey const& key) -> size_t { return tbk::BlockKeyHasher{}(key); }); nb::class_(m, "BlockKeyHasher") .def_static("hash", &tbk::BlockKeyHasher::hash, nb::arg("block_key"), nb::arg("parent_hash") = 0); diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 2fe1f61928f0..10a33e36037f 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -1278,20 +1278,6 @@ def _prefill_contributed_blocks(self, active_requests: RequestList) -> tuple[set return newly_contributed_context_blocks, newly_contributed_cross_context_blocks - def _one_manager_beneficial_to_skip( - self, kv_cache_manager, unique_tokens, req: LlmRequest, newly_contributed_blocks: set - ) -> bool: - """ - Check if skipping is beneficial for one KV cache manager. - C++ reference: inlined skip-check logic in capacityScheduler.cpp - """ - summary = kv_cache_manager.analyze_prefix_reuse(unique_tokens, req) - if summary.first_new_block is not None: - if summary.first_new_block in newly_contributed_blocks: - return True - newly_contributed_blocks.add(summary.first_new_block) - return False - def _beneficial_to_skip( self, req: LlmRequest, @@ -1303,17 +1289,24 @@ def _beneficial_to_skip( A request should be skipped if it can reuse blocks contributed by already scheduled context requests. - C++ reference: capacityScheduler.cpp:97-123 (beneficialToSkip) + When the request is NOT skipped, its firstNewBlock contributions are + registered so that subsequent duplicate requests can be deferred. + + C++ reference: capacityScheduler.cpp (beneficialToSkip / oneManagerBeneficialToSkip) """ if not (req.is_context_init_state and req.is_first_context_chunk): return False + ctx_new_block = None + cross_new_block = None + if self.kv_cache_manager is not None and self.kv_cache_manager.enable_block_reuse: unique_tokens = req.get_unique_tokens(0) - if self._one_manager_beneficial_to_skip( - self.kv_cache_manager, unique_tokens, req, newly_contributed_context_blocks - ): - return True + summary = self.kv_cache_manager.analyze_prefix_reuse(unique_tokens, req) + if summary.first_new_block is not None: + if summary.first_new_block in newly_contributed_context_blocks: + return True + ctx_new_block = summary.first_new_block if ( self.cross_kv_cache_manager is not None @@ -1321,13 +1314,20 @@ def _beneficial_to_skip( ): encoder_unique_tokens = req.get_encoder_unique_tokens() if encoder_unique_tokens is not None: - if self._one_manager_beneficial_to_skip( - self.cross_kv_cache_manager, - encoder_unique_tokens, - req, - newly_contributed_cross_context_blocks, - ): - return True + summary = self.cross_kv_cache_manager.analyze_prefix_reuse( + encoder_unique_tokens, req + ) + if summary.first_new_block is not None: + if summary.first_new_block in newly_contributed_cross_context_blocks: + return True + cross_new_block = summary.first_new_block + + # Request is NOT skipped — register contributions so subsequent duplicate + # requests can be deferred correctly. + if ctx_new_block is not None: + newly_contributed_context_blocks.add(ctx_new_block) + if cross_new_block is not None: + newly_contributed_cross_context_blocks.add(cross_new_block) return False diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index a0dee5aa8d16..5e500d44292b 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -176,10 +176,13 @@ def scheduling_remove_sequence(self, req_id: int): def analyze_prefix_reuse(self, unique_tokens, req): if self.enable_block_reuse and unique_tokens: - # Derive a deterministic block key from the tokens so that - # duplicate requests produce the same first_new_block and - # trigger the beneficial-to-skip logic. - block_key = tuple(t if isinstance(t, int) else t for t in unique_tokens) + # Derive a deterministic, hashable block key from the tokens so that + # duplicate requests produce equal first_new_block values and trigger + # the beneficial-to-skip logic. UniqueToken objects (nanobind-wrapped + # C++ structs) have no __hash__/__eq__, so extract the primitive fields. + block_key = tuple( + t if isinstance(t, int) else (t.token_id, t.token_extra_id) for t in unique_tokens + ) return MockPrefixReuseSummary(first_new_block=block_key) return MockPrefixReuseSummary() @@ -2195,10 +2198,11 @@ def test_delay_duplicate_request(self): r1 = _make_request(1, prompt_len=21, input_tokens=tokens) r2 = _make_request(2, prompt_len=21, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1, r2]) - # With reuse enabled, r1 and r2 share the same prefix as r0 and are delayed + # With reuse enabled, r1 and r2 share the same prefix as r0 and are delayed. + # GUARANTEED_NO_EVICT skips them via continue — they are not in paused. assert len(fitting) == 1 assert fitting[0].request_id == 0 - assert len(paused) == 2 + assert len(paused) == 0 def test_delay_duplicate_request_chunked(self): """C++ ref: DelayDuplicateRequestChunked""" @@ -2214,10 +2218,11 @@ def test_delay_duplicate_request_chunked(self): r1 = _make_request(1, prompt_len=50, input_tokens=tokens) r1.context_chunk_size = 20 fitting, disagg, paused = scheduler.schedule_request([r0, r1]) - # r1 shares the same prefix as r0 and is delayed + # r1 shares the same prefix as r0 and is delayed. + # GUARANTEED_NO_EVICT skips it via continue — it is not in paused. assert len(fitting) == 1 assert fitting[0].request_id == 0 - assert len(paused) == 1 + assert len(paused) == 0 def test_delay_five_requests_complicated(self): """C++ ref: DelayFiveRequestsComplicated""" @@ -2247,10 +2252,11 @@ def test_reuse_aware_allows_more_requests(self): r0 = _make_request(0, prompt_len=20, input_tokens=tokens) r1 = _make_request(1, prompt_len=20, input_tokens=tokens) fitting, disagg, paused = scheduler.schedule_request([r0, r1]) - # r1 shares the same prefix as r0 and is delayed + # r1 shares the same prefix as r0 and is delayed. + # GUARANTEED_NO_EVICT skips it via continue — it is not in paused. assert len(fitting) == 1 assert fitting[0].request_id == 0 - assert len(paused) == 1 + assert len(paused) == 0 def test_reuse_aware_partial_prefix_match(self): """C++ ref: ReuseAwareSchedulingWithPartialPrefixMatch""" From 05ca2fd6b939cafe70905eb8ab3877e3c598b4f1 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Thu, 16 Apr 2026 17:21:27 -0700 Subject: [PATCH 3/3] Address coderabbit's comments, batch 2. Signed-off-by: Simeng Liu --- .../tensorrt_llm/batch_manager/kvCacheManager.h | 4 +++- .../batch_manager/capacityScheduler.cpp | 14 ++++++++++---- .../nanobind/batch_manager/kvCacheManager.cpp | 12 ++++++++---- 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index f37005d391e0..204ff813b3e2 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -574,7 +574,9 @@ struct PrefixReuseSummary /// Used by the token budget (NoEvict) since all cached tokens avoid recompute. SizeType32 reusableBlocksAll{0}; - /// First block key NOT found in the radix tree (nullopt = full prefix hit). + /// First block key NOT found in the radix tree. std::nullopt means either all full + /// prefix blocks matched (full prefix hit) or the request has no full block key to + /// probe yet; a concrete BlockKey identifies the first missing full block. /// Used by the capacity scheduler's skip-check logic to decide whether to defer a request. std::optional firstNewBlock{std::nullopt}; }; diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 0639787d0a16..be69cce666fb 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -287,12 +287,15 @@ std::tuple GuaranteedNoEvictScheduler::impl( std::optional crossSummary; if (isFirstChunkContext) { - if (kvCacheManager.isEnableBlockReuse()) + // analyzePrefixReuse asserts on variable-window managers; skip the walk there + // and let downstream callers fall back to their fresh tree-walk path. + if (kvCacheManager.isEnableBlockReuse() && !kvCacheManager.getBlockManager().isVariableWindow()) { auto uniqueTokens = req->getUniqueTokens(0); summary = kvCacheManager.analyzePrefixReuse(uniqueTokens, *req); } - if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse()) + if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse() + && !crossKvCacheManager->getBlockManager().isVariableWindow()) { auto uniqueTokens = *(req->getEncoderUniqueTokens().value()); crossSummary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req); @@ -412,7 +415,10 @@ std::tuple MaxUtilizationScheduler::operator()( bool const isFirstChunkContext = req->isContextInitState() && req->isFirstContextChunk() && !req->isDisaggGenerationInitState(); std::optional summary; - if (isFirstChunkContext && kvCacheManager.isEnableBlockReuse()) + // analyzePrefixReuse asserts on variable-window managers; skip the walk there + // and let downstream callers fall back to their fresh tree-walk path. + if (isFirstChunkContext && kvCacheManager.isEnableBlockReuse() + && !kvCacheManager.getBlockManager().isVariableWindow()) { auto uniqueTokens = req->getUniqueTokens(0); summary = kvCacheManager.analyzePrefixReuse(uniqueTokens, *req); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index b3a4db3f4d64..4ddeb2f73297 100755 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -371,10 +371,14 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) .def("get_kv_cache_stats", &BaseKVCacheManager::getKvCacheStats, nb::call_guard()) .def_prop_ro("max_blocks_per_seq", [](tbk::BaseKVCacheManager& self) { return self.getOffsetTableDimensions().maxBlocksPerSeq; }) - .def("get_needed_blocks_one_step", &BaseKVCacheManager::getNeededBlocksOneStep, - nb::call_guard()) - .def("get_remaining_blocks_to_completion", &BaseKVCacheManager::getRemainingBlocksToCompletion, + // nanobind does not inherit C++ default arguments from member-function pointers, so the + // cached_summary default must be spelled out here or Python callers that omit it raise + // TypeError. See kvCacheManager.h for the C++ signatures. + .def("get_needed_blocks_one_step", &BaseKVCacheManager::getNeededBlocksOneStep, nb::arg("req"), + nb::arg("two_steps_look_ahead"), nb::arg("window_size"), nb::arg("cached_summary") = std::nullopt, nb::call_guard()) + .def("get_remaining_blocks_to_completion", &BaseKVCacheManager::getRemainingBlocksToCompletion, nb::arg("req"), + nb::arg("window_size"), nb::arg("cached_summary") = std::nullopt, nb::call_guard()) .def("add_token", &BaseKVCacheManager::addToken, nb::call_guard()) .def("add_sequence", &BaseKVCacheManager::addSequence, nb::call_guard()) .def("remove_sequence", &BaseKVCacheManager::removeSequence, nb::call_guard())