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 d9ea41926ab1..881e2c706f37 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,62 @@ 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) +{ + // 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) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6564be1ba449..5a99bd20189e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2859,6 +2859,50 @@ 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: + 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, 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) + return num_scheduled_ctx_tokens + num_scheduled_gen_tokens + def _waiting_requests(self, context_requests: list[LlmRequest], generation_requests: list[LlmRequest]): """ @@ -2868,11 +2912,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/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_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 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" # ############################################################################