From 4655598b17726a4535c728fbc1e1cea3e294298d Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Mon, 13 Apr 2026 07:46:26 -0700 Subject: [PATCH 1/2] [None][fix] Fix compute token accounting for KV cache reuse with context chunking When KV cache reuse is combined with context chunking, setPrepopulatedPromptLen shifts the chunk window right by the reused amount rather than shrinking it. Non-last chunks still process ~chunkSize tokens in the forward pass. The old formula (chunkSize - reusable) underestimated compute cost for non-last chunks, causing the scheduler to over-pack batches beyond the token budget. This adds reuse_adjusted_compute() to correctly account for chunk-shift behavior, and updates _waiting_requests in PyExecutor to subtract reusable tokens from scheduled token counts. Signed-off-by: Lanyu Liao --- .../batch_manager/microBatchScheduler.cpp | 75 ++++++---- .../batch_manager/microBatchSchedulerTest.cpp | 52 +++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 50 ++++++- .../_torch/executor/test_py_executor.py | 138 ++++++++++++++++++ 4 files changed, 278 insertions(+), 37 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp index 40b760c3cb0a..3e8fca0be052 100644 --- a/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp @@ -24,6 +24,25 @@ namespace tensorrt_llm::batch_manager using SizeType32 = MicroBatchScheduler::SizeType32; +/// Return the forward-pass token cost for a context chunk with KV cache reuse. +/// +/// setPrepopulatedPromptLen shifts the chunk window right by the prepopulated +/// amount rather than shrinking it. For non-last chunks the model still +/// processes approximately @p chunkSize tokens; only for the last chunk is the +/// cost @p contextRemaining - @p reusable. +static SizeType32 reuse_adjusted_compute(SizeType32 chunkSize, SizeType32 reusable, SizeType32 contextRemaining) +{ + if (reusable <= 0) + { + return chunkSize; + } + if (reusable + chunkSize < contextRemaining) + { + return chunkSize; + } + return std::max(0, contextRemaining - reusable); +} + MicroBatchScheduler::MicroBatchScheduler(std::optional ctxChunkConfig, std::optional maxContextLength, LlmRequestState noScheduleUntilState, LlmRequestState noScheduleAfterState) @@ -38,16 +57,15 @@ void MicroBatchScheduler::fitDraftTokens(RequestVector& contextsToBeChunked, std::optional ctxTokensCapacity, SizeType32 const chunkUnitSize, std::optional const& maxContextLength) { - // How many compute tokens (chunk - reusable) are in this batch already? + // How many compute tokens are in this batch already? SizeType32 numCtxTokens{0}; for (auto const& llmReq : contextsToBeChunked) { 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)); + numCtxTokens += reuse_adjusted_compute(chunkSize, reusable, contextRemaining); } // Discard draft tokens that won't fit into the existing chunk unit, max @@ -125,14 +143,13 @@ void MicroBatchScheduler::setCtxRequestsChunkSizegetContextChunkSize(); SizeType32 actualIncrement = actualChunkSize - pastChunkSize; - // 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)); + // Compute-aware budget accounting for setPrepopulatedPromptLen's + // chunk-shift behaviour (non-last chunks keep their full size). + SizeType32 const contextRemaining = llmReq->getContextRemainingLength(); + SizeType32 const reusable + = llmReq->isFirstContextChunk() ? std::min(llmReq->getEstimatedReusableTokens(), contextRemaining) : 0; + SizeType32 const pastCompute = reuse_adjusted_compute(pastChunkSize, reusable, contextRemaining); + SizeType32 const actualCompute = reuse_adjusted_compute(actualChunkSize, reusable, contextRemaining); SizeType32 const computeIncrement = actualCompute - pastCompute; if ((ctxTokensCapacity && numCtxTokens + computeIncrement > ctxTokensCapacity.value()) @@ -181,24 +198,21 @@ 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 const computeCost = reuse_adjusted_compute(suggestedChunkSize, reusable, suggestedChunkSize); SizeType32 actualChunkSize = suggestedChunkSize; if (ctxTokensCapacity && computeCost > ctxTokensCapacity.value()) { - // 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) { - // maxContextLength limits compute tokens, not total tokens. - SizeType32 const actualCompute = std::max(0, actualChunkSize - reusable); + SizeType32 const actualCompute = reuse_adjusted_compute(actualChunkSize, reusable, suggestedChunkSize); if (actualCompute > maxContextLength.value()) { - actualChunkSize = std::min(reusable + maxContextLength.value(), suggestedChunkSize); + actualChunkSize = maxContextLength.value(); + actualChunkSize = std::min(actualChunkSize, suggestedChunkSize); } } if (actualChunkSize != suggestedChunkSize) @@ -208,10 +222,7 @@ void MicroBatchScheduler::setCtxRequestsChunkSizesetContextChunkSize(actualChunkSize); if (ctxTokensCapacity) { - // 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)); + SizeType32 const modelCost = reuse_adjusted_compute(actualChunkSize, reusable, suggestedChunkSize); ctxTokensCapacity = ctxTokensCapacity.value() - modelCost; } } @@ -349,10 +360,12 @@ std::tuple MicroBatchScheduler::operator()(Request if (!mCtxChunkConfig) // skip chunking { constexpr SizeType32 beam{0}; - reqNumTokens - = llmReq->getNumTokens(beam) + (llmReq->hasDraftTokens() ? llmReq->getNumDraftTokens() : 0); - // Compute tokens = total - reusable (at least 1 to make progress) - SizeType32 const computeTokens = std::max(1, reqNumTokens - reusable); + SizeType32 const contextTokens = llmReq->getNumTokens(beam); + SizeType32 const draftTokens = llmReq->hasDraftTokens() ? llmReq->getNumDraftTokens() : 0; + reqNumTokens = contextTokens + draftTokens; + SizeType32 const contextCompute + = reuse_adjusted_compute(contextTokens, reusable, llmReq->getContextRemainingLength()); + SizeType32 const computeTokens = std::max(1, contextCompute + draftTokens); TLLM_CHECK_WITH_INFO(!mMaxContextLength || computeTokens <= mMaxContextLength.value(), "Context compute tokens (%d) exceeds the limit value (%d)", computeTokens, mMaxContextLength.value()); @@ -369,9 +382,8 @@ std::tuple MicroBatchScheduler::operator()(Request llmReq->setContextChunkSize(llmReq->getContextRemainingLength()); auto const draftTokens = (llmReq->isLastContextChunk() && llmReq->hasDraftTokens()) ? llmReq->getNumDraftTokens() : 0; - // 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 const contextCompute = reuse_adjusted_compute( + llmReq->getContextChunkSize(), reusable, llmReq->getContextRemainingLength()); SizeType32 computeTokens = contextCompute + draftTokens; if (mMaxContextLength) @@ -444,10 +456,9 @@ std::tuple MicroBatchScheduler::operator()(Request if (llmReq->getContextChunkSize() > 0) { contextRequests.emplace_back(llmReq); - // 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() - reusable); + SizeType32 const computeTokens + = reuse_adjusted_compute(llmReq->getContextChunkSize(), reusable, llmReq->getContextRemainingLength()); 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/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp index fd64a9f67e91..be2637c02203 100644 --- a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp @@ -998,6 +998,58 @@ TEST_F(MicroBatchSchedulerTest, ReusableTokensWithChunkedContextEqualProgress) EXPECT_EQ(req1->getContextChunkSize(), 11) << "req1: reusable(10) + 1 compute token = chunk 11"; } +TEST_F(MicroBatchSchedulerTest, ReusableTokensChunkShiftNonLastChunk) +{ + // Test that reuse_adjusted_compute returns chunkSize (not chunkSize - reusable) + // for non-last chunks where reusable + chunkSize < contextRemaining. + // + // setPrepopulatedPromptLen shifts the chunk window right by the reused amount + // rather than shrinking it, so non-last chunks still process ~chunkSize tokens. + // + // Setup: 1 request, promptLen=100, reusable=30, FCFS, chunkUnit=1. + // Budget = 25 (maxNumTokens). + // + // Old (wrong) formula: compute = max(0, chunkSize - reusable) + // → chunk=100, compute = max(0, 100-30) = 70 > 25 → chunked to 25, + // compute = max(0, 25-30) = 0 → budget barely touched. + // + // New (correct) formula: reuse_adjusted_compute(chunkSize=25, reusable=30, remaining=100) + // → reusable(30) + chunkSize(25) = 55 < remaining(100) → non-last chunk + // → compute = chunkSize = 25 = budget → correct accounting. + // + // With 2 requests: budget=25 should only fit one non-last chunk of 25 tokens, + // not two (which the old formula would allow by underestimating compute to 0). + constexpr SizeType32 maxNumTokens = 25; + constexpr SizeType32 maxBatchSize = 4; + constexpr SizeType32 chunkUnitSize = 1; + constexpr SizeType32 reusableTokens = 30; + constexpr SizeType32 promptLen = 100; + constexpr SizeType32 maxNewTokens = 5; + constexpr ContextChunkingPolicy ctxChunkPolicy{ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED}; + + 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); + + // req0: chunk=25, reuse_adjusted_compute(25, 30, 100) = 25 (non-last: 30+25<100) + // Budget fully consumed → req1 gets chunk=0. + EXPECT_EQ(ctx.size(), 1u) << "Only req0 fits; non-last chunk costs full chunkSize"; + EXPECT_EQ(req0->getContextChunkSize(), 25) << "req0: chunk=25 (budget fully consumed by non-last chunk compute)"; + EXPECT_EQ(req1->getContextChunkSize(), 0) << "req1: no budget remaining"; +} + TEST_F(MicroBatchSchedulerTest, ReusableTokensZeroHasNoEffect) { // Verify that zero reusable tokens (the default) produces identical scheduling diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 37e80f71f7a4..a14ea0ea198b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2859,6 +2859,49 @@ def _balance_adp_requests(self, context_requests: list[LlmRequest], balanced_context_requests = context_requests return balanced_context_requests + @staticmethod + def _compute_scheduled_tokens(context_requests, generation_requests): + """Compute the total number of scheduled tokens for batch waiting decisions. + + For context requests, we estimate the actual compute tokens for this + iteration (excluding tokens served from KV cache). + + For generation requests, each contributes 1 + num_draft_tokens. + + Note on reusable token handling: + estimated_reusable_tokens is an absolute count from position 0. + Depending on the scheduler, context_current_position may or may not + have been advanced past the reusable prefix by the time this method + is called: + - V1 scheduler: prepare_context runs after scheduling, so + context_current_position is still 0. + - V2 scheduler: prepare_context runs during scheduling, so + context_current_position is already advanced to the reused offset. + To handle both correctly, the reusable credit applied to the current + chunk is max(0, reusable - context_current_position), i.e. only the + portion of the reusable range that falls within this chunk's span. + """ + num_scheduled_ctx_tokens = 0 + for ctx_req in context_requests: + reusable = (ctx_req.estimated_reusable_tokens + if ctx_req.is_first_context_chunk else 0) + # Credit only the reusable tokens that overlap with the current + # chunk: if context_current_position has already been advanced past + # the reusable prefix (V2), the credit is 0; if not (V1), the full + # reusable count is subtracted. + reusable_in_chunk = max(0, + reusable - ctx_req.context_current_position) + remaining = ctx_req.context_remaining_length + if (reusable_in_chunk > 0 and + reusable_in_chunk + ctx_req.context_chunk_size < remaining): + compute = ctx_req.context_chunk_size + else: + compute = max(1, ctx_req.context_chunk_size - reusable_in_chunk) + num_scheduled_ctx_tokens += compute + num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens + for gen_req in generation_requests) + return num_scheduled_ctx_tokens + num_scheduled_gen_tokens + def _waiting_requests(self, context_requests: list[LlmRequest], generation_requests: list[LlmRequest]): """ @@ -2868,11 +2911,8 @@ def _waiting_requests(self, context_requests: list[LlmRequest], - The number of waiting iterations is smaller than `self.batch_wait_timeout_iters`. """ - num_scheduled_ctx_tokens = sum( - len(ctx_req.get_tokens(0)) for ctx_req in context_requests) - num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens - for gen_req in generation_requests) - num_scheduled_tokens = num_scheduled_ctx_tokens + num_scheduled_gen_tokens + num_scheduled_tokens = self._compute_scheduled_tokens( + context_requests, generation_requests) should_waiting = self.batch_wait_iters_count < self.batch_wait_timeout_iters and num_scheduled_tokens < self.batch_wait_max_tokens_ratio * self.max_num_tokens if should_waiting: diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 362d904f0f0d..25ef3f9afde2 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -17,6 +17,7 @@ SHUTDOWN_REQUEST_ID, RequestQueueItem, ) +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.scheduler import FCFSWaitingQueue @@ -245,3 +246,140 @@ def test_partial_reuse_disabled_by_pp(self): """PP > 1 disables partial reuse path, falling back to normal logic.""" req = _make_request(complete_state=True, transmission_state=False) assert _classify_termination(req, True, False, 2) == "stats_only" + + +# --------------------------------------------------------------------------- +# Tests for _compute_scheduled_tokens with KV cache reuse chunk-shift logic +# --------------------------------------------------------------------------- + + +def _make_ctx_request( + context_chunk_size, + context_remaining_length, + estimated_reusable_tokens=0, + is_first_context_chunk=True, + context_current_position=0, +): + """Helper to create a mock context request for token computation tests.""" + req = Mock() + req.context_chunk_size = context_chunk_size + req.context_remaining_length = context_remaining_length + req.estimated_reusable_tokens = estimated_reusable_tokens + req.is_first_context_chunk = is_first_context_chunk + req.context_current_position = context_current_position + return req + + +def _make_gen_request(num_draft_tokens=0): + """Helper to create a mock generation request.""" + req = Mock() + req.num_draft_tokens = num_draft_tokens + return req + + +class TestComputeScheduledTokens: + """Tests for PyExecutor._compute_scheduled_tokens. + + Validates the chunk-shift aware token accounting: setPrepopulatedPromptLen + shifts the chunk window right by the reused amount rather than shrinking it. + Non-last chunks cost chunkSize; only last chunks cost remaining - reusable. + """ + + def test_no_reuse(self): + """Without reuse, compute = chunk_size.""" + ctx = [_make_ctx_request(context_chunk_size=100, context_remaining_length=100)] + assert PyExecutor._compute_scheduled_tokens(ctx, []) == 100 + + def test_last_chunk_with_reuse(self): + """Last chunk (reusable + chunk >= remaining): compute = chunk - reusable.""" + # promptLen=100, reusable=60, chunk=100 (full context) + # 60 + 100 >= 100 → last chunk → compute = max(1, 100 - 60) = 40 + ctx = [ + _make_ctx_request( + context_chunk_size=100, context_remaining_length=100, estimated_reusable_tokens=60 + ) + ] + assert PyExecutor._compute_scheduled_tokens(ctx, []) == 40 + + def test_non_last_chunk_with_reuse(self): + """Non-last chunk (reusable + chunk < remaining): compute = chunk_size. + + This is the core chunk-shift scenario. The old formula would compute + max(0, 25 - 30) = 0, but the correct cost is 25 because the chunk + window shifts right rather than shrinking. + """ + # promptLen=100, reusable=30, chunk=25 + # 30 + 25 = 55 < 100 → non-last chunk → compute = 25 + ctx = [ + _make_ctx_request( + context_chunk_size=25, context_remaining_length=100, estimated_reusable_tokens=30 + ) + ] + assert PyExecutor._compute_scheduled_tokens(ctx, []) == 25 + + def test_non_first_chunk_ignores_reuse(self): + """Reusable tokens only apply to the first context chunk.""" + ctx = [ + _make_ctx_request( + context_chunk_size=50, + context_remaining_length=50, + estimated_reusable_tokens=30, + is_first_context_chunk=False, + ) + ] + assert PyExecutor._compute_scheduled_tokens(ctx, []) == 50 + + def test_v2_scheduler_position_advanced(self): + """V2 scheduler: context_current_position already advanced past reuse. + + reusable_in_chunk = max(0, 30 - 30) = 0 → no credit → compute = chunk. + """ + ctx = [ + _make_ctx_request( + context_chunk_size=50, + context_remaining_length=70, + estimated_reusable_tokens=30, + context_current_position=30, + ) + ] + assert PyExecutor._compute_scheduled_tokens(ctx, []) == 50 + + def test_min_compute_is_one(self): + """Compute cost is floored at 1 even when reusable >= chunk_size.""" + # chunk=10, remaining=10, reusable=15 → last chunk → max(1, 10-15) = 1 + ctx = [ + _make_ctx_request( + context_chunk_size=10, context_remaining_length=10, estimated_reusable_tokens=15 + ) + ] + assert PyExecutor._compute_scheduled_tokens(ctx, []) == 1 + + def test_generation_tokens(self): + """Generation requests contribute 1 + num_draft_tokens each.""" + gen = [_make_gen_request(3), _make_gen_request(0)] + assert PyExecutor._compute_scheduled_tokens([], gen) == (1 + 3) + (1 + 0) + + def test_mixed_context_and_generation(self): + """Combined context (with chunk-shift) and generation tokens.""" + # Non-last chunk: compute = 25 + ctx = [ + _make_ctx_request( + context_chunk_size=25, context_remaining_length=100, estimated_reusable_tokens=30 + ) + ] + gen = [_make_gen_request(2)] + # 25 ctx + (1 + 2) gen = 28 + assert PyExecutor._compute_scheduled_tokens(ctx, gen) == 28 + + def test_multiple_ctx_requests_mixed_chunks(self): + """Multiple context requests: one non-last chunk, one last chunk.""" + # req0: non-last chunk → compute = 20 + req0 = _make_ctx_request( + context_chunk_size=20, context_remaining_length=100, estimated_reusable_tokens=30 + ) + # req1: last chunk (reuse=10, chunk=50, remaining=50) → 10+50>=50 + # → compute = max(1, 50-10) = 40 + req1 = _make_ctx_request( + context_chunk_size=50, context_remaining_length=50, estimated_reusable_tokens=10 + ) + assert PyExecutor._compute_scheduled_tokens([req0, req1], []) == 20 + 40 From 2f1840300f214ef8feae2ac044fae33e58b76969 Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Tue, 14 Apr 2026 08:04:59 -0700 Subject: [PATCH 2/2] [None][fix] Fix compute token accounting for KV cache reuse with context chunking Introduce _reuse_adjusted_compute() to correctly calculate forward-pass cost when setPrepopulatedPromptLen shifts the chunk window. Non-last chunks cost chunkSize; last chunks cost contextRemaining - reusable. Update C++ and Python tests to match the corrected behavior. Signed-off-by: Lanyu Liao --- .../batch_manager/microBatchSchedulerTest.cpp | 34 +++++----- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 +- .../_torch/pyexecutor/scheduler/scheduler.py | 66 +++++++++++++------ .../_torch/executor/test_py_scheduler.py | 50 +++++++------- 4 files changed, 91 insertions(+), 66 deletions(-) diff --git a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp index be2637c02203..714efdb680cd 100644 --- a/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp @@ -953,21 +953,24 @@ 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. + // Setup: 2 requests, promptLen=15, reusable=10, compute budget=12, 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. + // setPrepopulatedPromptLen shifts the chunk window right by the reused amount. + // For non-last chunks (reusable + chunkSize < contextRemaining), cost = chunkSize. + // For last chunks (reusable + chunkSize >= contextRemaining), cost = contextRemaining - reusable. // - // 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). + // With reusable=10, contextRemaining=15: + // chunkSize 1-4: non-last (10+4=14 < 15), cost = chunkSize. + // chunkSize >= 5: last (10+5=15 >= 15), cost = max(0, 15-10) = 5. // - // 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; + // EQUAL_PROGRESS grows chunks in lock-step: + // Iters 1-4: both grow to chunk=4, total compute = 2*4 = 8. + // Iter 5: both reach chunk=5 (last-chunk threshold), cost=5, increment=1 each. total=10. + // Iters 6+: cost stays 5, compute increment=0 — chunks grow for free to full context. + // Result: both complete full context (chunk=15), total compute = 10 < 12. + // + // Without reusable tokens, budget=12 would only allow both to reach chunk=6 (total=12). + constexpr SizeType32 maxNumTokens = 12; constexpr SizeType32 maxBatchSize = 4; constexpr SizeType32 chunkUnitSize = 1; constexpr SizeType32 reusableTokens = 10; @@ -993,9 +996,10 @@ TEST_F(MicroBatchSchedulerTest, ReusableTokensWithChunkedContextEqualProgress) 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"; + // Both complete full context: once past the last-chunk threshold (chunkSize=5), + // additional tokens cost 0 compute, so chunks grow to full promptLen. + EXPECT_EQ(req0->getContextChunkSize(), 15) << "req0: full context (last-chunk cost capped at 5)"; + EXPECT_EQ(req1->getContextChunkSize(), 15) << "req1: full context (last-chunk cost capped at 5)"; } TEST_F(MicroBatchSchedulerTest, ReusableTokensChunkShiftNonLastChunk) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a14ea0ea198b..f5f8b63eef10 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2892,11 +2892,12 @@ def _compute_scheduled_tokens(context_requests, generation_requests): reusable_in_chunk = max(0, reusable - ctx_req.context_current_position) remaining = ctx_req.context_remaining_length - if (reusable_in_chunk > 0 and - reusable_in_chunk + ctx_req.context_chunk_size < remaining): + if reusable_in_chunk <= 0: + compute = ctx_req.context_chunk_size + elif reusable_in_chunk + ctx_req.context_chunk_size < remaining: compute = ctx_req.context_chunk_size else: - compute = max(1, ctx_req.context_chunk_size - reusable_in_chunk) + compute = max(1, remaining - reusable_in_chunk) num_scheduled_ctx_tokens += compute num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens for gen_req in generation_requests) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 4f2d56c657be..c99a1d205360 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -318,6 +318,21 @@ class MicroBatchScheduler: """Base class to match structure.""" +def _reuse_adjusted_compute(chunk_size: int, reusable: int, context_remaining: int) -> int: + """Return the forward-pass token cost for a context chunk with KV cache reuse. + + setPrepopulatedPromptLen shifts the chunk window right by the prepopulated + amount rather than shrinking it. For non-last chunks the model still + processes approximately *chunk_size* tokens; only for the last chunk is the + cost *context_remaining - reusable*. + """ + if reusable <= 0: + return chunk_size + if reusable + chunk_size < context_remaining: + return chunk_size + return max(0, context_remaining - reusable) + + class PyMicroBatchScheduler(MicroBatchScheduler): def __init__( self, @@ -423,7 +438,10 @@ def schedule( draft_tokens = req.num_draft_tokens if req.has_draft_tokens else 0 req_num_tokens = base_tokens + draft_tokens - compute_tokens = max(1, req_num_tokens - reusable) + context_compute = _reuse_adjusted_compute( + base_tokens, reusable, req.context_remaining_length + ) + compute_tokens = max(1, context_compute + draft_tokens) assert max_context_length is None or compute_tokens <= max_context_length, ( f"Context compute tokens ({compute_tokens}) exceeds limit ({max_context_length})" @@ -451,7 +469,9 @@ def schedule( ) # Compute cost: context compute + draft tokens # (reusable tokens only offset context tokens, not draft tokens) - context_compute = max(0, req.context_chunk_size - reusable) + context_compute = _reuse_adjusted_compute( + req.context_chunk_size, reusable, req.context_remaining_length + ) compute_tokens = context_compute + draft_tokens if max_context_length is not None: @@ -522,7 +542,9 @@ 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 - reusable) + compute_tokens = _reuse_adjusted_compute( + req.context_chunk_size, reusable, req.context_remaining_length + ) batch_num_tokens += compute_tokens logger.debug( f"context request scheduled: ID {req.request_id}, " @@ -644,10 +666,15 @@ def _chunk_equal_progress(self, requests: RequestList, capacity: Optional[int], # 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-aware budget accounting for setPrepopulatedPromptLen's + # chunk-shift behaviour (non-last chunks keep their full size). + context_remaining = req.context_remaining_length + reusable = min( + req.estimated_reusable_tokens if req.is_first_context_chunk else 0, + context_remaining, + ) + past_compute = _reuse_adjusted_compute(past_size, reusable, context_remaining) + actual_compute = _reuse_adjusted_compute(actual_size, reusable, context_remaining) compute_increment = actual_compute - past_compute # Check Constraints — mirrors the if-block guard in C++ before numCtxTokens += @@ -690,7 +717,7 @@ def _chunk_fcfs( req.estimated_reusable_tokens if req.is_first_context_chunk else 0, suggested_size, ) - compute_cost = suggested_size - reusable + compute_cost = _reuse_adjusted_compute(suggested_size, reusable, suggested_size) # Start with full context as the allocation target. actual_size = suggested_size @@ -704,9 +731,9 @@ def _chunk_fcfs( # 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) + actual_compute = _reuse_adjusted_compute(actual_size, reusable, suggested_size) if actual_compute > self.max_context_length: - actual_size = reusable + self.max_context_length + actual_size = self.max_context_length actual_size = min(actual_size, suggested_size) # Align down to unit_size when either constraint trimmed the chunk, to avoid @@ -716,9 +743,9 @@ def _chunk_fcfs( req.context_chunk_size = int(actual_size) - # 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)) + actual_model_cost = _reuse_adjusted_compute( + req.context_chunk_size, reusable, suggested_size + ) if capacity is not None: current_compute_capacity -= actual_model_cost @@ -749,16 +776,13 @@ def _fit_draft_tokens(self, requests: RequestList, capacity: Optional[int], unit # min(chunk_size, P - reusable), where P = context_remaining_length # (the full prompt length on the first context chunk). num_ctx_tokens = sum( - min( + _reuse_adjusted_compute( 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, - ), + min( + req.estimated_reusable_tokens if req.is_first_context_chunk else 0, + req.context_remaining_length, ), + req.context_remaining_length, ) for req in requests ) diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index d1f78827e1fd..e11abac0c9f2 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -1019,46 +1019,42 @@ def test_reusable_tokens_equal_progress(self): 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. + individual chunk sizes are capped at max_num_tokens. - 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. + setPrepopulatedPromptLen shifts the chunk window right by the reused amount. + For non-last chunks (reusable + chunkSize < contextRemaining), cost = chunkSize. + For last chunks (reusable + chunkSize >= contextRemaining), + cost = contextRemaining - reusable. - Reusable tokens 0-4 are "free" (compute_increment = 0 until chunk > 5). - Budget is only consumed for tokens beyond the reusable prefix. + Setup: 2 requests, prompt_len=15, max_num_tokens=15, chunk_unit=1. + req0: reusable=10 → last-chunk threshold at chunk=5 (10+5=15). + req1: reusable=0 → all tokens cost budget. + Tentative compute: 5 + 15 = 20 > 15 → chunking exercised. - 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. + EQUAL_PROGRESS trace: + iter 1-4: both non-last, cost=chunkSize each. total=8. + iter 5: req0 crosses threshold (last-chunk, cost=5, increment=1). + req1 still non-last (cost=5, increment=1). total=10. + iter 6-10: req0 increment=0 (free growth), req1 increment=1. total=15. + Result: req0=10, req1=10, total compute=15=budget. - 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) + Without reusable tokens, budget=15 yields req0=8, req1=7. """ config = ContextChunkingConfig(ChunkingPolicy.EQUAL_PROGRESS, chunk_unit_size=1) scheduler = PyMicroBatchScheduler( - max_batch_size=4, max_num_tokens=8, ctx_chunk_config=config + max_batch_size=4, max_num_tokens=15, 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 + req0.estimated_reusable_tokens = 10 + req1.estimated_reusable_tokens = 0 - ctx, gen = scheduler.schedule([req0, req1, req2], set()) + ctx, gen = scheduler.schedule([req0, req1], 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" + assert len(ctx) == 2, "Both requests should be scheduled" + assert chunks[0] == 10, "req0: free growth past threshold, capped by budget" + assert chunks[1] == 10, "req1: benefits from req0's free growth" # ############################################################################