From 2f9d14db7b84b9e61d3134ac24e8fea7b2e2b136 Mon Sep 17 00:00:00 2001 From: SimengLiu-nv Date: Thu, 19 Feb 2026 14:54:55 -0800 Subject: [PATCH 1/4] [TRTLLM-11070][feat] Account for reusable KV cache blocks in micro batch scheduler capacity scheduling Follow up on https://github.com/NVIDIA/TensorRT-LLM/pull/11490 which enable reuse across requests with shared context, allowing better resource utilization and increased batch capacity. This PR enables mirco batch scheduler to account for cached KV blocks from shared prefixes when making scheduling decisions. The capacity scheduler populates estimatedReusableTokens on requests (via radix tree lookup), and the micro batch scheduler subtracts these from the compute budget, allowing more requests to fit in. Signed-off-by: SimengLiu-nv --- .../tensorrt_llm/batch_manager/llmRequest.h | 27 ++ .../batch_manager/kvCacheManager.cpp | 37 +- .../batch_manager/microBatchScheduler.cpp | 46 +- .../nanobind/batch_manager/bindings.cpp | 1 + .../nanobind/batch_manager/kvCacheManager.cpp | 2 + .../batch_manager/kvCacheManagerTest.cpp | 33 +- .../batch_manager/microBatchSchedulerTest.cpp | 421 ++++++++++++++++++ .../_torch/pyexecutor/resource_manager.py | 3 +- .../_torch/pyexecutor/scheduler/scheduler.py | 108 +++-- 9 files changed, 614 insertions(+), 64 deletions(-) mode change 100644 => 100755 cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp mode change 100644 => 100755 cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp mode change 100644 => 100755 cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index 12b5c360741a..11bbc6bd352c 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -866,6 +866,7 @@ class GenericLlmRequest mPrepopulatedPromptLenDraft = 0; mContextChunkSizeTarget = mPromptLen; mContextChunkSizeDraft = mPromptLen; + mEstimatedReusableTokens = 0; mSeqSlot.reset(); } @@ -1135,6 +1136,23 @@ class GenericLlmRequest } } + /// @brief Get the estimated number of reusable tokens from the KV cache. + /// @details Set by the capacity scheduler so the micro batch scheduler can + /// account for cached tokens in its token budget. For subsequent + /// context chunks, this returns 0 because contextRemainingLength + /// already reflects the advancement from setPrepopulatedPromptLen. + [[nodiscard]] SizeType32 getEstimatedReusableTokens() const noexcept + { + return mEstimatedReusableTokens; + } + + /// @brief Set the estimated number of reusable tokens. Const because + /// the field is mutable (it's a scheduling cache, not request state). + void setEstimatedReusableTokens(SizeType32 estimatedReusableTokens) const noexcept + { + mEstimatedReusableTokens = estimatedReusableTokens; + } + void setDraftTokens(std::shared_ptr const& draftTokens) { mDraftTokens = draftTokens; @@ -1964,6 +1982,15 @@ class GenericLlmRequest SizeType32 mPrepopulatedPromptLenTarget{0}; SizeType32 mPrepopulatedPromptLenDraft{0}; + // Estimated number of reusable tokens from the KV cache radix tree. + // Set by the capacity scheduler (during getNeededBlocksOneStep / + // getRemainingBlocksToCompletion) so that the micro batch scheduler + // can account for cached tokens when computing the token budget. + // Marked mutable because it is a cache/estimate set during const + // capacity-scheduler queries. Reset to 0 after addSequence sets + // the authoritative mPrepopulatedPromptLen and advances context position. + mutable SizeType32 mEstimatedReusableTokens{0}; + SizeType32 mMaxSentTokenLen; std::optional mEmbeddingBias{std::nullopt}; diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp old mode 100644 new mode 100755 index 7db8a320dbda..34a58ea77a9d --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -788,14 +788,10 @@ void BlockManager::storeContextBlocks(GenerationRequest& sequence, LlmRequest co constexpr int beamIdx = 0; // no need to consider more than one beam for input tokens for (auto const& [windowSize, _] : mWindowBlockManagers) { - if (mWindowBlockManagers.at(windowSize).isSWA()) - { - // SWA cannot store new blocks on the fly because the block stored - // may go OOW and be reused by another sequence. - continue; - } auto cacheBlockIds = sequence.getCacheBlockIds(windowSize); auto const& uniqueTokens = llmRequest.getUniqueTokens(beamIdx); + TLLM_LOG_DEBUG("storeContextBlocks for request %lu on window %d with %d unique tokens", llmRequest.mRequestId, + windowSize, uniqueTokens.size()); auto blockedUniqueTokens = chopVectorIntoBlocks(uniqueTokens, uniqueTokens.size() - 1, getTokensPerBlock(), false); @@ -1298,8 +1294,8 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& *blockItr, blockItr->uniqueTokens.size() == static_cast(mTokensPerBlock)); } matchingBlock->setHash(); - TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Copied partially filled block %d", mLogPrefix.c_str(), - matchingBlockId); + TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks for request %lu - Copied partially filled block %d", + mLogPrefix.c_str(), sequence.getRequestId(), matchingBlockId); } else { @@ -1307,8 +1303,8 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& freeLeafBlock(matchingBlock); mEvictionPolicy->claimBlock( matchingBlock, perBlockRetentions[bi].retentionPriority, perBlockRetentions[bi].durationMs); - TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Reused partially filled block %d", mLogPrefix.c_str(), - matchingBlockId); + TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks for request %lu - Reused partially filled block %d", + mLogPrefix.c_str(), sequence.getRequestId(), matchingBlockId); } searchRoot = nullptr; // no matching needed for following blocks } @@ -1317,7 +1313,8 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& // Recover block and reuse mEvictionPolicy->claimBlock( matchingBlock, perBlockRetentions[bi].retentionPriority, perBlockRetentions[bi].durationMs); - TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Matched full block %d", mLogPrefix.c_str(), matchingBlockId); + TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks for request %lu - Matched full block %d", mLogPrefix.c_str(), + sequence.getRequestId(), matchingBlockId); searchRoot = matchingBlock; } onboardBlock(sequence, matchingBlock, mode, directory); @@ -1383,7 +1380,6 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& } } - sequence.setCurrentPrepopulatedPromptLen(numMatchedTokens); return numMatchedTokens; } @@ -2246,19 +2242,17 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep( = tc::ceilDiv(numUnSharedTokens, getTokensPerBlock()) * req.mSamplingConfig.beamWidth; auto numRequiredBlocks = numSharedBlocks + numUnSharedBlocks; - // Subtract only reusable blocks that are already allocated (have active refs from another - // sequence). Sharing an allocated block doesn't consume from the free pool. We must NOT - // subtract free reusable blocks (blocks cached in the radix tree but without active refs), - // because those are already counted in the eviction policy's free block count — subtracting - // them would double-count (once as reducing demand, once as inflating supply), causing the - // scheduler to over-admit requests and leading to "No free block found" errors. - if (mEnableBlockReuse && !mBlockManager.isVariableWindow() && !isCrossKv()) + // Subtract reusable blocks if block reuse is enabled and we're not using variable window attention + if (mEnableBlockReuse && !mBlockManager.isVariableWindow() && !isCrossKv() + && !req.isDisaggGenerationInitState()) { auto const uniqueTokens = req.getUniqueTokens(0); auto const numReusableBlocks = countReusableBlocks(uniqueTokens, req, /*onlyAllocated=*/true); // Only subtract from shared blocks (reusable blocks are always shared) auto const reusableSharedBlocks = std::min(numReusableBlocks, numSharedBlocks); numRequiredBlocks -= reusableSharedBlocks; + // Store on request so the micro batch scheduler can use it for token budget + req.setEstimatedReusableTokens(reusableSharedBlocks * getTokensPerBlock()); } return numRequiredBlocks; } @@ -2336,6 +2330,8 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion(LlmRequest const& req, auto const uniqueTokens = req.getUniqueTokens(0); auto const numReusableBlocks = countReusableBlocks(uniqueTokens, req, /*onlyAllocated=*/true); numReusableContextBlocks = std::min(numReusableBlocks, numContextBlocks); + // Store on request so the micro batch scheduler can use it for token budget + req.setEstimatedReusableTokens(numReusableContextBlocks * getTokensPerBlock()); TLLM_LOG_DEBUG( "getRemainingBlocksToCompletion: request ID %lu, numContextBlocks=%d, numReusableBlocks=%d, " "numReusableContextBlocks=%d, numGenBlocksPerBeam=%d", @@ -2558,6 +2554,9 @@ void KVCacheManager::addSequence( { TLLM_LOG_DEBUG("KVCacheManager::addSequence: Setting prepopulatedPromptLen to %d", minPrepopulatedPromptLen); llmRequest->setPrepopulatedPromptLen(minPrepopulatedPromptLen, getTokensPerBlock()); + // Clear the scheduling estimate now that the authoritative value is set. + // This prevents subsequent chunks from double-counting reusable tokens. + llmRequest->setEstimatedReusableTokens(0); } if (llmRequest) diff --git a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp index 6a2dc46d530d..cb380eba37f8 100644 --- a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp @@ -176,12 +176,13 @@ std::tuple MicroBatchScheduler::operator()(Request NVTX3_SCOPED_RANGE(microBatcherScheduleRequests); RequestVector contextRequests, generationRequests; + // batchNumTokens tracks COMPUTE tokens only (excluding reusable cached tokens) SizeType32 batchNumTokens{0}; SizeType32 scheduledReqSize{0}; SizeType32 scheduledBeamWidth{0}; // 0 means no request is scheduled RequestVector contextsToBeChunked; - SizeType32 numChunkedTokens{0}; + SizeType32 numChunkedComputeTokens{0}; bool allContextRequestsFit{true}; // 1. Select the generation phase requests that meet the criteria of total token size. @@ -217,41 +218,51 @@ std::tuple MicroBatchScheduler::operator()(Request } else if (llmReq->isContextInitState()) { + // Reusable tokens set by capacity scheduler (from radix tree lookup). + // Only valid for the first context chunk; subsequent chunks must compute all remaining tokens. + SizeType32 const reusable = llmReq->isFirstContextChunk() ? llmReq->getEstimatedReusableTokens() : 0; + if (!mCtxChunkConfig) // skip chunking { constexpr SizeType32 beam{0}; reqNumTokens = llmReq->getNumTokens(beam) + (llmReq->hasDraftTokens() ? llmReq->getNumDraftTokens() : 0); - TLLM_CHECK_WITH_INFO(!mMaxContextLength || reqNumTokens <= mMaxContextLength.value(), - "The number of context tokens (%d) exceeds the limit value (%d)", reqNumTokens, + // Compute tokens = total - reusable (at least 1 to make progress) + SizeType32 const computeTokens = std::max(1, reqNumTokens - reusable); + TLLM_CHECK_WITH_INFO(!mMaxContextLength || computeTokens <= mMaxContextLength.value(), + "Context compute tokens (%d) exceeds the limit value (%d)", computeTokens, mMaxContextLength.value()); - if (maxNumTokensRuntime && batchNumTokens + reqNumTokens > maxNumTokensRuntime.value()) + if (maxNumTokensRuntime && batchNumTokens + computeTokens > maxNumTokensRuntime.value()) { break; } - TLLM_LOG_DEBUG("context request scheduled: ID %u", llmReq->mRequestId); + TLLM_LOG_DEBUG("context request scheduled: ID %u (reusable %d)", llmReq->mRequestId, reusable); contextRequests.emplace_back(llmReq); - batchNumTokens += reqNumTokens; + batchNumTokens += computeTokens; } else { llmReq->setContextChunkSize(llmReq->getContextRemainingLength()); auto const draftTokens = (llmReq->isLastContextChunk() && llmReq->hasDraftTokens()) ? llmReq->getNumDraftTokens() : 0; - reqNumTokens = llmReq->getContextChunkSize() + draftTokens; + // Compute cost: context compute + draft tokens + // (reusable tokens only offset context tokens, not draft tokens) + SizeType32 const contextCompute = std::max(0, llmReq->getContextChunkSize() - reusable); + SizeType32 computeTokens = contextCompute + draftTokens; if (mMaxContextLength) { - if (mMaxContextLength.value() < reqNumTokens) + if (mMaxContextLength.value() < computeTokens) { // The context exceeds the length limit, we need to try chunking later. - reqNumTokens = mMaxContextLength.value(); + computeTokens = mMaxContextLength.value(); allContextRequestsFit = false; } } contextsToBeChunked.emplace_back(llmReq); - numChunkedTokens += reqNumTokens; - TLLM_LOG_DEBUG("contexts-to-be-chunked request scheduled: ID %u", llmReq->mRequestId); + numChunkedComputeTokens += computeTokens; + TLLM_LOG_DEBUG( + "contexts-to-be-chunked request scheduled: ID %u (reusable %d)", llmReq->mRequestId, reusable); } } else // (llmReq->isGenerationInProgressState()) @@ -284,7 +295,7 @@ std::tuple MicroBatchScheduler::operator()(Request } } - if (maxNumTokensRuntime && numChunkedTokens > maxNumTokensRuntime.value() - batchNumTokens) + if (maxNumTokensRuntime && numChunkedComputeTokens > maxNumTokensRuntime.value() - batchNumTokens) { allContextRequestsFit = false; } @@ -303,9 +314,14 @@ std::tuple MicroBatchScheduler::operator()(Request if (llmReq->getContextChunkSize() > 0) { contextRequests.emplace_back(llmReq); - batchNumTokens += llmReq->getContextChunkSize(); - TLLM_LOG_DEBUG( - "context request scheduled: ID %lu, chunk size %d", llmReq->mRequestId, llmReq->getContextChunkSize()); + // Only count compute tokens (total - reusable). + // Reusable credit only applies to the first context chunk. + SizeType32 const reusable = llmReq->isFirstContextChunk() ? llmReq->getEstimatedReusableTokens() : 0; + SizeType32 const computeTokens + = std::max(0, llmReq->getContextChunkSize() - std::min(reusable, llmReq->getContextChunkSize())); + batchNumTokens += computeTokens; + TLLM_LOG_DEBUG("context request scheduled: ID %lu, chunk size %d%s", llmReq->mRequestId, + llmReq->getContextChunkSize(), reusable > 0 ? (", reusable " + std::to_string(reusable)).c_str() : ""); } } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 5e87afc8439c..8b9ad85360cd 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -169,6 +169,7 @@ void initBindings(nb::module_& m) .def_prop_ro("prepopulated_prompt_len", &GenLlmReq::getPrepopulatedPromptLen) .def("set_prepopulated_prompt_len", &GenLlmReq::setPrepopulatedPromptLen, nb::arg("prepopulated_prompt_len"), nb::arg("kv_tokens_per_block")) + .def_prop_ro("estimated_reusable_tokens", &GenLlmReq::getEstimatedReusableTokens) .def_prop_rw("guided_decoding_params", &GenLlmReq::getGuidedDecodingParams, &GenLlmReq::setGuidedDecodingParams) .def_prop_rw("context_phase_params", &GenLlmReq::getContextPhaseParams, &GenLlmReq::setContextPhaseParams) .def_prop_ro("is_context_only_request", &GenLlmReq::isContextOnlyRequest) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp old mode 100644 new mode 100755 index ad70994cc21c..4c09d8619d9e --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -493,6 +493,8 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::call_guard()) .def("find_new_context_block", &BaseKVCacheManager::findNewContextBlock, 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 old mode 100644 new mode 100755 index 850ce7801dc8..117dd0731e02 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -4339,8 +4339,10 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamWindowSize) events = getEvents(kvCacheManager); - // Expecting only 1 event, storeContextBlock is not called for sliding window. - EXPECT_EQ(events.size(), 1); + EXPECT_EQ(events.size(), 2); + + EXPECT_EQ(events.front().windowSize, slidingWindow); + EXPECT_TRUE(std::holds_alternative(events.front().data)); EXPECT_EQ(events.back().windowSize, maxAttentionWindow); EXPECT_TRUE(std::holds_alternative(events.back().data)); @@ -5819,9 +5821,16 @@ TEST(KVCacheManagerReuseAccountingTest, ReuseAwareBlockEstimatesStayConsistentAf auto const remainingBeforeAdd = kvCacheManager->getRemainingBlocksToCompletion(req1, onlyWindowSize); EXPECT_EQ(remainingBeforeAdd, numContextBlocks + numGenBlocks); + // Verify estimatedReusableTokens is still set after getRemainingBlocksToCompletion + EXPECT_EQ(req1.getEstimatedReusableTokens(), expectedReusableBlocks * tokensPerBlock); + // After addSequence, context blocks are allocated (reuse already applied during allocation) // Only generation blocks remain to be allocated kvCacheManager->addSequence(req1.mRequestId, req1.getPromptLen(), maxBeamWidth, req1); + + // Verify estimatedReusableTokens is cleared to 0 after addSequence + EXPECT_EQ(req1.getEstimatedReusableTokens(), 0); + auto const remainingAfterContextAlloc = kvCacheManager->getRemainingBlocksToCompletion(req1, onlyWindowSize); EXPECT_EQ(remainingAfterContextAlloc, maxNewTokens / tokensPerBlock); } @@ -5934,6 +5943,9 @@ TEST(KVCacheManagerReuseAccountingTest, CountReusableBlocksPartialMatch) auto const neededOneStep = kvCacheManager->getNeededBlocksOneStep(req1, /*twoStepsLookAhead=*/false, onlyWindowSize); EXPECT_EQ(neededOneStep, promptLength / tokensPerBlock); // All 4 context blocks + + // Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks. + EXPECT_EQ(req1.getEstimatedReusableTokens(), 0); } TEST(KVCacheManagerReuseAccountingTest, GetRemainingBlocksToCompletionWithPartialReuse) @@ -5995,6 +6007,9 @@ TEST(KVCacheManagerReuseAccountingTest, GetRemainingBlocksToCompletionWithPartia auto const numContextBlocks = promptLength / tokensPerBlock; // 5 blocks auto const numGenBlocks = maxNewTokens / tokensPerBlock; // 3 blocks EXPECT_EQ(remaining, numContextBlocks + numGenBlocks); // 5 context + 3 generation = 8 + + // Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks. + EXPECT_EQ(req1.getEstimatedReusableTokens(), 0); } TEST(KVCacheManagerReuseAccountingTest, GetNeededBlocksOneStepWithFullReuse) @@ -6055,6 +6070,9 @@ TEST(KVCacheManagerReuseAccountingTest, GetNeededBlocksOneStepWithFullReuse) = kvCacheManager->getNeededBlocksOneStep(req1, /*twoStepsLookAhead=*/false, onlyWindowSize); auto const numSharedBlocks = promptLength / tokensPerBlock; // 3 blocks EXPECT_EQ(neededOneStep, numSharedBlocks); // All 3 context blocks + + // Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks. + EXPECT_EQ(req1.getEstimatedReusableTokens(), 0); } TEST(KVCacheManagerReuseAccountingTest, ReuseDisabledReturnsFullBlockCount) @@ -6100,11 +6118,17 @@ TEST(KVCacheManagerReuseAccountingTest, ReuseDisabledReturnsFullBlockCount) auto const neededOneStep = kvCacheManager->getNeededBlocksOneStep(req, /*twoStepsLookAhead=*/false, onlyWindowSize); EXPECT_EQ(neededOneStep, promptLength / tokensPerBlock); // All 4 context blocks + // Verify estimatedReusableTokens stays 0 when reuse is disabled + EXPECT_EQ(req.getEstimatedReusableTokens(), 0); + // getRemainingBlocksToCompletion should include both context and generation blocks auto const remaining = kvCacheManager->getRemainingBlocksToCompletion(req, onlyWindowSize); auto const expectedContextBlocks = promptLength / tokensPerBlock; auto const expectedGenBlocks = maxNewTokens / tokensPerBlock; EXPECT_EQ(remaining, expectedContextBlocks + expectedGenBlocks); // 4 + 2 = 6 blocks + + // Verify estimatedReusableTokens still 0 after getRemainingBlocksToCompletion + EXPECT_EQ(req.getEstimatedReusableTokens(), 0); } TEST(KVCacheManagerReuseAccountingTest, MultipleRequestsWithSharedPrefix) @@ -6184,7 +6208,10 @@ TEST(KVCacheManagerReuseAccountingTest, MultipleRequestsWithSharedPrefix) = kvCacheManager->getNeededBlocksOneStep(req1, /*twoStepsLookAhead=*/false, onlyWindowSize); EXPECT_EQ(neededOneStep, promptLength / tokensPerBlock); // All 4 context blocks - // getRemainingBlocksToCompletion: 4 context + 1 gen = 5 blocks (no subtraction) + // Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks. + EXPECT_EQ(req1.getEstimatedReusableTokens(), 0); + + // getRemainingBlocksToCompletion: 4 context + 1 gen = 5 blocks (no subtraction; blocks are free) auto const remaining = kvCacheManager->getRemainingBlocksToCompletion(req1, onlyWindowSize); EXPECT_EQ(remaining, (promptLength / tokensPerBlock) + (maxNewTokens / tokensPerBlock)); } diff --git a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp index d727bf31af57..8436f08bb43c 100644 --- a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp @@ -18,10 +18,13 @@ #include #include +#include "tensorrt_llm/batch_manager/capacityScheduler.h" #include "tensorrt_llm/batch_manager/common.h" +#include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/llmRequest.h" #include "tensorrt_llm/batch_manager/microBatchScheduler.h" +#include #include using namespace tensorrt_llm::runtime; @@ -796,6 +799,424 @@ TEST_F(MicroBatchSchedulerTest, DraftTokensGreaterThanChunkSize) } } +TEST_F(MicroBatchSchedulerTest, ReusableTokensReduceComputeBudget) +{ + // Test that reusable tokens (set by the capacity scheduler as a side effect of + // getNeededBlocksOneStep / getRemainingBlocksToCompletion) allow more context + // requests to be scheduled within a tight maxNumTokens budget. + constexpr SizeType32 maxNumTokens = 20; + constexpr SizeType32 maxBatchSize = 4; + + mNumContexts = 1; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler = std::make_shared(); + + constexpr int32_t maxNewTokens = 5; + constexpr int32_t promptLen = 20; + + // Test 1: Without reusable tokens, only 1 request fits in the budget + // (20 context tokens == maxNumTokens, leaving no room for request 1) + { + RequestVector activeRequests; + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 0)); + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 1)); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + EXPECT_EQ(ctx.size(), 1) << "Without reuse, only 1 request fits (20 tokens = budget)"; + EXPECT_EQ(gen.size(), 0); + } + + // Test 2: With reusable tokens set on both requests, both fit. + // Each request has 15 reusable tokens -> compute cost = max(1, 20-15) = 5 per request. + // Total compute = 5 + 5 = 10 < 20 budget. + { + RequestVector activeRequests; + auto req0 = createRequest(promptLen, maxNewTokens, 0); + auto req1 = createRequest(promptLen, maxNewTokens, 1); + req0->setEstimatedReusableTokens(15); + req1->setEstimatedReusableTokens(15); + activeRequests.push_back(req0); + activeRequests.push_back(req1); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + EXPECT_EQ(ctx.size(), 2) << "With reuse (15 each), both fit: 5 + 5 = 10 compute < 20 budget"; + EXPECT_EQ(gen.size(), 0); + } +} + +TEST_F(MicroBatchSchedulerTest, ReusableTokensWithChunkedContextFCFS) +{ + // Test that reusable tokens are correctly accounted for in FCFS chunking: + // the reusable portion is "free" and doesn't consume the forward-pass compute budget. + constexpr SizeType32 maxNumTokens = 15; + constexpr SizeType32 maxBatchSize = 4; + constexpr SizeType32 chunkUnitSize = 5; + constexpr ContextChunkingPolicy ctxChunkPolicy{ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED}; + + mNumContexts = 1; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler + = std::make_shared(ContextChunkingConfig{ctxChunkPolicy, chunkUnitSize}, std::nullopt); + + constexpr int32_t maxNewTokens = 5; + constexpr int32_t promptLen = 20; + + // Without reuse: 20 context tokens > 15 budget -> chunked to 15 tokens + { + RequestVector activeRequests; + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 0)); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + EXPECT_EQ(ctx.size(), 1); + EXPECT_EQ(ctx.at(0)->getContextChunkSize(), 15) << "Without reuse, chunked to 15 tokens"; + } + + // With 10 reusable tokens: compute = 20 - 10 = 10 < 15 budget -> full context fits + { + RequestVector activeRequests; + auto req0 = createRequest(promptLen, maxNewTokens, 0); + req0->setEstimatedReusableTokens(10); + activeRequests.push_back(req0); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + EXPECT_EQ(ctx.size(), 1); + EXPECT_EQ(ctx.at(0)->getContextChunkSize(), promptLen) + << "With 10 reusable tokens, full context fits (compute = 10 < 15)"; + } +} + +TEST_F(MicroBatchSchedulerTest, ReusableTokensZeroHasNoEffect) +{ + // Verify that zero reusable tokens (the default) produces identical scheduling + // to the original behavior — this guards against regressions. + constexpr SizeType32 maxNumTokens = 12; + constexpr SizeType32 maxBatchSize = 4; + + mNumContexts = 1; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler = std::make_shared(); + + constexpr int32_t maxNewTokens = 10; + constexpr int32_t promptLen = 10; + + RequestVector activeRequests; + auto req0 = createRequest(promptLen, maxNewTokens, 0); + auto req1 = createRequest(promptLen, maxNewTokens, 1); + // Explicitly set to 0 — same as default + req0->setEstimatedReusableTokens(0); + req1->setEstimatedReusableTokens(0); + activeRequests.push_back(req0); + activeRequests.push_back(req1); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + // 10 tokens fits, but 10 + 10 = 20 > 12, so only 1 context request scheduled + EXPECT_EQ(ctx.size(), 1); + EXPECT_EQ(gen.size(), 0); +} + +//// +// Combined Capacity Scheduler + Micro Batch Scheduler tests. +// These verify the end-to-end data flow: the capacity scheduler populates +// estimatedReusableTokens on requests (via getNeededBlocksOneStep), then +// the micro batch scheduler reads that value for compute budget decisions. +//// + +class CombinedSchedulerTest : public ::testing::Test +{ +protected: + static std::shared_ptr createKvCacheManager(SizeType32 maxNumRequests, + SizeType32 tokensPerBlock, SizeType32 maxNumTokens, SizeType32 maxNumTokensPerSeq, bool enableReuse) + { + auto const maxNumBlocks = (maxNumTokens + tokensPerBlock - 1) / tokensPerBlock; + auto stream = std::make_shared(); + + using BlocksPerWindow = std::map>; + auto const blocksPerWindow = BlocksPerWindow{{maxNumTokensPerSeq, {maxNumBlocks, 0}}}; + + return std::make_shared( + /*numLayers=*/10, /*nbKvHeads=*/10, /*sizePerHead=*/1, tokensPerBlock, blocksPerWindow, maxNumRequests, + /*maxBeamWidth=*/1, std::vector{maxNumTokensPerSeq}, std::nullopt, nvinfer1::DataType::kHALF, + /*sinkTokenLength=*/0, stream, maxNumTokensPerSeq, enableReuse, /*onboardBlocks=*/true); + } + + static std::shared_ptr createRequestWithTokens( + std::shared_ptr> inputTokens, int32_t maxNewTokens, uint64_t reqId) + { + tensorrt_llm::runtime::SamplingConfig samplingConfig; + return std::make_shared(reqId, maxNewTokens, inputTokens, samplingConfig, /*isStreaming=*/false); + } +}; + +TEST_F(CombinedSchedulerTest, CapacitySchedulerSetsReusableTokensForMicroBatch) +{ + // End-to-end flow: + // 1. Request 0 completes context → blocks stored in radix tree + // 2. Capacity scheduler (MAX_UTILIZATION) calls getNeededBlocksOneStep for request 1 + // → traverses radix tree, finds reusable blocks, sets estimatedReusableTokens + // 3. Micro batch scheduler reads estimatedReusableTokens → reduces compute cost → + // request 1 fits in a tight token budget that wouldn't fit without reuse + + constexpr SizeType32 tokensPerBlock = 10; + constexpr SizeType32 kvCacheMaxNumTokens = 100; + constexpr SizeType32 kvCacheMaxNumTokensPerSeq = 50; + constexpr SizeType32 maxNumRequests = 4; + + auto kvCacheManager = createKvCacheManager( + maxNumRequests, tokensPerBlock, kvCacheMaxNumTokens, kvCacheMaxNumTokensPerSeq, /*enableReuse=*/true); + + auto capacityScheduler + = CapacityScheduler(maxNumRequests, CapacitySchedulerPolicy::kMAX_UTILIZATION, /*hasKvCacheManager=*/true); + + auto microBatchScheduler = MicroBatchScheduler(); + + // promptLen = 30 → 3 blocks, 2 full blocks stored by storeContextBlocks + // (formula: (promptLen-1)/tokensPerBlock = 29/10 = 2) → estimatedReusableTokens = 20 + constexpr int32_t promptLen = 30; + constexpr int32_t maxNewTokens = 5; + auto inputTokens = std::make_shared>(promptLen); + std::iota(inputTokens->begin(), inputTokens->end(), 0); + + auto req0 = createRequestWithTokens(inputTokens, maxNewTokens, 0); + auto req1 = createRequestWithTokens(inputTokens, maxNewTokens, 1); + + // === Iteration 0: Schedule request 0, populate cache === + RequestList activeList; + activeList.push_back(req0); + activeList.push_back(req1); + + auto [scheduled0, disaggInit0, paused0] + = capacityScheduler(activeList, *kvCacheManager, /*peftCacheManager=*/std::nullopt); + + // Request 0 should be scheduled + ASSERT_GE(scheduled0.size(), 1u); + + // Process request 0: addSequence → complete context → store blocks + kvCacheManager->addSequence(req0->mRequestId, promptLen, /*beamWidth=*/1, req0); + req0->moveToNextContextChunk(); + kvCacheManager->storeContextBlocks(*req0); + req0->addNewTokens({0}); + req0->setState(LlmRequestState::kGENERATION_IN_PROGRESS); + kvCacheManager->addToken(req0->mRequestId); + + // Handle paused requests (req1 deferred for reuse benefit) + for (auto const& req : paused0) + { + if (req->mRequestId != 0) + { + req->pause(/*maxInputLen=*/1000); + } + } + + // === Iteration 1: Capacity scheduler sets estimatedReusableTokens on req1 === + auto [scheduled1, disaggInit1, paused1] + = capacityScheduler(activeList, *kvCacheManager, /*peftCacheManager=*/std::nullopt); + + // Verify estimatedReusableTokens was set on req1 by getNeededBlocksOneStep + EXPECT_EQ(req1->getEstimatedReusableTokens(), 20) << "2 full blocks (tokens 0-9, 10-19) reusable = 20 tokens"; + + // === Micro batch scheduler with tight budget === + // Budget = 15. Without reuse: req1 = 30 tokens → doesn't fit alongside gen req0. + // With reuse: req1 compute = max(1, 30-20) = 10, plus req0 gen = 1, total = 11 < 15. + RequestVector microBatchActive; + for (auto& req : scheduled1) + { + microBatchActive.push_back(req); + } + + constexpr SizeType32 maxNumTokensBudget = 15; + constexpr SizeType32 maxBatchSize = 4; + ReqIdsSet inflightReqIds; + auto [ctx, gen] = microBatchScheduler(microBatchActive, inflightReqIds, maxBatchSize, maxNumTokensBudget); + + // Request 1 should be scheduled (context) thanks to reuse + bool req1InCtx = std::any_of(ctx.begin(), ctx.end(), [](auto const& r) { return r->mRequestId == 1; }); + EXPECT_TRUE(req1InCtx) << "With 20 reusable tokens, req1 compute cost = 10 fits in 15 budget"; + + // Request 0 should be scheduled (generation) + bool req0InGen = std::any_of(gen.begin(), gen.end(), [](auto const& r) { return r->mRequestId == 0; }); + EXPECT_TRUE(req0InGen) << "Request 0 (gen, 1 token) should be scheduled"; + + // === Compare: without reuse, req1 would NOT fit === + auto req2 = createRequestWithTokens(inputTokens, maxNewTokens, 2); + req2->setEstimatedReusableTokens(0); + + RequestVector noReuseActive; + noReuseActive.push_back(req0); // gen (1 token) + noReuseActive.push_back(req2); // ctx (30 tokens, no reuse) + + ReqIdsSet inflightReqIds2; + auto [ctx2, gen2] = microBatchScheduler(noReuseActive, inflightReqIds2, maxBatchSize, maxNumTokensBudget); + + // Without reuse: gen(1) + ctx(30) = 31 > 15 → req2 doesn't fit + bool req2InCtx = std::any_of(ctx2.begin(), ctx2.end(), [](auto const& r) { return r->mRequestId == 2; }); + EXPECT_FALSE(req2InCtx) << "Without reuse, 30 context tokens + 1 gen token exceeds 15 budget"; + + kvCacheManager->removeSequence(req0->mRequestId, req0); +} + +TEST_F(CombinedSchedulerTest, CapacitySchedulerReusableTokensWithChunkedMicroBatch) +{ + // Same pipeline but with chunked prefill. + // Without reuse the context is chunked; with reuse it fits without chunking. + + constexpr SizeType32 tokensPerBlock = 10; + constexpr SizeType32 kvCacheMaxNumTokens = 100; + constexpr SizeType32 kvCacheMaxNumTokensPerSeq = 50; + constexpr SizeType32 maxNumRequests = 4; + + auto kvCacheManager = createKvCacheManager( + maxNumRequests, tokensPerBlock, kvCacheMaxNumTokens, kvCacheMaxNumTokensPerSeq, /*enableReuse=*/true); + + auto capacityScheduler + = CapacityScheduler(maxNumRequests, CapacitySchedulerPolicy::kMAX_UTILIZATION, /*hasKvCacheManager=*/true); + + // FCFS chunking with chunkUnitSize = 5 + constexpr SizeType32 chunkUnitSize = 5; + auto microBatchScheduler = MicroBatchScheduler( + ContextChunkingConfig{ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED, chunkUnitSize}, std::nullopt); + + constexpr int32_t promptLen = 30; + constexpr int32_t maxNewTokens = 5; + auto inputTokens = std::make_shared>(promptLen); + std::iota(inputTokens->begin(), inputTokens->end(), 0); + + auto req0 = createRequestWithTokens(inputTokens, maxNewTokens, 0); + auto req1 = createRequestWithTokens(inputTokens, maxNewTokens, 1); + + // Iteration 0: Schedule and process request 0 + RequestList activeList; + activeList.push_back(req0); + activeList.push_back(req1); + + auto [scheduled0, disaggInit0, paused0] + = capacityScheduler(activeList, *kvCacheManager, /*peftCacheManager=*/std::nullopt); + + kvCacheManager->addSequence(req0->mRequestId, promptLen, /*beamWidth=*/1, req0); + req0->moveToNextContextChunk(); + kvCacheManager->storeContextBlocks(*req0); + req0->addNewTokens({0}); + req0->setState(LlmRequestState::kGENERATION_IN_PROGRESS); + kvCacheManager->addToken(req0->mRequestId); + + for (auto const& req : paused0) + { + if (req->mRequestId != 0) + { + req->pause(/*maxInputLen=*/1000); + } + } + + // Iteration 1: Capacity scheduler sets estimatedReusableTokens on req1 + auto [scheduled1, disaggInit1, paused1] + = capacityScheduler(activeList, *kvCacheManager, /*peftCacheManager=*/std::nullopt); + + EXPECT_EQ(req1->getEstimatedReusableTokens(), 20); + + // Micro batch scheduler with tight budget (15 tokens) + // With reuse: compute = max(0, 30-20) = 10 < 15 → full context fits, no chunking needed + RequestVector microBatchActive; + for (auto& req : scheduled1) + { + microBatchActive.push_back(req); + } + + constexpr SizeType32 maxNumTokensBudget = 15; + constexpr SizeType32 maxBatchSize = 4; + ReqIdsSet inflightReqIds; + auto [ctx, gen] = microBatchScheduler(microBatchActive, inflightReqIds, maxBatchSize, maxNumTokensBudget); + + // With reuse, full context should fit without chunking + for (auto const& req : ctx) + { + if (req->mRequestId == 1) + { + EXPECT_EQ(req->getContextChunkSize(), promptLen) + << "With 20 reusable tokens, full 30-token context fits (compute = 10 < 15)"; + } + } + + // === Compare: without reuse, context gets chunked === + auto req2 = createRequestWithTokens(inputTokens, maxNewTokens, 2); + req2->setEstimatedReusableTokens(0); + + RequestVector noReuseActive; + noReuseActive.push_back(req2); // ctx only, no gen + + ReqIdsSet inflightReqIds2; + auto [ctx2, gen2] = microBatchScheduler(noReuseActive, inflightReqIds2, maxBatchSize, maxNumTokensBudget); + + ASSERT_EQ(ctx2.size(), 1u); + EXPECT_EQ(ctx2.at(0)->getContextChunkSize(), 15) + << "Without reuse, 30-token context chunked to 15 (budget limit, aligned to chunkUnitSize=5)"; + + kvCacheManager->removeSequence(req0->mRequestId, req0); +} + +TEST_F(CombinedSchedulerTest, NoReuseMicroBatchUnchanged) +{ + // When block reuse is disabled, estimatedReusableTokens stays 0 + // and the micro batch scheduler behaves identically to before. + + constexpr SizeType32 tokensPerBlock = 10; + constexpr SizeType32 kvCacheMaxNumTokens = 100; + constexpr SizeType32 kvCacheMaxNumTokensPerSeq = 50; + constexpr SizeType32 maxNumRequests = 4; + + // Reuse DISABLED + auto kvCacheManager = createKvCacheManager( + maxNumRequests, tokensPerBlock, kvCacheMaxNumTokens, kvCacheMaxNumTokensPerSeq, /*enableReuse=*/false); + + auto capacityScheduler + = CapacityScheduler(maxNumRequests, CapacitySchedulerPolicy::kMAX_UTILIZATION, /*hasKvCacheManager=*/true); + + auto microBatchScheduler = MicroBatchScheduler(); + + constexpr int32_t promptLen = 20; + constexpr int32_t maxNewTokens = 5; + auto inputTokens = std::make_shared>(promptLen); + std::iota(inputTokens->begin(), inputTokens->end(), 0); + + auto req0 = createRequestWithTokens(inputTokens, maxNewTokens, 0); + auto req1 = createRequestWithTokens(inputTokens, maxNewTokens, 1); + + // With reuse disabled, both requests are scheduled in iteration 0 + RequestList activeList; + activeList.push_back(req0); + activeList.push_back(req1); + + auto [scheduled0, disaggInit0, paused0] + = capacityScheduler(activeList, *kvCacheManager, /*peftCacheManager=*/std::nullopt); + + // Both should be scheduled (no beneficialToSkip without reuse) + EXPECT_EQ(scheduled0.size(), 2u); + + // estimatedReusableTokens should be 0 for both + EXPECT_EQ(req0->getEstimatedReusableTokens(), 0); + EXPECT_EQ(req1->getEstimatedReusableTokens(), 0); + + // Micro batch scheduler with budget that fits only 1 context request + // Budget = 25: first req = 20 tokens, second req = 20 tokens → 40 > 25 + RequestVector microBatchActive; + for (auto& req : scheduled0) + { + microBatchActive.push_back(req); + } + + constexpr SizeType32 maxNumTokensBudget = 25; + constexpr SizeType32 maxBatchSize = 4; + ReqIdsSet inflightReqIds; + auto [ctx, gen] = microBatchScheduler(microBatchActive, inflightReqIds, maxBatchSize, maxNumTokensBudget); + + // Without reuse: each request costs 20 tokens, only 1 fits in budget of 25 + EXPECT_EQ(ctx.size(), 1u) << "Without reuse, only 1 of 2 context requests fits in 25-token budget"; + EXPECT_EQ(gen.size(), 0u); +} + class ContextChunkingTest : public MicroBatchSchedulerTest { protected: diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 8168aaee53a3..7728c2da24b7 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -787,7 +787,8 @@ def update_resources(self, # For context requests, we store the blocks for reuse. for request in scheduled_batch.context_requests: - self.impl.store_context_blocks(request) + if request.context_remaining_length == 0: # only after all chunks done + self.impl.store_context_blocks(request) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): return self.impl.remove_sequence(request.py_request_id, request, diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 540a0788e2ff..e1a9a3fa63a1 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -360,14 +360,14 @@ def schedule( context_requests: RequestList = [] generation_requests: RequestList = [] - # Current total tokens in the scheduled batch (Generation + Context) + # Current compute tokens in the scheduled batch (non-reusable only) batch_num_tokens = 0 scheduled_req_size = 0 scheduled_beam_width = 0 contexts_to_be_chunked: RequestList = [] - # Total tokens required by chunked requests (calculated tentatively) - num_chunked_tokens = 0 + # Compute tokens required by chunked requests (non-reusable only) + num_chunked_compute_tokens = 0 all_context_requests_fit = True # Cache instance attributes as locals for faster access in loop @@ -375,7 +375,6 @@ def schedule( max_num_tokens = self.max_num_tokens max_context_length = self.max_context_length ctx_chunk_config = self.ctx_chunk_config - # 1. Main Scheduling Loop for req in active_requests: req_state_value = req.state_value @@ -419,37 +418,55 @@ def schedule( draft_tokens = req.num_draft_tokens if req.has_draft_tokens else 0 req_num_tokens = base_tokens + draft_tokens - assert max_context_length is None or req_num_tokens <= max_context_length, ( - f"Context tokens ({req_num_tokens}) exceeds limit ({max_context_length})" + # Reusable tokens set by capacity scheduler (from radix tree lookup). + # Only valid for the first context chunk; subsequent chunks must compute all remaining tokens. + reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 + compute_tokens = max(1, req_num_tokens - reusable) + + assert max_context_length is None or compute_tokens <= max_context_length, ( + f"Context compute tokens ({compute_tokens}) exceeds limit ({max_context_length})" ) if max_num_tokens is not None and ( - batch_num_tokens + req_num_tokens > max_num_tokens + batch_num_tokens + compute_tokens > max_num_tokens ): break - logger.debug(f"context request scheduled: ID {req.request_id}") + logger.debug( + f"context request scheduled: ID {req.request_id}" + + (f" (reusable {reusable})" if reusable > 0 else "") + ) context_requests.append(req) - batch_num_tokens += req_num_tokens + batch_num_tokens += compute_tokens else: # Chunking Enabled: Tentative schedule req.context_chunk_size = req.context_remaining_length + # Reusable tokens set by capacity scheduler (from radix tree lookup). + # Only valid for the first context chunk; subsequent chunks must compute all remaining tokens. + reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 + draft_tokens = ( req.num_draft_tokens if (req.is_last_context_chunk and req.has_draft_tokens) else 0 ) - req_num_tokens = req.context_chunk_size + draft_tokens + # Compute cost: context compute + draft tokens + # (reusable tokens only offset context tokens, not draft tokens) + context_compute = max(0, req.context_chunk_size - reusable) + compute_tokens = context_compute + draft_tokens if max_context_length is not None: - if max_context_length < req_num_tokens: - req_num_tokens = max_context_length + if max_context_length < compute_tokens: + compute_tokens = max_context_length all_context_requests_fit = False - logger.debug(f"contexts-to-be-chunked request scheduled: ID {req.request_id}") + logger.debug( + f"contexts-to-be-chunked request scheduled: ID {req.request_id}" + + (f" (reusable {reusable})" if reusable > 0 else "") + ) contexts_to_be_chunked.append(req) - num_chunked_tokens += req_num_tokens + num_chunked_compute_tokens += compute_tokens # --- C. Generation Request Handling --- else: @@ -481,8 +498,10 @@ def schedule( if scheduled_req_size >= max_batch_size: break - # 2. Verify Chunking Fits - if max_num_tokens is not None and num_chunked_tokens > (max_num_tokens - batch_num_tokens): + # 2. Verify Chunking Fits (using compute tokens, not total) + if max_num_tokens is not None and num_chunked_compute_tokens > ( + max_num_tokens - batch_num_tokens + ): all_context_requests_fit = False # 3. Apply Chunking Strategy if needed @@ -500,10 +519,16 @@ def schedule( for req in contexts_to_be_chunked: if req.context_chunk_size > 0: context_requests.append(req) - batch_num_tokens += req.context_chunk_size + # Reusable credit only applies to the first context chunk. + reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 + compute_tokens = max( + 0, req.context_chunk_size - min(reusable, req.context_chunk_size) + ) + batch_num_tokens += compute_tokens logger.debug( f"context request scheduled: ID {req.request_id}, " f"chunk size {req.context_chunk_size}" + + (f", reusable {reusable}" if reusable > 0 else "") ) # Sort requests for consistency with C++ @@ -515,7 +540,10 @@ def schedule( f"batchSize (num ctx/enc requests + num gen requests): " f"{len(context_requests) + len(generation_requests)}" ) - logger.debug(f"batchNumTokens / maxNumTokens: {batch_num_tokens} / {max_num_tokens or 0}") + logger.debug( + f"batchNumTokens (num ctx/enc input tokens + num gen input tokens) " + f"/ maxNumTokens: {batch_num_tokens} / {max_num_tokens or 0}" + ) return context_requests, generation_requests @@ -555,7 +583,11 @@ def get_lora_task_id(req: LlmRequest): generation_requests.sort(key=get_lora_task_id) - def _set_ctx_requests_chunk_size(self, requests: RequestList, capacity: Optional[int]): + def _set_ctx_requests_chunk_size( + self, + requests: RequestList, + capacity: Optional[int], + ): # C++: Resets all chunk sizes to 0 at start for req in requests: req.context_chunk_size = 0 @@ -608,18 +640,41 @@ def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], num_ctx_tokens += actual_increment num_tokens_single_loop += actual_increment - def _chunk_fcfs(self, requests: RequestList, capacity: Optional[int], unit_size: int): - current_capacity = capacity if capacity is not None else float("inf") + def _chunk_fcfs( + self, + requests: RequestList, + capacity: Optional[int], + unit_size: int, + ): + # capacity represents the COMPUTE budget (non-reusable tokens). + # Reusable tokens are "free" — they don't consume forward-pass compute. + current_compute_capacity = capacity if capacity is not None else float("inf") for req in requests: suggested_size = req.context_remaining_length + # Only the first context chunk can reuse cached KV blocks + # subsequent chunks must compute all remaining tokens. + reusable = min( + req.estimated_reusable_tokens if req.is_first_context_chunk else 0, + suggested_size, + ) + compute_cost = suggested_size - reusable + + # Start with full context as the allocation target actual_size = suggested_size - if current_capacity < actual_size: - actual_size = current_capacity + # Limit by compute capacity (reusable tokens are free) + if current_compute_capacity < compute_cost: + # Can't fit all new tokens; allow all reusable + whatever compute fits + actual_size = reusable + current_compute_capacity + actual_size = min(actual_size, suggested_size) + # Limit by max_context_length (applies to compute portion only) if self.max_context_length is not None: - actual_size = min(self.max_context_length, actual_size) + actual_compute = max(0, actual_size - reusable) + if actual_compute > self.max_context_length: + actual_size = reusable + self.max_context_length + actual_size = min(actual_size, suggested_size) # Round down to unit size if we had to truncate if actual_size < suggested_size: @@ -627,9 +682,10 @@ def _chunk_fcfs(self, requests: RequestList, capacity: Optional[int], unit_size: req.context_chunk_size = int(actual_size) - # C++: ctxTokensCapacity = ctxTokensCapacity - actualChunkSize + # Subtract only the compute portion from capacity + actual_compute = max(0, req.context_chunk_size - min(reusable, req.context_chunk_size)) if capacity is not None: - current_capacity -= req.context_chunk_size + current_compute_capacity -= actual_compute def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit_size: int): # Calculate tokens already taken by the batch so far From 97e08c9949200eb5e8e6efeccca721526df584a2 Mon Sep 17 00:00:00 2001 From: SimengLiu-nv Date: Thu, 19 Mar 2026 20:49:33 -0700 Subject: [PATCH 2/4] Add python tests and resolve kvCacheManagerTest failures. Signed-off-by: SimengLiu-nv --- .../batch_manager/kvCacheManager.cpp | 31 ++-- .../batch_manager/microBatchScheduler.cpp | 3 +- .../nanobind/batch_manager/bindings.cpp | 3 +- .../batch_manager/kvCacheManagerTest.cpp | 15 +- .../_torch/pyexecutor/resource_manager.py | 7 +- .../_torch/pyexecutor/scheduler/scheduler.py | 20 +-- .../_torch/executor/test_py_scheduler.py | 148 ++++++++++++++++++ 7 files changed, 195 insertions(+), 32 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 34a58ea77a9d..2ca65be219dc 100755 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -786,8 +786,14 @@ bool WindowBlockManager::verifyQueueIntegrity() void BlockManager::storeContextBlocks(GenerationRequest& sequence, LlmRequest const& llmRequest) { constexpr int beamIdx = 0; // no need to consider more than one beam for input tokens - for (auto const& [windowSize, _] : mWindowBlockManagers) - { + // Iterate in descending window-size order (largest/full-attention windows first). + // This guarantees that the Stored event for the full-attention window is committed + // before flushRemovedEvents fires for SWA windows, preserving the per-window + // ordering guarantee: Removed events precede the Stored event for the same window, + // and Stored(full) is not interleaved with Removed(SWA). + for (auto it = mWindowBlockManagers.rbegin(); it != mWindowBlockManagers.rend(); ++it) + { + auto& [windowSize, manager] = *it; auto cacheBlockIds = sequence.getCacheBlockIds(windowSize); auto const& uniqueTokens = llmRequest.getUniqueTokens(beamIdx); TLLM_LOG_DEBUG("storeContextBlocks for request %lu on window %d with %d unique tokens", llmRequest.mRequestId, @@ -796,7 +802,7 @@ void BlockManager::storeContextBlocks(GenerationRequest& sequence, LlmRequest co auto blockedUniqueTokens = chopVectorIntoBlocks(uniqueTokens, uniqueTokens.size() - 1, getTokensPerBlock(), false); auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); - (void) mWindowBlockManagers.at(windowSize).storeBlocks(std::move(blockKeys), cacheBlockIds[beamIdx]); + (void) manager.storeBlocks(std::move(blockKeys), cacheBlockIds[beamIdx]); } } @@ -2328,14 +2334,21 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion(LlmRequest const& req, && numAllocBlocksPerBeam == 0) { auto const uniqueTokens = req.getUniqueTokens(0); - auto const numReusableBlocks = countReusableBlocks(uniqueTokens, req, /*onlyAllocated=*/true); - numReusableContextBlocks = std::min(numReusableBlocks, numContextBlocks); - // Store on request so the micro batch scheduler can use it for token budget - req.setEstimatedReusableTokens(numReusableContextBlocks * getTokensPerBlock()); + // 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); + // 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()); TLLM_LOG_DEBUG( - "getRemainingBlocksToCompletion: request ID %lu, numContextBlocks=%d, numReusableBlocks=%d, " + "getRemainingBlocksToCompletion: request ID %lu, numContextBlocks=%d, " + "numReusableBlocksAllocated=%d, numReusableBlocksAll=%d, " "numReusableContextBlocks=%d, numGenBlocksPerBeam=%d", - req.mRequestId, numContextBlocks, numReusableBlocks, numReusableContextBlocks, numGenBlocksPerBeam); + req.mRequestId, numContextBlocks, numReusableBlocksAllocated, numReusableBlocksAll, + numReusableContextBlocks, numGenBlocksPerBeam); } // In case of sliding window attention, a new block is allocated when the diff --git a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp index cb380eba37f8..8d926d56adc6 100644 --- a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp @@ -317,8 +317,7 @@ std::tuple MicroBatchScheduler::operator()(Request // Only count compute tokens (total - reusable). // Reusable credit only applies to the first context chunk. SizeType32 const reusable = llmReq->isFirstContextChunk() ? llmReq->getEstimatedReusableTokens() : 0; - SizeType32 const computeTokens - = std::max(0, llmReq->getContextChunkSize() - std::min(reusable, llmReq->getContextChunkSize())); + SizeType32 const computeTokens = std::max(0, llmReq->getContextChunkSize() - reusable); batchNumTokens += computeTokens; TLLM_LOG_DEBUG("context request scheduled: ID %lu, chunk size %d%s", llmReq->mRequestId, llmReq->getContextChunkSize(), reusable > 0 ? (", reusable " + std::to_string(reusable)).c_str() : ""); diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 8b9ad85360cd..dfd73320ecdf 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -169,7 +169,8 @@ void initBindings(nb::module_& m) .def_prop_ro("prepopulated_prompt_len", &GenLlmReq::getPrepopulatedPromptLen) .def("set_prepopulated_prompt_len", &GenLlmReq::setPrepopulatedPromptLen, nb::arg("prepopulated_prompt_len"), nb::arg("kv_tokens_per_block")) - .def_prop_ro("estimated_reusable_tokens", &GenLlmReq::getEstimatedReusableTokens) + .def_prop_rw( + "estimated_reusable_tokens", &GenLlmReq::getEstimatedReusableTokens, &GenLlmReq::setEstimatedReusableTokens) .def_prop_rw("guided_decoding_params", &GenLlmReq::getGuidedDecodingParams, &GenLlmReq::setGuidedDecodingParams) .def_prop_rw("context_phase_params", &GenLlmReq::getContextPhaseParams, &GenLlmReq::setContextPhaseParams) .def_prop_ro("is_context_only_request", &GenLlmReq::isContextOnlyRequest) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 117dd0731e02..995453fd669d 100755 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -4341,10 +4341,12 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamWindowSize) EXPECT_EQ(events.size(), 2); - EXPECT_EQ(events.front().windowSize, slidingWindow); + // storeContextBlocks iterates windows in descending order (largest first) to preserve + // per-window event ordering, so the full-attention window is emitted before the SWA window. + EXPECT_EQ(events.front().windowSize, maxAttentionWindow); EXPECT_TRUE(std::holds_alternative(events.front().data)); - EXPECT_EQ(events.back().windowSize, maxAttentionWindow); + EXPECT_EQ(events.back().windowSize, slidingWindow); EXPECT_TRUE(std::holds_alternative(events.back().data)); } @@ -6001,15 +6003,18 @@ TEST(KVCacheManagerReuseAccountingTest, GetRemainingBlocksToCompletionWithPartia }; // After removeSequence, reusable blocks are free (no active refs). - // getRemainingBlocksToCompletion must NOT subtract free reusable blocks. + // getRemainingBlocksToCompletion must NOT subtract free reusable blocks from the BLOCK budget. // Needs all 5 context + 3 generation = 8 blocks. auto const remaining = kvCacheManager->getRemainingBlocksToCompletion(req1, onlyWindowSize); auto const numContextBlocks = promptLength / tokensPerBlock; // 5 blocks auto const numGenBlocks = maxNewTokens / tokensPerBlock; // 3 blocks EXPECT_EQ(remaining, numContextBlocks + numGenBlocks); // 5 context + 3 generation = 8 - // Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks. - EXPECT_EQ(req1.getEstimatedReusableTokens(), 0); + // storeContextBlocks stores (promptLength - 1) / tokensPerBlock = 4 full blocks. + // getRemainingBlocksToCompletion counts ALL reusable blocks (free or allocated) for the + // TOKEN budget, so estimatedReusableTokens = min(4, 5) * tokensPerBlock = 64. + auto const numStoredBlocks = (promptLength - 1) / tokensPerBlock; // 4 + EXPECT_EQ(req1.getEstimatedReusableTokens(), std::min(numStoredBlocks, numContextBlocks) * tokensPerBlock); } TEST(KVCacheManagerReuseAccountingTest, GetNeededBlocksOneStepWithFullReuse) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 7728c2da24b7..0fefec42a565 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -785,9 +785,12 @@ def update_resources(self, if request.py_rewind_len > 0: self.rewind_kv_cache(request, request.py_rewind_len) - # For context requests, we store the blocks for reuse. + # For context requests, store completed context blocks for KV cache reuse. + # We wait until context_remaining_length == 0 (all chunks processed) before + # storing, so that SWA windows are safe to store — blocks won't go out-of-window + # and be evicted while the context is still in-flight. for request in scheduled_batch.context_requests: - if request.context_remaining_length == 0: # only after all chunks done + if request.context_remaining_length == 0: self.impl.store_context_blocks(request) def free_resources(self, request: LlmRequest, pin_on_release: bool = False): diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index e1a9a3fa63a1..6fb322346b83 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -411,6 +411,10 @@ def schedule( # --- B. Context Request Handling --- elif req_state_value == self._context_init_state_value: + # Reusable tokens set by capacity scheduler (from radix tree lookup). + # Only valid for the first context chunk; subsequent chunks must compute all remaining tokens. + reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 + if not ctx_chunk_config: # No Chunking: Schedule full context # C++ uses getNumTokens(beam=0) which is tokens.size() - numPreDecodedTokens @@ -418,9 +422,6 @@ def schedule( draft_tokens = req.num_draft_tokens if req.has_draft_tokens else 0 req_num_tokens = base_tokens + draft_tokens - # Reusable tokens set by capacity scheduler (from radix tree lookup). - # Only valid for the first context chunk; subsequent chunks must compute all remaining tokens. - reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 compute_tokens = max(1, req_num_tokens - reusable) assert max_context_length is None or compute_tokens <= max_context_length, ( @@ -442,10 +443,6 @@ def schedule( # Chunking Enabled: Tentative schedule req.context_chunk_size = req.context_remaining_length - # Reusable tokens set by capacity scheduler (from radix tree lookup). - # Only valid for the first context chunk; subsequent chunks must compute all remaining tokens. - reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 - draft_tokens = ( req.num_draft_tokens if (req.is_last_context_chunk and req.has_draft_tokens) @@ -521,9 +518,7 @@ def schedule( context_requests.append(req) # Reusable credit only applies to the first context chunk. reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 - compute_tokens = max( - 0, req.context_chunk_size - min(reusable, req.context_chunk_size) - ) + compute_tokens = max(0, req.context_chunk_size - reusable) batch_num_tokens += compute_tokens logger.debug( f"context request scheduled: ID {req.request_id}, " @@ -666,8 +661,7 @@ def _chunk_fcfs( # Limit by compute capacity (reusable tokens are free) if current_compute_capacity < compute_cost: # Can't fit all new tokens; allow all reusable + whatever compute fits - actual_size = reusable + current_compute_capacity - actual_size = min(actual_size, suggested_size) + actual_size = min(reusable + current_compute_capacity, suggested_size) # Limit by max_context_length (applies to compute portion only) if self.max_context_length is not None: @@ -683,7 +677,7 @@ def _chunk_fcfs( req.context_chunk_size = int(actual_size) # Subtract only the compute portion from capacity - actual_compute = max(0, req.context_chunk_size - min(reusable, req.context_chunk_size)) + actual_compute = max(0, req.context_chunk_size - reusable) if capacity is not None: current_compute_capacity -= actual_compute diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index 493a92a1cb9f..3ebf53bbc17e 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -781,6 +781,154 @@ def test_sort_by_lora_task_id(self): assert gen[2].request_id == 0 +# ############################################################################ +# +# Reusable KV Cache Token Tests +# +# ############################################################################ + + +class TestPyMicroBatchSchedulerReusableTokens: + """ + Tests for estimated_reusable_tokens logic in PyMicroBatchScheduler. + Mirrors C++ ReusableTokens tests in microBatchSchedulerTest.cpp. + + The reusable credit reduces the compute cost of a request when the KV + cache radix tree has already computed some prefix blocks. + """ + + def test_reusable_tokens_reduce_compute_budget(self): + """ + Reusable tokens shrink the compute cost so more requests fit. + C++ ref: ReusableTokensReduceComputeBudget + + max_num_tokens=20, two requests of prompt_len=20. + Without reuse: req0 costs 20 == budget → req1 doesn't fit. + With reusable=15 each: compute = max(1, 20-15) = 5 → total 10 < 20 + → both fit. + """ + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=20, ctx_chunk_config=None + ) + req0 = make_context_request(0, prompt_len=20) + req1 = make_context_request(1, prompt_len=20) + req0.estimated_reusable_tokens = 15 + req1.estimated_reusable_tokens = 15 + + ctx, gen = scheduler.schedule([req0, req1], set()) + assert len(ctx) == 2 + + def test_reusable_tokens_zero_has_no_effect(self): + """ + Explicitly setting reusable=0 is identical to the default (no reuse). + C++ ref: ReusableTokensZeroHasNoEffect + + max_num_tokens=12, two requests of prompt_len=10. + 10+10=20 > 12 → only the first request fits. + """ + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=12, ctx_chunk_config=None + ) + req0 = make_context_request(0, prompt_len=10) + req1 = make_context_request(1, prompt_len=10) + req0.estimated_reusable_tokens = 0 + req1.estimated_reusable_tokens = 0 + + ctx, gen = scheduler.schedule([req0, req1], set()) + assert len(ctx) == 1 + assert ctx[0].request_id == 0 + + def test_reusable_tokens_chunked_context_fcfs_full_context_fits(self): + """ + Reusable tokens can allow the full context to fit without chunking. + C++ ref: ReusableTokensWithChunkedContextFCFS + + max_num_tokens=15, prompt_len=20, FCFS, chunk_unit=5. + Without reuse: tentative compute = 20, > 15 → chunked to 15. + With reusable=10: context_compute = max(0, 20-10) = 10 < 15 + → all_context_requests_fit stays True → full 20-token context scheduled. + """ + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=5) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=15, ctx_chunk_config=config + ) + req = make_context_request(0, prompt_len=20) + req.estimated_reusable_tokens = 10 + + ctx, gen = scheduler.schedule([req], set()) + assert len(ctx) == 1 + # Full context fits — chunk_size equals the full prompt length. + assert ctx[0].context_chunk_size == 20 + + def test_reusable_tokens_only_on_first_chunk(self): + """ + The reusable credit is only applied on the first context chunk. + On subsequent chunks is_first_context_chunk == False, so reusable = 0. + + Use a request already past its first chunk (context_position=10) with + prompt_len=30 and reusable=20 — the reuse should be ignored. + Without reuse the remaining 20 tokens fill the budget exactly. + With another fresh request (reusable=20, first chunk) the reuse IS + applied: compute = max(1, 30-20) = 10 → both fit in budget=30. + """ + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=30, ctx_chunk_config=None + ) + # req0: already past first chunk — reusable credit must be ignored. + req0 = make_context_request(0, prompt_len=30, context_position=10) + req0.estimated_reusable_tokens = 20 + assert not req0.is_first_context_chunk + + ctx, gen = scheduler.schedule([req0], set()) + assert len(ctx) == 1 + # Remaining tokens = 30 - 10 = 20; reuse ignored → compute = 20 + # Budget = 30, 20 <= 30, so it fits. + # Now add a second request — without reuse it would push total to 40 > 30. + req1 = make_context_request(1, prompt_len=20) + # No reuse → compute = 20; total = 20 + 20 = 40 > 30 → req1 doesn't fit. + scheduler2 = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=30, ctx_chunk_config=None + ) + ctx2, _ = scheduler2.schedule([req0, req1], set()) + req0_again = make_context_request(0, prompt_len=30, context_position=10) + req0_again.estimated_reusable_tokens = 20 + # Confirm: fresh req0 (non-first chunk) + req1 → only req0 fits + assert len(ctx2) == 1 + + # Contrast: first-chunk request with reuse DOES get the credit. + req2 = make_context_request(2, prompt_len=30) + req2.estimated_reusable_tokens = 20 + assert req2.is_first_context_chunk + req3 = make_context_request(3, prompt_len=20) + scheduler3 = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=30, ctx_chunk_config=None + ) + ctx3, _ = scheduler3.schedule([req2, req3], set()) + # req2 compute = max(1, 30-20) = 10; req3 compute = 20; total = 30 → both fit. + assert len(ctx3) == 2 + + def test_reusable_tokens_no_chunking_min_cost_is_one(self): + """ + The no-chunking path floors compute cost at 1 even if reusable > prompt_len. + Python-specific floor test. + + prompt_len=10, reusable=15: compute = max(1, 10-15) = max(1,-5) = 1. + max_num_tokens=10: req0 costs 1, req1 (no reuse, 10 tokens) costs 1+10=11 > 10 + → req1 doesn't fit. + """ + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=10, ctx_chunk_config=None + ) + req0 = make_context_request(0, prompt_len=10) + req0.estimated_reusable_tokens = 15 # exceeds prompt length + req1 = make_context_request(1, prompt_len=10) + + ctx, gen = scheduler.schedule([req0, req1], set()) + # req0 compute = 1; req1 compute = 10; 1 + 10 = 11 > 10 → only req0 fits. + assert len(ctx) == 1 + assert ctx[0].request_id == 0 + + # ############################################################################ # # Part 3: Direct Context Chunking Tests (mirrors C++ ContextChunkingTest) From d52c859c03944521fdd0a381d31a1a8d508ef56d Mon Sep 17 00:00:00 2001 From: SimengLiu-nv Date: Sat, 21 Mar 2026 21:48:42 -0700 Subject: [PATCH 3/4] Address coderaddit's comments Signed-off-by: SimengLiu-nv --- .../batch_manager/kvCacheManager.cpp | 5 ++ .../batch_manager/microBatchScheduler.cpp | 32 ++++++++-- .../batch_manager/microBatchSchedulerTest.cpp | 59 +++++++++++++++++++ .../_torch/pyexecutor/scheduler/scheduler.py | 48 ++++++++++----- 4 files changed, 124 insertions(+), 20 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 2ca65be219dc..107c703b0b79 100755 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -2234,6 +2234,9 @@ void KVCacheManager::startScheduling() SizeType32 KVCacheManager::getNeededBlocksOneStep( LlmRequest const& req, bool twoStepsLookAhead, SizeType32 windowSize) const { + // Default to zero; overwritten below when block reuse is active for a first-chunk context request. + req.setEstimatedReusableTokens(0); + if ((req.isContextInitState() && req.isFirstContextChunk()) || req.isDisaggGenerationInitState()) { auto const chunkSize = req.mMaxNewTokens; @@ -2293,6 +2296,8 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep( SizeType32 KVCacheManager::getRemainingBlocksToCompletion(LlmRequest const& req, SizeType32 windowSize) const { + // Default to zero; overwritten below when block reuse is active for a first-chunk context request. + req.setEstimatedReusableTokens(0); if (isCrossKv()) { diff --git a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp index 8d926d56adc6..01c5555797a3 100644 --- a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp @@ -38,11 +38,16 @@ void MicroBatchScheduler::fitDraftTokens(RequestVector& contextsToBeChunked, std::optional ctxTokensCapacity, SizeType32 const chunkUnitSize, std::optional const& maxContextLength) { - // How many context tokens are in this batch already? + // How many compute tokens (chunk - reusable) are in this batch already? SizeType32 numCtxTokens{0}; for (auto const& llmReq : contextsToBeChunked) { - numCtxTokens += llmReq->getContextChunkSize(); + SizeType32 const chunkSize = llmReq->getContextChunkSize(); + // contextRemaining = P for first chunk; used to compute actual model token count. + SizeType32 const contextRemaining = llmReq->getContextRemainingLength(); + SizeType32 const reusable + = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), contextRemaining) : 0; + numCtxTokens += std::min(chunkSize, std::max(0, contextRemaining - reusable)); } // Discard draft tokens that won't fit into the existing chunk unit, max @@ -122,14 +127,25 @@ void MicroBatchScheduler::setCtxRequestsChunkSizegetContextRemainingLength(); + // Reusable tokens are "free" — they don't consume forward-pass compute budget. + SizeType32 const reusable + = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), suggestedChunkSize) : 0; + SizeType32 const computeCost = suggestedChunkSize - reusable; SizeType32 actualChunkSize = suggestedChunkSize; - if (ctxTokensCapacity) + if (ctxTokensCapacity && computeCost > ctxTokensCapacity.value()) { - actualChunkSize = std::min(ctxTokensCapacity.value(), actualChunkSize); + // Model processes min(chunk_size, P - reusable) tokens starting from position reusable. + // To keep model tokens within budget: chunk_size <= capacity (not reusable + capacity). + actualChunkSize = ctxTokensCapacity.value(); } if (maxContextLength) { - actualChunkSize = std::min(maxContextLength.value(), actualChunkSize); + // maxContextLength limits compute tokens, not total tokens. + SizeType32 const actualCompute = std::max(0, actualChunkSize - reusable); + if (actualCompute > maxContextLength.value()) + { + actualChunkSize = std::min(reusable + maxContextLength.value(), suggestedChunkSize); + } } if (actualChunkSize != suggestedChunkSize) { @@ -138,7 +154,11 @@ void MicroBatchScheduler::setCtxRequestsChunkSizesetContextChunkSize(actualChunkSize); if (ctxTokensCapacity) { - ctxTokensCapacity = ctxTokensCapacity.value() - actualChunkSize; + // Decrement by actual model token count: min(chunk_size, P - reusable). + // This equals min(actualChunkSize, computeCost) since computeCost = suggestedChunkSize - reusable. + SizeType32 const modelCost + = std::min(actualChunkSize, std::max(0, suggestedChunkSize - reusable)); + ctxTokensCapacity = ctxTokensCapacity.value() - modelCost; } } } diff --git a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp index 8436f08bb43c..32b75309c6c6 100644 --- a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp @@ -889,6 +889,65 @@ TEST_F(MicroBatchSchedulerTest, ReusableTokensWithChunkedContextFCFS) } } +TEST_F(MicroBatchSchedulerTest, ReusableTokensWithChunkedContextFCFS_OverBudgetMultiRequest) +{ + // Regression test for FCFS compute-aware scheduling with multiple requests and reusable tokens. + // + // Setup: 3 requests, each promptLen=20, reusable=15. + // Compute budget = 7 (maxNumTokens=7, no generation requests). + // Total compute if all fit: 3 * (20 - 15) = 15 > 7 → allContextRequestsFit = false, + // FCFS chunk allocation is exercised. + // + // The model processes min(chunk_size, P - reusable) tokens per request starting from + // context_current_position = reusable. chunk_size must stay <= capacity so that model + // tokens do not exceed max_num_tokens. + // + // Expected (correct compute-aware FCFS): + // req0: chunk=20 (full context fits; model processes min(20,5)=5 tokens; budget 7→2) + // req1: chunk=2 (capacity=2 remaining; model processes min(2,5)=2 tokens; budget 2→0) + // req2: chunk=0 (no budget remaining; not scheduled) + // + // Bug (raw-token FCFS): req0 gets chunk=7 (limited by raw budget), req1 and req2 starved. + constexpr SizeType32 maxNumTokens = 7; + constexpr SizeType32 maxBatchSize = 4; + constexpr SizeType32 chunkUnitSize = 1; + constexpr SizeType32 reusableTokens = 15; + constexpr SizeType32 promptLen = 20; + constexpr SizeType32 maxNewTokens = 5; + constexpr ContextChunkingPolicy ctxChunkPolicy{ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED}; + + mNumContexts = 3; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler + = std::make_shared(ContextChunkingConfig{ctxChunkPolicy, chunkUnitSize}, std::nullopt); + + RequestVector activeRequests; + auto req0 = createRequest(promptLen, maxNewTokens, 0); + auto req1 = createRequest(promptLen, maxNewTokens, 1); + auto req2 = createRequest(promptLen, maxNewTokens, 2); + req0->setEstimatedReusableTokens(reusableTokens); + req1->setEstimatedReusableTokens(reusableTokens); + req2->setEstimatedReusableTokens(reusableTokens); + activeRequests.push_back(req0); + activeRequests.push_back(req1); + activeRequests.push_back(req2); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + + // req0 and req1 are scheduled; req2 gets chunk=0 (no budget) and is filtered out. + EXPECT_EQ(ctx.size(), 2u) << "req0 and req1 fit; req2 has no budget and is not scheduled"; + + // req0: full context fits (compute cost = 5 <= budget 7); chunk = promptLen. + EXPECT_EQ(req0->getContextChunkSize(), promptLen) << "req0: full context fits (model 5 tokens out of budget 7)"; + + // req1: 2 compute tokens remain; chunk = 2 (the remaining capacity, not reusable + capacity). + EXPECT_EQ(req1->getContextChunkSize(), 2) << "req1: chunk = remaining capacity (2), model processes 2 tokens"; + + // req2: no budget remaining; chunk = 0, not scheduled. + EXPECT_EQ(req2->getContextChunkSize(), 0) << "req2: budget exhausted, chunk = 0"; +} + TEST_F(MicroBatchSchedulerTest, ReusableTokensZeroHasNoEffect) { // Verify that zero reusable tokens (the default) produces identical scheduling diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 6fb322346b83..6291a57df670 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -600,11 +600,12 @@ def _set_ctx_requests_chunk_size( self._fit_draft_tokens(requests, capacity, unit_size) def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], unit_size: int): - num_ctx_tokens = 0 + # capacity is a compute-token budget (non-reusable only). Track compute tokens, not including reusable tokens. + num_compute_tokens = 0 num_tokens_single_loop = 1 # C++ Loop: while ((!capacity || numCtxTokens < capacity) && numTokensSingleLoop) - while (capacity is None or num_ctx_tokens < capacity) and num_tokens_single_loop > 0: + while (capacity is None or num_compute_tokens < capacity) and num_tokens_single_loop > 0: num_tokens_single_loop = 0 for req in requests: past_size = req.context_chunk_size @@ -619,11 +620,15 @@ def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], req.context_chunk_size = suggested_size actual_size = req.context_chunk_size - actual_increment = actual_size - past_size + # Compute the compute-token increment (reusable tokens don't consume budget). + reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 + past_compute = max(0, past_size - min(reusable, past_size)) + actual_compute = max(0, actual_size - min(reusable, actual_size)) + compute_increment = actual_compute - past_compute # Check Constraints - # 1. Capacity - if capacity is not None and (num_ctx_tokens + actual_increment > capacity): + # 1. Capacity (in compute tokens) + if capacity is not None and (num_compute_tokens + compute_increment > capacity): req.context_chunk_size = past_size # Revert continue @@ -632,8 +637,8 @@ def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], req.context_chunk_size = past_size # Revert continue - num_ctx_tokens += actual_increment - num_tokens_single_loop += actual_increment + num_compute_tokens += compute_increment + num_tokens_single_loop += compute_increment def _chunk_fcfs( self, @@ -660,8 +665,9 @@ def _chunk_fcfs( # Limit by compute capacity (reusable tokens are free) if current_compute_capacity < compute_cost: - # Can't fit all new tokens; allow all reusable + whatever compute fits - actual_size = min(reusable + current_compute_capacity, suggested_size) + # Model processes min(chunk_size, P - reusable) tokens from position reusable. + # To keep model tokens within budget: chunk_size = capacity (not reusable + capacity). + actual_size = int(current_compute_capacity) # Limit by max_context_length (applies to compute portion only) if self.max_context_length is not None: @@ -676,14 +682,28 @@ def _chunk_fcfs( req.context_chunk_size = int(actual_size) - # Subtract only the compute portion from capacity - actual_compute = max(0, req.context_chunk_size - reusable) + # Subtract actual model token count: min(chunk_size, P - reusable). + actual_model_cost = min(req.context_chunk_size, max(0, int(suggested_size) - reusable)) if capacity is not None: - current_compute_capacity -= actual_compute + current_compute_capacity -= actual_model_cost def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit_size: int): - # Calculate tokens already taken by the batch so far - num_ctx_tokens = sum(req.context_chunk_size for req in requests) + # capacity is a compute-token budget. Sum actual model tokens per request: + # min(chunk_size, P - reusable), where P = context_remaining_length for first chunk. + num_ctx_tokens = sum( + min( + req.context_chunk_size, + max( + 0, + req.context_remaining_length + - min( + req.estimated_reusable_tokens if req.is_first_context_chunk else 0, + req.context_remaining_length, + ), + ), + ) + for req in requests + ) for req in requests: if req.is_last_context_chunk and req.has_draft_tokens: From e0d5deb099e4f40fbc1a7ff6c81f672527b106bd Mon Sep 17 00:00:00 2001 From: SimengLiu-nv Date: Sun, 22 Mar 2026 10:27:47 -0700 Subject: [PATCH 4/4] Address lancelly's comment and sync up the code changes in C++ and python implementation of microBatchSchedulerTest. Added tests with full coverage in both C++ and python tests. Signed-off-by: SimengLiu-nv --- .../batch_manager/microBatchScheduler.cpp | 68 ++++++- .../batch_manager/microBatchSchedulerTest.cpp | 190 ++++++++++++++++++ .../_torch/pyexecutor/scheduler/scheduler.py | 62 ++++-- .../_torch/executor/test_py_scheduler.py | 132 ++++++++++++ 4 files changed, 434 insertions(+), 18 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp index 01c5555797a3..92d4589e0b66 100644 --- a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp @@ -86,6 +86,24 @@ void MicroBatchScheduler::fitDraftTokens(RequestVector& contextsToBeChunked, } } +// Assigns chunk sizes to context requests under the kEQUAL_PROGRESS policy. +// +// All requests advance together in lock-step: each iteration of the outer while-loop +// offers every request one additional chunkUnitSize of tokens. This continues until +// the compute budget (ctxTokensCapacity) is exhausted or no request can advance further. +// +// Budget accounting is compute-aware: tokens covered by the reusable KV-cache prefix +// (getEstimatedReusableTokens) are served from cache and do not consume forward-pass +// capacity. Only tokens beyond that prefix count against ctxTokensCapacity. +// The reusable prefix is only considered on the very first chunk of a request +// (isFirstContextChunk), since subsequent chunks start past the cached prefix. +// +// A request is skipped for this iteration if adding chunkUnitSize would exceed +// ctxTokensCapacity or maxContextLength; its chunk size is left unchanged. +// +// Loop-termination uses the raw token increment (actualIncrement), not the +// compute-adjusted one, so a request whose entire new chunk is reusable still +// counts as progress and does not cause a premature exit. template <> void MicroBatchScheduler::setCtxRequestsChunkSize( RequestVector& contextsToBeChunked, std::optional ctxTokensCapacity, SizeType32 const chunkUnitSize, @@ -107,18 +125,54 @@ void MicroBatchScheduler::setCtxRequestsChunkSizegetContextChunkSize(); SizeType32 actualIncrement = actualChunkSize - pastChunkSize; - if ((ctxTokensCapacity && numCtxTokens + actualIncrement > ctxTokensCapacity.value()) + // Compute-aware budget: reusable tokens are served from cache and do not + // consume forward-pass capacity. Only the tokens beyond the reusable prefix count. + SizeType32 const reusable = llmReq->isFirstContextChunk() + ? std::min(llmReq->getEstimatedReusableTokens(), llmReq->getContextRemainingLength()) + : 0; + SizeType32 const pastCompute = std::max(0, pastChunkSize - std::min(reusable, pastChunkSize)); + SizeType32 const actualCompute + = std::max(0, actualChunkSize - std::min(reusable, actualChunkSize)); + SizeType32 const computeIncrement = actualCompute - pastCompute; + + if ((ctxTokensCapacity && numCtxTokens + computeIncrement > ctxTokensCapacity.value()) || (maxContextLength && actualChunkSize > maxContextLength.value())) { llmReq->setContextChunkSize(pastChunkSize); continue; } - numCtxTokens += actualIncrement; + numCtxTokens += computeIncrement; + // Keep raw actualIncrement for loop-termination detection (not compute-aware). numTokensSingleLoop += actualIncrement; } } } +// Assigns chunk sizes to context requests under the kFIRST_COME_FIRST_SERVED policy. +// +// Requests are processed in order. Each request greedily claims as many tokens as +// possible from the remaining compute budget (ctxTokensCapacity) before the next +// request is considered — hence "first come, first served". +// +// For each request the desired chunk is its full remaining context length. The actual +// chunk size is then reduced by two independent constraints (applied in order): +// +// 1. ctxTokensCapacity: the available forward-pass compute budget. +// Reusable tokens (cached KV prefix, first chunk only) are free; only +// non-reusable tokens count. If the non-reusable portion exceeds the budget, +// the chunk is capped at ctxTokensCapacity (not reusable + ctxTokensCapacity, +// because the model processes tokens starting from position 0 of the chunk). +// +// 2. maxContextLength: an upper bound on the number of compute tokens per chunk. +// If the non-reusable portion exceeds this, the chunk is capped at +// reusable + maxContextLength (clamped to suggestedChunkSize). +// +// When either constraint trims the chunk, the result is aligned down to the nearest +// chunkUnitSize boundary to avoid KV-cache fragmentation. +// +// After assigning the chunk, ctxTokensCapacity is decremented by the actual model +// cost: min(actualChunkSize, non-reusable tokens), so the budget available to +// subsequent requests reflects only the compute consumed by this one. template <> void MicroBatchScheduler::setCtxRequestsChunkSize( RequestVector& contextsToBeChunked, std::optional ctxTokensCapacity, SizeType32 const chunkUnitSize, @@ -163,6 +217,16 @@ void MicroBatchScheduler::setCtxRequestsChunkSize ctxTokensCapacity, SizeType32 const chunkUnitSize, std::optional const& maxContextLength) diff --git a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp index 32b75309c6c6..4edfda5f1ebf 100644 --- a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp @@ -948,6 +948,55 @@ TEST_F(MicroBatchSchedulerTest, ReusableTokensWithChunkedContextFCFS_OverBudgetM EXPECT_EQ(req2->getContextChunkSize(), 0) << "req2: budget exhausted, chunk = 0"; } +TEST_F(MicroBatchSchedulerTest, ReusableTokensWithChunkedContextEqualProgress) +{ + // Test compute-aware budget tracking in EQUAL_PROGRESS with reusable tokens. + // + // Setup: 2 requests, promptLen=15, reusable=10, compute budget=3, chunkUnit=1. + // + // In EQUAL_PROGRESS, chunks grow by 1 per request per iteration. Reusable tokens + // do not consume forward-pass capacity, so the budget is only charged for tokens + // beyond the reusable prefix. + // + // Expected (compute-aware EQUAL_PROGRESS): + // Tokens 0-10 are "free" (all cached). Budget is only consumed for tokens > 10. + // req0: chunk=12 (model cost = 12-10 = 2) + // req1: chunk=11 (model cost = 11-10 = 1) + // total compute = 3 = budget (fully utilised). + // + // Bug (raw-token EQUAL_PROGRESS): chunks capped at 2 and 1 (budget fully consumed + // after the very first two tokens, ignoring reusable credits). + constexpr SizeType32 maxNumTokens = 3; + constexpr SizeType32 maxBatchSize = 4; + constexpr SizeType32 chunkUnitSize = 1; + constexpr SizeType32 reusableTokens = 10; + constexpr SizeType32 promptLen = 15; + constexpr SizeType32 maxNewTokens = 5; + constexpr ContextChunkingPolicy ctxChunkPolicy{ContextChunkingPolicy::kEQUAL_PROGRESS}; + + mNumContexts = 2; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler + = std::make_shared(ContextChunkingConfig{ctxChunkPolicy, chunkUnitSize}, std::nullopt); + + RequestVector activeRequests; + auto req0 = createRequest(promptLen, maxNewTokens, 0); + auto req1 = createRequest(promptLen, maxNewTokens, 1); + req0->setEstimatedReusableTokens(reusableTokens); + req1->setEstimatedReusableTokens(reusableTokens); + activeRequests.push_back(req0); + activeRequests.push_back(req1); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + + EXPECT_EQ(ctx.size(), 2u) << "Both requests should be scheduled"; + + // req0 gets one extra unit over req1 due to equal-progress ordering. + EXPECT_EQ(req0->getContextChunkSize(), 12) << "req0: reusable(10) + 2 compute tokens = chunk 12"; + EXPECT_EQ(req1->getContextChunkSize(), 11) << "req1: reusable(10) + 1 compute token = chunk 11"; +} + TEST_F(MicroBatchSchedulerTest, ReusableTokensZeroHasNoEffect) { // Verify that zero reusable tokens (the default) produces identical scheduling @@ -978,6 +1027,147 @@ TEST_F(MicroBatchSchedulerTest, ReusableTokensZeroHasNoEffect) EXPECT_EQ(gen.size(), 0); } +TEST_F(MicroBatchSchedulerTest, ReusableTokensNoChunkingMinCostIsOne) +{ + // Test (no-chunking path) that the compute cost of a context request is floored at 1 + // even when the number of reusable tokens exceeds the prompt length. + // + // req0: promptLen=10, reusable=15 → compute = max(1, 10-15) = 1 + // req1: promptLen=10, reusable=0 → compute = 10 + // maxNumTokens=10: req0 costs 1 token of compute; adding req1 would bring the total + // to 11 which exceeds the budget → req1 is not scheduled. + // + // Python ref: test_reusable_tokens_no_chunking_min_cost_is_one + constexpr SizeType32 maxNumTokens = 10; + constexpr SizeType32 maxBatchSize = 4; + constexpr int32_t promptLen = 10; + constexpr int32_t maxNewTokens = 5; + + mNumContexts = 1; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler = std::make_shared(); // no chunking + + auto req0 = createRequest(promptLen, maxNewTokens, 0); + auto req1 = createRequest(promptLen, maxNewTokens, 1); + req0->setEstimatedReusableTokens(15); // exceeds promptLen + + RequestVector activeRequests = {req0, req1}; + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + + EXPECT_EQ(ctx.size(), 1u) << "req0 costs 1 compute (floored), req1 costs 10; only req0 fits"; + EXPECT_EQ(ctx.at(0)->mRequestId, 0u) << "req0 should be the scheduled request"; + EXPECT_EQ(gen.size(), 0u); +} + +TEST_F(MicroBatchSchedulerTest, BeamWidthMismatchSkipped) +{ + // Test that generation requests whose beam width differs from the first scheduled + // generation request are skipped for that step. + // + // req0: BW=1 (scheduled), req1: BW=4 (skipped), req2: BW=1 (scheduled alongside req0). + // + // Python ref: test_beam_width_mismatch_skipped + constexpr SizeType32 maxBatchSize = 4; + constexpr int32_t promptLen = 2; + constexpr int32_t maxNewTokens = 5; + + mNumContexts = 1; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler = std::make_shared(); + + // Helper: manually advance a context request to generation state. + auto advanceToGeneration = [](std::shared_ptr req) + { + req->setContextChunkSize(req->getContextRemainingLength()); + req->moveToNextContextChunk(); + req->setState(LlmRequestState::kGENERATION_IN_PROGRESS); + req->addNewTokens({42}); + }; + + auto req0 = createRequest(promptLen, maxNewTokens, 0, /*beamWidth=*/1); + auto req1 = createRequest(promptLen, maxNewTokens, 1, /*beamWidth=*/4); + auto req2 = createRequest(promptLen, maxNewTokens, 2, /*beamWidth=*/1); + advanceToGeneration(req0); + advanceToGeneration(req1); + advanceToGeneration(req2); + + RequestVector activeRequests = {req0, req1, req2}; + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, std::nullopt); + + EXPECT_EQ(gen.size(), 2u) << "req0 and req2 (BW=1) scheduled; req1 (BW=4) skipped"; + EXPECT_EQ(gen.at(0)->mRequestId, 0u); + EXPECT_EQ(gen.at(1)->mRequestId, 2u); + EXPECT_EQ(ctx.size(), 0u); +} + +TEST_F(MicroBatchSchedulerTest, InflightRequestsExcluded) +{ + // Test that requests already tracked in inflightReqIds are not re-scheduled. + // This guards the 2-micro-batch overlap mechanism: when requests from micro-batch + // slot 0 are inflight, the scheduler must not pick them up for slot 1. + // + // Python ref: test_inflight_requests_excluded + constexpr SizeType32 maxBatchSize = 4; + constexpr int32_t promptLen = 10; + constexpr int32_t maxNewTokens = 5; + + mNumContexts = 1; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler = std::make_shared(); + + RequestVector activeRequests; + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 0)); + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 1)); + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 2)); + + // Mark requests 0 and 2 as inflight — only request 1 should be scheduled. + ReqIdsSet inflightReqIds = {0, 2}; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, std::nullopt); + + EXPECT_EQ(ctx.size(), 1u) << "Only req1 (not inflight) should be scheduled"; + EXPECT_EQ(ctx.at(0)->mRequestId, 1u); + EXPECT_EQ(gen.size(), 0u); +} + +TEST_F(MicroBatchSchedulerTest, ChunkSizeZeroNotScheduled) +{ + // Test that context requests whose chunk_size remains 0 after the chunking step + // are excluded from the scheduled output. + // + // Setup: budget=5, chunkUnit=5, EQUAL_PROGRESS, 2 requests with promptLen=20. + // On the first scheduling step EQUAL_PROGRESS distributes one unit at a time; + // with budget=5 only one request receives chunk=5 while the other stays at 0. + // The zero-chunk request must not appear in the context output. + // + // Python ref: test_chunk_size_zero_not_scheduled + constexpr SizeType32 maxNumTokens = 5; + constexpr SizeType32 maxBatchSize = 4; + constexpr SizeType32 chunkUnitSize = 5; + constexpr ContextChunkingPolicy ctxChunkPolicy{ContextChunkingPolicy::kEQUAL_PROGRESS}; + constexpr int32_t promptLen = 20; + constexpr int32_t maxNewTokens = 5; + + mNumContexts = 1; + mContextRequests.resize(mNumContexts); + mMicroBatchScheduler + = std::make_shared(ContextChunkingConfig{ctxChunkPolicy, chunkUnitSize}, std::nullopt); + + RequestVector activeRequests; + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 0)); + activeRequests.push_back(createRequest(promptLen, maxNewTokens, 1)); + + ReqIdsSet inflightReqIds; + auto const [ctx, gen] = (*mMicroBatchScheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens); + + for (auto const& req : ctx) + { + EXPECT_GT(req->getContextChunkSize(), 0) << "Requests with chunk_size=0 must not appear in output"; + } + EXPECT_EQ(ctx.size(), 1u) << "Only the request that received chunk=5 should be scheduled"; +} + //// // Combined Capacity Scheduler + Micro Batch Scheduler tests. // These verify the end-to-end data flow: the capacity scheduler populates diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 6291a57df670..82555e17e28f 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -583,7 +583,11 @@ def _set_ctx_requests_chunk_size( requests: RequestList, capacity: Optional[int], ): - # C++: Resets all chunk sizes to 0 at start + """ + Mirrors MicroBatchScheduler::setCtxRequestsChunkSize (microBatchScheduler.cpp). + + Structure: reset all chunk sizes to 0, dispatch to policy, then trim draft tokens. + """ for req in requests: req.context_chunk_size = 0 @@ -600,7 +604,16 @@ def _set_ctx_requests_chunk_size( self._fit_draft_tokens(requests, capacity, unit_size) def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], unit_size: int): - # capacity is a compute-token budget (non-reusable only). Track compute tokens, not including reusable tokens. + """ + Mirrors the kEQUAL_PROGRESS specialization of setCtxRequestsChunkSize (microBatchScheduler.cpp). + + All requests advance in lock-step: each while-loop iteration offers every request one + additional unit_size worth of tokens, until the compute budget (capacity) is exhausted + or no request can make further progress. + + capacity is a compute-token budget (non-reusable only). Tokens covered by the reusable + KV-cache prefix are not charged against it. + """ num_compute_tokens = 0 num_tokens_single_loop = 1 @@ -613,20 +626,25 @@ def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], # C++ logic: suggested = past + unit suggested_size = past_size + unit_size - # Ensure we don't exceed what the request actually needs + # In C++, setContextChunkSize() clamps to getContextRemainingLength() internally. + # Python's setter does not clamp, so we replicate that clamp here explicitly. remaining_total = req.context_remaining_length suggested_size = min(suggested_size, remaining_total) req.context_chunk_size = suggested_size + # In C++, getContextChunkSize() is called after the setter to observe the + # clamped value. In Python the setter does not re-clamp, so actual_size == + # suggested_size; the readback is kept for structural symmetry with C++. actual_size = req.context_chunk_size + # Compute the compute-token increment (reusable tokens don't consume budget). reusable = req.estimated_reusable_tokens if req.is_first_context_chunk else 0 past_compute = max(0, past_size - min(reusable, past_size)) actual_compute = max(0, actual_size - min(reusable, actual_size)) compute_increment = actual_compute - past_compute - # Check Constraints + # Check Constraints — mirrors the if-block guard in C++ before numCtxTokens += # 1. Capacity (in compute tokens) if capacity is not None and (num_compute_tokens + compute_increment > capacity): req.context_chunk_size = past_size # Revert @@ -638,7 +656,9 @@ def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], continue num_compute_tokens += compute_increment - num_tokens_single_loop += compute_increment + # Keep raw increment (actual_size - past_size) including reusable tokens for loop-termination detection, + # not compute_increment. + num_tokens_single_loop += actual_size - past_size def _chunk_fcfs( self, @@ -646,13 +666,19 @@ def _chunk_fcfs( capacity: Optional[int], unit_size: int, ): - # capacity represents the COMPUTE budget (non-reusable tokens). - # Reusable tokens are "free" — they don't consume forward-pass compute. + """Mirrors the kFIRST_COME_FIRST_SERVED specialization of setCtxRequestsChunkSize (microBatchScheduler.cpp). + + Requests are processed in order; each greedily claims as much of the remaining compute + budget as possible before the next request is considered — hence "first come, first served". + + capacity represents the COMPUTE budget (non-reusable tokens). + Reusable tokens are "free" — they don't consume forward-pass compute. + """ current_compute_capacity = capacity if capacity is not None else float("inf") for req in requests: suggested_size = req.context_remaining_length - # Only the first context chunk can reuse cached KV blocks + # Only the first context chunk can reuse cached KV blocks; # subsequent chunks must compute all remaining tokens. reusable = min( req.estimated_reusable_tokens if req.is_first_context_chunk else 0, @@ -660,36 +686,40 @@ def _chunk_fcfs( ) compute_cost = suggested_size - reusable - # Start with full context as the allocation target + # Start with full context as the allocation target. actual_size = suggested_size - # Limit by compute capacity (reusable tokens are free) + # Constraint 1: compute capacity. if current_compute_capacity < compute_cost: - # Model processes min(chunk_size, P - reusable) tokens from position reusable. - # To keep model tokens within budget: chunk_size = capacity (not reusable + capacity). + # Model processes min(chunk_size, P - reusable) tokens from position reusable, + # where P = suggested_size = context_remaining_length (full prompt on first chunk). + # To keep model tokens within budget: chunk_size = capacity (not reusable + capacity), actual_size = int(current_compute_capacity) - # Limit by max_context_length (applies to compute portion only) + # Constraint 2: max_context_length (applies to compute portion only). if self.max_context_length is not None: actual_compute = max(0, actual_size - reusable) if actual_compute > self.max_context_length: actual_size = reusable + self.max_context_length actual_size = min(actual_size, suggested_size) - # Round down to unit size if we had to truncate + # Align down to unit_size when either constraint trimmed the chunk, to avoid + # KV-cache fragmentation. if actual_size < suggested_size: actual_size = (int(actual_size) // unit_size) * unit_size req.context_chunk_size = int(actual_size) - # Subtract actual model token count: min(chunk_size, P - reusable). + # Decrement budget by actual model token count: min(chunk_size, P - reusable). + # where P = suggested_size =context_remaining_length (full prompt on first chunk) actual_model_cost = min(req.context_chunk_size, max(0, int(suggested_size) - reusable)) if capacity is not None: current_compute_capacity -= actual_model_cost def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit_size: int): # capacity is a compute-token budget. Sum actual model tokens per request: - # min(chunk_size, P - reusable), where P = context_remaining_length for first chunk. + # min(chunk_size, P - reusable), where P = context_remaining_length + # (the full prompt length on the first context chunk). num_ctx_tokens = sum( min( req.context_chunk_size, diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index 3ebf53bbc17e..01b5c90edd30 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -493,6 +493,52 @@ def test_simple_no_overlap_max_context_length(self): assert len(ctx2) == 1 assert ctx2[0].context_chunk_size <= 12 + def test_simple_with_overlap(self): + """ + 2-micro-batch overlap: with max_batch_size=2 and 4 context requests, + the inflight mechanism ensures alternating pairs of requests are scheduled + in successive steps, mimicking the 2-slot pipeline in C++. + + Step 1: no inflight → req0, req1 scheduled (slot 0) + Step 2: {0,1} inflight → req2, req3 scheduled (slot 1) + Step 3: {2,3} inflight → req0, req1 scheduled again (slot 0 freed) + + C++ ref: SimpleWithOverlap + """ + scheduler = PyMicroBatchScheduler(max_batch_size=2, max_num_tokens=None) + requests = [make_context_request(i, prompt_len=10) for i in range(4)] + + # Step 1: slot 0 — req0, req1 scheduled + ctx0, _ = scheduler.schedule(requests, set()) + assert {r.request_id for r in ctx0} == {0, 1} + slot0_inflight = {r.request_id for r in ctx0} + + # Step 2: slot 1 — req0/req1 still inflight, req2/req3 scheduled + ctx1, _ = scheduler.schedule(requests, slot0_inflight) + assert {r.request_id for r in ctx1} == {2, 3} + slot1_inflight = {r.request_id for r in ctx1} + + # Step 3: slot 0 freed (inflight = slot1 only) — req0/req1 scheduled again + ctx2, _ = scheduler.schedule(requests, slot1_inflight) + assert {r.request_id for r in ctx2} == {0, 1} + + def test_gen_draft_tokens_max_num_tokens(self): + """ + Draft tokens in generation phase reduce batch size when maxNumTokens is tight. + Each gen request costs beam_width + num_draft_tokens = 1 + 63 = 64 tokens. + With maxNumTokens=128, only 2 of the 4 gen requests fit (2*64=128). + + C++ ref: GenDraftTokensMaxNumTokens + """ + scheduler = PyMicroBatchScheduler(max_batch_size=64, max_num_tokens=128) + requests = [make_generation_request(i, beam_width=1, draft_tokens_len=63) for i in range(4)] + ctx, gen = scheduler.schedule(requests, set()) + # Each request costs 1 + 63 = 64 tokens; 2 fit (128 = budget), 3 don't (192 > 128). + assert len(gen) == 2 + assert gen[0].request_id == 0 + assert gen[1].request_id == 1 + assert len(ctx) == 0 + # ############################################################################ # @@ -928,6 +974,92 @@ def test_reusable_tokens_no_chunking_min_cost_is_one(self): assert len(ctx) == 1 assert ctx[0].request_id == 0 + def test_reusable_tokens_fcfs_over_budget_multi_request(self): + """ + FCFS compute-aware scheduling with multiple requests exceeding the compute budget. + C++ ref: ReusableTokensWithChunkedContextFCFS_OverBudgetMultiRequest + + Setup: 3 requests, prompt_len=15, reusable=8, compute budget=20, chunk_unit=1. + Note: In PyMicroBatchScheduler, max_context_length = max_num_tokens = 20, so + prompt_len=15 keeps chunk sizes within the max_context_length limit. + Total compute if all scheduled in full: 3 * (15 - 8) = 21 > 20 → chunking exercised. + + The model processes min(chunk_size, P - reusable) tokens from position reusable, + where P = context_remaining_length = 15 (full prompt on first chunk). + + Expected (compute-aware FCFS): + req0: full context fits (compute=7 <= 20); chunk=15; budget 20→13 + req1: full context fits (compute=7 <= 13); chunk=15; budget 13→6 + req2: compute=7 > 6 → chunk=6 (remaining capacity); budget→0 + """ + config = ContextChunkingConfig(ChunkingPolicy.FIRST_COME_FIRST_SERVED, chunk_unit_size=1) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=20, ctx_chunk_config=config + ) + req0 = make_context_request(0, prompt_len=15) + req1 = make_context_request(1, prompt_len=15) + req2 = make_context_request(2, prompt_len=15) + req0.estimated_reusable_tokens = 8 + req1.estimated_reusable_tokens = 8 + req2.estimated_reusable_tokens = 8 + + ctx, gen = scheduler.schedule([req0, req1, req2], set()) + + # Note: ctx is sorted — partially-chunked requests come before full-context ones. + # Look up by request_id rather than by position. + chunks = {r.request_id: r.context_chunk_size for r in ctx} + assert len(ctx) == 3, "All three requests should be scheduled" + assert chunks[0] == 15, "req0: full context (compute=7, budget 20→13)" + assert chunks[1] == 15, "req1: full context (compute=7, budget 13→6)" + assert chunks[2] == 6, "req2: chunk=6 (remaining capacity)" + + def test_reusable_tokens_equal_progress(self): + """ + EQUAL_PROGRESS compute-aware budget tracking with reusable tokens. + C++ ref: ReusableTokensWithChunkedContextEqualProgress + + Note: In PyMicroBatchScheduler, max_context_length = max_num_tokens, so + individual chunk sizes are capped at max_num_tokens. Parameters are chosen + so the reusable prefix is smaller than max_num_tokens, allowing the + compute-aware path to produce noticeably larger chunks than the raw path. + + Setup: 3 requests, prompt_len=15, reusable=5, compute budget=8, chunk_unit=1. + max_context_length = max_num_tokens = 8 (caps individual chunk sizes at 8). + Total compute if all scheduled in full: 3 * (15 - 5) = 30 > 8 → chunking exercised. + + Reusable tokens 0-4 are "free" (compute_increment = 0 until chunk > 5). + Budget is only consumed for tokens beyond the reusable prefix. + + This test validates the loop-termination fix: num_tokens_single_loop must use + the raw chunk increment (not compute_increment). With the bug, the loop exits + on the very first iteration because compute_increment=0 for all reqs → loops + terminates at chunk=1. With the fix, the loop continues until the compute + budget is exhausted or max_context_length is reached. + + Expected (compute-aware EQUAL_PROGRESS): + req0: chunk=8 (model cost = 8 - 5 = 3) + req1: chunk=8 (model cost = 3) + req2: chunk=7 (model cost = 7 - 5 = 2; total compute 3+3+2=8 = budget) + """ + config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=1) + scheduler = PyMicroBatchScheduler( + max_batch_size=4, max_num_tokens=8, ctx_chunk_config=config + ) + req0 = make_context_request(0, prompt_len=15) + req1 = make_context_request(1, prompt_len=15) + req2 = make_context_request(2, prompt_len=15) + req0.estimated_reusable_tokens = 5 + req1.estimated_reusable_tokens = 5 + req2.estimated_reusable_tokens = 5 + + ctx, gen = scheduler.schedule([req0, req1, req2], set()) + + chunks = {r.request_id: r.context_chunk_size for r in ctx} + assert len(ctx) == 3, "All three requests should be scheduled" + assert chunks[0] == 8, "req0: reusable(5) + 3 compute tokens = 8" + assert chunks[1] == 8, "req1: reusable(5) + 3 compute tokens = 8" + assert chunks[2] == 7, "req2: reusable(5) + 2 compute tokens = 7" + # ############################################################################ #