From f0a16e89b48d59fcc80f37702ea44b3929736c72 Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Wed, 29 Apr 2026 00:20:34 -0700 Subject: [PATCH 1/3] [None][fix] Add Python-side remaining_budget guard for KV reuse budget overflow Cherry-pick of Python portions from feat/bench_y PRs #12682 and #12806 that were not included in #12976 (which ported only the C++ fix to main). Adds a remaining_budget re-validation guard in KVCacheManager.prepare_resources() that re-probes the radix tree for actual reusable blocks after KV cache allocation and skips requests whose forward cost exceeds the remaining budget. This catches the estimation-vs-reality gap when cache eviction between scheduling and prepare_resources() reduces actual reuse below the scheduler's estimate. Original authors: Liao Lanyu (@lancelly), Jin Li (@liji-nv) Signed-off-by: Yuewei Na --- .../_torch/pyexecutor/model_engine.py | 18 +++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 77 ++++++++++++++++++- .../_torch/pyexecutor/resource_manager.py | 69 ++++++++++++++++- 3 files changed, 161 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index dd6c7b99b518..af0b8f5ce63c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2669,6 +2669,24 @@ def previous_seq_slots_device(): num_tokens = len(input_ids) num_draft_tokens = len(draft_tokens) total_num_tokens = len(position_ids) + if total_num_tokens > self.max_num_tokens: + ctx_details = [] + for r in scheduled_requests.context_requests: + pos = r.context_current_position + csz = r.context_chunk_size + full = len(r.get_tokens(0)) + tokens = min(csz, max(0, full - pos)) + ctx_details.append( + f"rid={r.py_request_id} pos={pos} chunk={csz} " + f"full={full} tokens={tokens}") + gen_count = len(scheduled_requests.generation_requests) + from tensorrt_llm.logger import logger as _mnt_logger + _mnt_logger.error( + f"MNT overflow: total={total_num_tokens} " + f"max={self.max_num_tokens} " + f"ctx_reqs={len(scheduled_requests.context_requests)} " + f"gen_reqs={gen_count} " + f"ctx_breakdown=[{'; '.join(ctx_details)}]") assert total_num_tokens <= self.max_num_tokens, ( f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self.max_num_tokens})" ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b8b09bca9c9e..ff2b6f86f122 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3049,6 +3049,73 @@ def _compute_scheduled_tokens(context_requests, generation_requests): for gen_req in generation_requests) return num_scheduled_ctx_tokens + num_scheduled_gen_tokens + def _maybe_log_batch_wait_decision( + self, + context_requests: list[LlmRequest], + generation_requests: list[LlmRequest], + num_scheduled_tokens: int, + wait_threshold: float, + should_waiting: bool, + ) -> None: + """Diagnostics for batch_wait: set TLLM_LOG_BATCH_WAIT=1 (rank 0 only).""" + if self.dist.rank != 0: + return + + num_scheduled_gen_tokens = sum(1 + gen_req.num_draft_tokens + for gen_req in generation_requests) + num_scheduled_ctx_formula = num_scheduled_tokens - num_scheduled_gen_tokens + + chunk_ctx_sum = 0 + ctx_summaries: List[str] = [] + max_detail = 4 + for i, ctx_req in enumerate(context_requests): + full_len = len(ctx_req.get_tokens(0)) + begin = ctx_req.context_current_position + chunk_sz = ctx_req.context_chunk_size + this_chunk = min(chunk_sz, max(0, full_len - begin)) + chunk_ctx_sum += this_chunk + reusable = (ctx_req.estimated_reusable_tokens + if ctx_req.is_first_context_chunk else 0) + reusable_in_chunk = max(0, reusable - begin) + remaining = ctx_req.context_remaining_length + if (reusable_in_chunk > 0 + and reusable_in_chunk + chunk_sz < remaining): + formula_contrib = chunk_sz + else: + formula_contrib = max(1, chunk_sz - reusable_in_chunk) + if i < max_detail: + ctx_summaries.append( + f"rid={ctx_req.py_request_id} full={full_len} pos={begin} " + f"chunk_sz={chunk_sz} this_chunk={this_chunk} " + f"reusable={reusable} formula_contrib={formula_contrib}") + n_ctx = len(context_requests) + if n_ctx > max_detail: + ctx_summaries.append(f"... +{n_ctx - max_detail} more ctx req(s)") + + logger.info( + "batch_wait: formula_total=", + num_scheduled_tokens, + " formula_ctx=", + num_scheduled_ctx_formula, + " formula_gen=", + num_scheduled_gen_tokens, + " chunk_ctx_sum=", + chunk_ctx_sum, + " threshold=", + wait_threshold, + " wait_iter=", + self.batch_wait_iters_count, + "/", + self.batch_wait_timeout_iters, + " should_defer_ctx=", + should_waiting, + " num_gen=", + len(generation_requests), + " ctx_detail=[", + "; ".join(ctx_summaries), + "]", + ) + def _waiting_requests(self, context_requests: list[LlmRequest], generation_requests: list[LlmRequest]): """ @@ -3060,8 +3127,14 @@ def _waiting_requests(self, context_requests: list[LlmRequest], 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 + wait_threshold = (self.batch_wait_max_tokens_ratio * + self.max_num_tokens) + + should_waiting = self.batch_wait_iters_count < self.batch_wait_timeout_iters and num_scheduled_tokens < wait_threshold + self._maybe_log_batch_wait_decision(context_requests, + generation_requests, + num_scheduled_tokens, + wait_threshold, should_waiting) if should_waiting: self.batch_wait_iters_count += 1 return [] diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index d7b744c059df..fb90919f33ce 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -661,6 +661,30 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # wait for all pending work to finish before launching offload/onboarding/partial copy self.impl.sync_transfer_manager_with_buffer_manager() + # Pre-addSequence budget re-validation. The C++ scheduler + # should already account for the chunk-shift cost, but under + # heavy KV-cache eviction the actual reuse may be lower than + # estimated. We re-probe the radix tree and estimate the + # true forward cost; if it exceeds the remaining budget the + # request is skipped (re-scheduled next iteration). + remaining_budget = None + if self.enable_block_reuse and not self.is_draft: + gen_tokens = sum( + req.get_beam_width_by_iter(for_next_iteration=False) + + get_draft_token_length(req) + for req in scheduled_batch.generation_requests) + remaining_budget = self.max_num_tokens - gen_tokens + + # Pre-subtract the fixed cost of non-first-chunk context + # requests. These have no reuse to re-validate and their + # compute cost is committed, so first-chunk budget checks + # must see the budget with these costs already removed. + for req in scheduled_batch.context_requests: + if not req.is_first_context_chunk: + remaining_budget -= req.context_chunk_size + + accepted_ctx_requests = [] + # Collect first-chunk requests eligible for add_sequence_batch. # When block reuse is enabled, addSequenceBatch uses a two-phase # claim-then-onboard strategy that prevents host offloading from @@ -689,11 +713,37 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): else: if req.is_first_context_chunk and self._kv_connector_should_add_sequence( req): + if remaining_budget is not None: + unique_tokens = req.get_unique_tokens(0) + reusable_blocks = self.impl.count_reusable_blocks( + unique_tokens, req, False) + actual_reuse = (reusable_blocks * + self.tokens_per_block) + req_compute = self._estimate_post_reuse_compute( + actual_reuse, req.context_chunk_size, + req.prompt_len) + if req_compute > remaining_budget: + logger.warning( + f"Reuse budget: skip req " + f"{req.py_request_id} " + f"(compute={req_compute}, " + f"chunk={req.context_chunk_size}, " + f"reuse={actual_reuse}, " + f"remaining={remaining_budget})") + continue + remaining_budget -= req_compute + # Batch path: two-phase claim-then-onboard batch_request_infos.append( (req.py_request_id, req.prompt_len, req_beam_width)) batch_llm_requests.append(req) batch_ctx_requests.append(req) + elif remaining_budget is not None and req.is_first_context_chunk: + reusable = req.estimated_reusable_tokens + remaining_budget -= self._estimate_post_reuse_compute( + reusable, req.context_chunk_size, req.prompt_len) + + accepted_ctx_requests.append(req) if batch_request_infos: self.impl.add_sequence_batch(batch_request_infos, @@ -711,7 +761,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # A request may change from `context_requests_chunking` to `context_requests_last_chunk` in # `add_sequence_batch` due to KV cache reuse, so we rebuild the context request lists here. - scheduled_batch.reset_context_requests() + scheduled_batch.reset_context_requests(accepted_ctx_requests) for req in scheduled_batch.generation_requests: if self.mapping.has_cp_helix(): @@ -737,6 +787,23 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self.kv_connector_manager.build_scheduler_output( scheduled_batch, self) + def _estimate_post_reuse_compute(self, reuse_tokens: int, chunk_size: int, + prompt_len: int) -> int: + """Estimate forward compute tokens after setPrepopulatedPromptLen. + + For non-last chunks the chunk window shifts right by the reused + amount and the forward cost is approximately chunk_size. For + last chunks the cost is prompt_len - reuse (original formula). + """ + P = reuse_tokens + if P > 0 and P < prompt_len: + if P + chunk_size < prompt_len: + aligned_end = ((P + chunk_size) // self.tokens_per_block * + self.tokens_per_block) + return max(1, aligned_end - P) + return max(1, prompt_len - P) + return chunk_size + def extend_capacity_for_tokens(self, request: LlmRequest) -> None: """No-op for V1; interface kept consistent with V2.""" From e06999eadb5301b745e81554a7548ff383f5adc3 Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Wed, 29 Apr 2026 16:39:10 -0700 Subject: [PATCH 2/3] [None][test] Add unit tests for KV reuse budget guard (DYN-2868) Adds three pure-Python test files (no GPU, no model weights) covering the budget guard introduced in the previous commit: - test_post_reuse_compute.py (Tier 1): Branch-coverage table for KVCacheManager._estimate_post_reuse_compute() plus parametrized "production-relevant divergence" cases proving the helper undercharges when chunk_size is not block-aligned (legal per test_kv_cache_v2_scheduler.py:243-247). Includes a precise property test asserting helper <= actual in the reuse branch (0 < P < prompt_len) and documents the short-circuit overcharge case. - test_resource_manager_v1_budget.py (Tier 2 + Tier 3): Mock-based integration tests driving the real prepare_resources() loop via KVCacheManager.__new__ to skip C++ init. Covers skip path, no-skip path, non-first-chunk pre-subtraction, gen+draft tokens consuming budget, no-op when reuse off / is_draft, reset_context_requests with is_last_context_chunk mutation, kv_connector elif branch, connector callbacks (update_state_after_alloc, build_scheduler_output), count_reusable_blocks call shape, and non-first-chunk fallthrough. Tier 3 contains the deterministic DYN-2868 regression pair: - test_m: post-prepare total <= max_num_tokens with the guard - test_n: negative control patches the helper to 0 and verifies overshoot via an INDEPENDENT helper (model_engine_total_tokens) that mirrors model_engine.py:2293-2297 + 2615-2616 token counting. - test_batch_wait_log.py (Tier 4): Smoke test for PyExecutor._maybe_log_batch_wait_decision rank-0 gate. Local verification: - Tier 1 arithmetic verified standalone (no TRT-LLM deps). - Tier 3 logic simulated independently (Test M = 512, Test N = 1024). - All three files pass py_compile. - Full pytest run requires tensorrt_llm.bindings (C++ build) and was not run locally; intended for CI. Signed-off-by: Yuewei Na --- .../_torch/executor/test_batch_wait_log.py | 40 ++ .../executor/test_post_reuse_compute.py | 155 +++++ .../test_resource_manager_v1_budget.py | 574 ++++++++++++++++++ 3 files changed, 769 insertions(+) create mode 100644 tests/unittest/_torch/executor/test_batch_wait_log.py create mode 100644 tests/unittest/_torch/executor/test_post_reuse_compute.py create mode 100644 tests/unittest/_torch/executor/test_resource_manager_v1_budget.py diff --git a/tests/unittest/_torch/executor/test_batch_wait_log.py b/tests/unittest/_torch/executor/test_batch_wait_log.py new file mode 100644 index 000000000000..d0352985a8e4 --- /dev/null +++ b/tests/unittest/_torch/executor/test_batch_wait_log.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Smoke test for PyExecutor._maybe_log_batch_wait_decision rank gating. + +Per ``py_executor.py:3060-3063`` the method early-returns when ``self.dist.rank +!= 0``. The docstring claims env-var gating via ``TLLM_LOG_BATCH_WAIT=1`` +which does NOT exist in the implementation; the only gate is rank. + +This test is a regression net — if a refactor removes the rank gate, the +function would crash at line 3107 (``self.batch_wait_iters_count``) since +we never set the wait counters on the fake executor. +""" +from unittest.mock import Mock + +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + +def test_log_decision_early_returns_on_nonzero_rank(): + """Non-rank-0 path should early-return without touching wait counters.""" + fake = PyExecutor.__new__(PyExecutor) + fake.dist = Mock() + fake.dist.rank = 1 # non-zero rank → early return at py_executor.py:3062 + + # If the gate is broken, the function would proceed to read + # self.batch_wait_iters_count (line 3107) which we never set → + # AttributeError. The successful no-crash path verifies the gate works. + PyExecutor._maybe_log_batch_wait_decision( + fake, + context_requests=[], + generation_requests=[], + num_scheduled_tokens=0, + wait_threshold=0.5, + should_waiting=False, + ) diff --git a/tests/unittest/_torch/executor/test_post_reuse_compute.py b/tests/unittest/_torch/executor/test_post_reuse_compute.py new file mode 100644 index 000000000000..e9d7c18a0cb4 --- /dev/null +++ b/tests/unittest/_torch/executor/test_post_reuse_compute.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Unit tests for KVCacheManager._estimate_post_reuse_compute() — DYN-2868. + +The helper is effectively pure: it depends only on the arguments and +``self.tokens_per_block``. Tests construct a stub holding tokens_per_block +and bind the unbound method, avoiding any real KVCacheManager initialization. + +Tier 1A: branch coverage of the helper. +Tier 1B: divergence vs ``_prepare_tp_inputs`` actual position-id count. +Tier 1C: precise property tests — undercharges in the reuse branch + (0 < P < prompt_len), short-circuits to chunk_size outside it. +""" +import pytest + +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager + + +class _Stub: + """Minimal stand-in for KVCacheManager exposing only tokens_per_block.""" + + def __init__(self, tokens_per_block): + self.tokens_per_block = tokens_per_block + + _estimate_post_reuse_compute = KVCacheManager._estimate_post_reuse_compute + + +def _model_engine_forward_count(begin_compute, chunk_size, prompt_len): + """Mirror model_engine.py:2293-2297 — len(prompt_tokens) per ctx req.""" + return max(0, min(chunk_size, prompt_len - begin_compute)) + + +# --------------------------------------------------------------------------- +# Tier 1A: branch coverage +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "P, chunk, prompt_len, block, expected, branch", + [ + # P <= 0 → return chunk_size + (0, 256, 1024, 64, 256, "no_reuse_zero"), + (-5, 256, 1024, 64, 256, "no_reuse_negative"), + # P >= prompt_len → return chunk_size + (1024, 256, 1024, 64, 256, "full_reuse_eq"), + (2048, 256, 1024, 64, 256, "full_reuse_gt"), + # last-chunk branch (P + chunk >= prompt_len): max(1, prompt_len - P) + (128, 1024, 200, 64, 72, "last_chunk_partial"), + (199, 4, 200, 64, 1, "last_chunk_just_fits"), + # non-last chunk: aligned_end = floor((P+chunk)/block)*block, max(1, aligned_end - P) + (64, 128, 1024, 64, 128, "nonlast_block_aligned"), + (70, 130, 1024, 64, 122, "nonlast_off_block"), + (100, 4, 1024, 64, 1, "nonlast_max_clamp"), + (50, 100, 1024, 1, 100, "block_size_one"), + ], +) +def test_branch_coverage(P, chunk, prompt_len, block, expected, branch): + stub = _Stub(tokens_per_block=block) + actual = stub._estimate_post_reuse_compute(P, chunk, prompt_len) + assert actual == expected, ( + f"branch={branch}: expected {expected}, got {actual} " + f"(P={P}, chunk={chunk}, prompt_len={prompt_len}, block={block})" + ) + + +# --------------------------------------------------------------------------- +# Tier 1B: helper vs model_engine — production-relevant divergence +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "P, chunk_size, prompt_len, tokens_per_block, expect_undercharge", + [ + # Block-aligned chunk + block-aligned reuse: helper == model_engine + (64, 128, 1024, 64, False), + (0, 256, 1024, 64, False), + (128, 256, 200, 64, False), # last chunk, exact match (helper=72, actual=72) + (192, 64, 1024, 64, False), + # Block-aligned reuse + UNALIGNED chunk: helper undercharges. This + # case is legal in production per + # tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py:243-247 + # (chunk_unit_size=100 with tokens_per_block=64). + (64, 100, 1000, 64, True), # aligned_end=128, helper=64, actual=100 + (128, 50, 1024, 64, True), # aligned_end=128, helper=1, actual=50 + # Block-unaligned reuse: also undercharges + (70, 130, 1024, 64, True), # aligned_end=192, helper=122, actual=130 + ], +) +def test_helper_vs_model_engine_per_case(P, chunk_size, prompt_len, + tokens_per_block, + expect_undercharge): + """Document where the helper diverges from _prepare_tp_inputs token + counting. The aligned cases agree; the unaligned cases undercharge — + which means the budget guard is conservative-enough to catch run-14's + +2 overshoot but not a strict ≤ max guarantee under all legal scheduler + outputs. Surfaced as a caveat in the PR description.""" + stub = _Stub(tokens_per_block) + helper = stub._estimate_post_reuse_compute(P, chunk_size, prompt_len) + actual = _model_engine_forward_count(P, chunk_size, prompt_len) + if expect_undercharge: + assert helper < actual, ( + f"Expected undercharge for unaligned (P={P}, chunk={chunk_size}, " + f"block={tokens_per_block}): helper={helper}, actual={actual}. " + f"If this fails, the helper has been tightened — update this " + f"test and review whether Tier 3 needs new scenarios." + ) + else: + assert helper == actual, ( + f"Helper/model_engine mismatch under aligned inputs: " + f"helper={helper}, actual={actual}, " + f"(P={P}, chunk={chunk_size}, prompt_len={prompt_len}, " + f"block={tokens_per_block})" + ) + + +# --------------------------------------------------------------------------- +# Tier 1C: precise property tests +# --------------------------------------------------------------------------- +def test_helper_undercharges_in_reuse_branch(): + """When 0 < P < prompt_len, helper <= actual model_engine forward count. + This is the invariant the budget guard relies on for the run-14 + scenario (block-aligned reuse + block-aligned chunk_size). Counts as + the load-bearing property — a regression that breaks this is a real bug.""" + counterexamples = [] + for P in [1, 32, 63, 64, 65, 70, 128, 199]: + for chunk in [1, 50, 64, 100, 128, 256]: + for prompt_len in [128, 200, 512, 1024]: + if P >= prompt_len: + continue # outside the reuse branch + for block in [1, 32, 64]: + stub = _Stub(block) + h = stub._estimate_post_reuse_compute(P, chunk, prompt_len) + a = _model_engine_forward_count(P, chunk, prompt_len) + if h > a: + counterexamples.append( + (P, chunk, prompt_len, block, h, a)) + assert not counterexamples, ( + f"helper > actual in reuse branch (first 5): {counterexamples[:5]}" + ) + + +def test_helper_short_circuits_outside_reuse_range(): + """When P <= 0 or P >= prompt_len, helper returns chunk_size verbatim + regardless of prompt_len — so the helper may OVERCHARGE actual forward + tokens here. This is intentional: outside the reuse branch the guard + falls back to the no-credit path, mirroring the pre-fix behavior.""" + stub = _Stub(tokens_per_block=64) + # P=0: no reuse credit → return chunk_size even when chunk > prompt_len + assert stub._estimate_post_reuse_compute(0, 256, 128) == 256 + # P negative: same path + assert stub._estimate_post_reuse_compute(-5, 256, 128) == 256 + # P beyond prompt_len: same path + assert stub._estimate_post_reuse_compute(2048, 256, 128) == 256 diff --git a/tests/unittest/_torch/executor/test_resource_manager_v1_budget.py b/tests/unittest/_torch/executor/test_resource_manager_v1_budget.py new file mode 100644 index 000000000000..2fae36ae2425 --- /dev/null +++ b/tests/unittest/_torch/executor/test_resource_manager_v1_budget.py @@ -0,0 +1,574 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Integration tests for KVCacheManager.prepare_resources() budget guard +(DYN-2868). Drives the real ``prepare_resources`` method against a mocked +manager (constructed via ``__new__`` to skip C++ deps) and a real +``ScheduledRequests``. + +Tier 2: branch coverage of the budget guard (skip / no-skip / non-first + pre-subtraction / draft tokens / connector branches / no-op when + reuse off / no-op when is_draft / reset_context_requests filter / + count_reusable_blocks call shape / non-first fallthrough). +Tier 3: deterministic DYN-2868 regression pair (Test M passes with guard, + Test N — negative control — overshoots when guard disabled). +""" +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests + + +# --------------------------------------------------------------------------- +# Mock factories — avoid Mock(spec=LlmRequest) because LlmRequest is a C++ +# binding; we just set the attributes prepare_resources reads. +# --------------------------------------------------------------------------- +def _make_first_chunk_req(rid, prompt_len, chunk_size, est_reuse=0, + num_draft_tokens=0, unique_tokens=None, + beam_width=1): + req = Mock() + req.request_id = rid + req.py_request_id = rid + req.prompt_len = prompt_len + req.context_chunk_size = chunk_size + req.context_current_position = 0 + req.is_first_context_chunk = True + req.is_last_context_chunk = (chunk_size >= prompt_len) + req.estimated_reusable_tokens = est_reuse + req.num_draft_tokens = num_draft_tokens + req.py_draft_tokens = list(range(num_draft_tokens)) + req.get_unique_tokens.return_value = unique_tokens or [0] * prompt_len + # Tier 3 helper uses len(req.get_tokens(0)) per model_engine.py:2291. + req.get_tokens.return_value = list(range(prompt_len)) + req.sampling_config = Mock() + req.sampling_config.beam_width = beam_width + return req + + +def _make_non_first_chunk_req(rid, chunk_size, prompt_len=None, + current_pos=128): + req = Mock() + req.request_id = rid + req.py_request_id = rid + pl = prompt_len if prompt_len is not None else ( + chunk_size + current_pos + 256) + req.prompt_len = pl + req.context_chunk_size = chunk_size + req.context_current_position = current_pos + req.is_first_context_chunk = False + req.is_last_context_chunk = False + req.num_draft_tokens = 0 + req.py_draft_tokens = [] + req.get_tokens.return_value = list(range(pl)) + req.sampling_config = Mock() + req.sampling_config.beam_width = 1 + return req + + +def _make_gen_req(rid, num_draft_tokens=0, beam_width=1): + req = Mock() + req.request_id = rid + req.py_request_id = rid + req.num_draft_tokens = num_draft_tokens + req.py_draft_tokens = list(range(num_draft_tokens)) + req.get_beam_width_by_iter.return_value = beam_width + req.sampling_config = Mock() + req.sampling_config.beam_width = beam_width + return req + + +@pytest.fixture +def fake_kv_manager(): + """KVCacheManager-shaped object with all attrs/methods that + ``prepare_resources`` touches before our budget code runs. Skips + ``__init__`` to avoid C++ deps. Each test gets a fresh fixture — + never share across tests because ``prepare_resources`` mutates + ``scheduled_batch`` and reads request fields.""" + mgr = KVCacheManager.__new__(KVCacheManager) + mgr.is_draft = False + mgr.enable_block_reuse = True + mgr.max_num_tokens = 512 + mgr.tokens_per_block = 64 + mgr.num_extra_kv_tokens = 0 + mgr.mapping = Mock() + mgr.mapping.cp_config = {} + mgr.mapping.has_cp_helix.return_value = False + mgr.kv_connector_manager = None # default; tests override for connector cases + mgr.impl = Mock() + mgr.impl.count_reusable_blocks.return_value = 0 + mgr.impl.add_sequence_batch = Mock() + mgr.impl.add_token = Mock() + mgr.impl.sync_transfer_manager_with_buffer_manager = Mock() + mgr.impl.refresh_blocks = Mock() + # For connector tests; harmless when kv_connector_manager is None. + mgr.get_cache_indices = Mock(return_value=[]) + return mgr + + +@pytest.fixture +def make_batch(): + """Construct a fresh ScheduledRequests using the public append API. + NEVER assign ``batch.context_requests`` directly — it's a property + derived from ``context_requests_chunking + context_requests_last_chunk``.""" + + def _factory(ctx_reqs, gen_reqs): + batch = ScheduledRequests() + for r in ctx_reqs: + batch.append_context_request(r) + for r in gen_reqs: + batch.append_generation_request(r) + return batch + + return _factory + + +# =========================================================================== +# Tier 2: branch coverage +# =========================================================================== +class TestSkipPath: + """Test A — skip when actual reuse < estimated reuse.""" + + def test_a_skip_when_actual_reuse_less_than_estimated( + self, fake_kv_manager, make_batch + ): + # 3 ctx reqs; scheduler thought each had ~768-token reuse, actual is 0. + # remaining_budget starts at 512. Each req: helper(0, 256, 1024) = 256. + # req0: 256 ≤ 512 → admit, remaining 256 + # req1: 256 ≤ 256 → admit, remaining 0 + # req2: 256 > 0 → SKIP via continue + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, + est_reuse=768) + for i in range(3) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + surviving = [r.py_request_id for r in batch.context_requests] + assert surviving == [0, 1], ( + f"req2 should be skipped via continue; got {surviving}") + + # The two admitted reqs went into add_sequence_batch. + fake_kv_manager.impl.add_sequence_batch.assert_called_once() + infos = fake_kv_manager.impl.add_sequence_batch.call_args[0][0] + assert {info[0] for info in infos} == {0, 1} + + +class TestNoSkipPath: + """Test B — no skip when estimate matches actual.""" + + def test_b_no_skip_when_estimate_matches_actual( + self, fake_kv_manager, make_batch + ): + # 3 ctx reqs; each has est_reuse=896 (14 blocks at block=64), and + # actual count_reusable_blocks=14 → actual_reuse=896. + # helper(896, 128, 1024): P+chunk=1024 == prompt_len → last-chunk + # → max(1, 1024-896)=128 + # 3 * 128 = 384 ≤ 512 → all admitted + fake_kv_manager.impl.count_reusable_blocks.return_value = 14 + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=128, + est_reuse=896) + for i in range(3) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + assert {r.py_request_id for r in batch.context_requests} == {0, 1, 2} + fake_kv_manager.impl.add_sequence_batch.assert_called_once() + + +class TestNonFirstChunkPreSubtract: + """Test C — non-first-chunk reqs reduce budget for first-chunk reqs.""" + + def test_c_non_first_chunk_pre_subtracted( + self, fake_kv_manager, make_batch + ): + # remaining_budget = 512 - 0(gen) - 300(non-first chunk_size) = 212 + # First-chunk req: helper(0, 256, 1024)=256 > 212 → SKIP + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + nf = _make_non_first_chunk_req(rid=10, chunk_size=300, + current_pos=256, prompt_len=1024) + fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=256, + est_reuse=0) + batch = make_batch([nf, fc], []) + + fake_kv_manager.prepare_resources(batch) + + surviving = [r.py_request_id for r in batch.context_requests] + # nf falls through both if/elif (not first_context_chunk) → kept. + # fc hits the if branch and is skipped via `continue`. + assert surviving == [10] + # No first-chunk admitted, so add_sequence_batch is never called. + fake_kv_manager.impl.add_sequence_batch.assert_not_called() + + +class TestGenTokensConsumeBudget: + """Test D — gen tokens (including draft) reduce remaining_budget.""" + + def test_d_gen_with_draft_consumes_budget( + self, fake_kv_manager, make_batch + ): + # remaining_budget = 512 - (1+5)(gen with 5 draft) = 506 + # First-chunk req: helper(0, 600, 1024)=600 > 506 → SKIP + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + gen = _make_gen_req(rid=99, num_draft_tokens=5) + fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=600, + est_reuse=0) + batch = make_batch([fc], [gen]) + + fake_kv_manager.prepare_resources(batch) + + # fc skipped via `continue`; only gen survives in context_requests + # (gen requests are not in context_requests anyway). + assert [r.py_request_id for r in batch.context_requests] == [] + + def test_d_gen_without_draft_admits( + self, fake_kv_manager, make_batch + ): + # remaining_budget = 512 - 1(gen, no draft) = 511 + # First-chunk req: helper(0, 256, 1024)=256 ≤ 511 → admit + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + gen = _make_gen_req(rid=99, num_draft_tokens=0) + fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=256, + est_reuse=0) + batch = make_batch([fc], [gen]) + + fake_kv_manager.prepare_resources(batch) + + assert [r.py_request_id for r in batch.context_requests] == [20] + + +class TestNoOpPaths: + """Tests E and F — guard is a no-op when reuse is off / draft manager.""" + + def test_e_no_op_when_reuse_disabled(self, fake_kv_manager, make_batch): + fake_kv_manager.enable_block_reuse = False + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, + est_reuse=768) + for i in range(4) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + # count_reusable_blocks must NEVER be called when the guard is off. + fake_kv_manager.impl.count_reusable_blocks.assert_not_called() + # All 4 admitted; no skip path engaged. + assert {r.py_request_id for r in batch.context_requests} == {0, 1, 2, 3} + + def test_f_no_op_when_is_draft(self, fake_kv_manager, make_batch): + fake_kv_manager.is_draft = True + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, + est_reuse=768) + for i in range(4) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + fake_kv_manager.impl.count_reusable_blocks.assert_not_called() + assert {r.py_request_id for r in batch.context_requests} == {0, 1, 2, 3} + + +class TestResetContextRequestsFiltered: + """Test G — reset_context_requests reclassifies the FILTERED list.""" + + def test_g_reset_with_chunk_type_mutation( + self, fake_kv_manager, make_batch + ): + """A request may flip is_last_context_chunk inside add_sequence_batch + (e.g., reuse covers most of the prompt, leaving only the last chunk). + Verify that after prepare_resources, the request lands in + context_requests_last_chunk via the reset_context_requests call.""" + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + # Start with a non-last chunk that gets flipped during add_sequence_batch. + fc = _make_first_chunk_req(rid=0, prompt_len=1024, chunk_size=256, + est_reuse=0) + fc.is_last_context_chunk = False # initial classification + + def _flip_to_last(infos, reqs): + for r in reqs: + r.is_last_context_chunk = True + + fake_kv_manager.impl.add_sequence_batch.side_effect = _flip_to_last + + batch = make_batch([fc], []) + # initial bucketing: non-last + assert len(batch.context_requests_chunking) == 1 + assert len(batch.context_requests_last_chunk) == 0 + + fake_kv_manager.prepare_resources(batch) + + # After prepare_resources + reset, the flipped req moves to last bucket. + assert len(batch.context_requests_chunking) == 0 + assert len(batch.context_requests_last_chunk) == 1 + assert batch.context_requests_last_chunk[0].py_request_id == 0 + + +class TestKvConnectorElifBranch: + """Test H — kv_connector returns False routes to elif branch.""" + + def test_h_connector_should_not_add_charges_estimated_reuse( + self, fake_kv_manager, make_batch + ): + """When kv_connector_manager.should_add_sequence(req) returns False + for a first-chunk req, resource_manager.py:741-744 charges the + budget using req.estimated_reusable_tokens (NOT actual reuse, since + we skip the count_reusable_blocks call for these). The req is also + NOT added to add_sequence_batch.""" + connector = Mock() + # req0: rejected by connector → elif branch + # req1, req2: accepted by connector → if branch + connector.should_add_sequence.side_effect = [False, True, True] + fake_kv_manager.kv_connector_manager = connector + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, + est_reuse=448) + for i in range(3) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + # Arithmetic walk: + # Initial budget = 512 - 0(gen) - 0(non-first) = 512 + # req0 elif: helper(448, 256, 512): P+chunk=704 >= 512 → max(1, 64) = 64 + # budget -= 64 → remaining 448 + # req1 if: count_reusable=0 → helper(0, 256, 512)=256 ≤ 448 → admit, budget 192 + # req2 if: helper=256 > 192 → SKIP + surviving = [r.py_request_id for r in batch.context_requests] + assert surviving == [0, 1], f"expected [0, 1], got {surviving}" + # Only req1 went into add_sequence_batch (req0 was rejected by connector). + infos = fake_kv_manager.impl.add_sequence_batch.call_args[0][0] + assert [info[0] for info in infos] == [1] + + +class TestKvConnectorCallbacks: + """Tests I, J — connector update_state_after_alloc and build_scheduler_output.""" + + def test_i_update_state_after_alloc_only_for_admitted( + self, fake_kv_manager, make_batch + ): + connector = Mock() + connector.should_add_sequence.return_value = True + fake_kv_manager.kv_connector_manager = connector + fake_kv_manager.get_cache_indices = Mock(return_value=[10, 20]) + + # 3 ctx reqs: budget allows 2 to be admitted, 1 skipped. + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, + est_reuse=768) + for i in range(3) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + # update_state_after_alloc should be called once per admitted req (2). + # Skipped req2 should NOT trigger the callback. + admitted_ids = [ + call.args[0].py_request_id + for call in connector.update_state_after_alloc.call_args_list + ] + assert sorted(admitted_ids) == [0, 1] + + def test_j_build_scheduler_output_after_filter( + self, fake_kv_manager, make_batch + ): + connector = Mock() + connector.should_add_sequence.return_value = True + fake_kv_manager.kv_connector_manager = connector + + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, + est_reuse=768) + for i in range(3) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + # build_scheduler_output should be called exactly once after the loop + # with the (already-filtered) batch. + connector.build_scheduler_output.assert_called_once() + passed_batch = connector.build_scheduler_output.call_args[0][0] + assert {r.py_request_id for r in passed_batch.context_requests} == {0, 1} + + +class TestContextDraftTokens: + """Test K — ctx request with draft tokens triggers extra add_token calls.""" + + def test_k_context_with_draft_tokens( + self, fake_kv_manager, make_batch + ): + # num_extra_kv_tokens=0, num_draft_tokens=3 → add_token called 3 times. + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + ctx = [ + _make_first_chunk_req(rid=0, prompt_len=512, chunk_size=256, + est_reuse=0, num_draft_tokens=3) + ] + batch = make_batch(ctx, []) + + fake_kv_manager.prepare_resources(batch) + + assert fake_kv_manager.impl.add_token.call_count == 3 + + +class TestCountReusableBlocksCallShape: + """Test L — the budget guard calls count_reusable_blocks with the + expected argument shape. A signature drift would silently break + re-probing.""" + + def test_l_count_reusable_blocks_call_args( + self, fake_kv_manager, make_batch + ): + sentinel_unique = [101, 102, 103, 104] + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + req = _make_first_chunk_req(rid=0, prompt_len=1024, chunk_size=256, + est_reuse=0, + unique_tokens=sentinel_unique) + batch = make_batch([req], []) + + fake_kv_manager.prepare_resources(batch) + + fake_kv_manager.impl.count_reusable_blocks.assert_called_once_with( + sentinel_unique, req, False) + + +class TestNonFirstChunkFallthrough: + """Test M (catalogue) — non-first-chunk reqs fall through both if/elif + branches and are kept by reset_context_requests, but NOT added to + add_sequence_batch.""" + + def test_m_non_first_chunk_falls_through( + self, fake_kv_manager, make_batch + ): + nf = _make_non_first_chunk_req(rid=0, chunk_size=128, + current_pos=128, prompt_len=1024) + batch = make_batch([nf], []) + + fake_kv_manager.prepare_resources(batch) + + assert [r.py_request_id for r in batch.context_requests] == [0] + fake_kv_manager.impl.add_sequence_batch.assert_not_called() + + +# =========================================================================== +# Tier 3: deterministic DYN-2868 regression +# =========================================================================== +def _model_engine_total_tokens(scheduled_batch): + """Reimplement model_engine.py:2291-2297 + 2615-2616 token counting, + INDEPENDENT of _estimate_post_reuse_compute. The negative control + (Test N) patches the helper to 0; the verifier MUST not go through + the same code path or it would self-invalidate.""" + total = 0 + for req in scheduled_batch.context_requests: + begin = req.context_current_position + all_tokens = req.get_tokens(0) + total += max(0, min(req.context_chunk_size, len(all_tokens) - begin)) + for req in scheduled_batch.generation_requests: + # Draft path (model_engine.py:2399, 2459): 1 + len(py_draft_tokens), + # NOT multiplied by beam. + # No-draft path (model_engine.py:2615-2616): one position per beam. + if len(req.py_draft_tokens) > 0: + total += 1 + len(req.py_draft_tokens) + else: + total += req.sampling_config.beam_width + return total + + +class TestDYN2868Regression: + """Tests M, N — the regression-net pair.""" + + def test_m_invariant_passes_with_guard( + self, fake_kv_manager, make_batch + ): + """The PR's purpose: the guard prevents post-prepare_resources + token count from exceeding max_num_tokens under run-14 conditions.""" + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 # eviction + + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, + est_reuse=448) + for i in range(4) + ] + batch = make_batch(ctx, []) + + # Simulate prepare_context (called inside add_sequence_batch on the + # real path): for admitted first-chunk reqs, set + # context_current_position to the actual reused offset (0 here). + def _set_pos(infos, reqs): + for r in reqs: + r.context_current_position = 0 + + fake_kv_manager.impl.add_sequence_batch.side_effect = _set_pos + + fake_kv_manager.prepare_resources(batch) + + total = _model_engine_total_tokens(batch) + assert total <= fake_kv_manager.max_num_tokens, ( + f"DYN-2868 regression: post-prepare total {total} > " + f"max_num_tokens {fake_kv_manager.max_num_tokens}. Surviving: " + f"{[r.py_request_id for r in batch.context_requests]}" + ) + + def test_n_invariant_breaks_when_guard_disabled( + self, fake_kv_manager, make_batch, monkeypatch + ): + """Negative control: with the guard disabled (helper patched → 0), + the same scenario admits ALL 4 reqs and overshoots. + + Trace with helper=0: + - Initial budget = 512 + - req_compute = 0 for all reqs + - line 725: `0 > 512` is False; no skips + - line 734: budget -= 0; remaining stays 512 + - All 4 admitted → total = 4 * 256 = 1024 > 512. + + Verifier uses _model_engine_total_tokens (NOT the patched helper), + so this is not self-invalidating.""" + monkeypatch.setattr( + KVCacheManager, "_estimate_post_reuse_compute", + lambda self, *a, **k: 0, + ) + fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + + ctx = [ + _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, + est_reuse=448) + for i in range(4) + ] + batch = make_batch(ctx, []) + + def _set_pos(infos, reqs): + for r in reqs: + r.context_current_position = 0 + + fake_kv_manager.impl.add_sequence_batch.side_effect = _set_pos + + fake_kv_manager.prepare_resources(batch) + + total = _model_engine_total_tokens(batch) + assert total > fake_kv_manager.max_num_tokens, ( + f"Negative control: expected overshoot, got total={total} ≤ " + f"max={fake_kv_manager.max_num_tokens}. The synthetic scenario " + f"no longer triggers DYN-2868 with the guard disabled — " + f"strengthen the workload (more requests / larger chunks)." + ) From 8322e01a4f07c58da56aeb6614e025553b55d3ac Mon Sep 17 00:00:00 2001 From: Yuewei Na Date: Mon, 4 May 2026 11:28:36 -0700 Subject: [PATCH 3/3] [None][fix] Use analyze_prefix_reuse for budget guard re-probe Replace the count_reusable_blocks() call in the prepare_resources budget guard with analyze_prefix_reuse(...).reusable_blocks_all, matching the public KVCacheManager binding API. Update tests to mock the new return shape via a small Mock(reusable_blocks_all=N) helper. Tighten comments and docstrings to describe current behavior without referencing PRs, issues, or specific line numbers; reformat with ruff format to satisfy the pre-commit hook. Signed-off-by: Yuewei Na --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 +- .../_torch/pyexecutor/resource_manager.py | 22 +- .../_torch/executor/test_batch_wait_log.py | 24 +- .../executor/test_post_reuse_compute.py | 80 ++-- .../test_resource_manager_v1_budget.py | 353 ++++++++---------- 5 files changed, 229 insertions(+), 257 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6fde1e4b23bc..890601f8131b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3102,7 +3102,12 @@ def _maybe_log_batch_wait_decision( wait_threshold: float, should_waiting: bool, ) -> None: - """Diagnostics for batch_wait: set TLLM_LOG_BATCH_WAIT=1 (rank 0 only).""" + """Emit per-iteration batch-wait diagnostics on rank 0. + + Logs the scheduling-formula token count next to the actual chunk-token + sum so that mismatches (e.g., from KV cache reuse estimation drift) are + visible in the rank-0 stream. Other ranks early-return. + """ if self.dist.rank != 0: return diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index e30d374e6d50..4032b7de135b 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -757,8 +757,13 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): req): if remaining_budget is not None: unique_tokens = req.get_unique_tokens(0) - reusable_blocks = self.impl.count_reusable_blocks( - unique_tokens, req, False) + # reusable_blocks_all counts every block that + # matches the prefix, including blocks not yet + # allocated to a request — the right basis for + # the budget re-probe since any of them could be + # claimed before add_sequence_batch runs. + reusable_blocks = self.impl.analyze_prefix_reuse( + unique_tokens, req).reusable_blocks_all actual_reuse = (reusable_blocks * self.tokens_per_block) req_compute = self._estimate_post_reuse_compute( @@ -831,11 +836,14 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): def _estimate_post_reuse_compute(self, reuse_tokens: int, chunk_size: int, prompt_len: int) -> int: - """Estimate forward compute tokens after setPrepopulatedPromptLen. - - For non-last chunks the chunk window shifts right by the reused - amount and the forward cost is approximately chunk_size. For - last chunks the cost is prompt_len - reuse (original formula). + """Estimate forward compute tokens for a context chunk given how many + prefix tokens will be served from KV cache reuse. + + For non-last chunks the chunk window shifts right by the reused amount + and the forward cost is approximately ``chunk_size`` (rounded down to + a ``tokens_per_block`` boundary). For last chunks the cost is + ``prompt_len - reuse``. Returns ``chunk_size`` verbatim when reuse is + out of range (``reuse <= 0`` or ``reuse >= prompt_len``). """ P = reuse_tokens if P > 0 and P < prompt_len: diff --git a/tests/unittest/_torch/executor/test_batch_wait_log.py b/tests/unittest/_torch/executor/test_batch_wait_log.py index d0352985a8e4..89382a55f009 100644 --- a/tests/unittest/_torch/executor/test_batch_wait_log.py +++ b/tests/unittest/_torch/executor/test_batch_wait_log.py @@ -6,30 +6,28 @@ # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 -"""Smoke test for PyExecutor._maybe_log_batch_wait_decision rank gating. +"""Smoke test for ``PyExecutor._maybe_log_batch_wait_decision`` rank gating. -Per ``py_executor.py:3060-3063`` the method early-returns when ``self.dist.rank -!= 0``. The docstring claims env-var gating via ``TLLM_LOG_BATCH_WAIT=1`` -which does NOT exist in the implementation; the only gate is rank. - -This test is a regression net — if a refactor removes the rank gate, the -function would crash at line 3107 (``self.batch_wait_iters_count``) since -we never set the wait counters on the fake executor. +The method early-returns when ``self.dist.rank != 0`` and otherwise reads +``self.batch_wait_iters_count`` and ``self.batch_wait_timeout_iters`` to +format diagnostic output. If a refactor removes the rank gate, the function +would crash on the wait counters since this test never sets them on the fake. """ + from unittest.mock import Mock from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor def test_log_decision_early_returns_on_nonzero_rank(): - """Non-rank-0 path should early-return without touching wait counters.""" + """Non-rank-0 path early-returns without touching wait counters.""" fake = PyExecutor.__new__(PyExecutor) fake.dist = Mock() - fake.dist.rank = 1 # non-zero rank → early return at py_executor.py:3062 + fake.dist.rank = 1 # non-zero rank → early return - # If the gate is broken, the function would proceed to read - # self.batch_wait_iters_count (line 3107) which we never set → - # AttributeError. The successful no-crash path verifies the gate works. + # No exception expected. If the gate is broken, the function would proceed + # to read attributes (``batch_wait_iters_count``, ``batch_wait_timeout_iters``) + # that we never set on the fake, raising AttributeError. PyExecutor._maybe_log_batch_wait_decision( fake, context_requests=[], diff --git a/tests/unittest/_torch/executor/test_post_reuse_compute.py b/tests/unittest/_torch/executor/test_post_reuse_compute.py index e9d7c18a0cb4..f6ce074d1297 100644 --- a/tests/unittest/_torch/executor/test_post_reuse_compute.py +++ b/tests/unittest/_torch/executor/test_post_reuse_compute.py @@ -6,24 +6,20 @@ # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 -"""Unit tests for KVCacheManager._estimate_post_reuse_compute() — DYN-2868. +"""Unit tests for ``KVCacheManager._estimate_post_reuse_compute``. The helper is effectively pure: it depends only on the arguments and -``self.tokens_per_block``. Tests construct a stub holding tokens_per_block -and bind the unbound method, avoiding any real KVCacheManager initialization. - -Tier 1A: branch coverage of the helper. -Tier 1B: divergence vs ``_prepare_tp_inputs`` actual position-id count. -Tier 1C: precise property tests — undercharges in the reuse branch - (0 < P < prompt_len), short-circuits to chunk_size outside it. +``self.tokens_per_block``. Tests construct a stub holding ``tokens_per_block`` +and bind the unbound method, avoiding any real ``KVCacheManager`` initialization. """ + import pytest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager class _Stub: - """Minimal stand-in for KVCacheManager exposing only tokens_per_block.""" + """Minimal stand-in for KVCacheManager exposing only ``tokens_per_block``.""" def __init__(self, tokens_per_block): self.tokens_per_block = tokens_per_block @@ -32,12 +28,12 @@ def __init__(self, tokens_per_block): def _model_engine_forward_count(begin_compute, chunk_size, prompt_len): - """Mirror model_engine.py:2293-2297 — len(prompt_tokens) per ctx req.""" + """Mirror ``_prepare_tp_inputs`` slicing: ``len(prompt_tokens[begin:end])``.""" return max(0, min(chunk_size, prompt_len - begin_compute)) # --------------------------------------------------------------------------- -# Tier 1A: branch coverage +# Branch coverage of the helper # --------------------------------------------------------------------------- @pytest.mark.parametrize( "P, chunk, prompt_len, block, expected, branch", @@ -68,34 +64,34 @@ def test_branch_coverage(P, chunk, prompt_len, block, expected, branch): # --------------------------------------------------------------------------- -# Tier 1B: helper vs model_engine — production-relevant divergence +# Helper vs model engine — production-relevant divergence # --------------------------------------------------------------------------- @pytest.mark.parametrize( "P, chunk_size, prompt_len, tokens_per_block, expect_undercharge", [ - # Block-aligned chunk + block-aligned reuse: helper == model_engine + # Block-aligned chunk + block-aligned reuse: helper == model engine count (64, 128, 1024, 64, False), (0, 256, 1024, 64, False), (128, 256, 200, 64, False), # last chunk, exact match (helper=72, actual=72) (192, 64, 1024, 64, False), - # Block-aligned reuse + UNALIGNED chunk: helper undercharges. This - # case is legal in production per - # tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py:243-247 - # (chunk_unit_size=100 with tokens_per_block=64). - (64, 100, 1000, 64, True), # aligned_end=128, helper=64, actual=100 - (128, 50, 1024, 64, True), # aligned_end=128, helper=1, actual=50 + # Block-aligned reuse + UNALIGNED chunk: helper undercharges. Unaligned + # chunk_unit_size is legal — V2 scheduler tests use chunk_unit_size=100 + # with tokens_per_block=64. + (64, 100, 1000, 64, True), # aligned_end=128, helper=64, actual=100 + (128, 50, 1024, 64, True), # aligned_end=128, helper=1, actual=50 # Block-unaligned reuse: also undercharges - (70, 130, 1024, 64, True), # aligned_end=192, helper=122, actual=130 + (70, 130, 1024, 64, True), # aligned_end=192, helper=122, actual=130 ], ) -def test_helper_vs_model_engine_per_case(P, chunk_size, prompt_len, - tokens_per_block, - expect_undercharge): - """Document where the helper diverges from _prepare_tp_inputs token - counting. The aligned cases agree; the unaligned cases undercharge — - which means the budget guard is conservative-enough to catch run-14's - +2 overshoot but not a strict ≤ max guarantee under all legal scheduler - outputs. Surfaced as a caveat in the PR description.""" +def test_helper_vs_model_engine_per_case( + P, chunk_size, prompt_len, tokens_per_block, expect_undercharge +): + """The helper short-circuits to ``chunk_size`` outside the reuse range and + rounds the chunk window down to a block boundary inside it. The aligned + cases agree with the actual ``_prepare_tp_inputs`` token count; the unaligned + cases undercharge — the budget guard is therefore a conservative-enough + backstop for small overshoots, not a strict ``≤ max_num_tokens`` guarantee + under all scheduler outputs.""" stub = _Stub(tokens_per_block) helper = stub._estimate_post_reuse_compute(P, chunk_size, prompt_len) actual = _model_engine_forward_count(P, chunk_size, prompt_len) @@ -103,8 +99,8 @@ def test_helper_vs_model_engine_per_case(P, chunk_size, prompt_len, assert helper < actual, ( f"Expected undercharge for unaligned (P={P}, chunk={chunk_size}, " f"block={tokens_per_block}): helper={helper}, actual={actual}. " - f"If this fails, the helper has been tightened — update this " - f"test and review whether Tier 3 needs new scenarios." + f"If this fails, the helper has been tightened — review whether " + f"the regression scenarios still trigger the bug condition." ) else: assert helper == actual, ( @@ -116,13 +112,13 @@ def test_helper_vs_model_engine_per_case(P, chunk_size, prompt_len, # --------------------------------------------------------------------------- -# Tier 1C: precise property tests +# Property tests # --------------------------------------------------------------------------- def test_helper_undercharges_in_reuse_branch(): - """When 0 < P < prompt_len, helper <= actual model_engine forward count. - This is the invariant the budget guard relies on for the run-14 - scenario (block-aligned reuse + block-aligned chunk_size). Counts as - the load-bearing property — a regression that breaks this is a real bug.""" + """When ``0 < P < prompt_len``, helper ≤ actual model-engine forward count. + This is the load-bearing invariant the budget guard relies on for + block-aligned reuse + block-aligned chunk_size — a regression that breaks + it can re-introduce overshoot.""" counterexamples = [] for P in [1, 32, 63, 64, 65, 70, 128, 199]: for chunk in [1, 50, 64, 100, 128, 256]: @@ -134,18 +130,14 @@ def test_helper_undercharges_in_reuse_branch(): h = stub._estimate_post_reuse_compute(P, chunk, prompt_len) a = _model_engine_forward_count(P, chunk, prompt_len) if h > a: - counterexamples.append( - (P, chunk, prompt_len, block, h, a)) - assert not counterexamples, ( - f"helper > actual in reuse branch (first 5): {counterexamples[:5]}" - ) + counterexamples.append((P, chunk, prompt_len, block, h, a)) + assert not counterexamples, f"helper > actual in reuse branch (first 5): {counterexamples[:5]}" def test_helper_short_circuits_outside_reuse_range(): - """When P <= 0 or P >= prompt_len, helper returns chunk_size verbatim - regardless of prompt_len — so the helper may OVERCHARGE actual forward - tokens here. This is intentional: outside the reuse branch the guard - falls back to the no-credit path, mirroring the pre-fix behavior.""" + """When ``P <= 0`` or ``P >= prompt_len``, helper returns ``chunk_size`` + verbatim regardless of ``prompt_len``. The helper may overcharge actual + forward tokens here, mirroring the no-reuse-credit fallback.""" stub = _Stub(tokens_per_block=64) # P=0: no reuse credit → return chunk_size even when chunk > prompt_len assert stub._estimate_post_reuse_compute(0, 256, 128) == 256 diff --git a/tests/unittest/_torch/executor/test_resource_manager_v1_budget.py b/tests/unittest/_torch/executor/test_resource_manager_v1_budget.py index 2fae36ae2425..283282185acc 100644 --- a/tests/unittest/_torch/executor/test_resource_manager_v1_budget.py +++ b/tests/unittest/_torch/executor/test_resource_manager_v1_budget.py @@ -6,18 +6,26 @@ # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 -"""Integration tests for KVCacheManager.prepare_resources() budget guard -(DYN-2868). Drives the real ``prepare_resources`` method against a mocked -manager (constructed via ``__new__`` to skip C++ deps) and a real -``ScheduledRequests``. - -Tier 2: branch coverage of the budget guard (skip / no-skip / non-first - pre-subtraction / draft tokens / connector branches / no-op when - reuse off / no-op when is_draft / reset_context_requests filter / - count_reusable_blocks call shape / non-first fallthrough). -Tier 3: deterministic DYN-2868 regression pair (Test M passes with guard, - Test N — negative control — overshoots when guard disabled). +"""Integration tests for the budget guard inside ``KVCacheManager.prepare_resources``. + +The guard re-probes the radix tree before admitting first-chunk context requests +and skips any whose post-reuse forward cost would exceed the remaining +``max_num_tokens`` budget. This file drives the real ``prepare_resources`` method +against a stand-in manager constructed via ``__new__`` (skips C++ initialization) +plus a real ``ScheduledRequests``. + +Branch coverage: skip / no-skip / non-first-chunk pre-subtraction / draft tokens / +connector should-add-sequence false branch / connector callbacks / no-op when +reuse is disabled / no-op when this is the draft manager / reset_context_requests +filtering / analyze_prefix_reuse call shape / non-first-chunk fallthrough. + +Regression: a deterministic pair simulating the eviction race that the guard +prevents — ``test_invariant_passes_with_guard`` admits to within budget; +``test_invariant_breaks_when_guard_disabled`` is the negative control that +patches the helper out and confirms the synthetic workload would otherwise +overshoot. """ + from unittest.mock import Mock import pytest @@ -30,9 +38,9 @@ # Mock factories — avoid Mock(spec=LlmRequest) because LlmRequest is a C++ # binding; we just set the attributes prepare_resources reads. # --------------------------------------------------------------------------- -def _make_first_chunk_req(rid, prompt_len, chunk_size, est_reuse=0, - num_draft_tokens=0, unique_tokens=None, - beam_width=1): +def _make_first_chunk_req( + rid, prompt_len, chunk_size, est_reuse=0, num_draft_tokens=0, unique_tokens=None, beam_width=1 +): req = Mock() req.request_id = rid req.py_request_id = rid @@ -40,25 +48,24 @@ def _make_first_chunk_req(rid, prompt_len, chunk_size, est_reuse=0, req.context_chunk_size = chunk_size req.context_current_position = 0 req.is_first_context_chunk = True - req.is_last_context_chunk = (chunk_size >= prompt_len) + req.is_last_context_chunk = chunk_size >= prompt_len req.estimated_reusable_tokens = est_reuse req.num_draft_tokens = num_draft_tokens req.py_draft_tokens = list(range(num_draft_tokens)) req.get_unique_tokens.return_value = unique_tokens or [0] * prompt_len - # Tier 3 helper uses len(req.get_tokens(0)) per model_engine.py:2291. + # The regression-net helper reads len(req.get_tokens(0)) to mirror the + # actual model-engine slicing; provide a concrete list of that length. req.get_tokens.return_value = list(range(prompt_len)) req.sampling_config = Mock() req.sampling_config.beam_width = beam_width return req -def _make_non_first_chunk_req(rid, chunk_size, prompt_len=None, - current_pos=128): +def _make_non_first_chunk_req(rid, chunk_size, prompt_len=None, current_pos=128): req = Mock() req.request_id = rid req.py_request_id = rid - pl = prompt_len if prompt_len is not None else ( - chunk_size + current_pos + 256) + pl = prompt_len if prompt_len is not None else (chunk_size + current_pos + 256) req.prompt_len = pl req.context_chunk_size = chunk_size req.context_current_position = current_pos @@ -72,6 +79,12 @@ def _make_non_first_chunk_req(rid, chunk_size, prompt_len=None, return req +def _summary(reusable_blocks): + """Stand-in for ``analyze_prefix_reuse``'s ``PrefixReuseSummary`` — + only ``reusable_blocks_all`` is read by the budget guard.""" + return Mock(reusable_blocks_all=reusable_blocks) + + def _make_gen_req(rid, num_draft_tokens=0, beam_width=1): req = Mock() req.request_id = rid @@ -86,11 +99,11 @@ def _make_gen_req(rid, num_draft_tokens=0, beam_width=1): @pytest.fixture def fake_kv_manager(): - """KVCacheManager-shaped object with all attrs/methods that - ``prepare_resources`` touches before our budget code runs. Skips - ``__init__`` to avoid C++ deps. Each test gets a fresh fixture — - never share across tests because ``prepare_resources`` mutates - ``scheduled_batch`` and reads request fields.""" + """``KVCacheManager``-shaped object holding every attribute and stubbed + method that ``prepare_resources`` touches before the budget code runs. + Skips ``__init__`` to avoid pulling in C++ resources. Each test gets a + fresh fixture — never share across tests because ``prepare_resources`` + mutates ``scheduled_batch`` and reads request fields.""" mgr = KVCacheManager.__new__(KVCacheManager) mgr.is_draft = False mgr.enable_block_reuse = True @@ -102,19 +115,19 @@ def fake_kv_manager(): mgr.mapping.has_cp_helix.return_value = False mgr.kv_connector_manager = None # default; tests override for connector cases mgr.impl = Mock() - mgr.impl.count_reusable_blocks.return_value = 0 + mgr.impl.analyze_prefix_reuse.return_value = _summary(0) mgr.impl.add_sequence_batch = Mock() mgr.impl.add_token = Mock() mgr.impl.sync_transfer_manager_with_buffer_manager = Mock() mgr.impl.refresh_blocks = Mock() - # For connector tests; harmless when kv_connector_manager is None. + # Stubbed for connector tests; harmless when kv_connector_manager is None. mgr.get_cache_indices = Mock(return_value=[]) return mgr @pytest.fixture def make_batch(): - """Construct a fresh ScheduledRequests using the public append API. + """Build a fresh ``ScheduledRequests`` via the public append API. NEVER assign ``batch.context_requests`` directly — it's a property derived from ``context_requests_chunking + context_requests_last_chunk``.""" @@ -130,23 +143,21 @@ def _factory(ctx_reqs, gen_reqs): # =========================================================================== -# Tier 2: branch coverage +# Branch coverage # =========================================================================== class TestSkipPath: - """Test A — skip when actual reuse < estimated reuse.""" + """The guard skips first-chunk reqs whose actual reuse is less than + the scheduler estimated, so admitting them would overshoot.""" - def test_a_skip_when_actual_reuse_less_than_estimated( - self, fake_kv_manager, make_batch - ): + def test_skip_when_actual_reuse_less_than_estimated(self, fake_kv_manager, make_batch): # 3 ctx reqs; scheduler thought each had ~768-token reuse, actual is 0. # remaining_budget starts at 512. Each req: helper(0, 256, 1024) = 256. # req0: 256 ≤ 512 → admit, remaining 256 # req1: 256 ≤ 256 → admit, remaining 0 # req2: 256 > 0 → SKIP via continue - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) ctx = [ - _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, - est_reuse=768) + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, est_reuse=768) for i in range(3) ] batch = make_batch(ctx, []) @@ -154,8 +165,7 @@ def test_a_skip_when_actual_reuse_less_than_estimated( fake_kv_manager.prepare_resources(batch) surviving = [r.py_request_id for r in batch.context_requests] - assert surviving == [0, 1], ( - f"req2 should be skipped via continue; got {surviving}") + assert surviving == [0, 1], f"req2 should be skipped via continue; got {surviving}" # The two admitted reqs went into add_sequence_batch. fake_kv_manager.impl.add_sequence_batch.assert_called_once() @@ -164,20 +174,17 @@ def test_a_skip_when_actual_reuse_less_than_estimated( class TestNoSkipPath: - """Test B — no skip when estimate matches actual.""" + """When actual reuse matches the scheduler estimate, no requests are skipped.""" - def test_b_no_skip_when_estimate_matches_actual( - self, fake_kv_manager, make_batch - ): + def test_no_skip_when_estimate_matches_actual(self, fake_kv_manager, make_batch): # 3 ctx reqs; each has est_reuse=896 (14 blocks at block=64), and - # actual count_reusable_blocks=14 → actual_reuse=896. + # analyze_prefix_reuse returns reusable_blocks_all=14 → actual_reuse=896. # helper(896, 128, 1024): P+chunk=1024 == prompt_len → last-chunk # → max(1, 1024-896)=128 # 3 * 128 = 384 ≤ 512 → all admitted - fake_kv_manager.impl.count_reusable_blocks.return_value = 14 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(14) ctx = [ - _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=128, - est_reuse=896) + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=128, est_reuse=896) for i in range(3) ] batch = make_batch(ctx, []) @@ -189,18 +196,14 @@ def test_b_no_skip_when_estimate_matches_actual( class TestNonFirstChunkPreSubtract: - """Test C — non-first-chunk reqs reduce budget for first-chunk reqs.""" + """Non-first-chunk reqs reduce the budget for first-chunk reqs.""" - def test_c_non_first_chunk_pre_subtracted( - self, fake_kv_manager, make_batch - ): + def test_non_first_chunk_pre_subtracted(self, fake_kv_manager, make_batch): # remaining_budget = 512 - 0(gen) - 300(non-first chunk_size) = 212 # First-chunk req: helper(0, 256, 1024)=256 > 212 → SKIP - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 - nf = _make_non_first_chunk_req(rid=10, chunk_size=300, - current_pos=256, prompt_len=1024) - fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=256, - est_reuse=0) + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) + nf = _make_non_first_chunk_req(rid=10, chunk_size=300, current_pos=256, prompt_len=1024) + fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=256, est_reuse=0) batch = make_batch([nf, fc], []) fake_kv_manager.prepare_resources(batch) @@ -214,34 +217,27 @@ def test_c_non_first_chunk_pre_subtracted( class TestGenTokensConsumeBudget: - """Test D — gen tokens (including draft) reduce remaining_budget.""" + """Generation-request tokens (including draft) reduce the remaining budget.""" - def test_d_gen_with_draft_consumes_budget( - self, fake_kv_manager, make_batch - ): + def test_gen_with_draft_consumes_budget(self, fake_kv_manager, make_batch): # remaining_budget = 512 - (1+5)(gen with 5 draft) = 506 # First-chunk req: helper(0, 600, 1024)=600 > 506 → SKIP - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) gen = _make_gen_req(rid=99, num_draft_tokens=5) - fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=600, - est_reuse=0) + fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=600, est_reuse=0) batch = make_batch([fc], [gen]) fake_kv_manager.prepare_resources(batch) - # fc skipped via `continue`; only gen survives in context_requests - # (gen requests are not in context_requests anyway). + # fc skipped via `continue`; gen requests are not in context_requests anyway. assert [r.py_request_id for r in batch.context_requests] == [] - def test_d_gen_without_draft_admits( - self, fake_kv_manager, make_batch - ): + def test_gen_without_draft_admits(self, fake_kv_manager, make_batch): # remaining_budget = 512 - 1(gen, no draft) = 511 # First-chunk req: helper(0, 256, 1024)=256 ≤ 511 → admit - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) gen = _make_gen_req(rid=99, num_draft_tokens=0) - fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=256, - est_reuse=0) + fc = _make_first_chunk_req(rid=20, prompt_len=1024, chunk_size=256, est_reuse=0) batch = make_batch([fc], [gen]) fake_kv_manager.prepare_resources(batch) @@ -250,53 +246,48 @@ def test_d_gen_without_draft_admits( class TestNoOpPaths: - """Tests E and F — guard is a no-op when reuse is off / draft manager.""" + """The guard is a no-op when reuse is disabled or this is the draft manager.""" - def test_e_no_op_when_reuse_disabled(self, fake_kv_manager, make_batch): + def test_no_op_when_reuse_disabled(self, fake_kv_manager, make_batch): fake_kv_manager.enable_block_reuse = False ctx = [ - _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, - est_reuse=768) + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, est_reuse=768) for i in range(4) ] batch = make_batch(ctx, []) fake_kv_manager.prepare_resources(batch) - # count_reusable_blocks must NEVER be called when the guard is off. - fake_kv_manager.impl.count_reusable_blocks.assert_not_called() + # analyze_prefix_reuse must NEVER be called when the guard is off. + fake_kv_manager.impl.analyze_prefix_reuse.assert_not_called() # All 4 admitted; no skip path engaged. assert {r.py_request_id for r in batch.context_requests} == {0, 1, 2, 3} - def test_f_no_op_when_is_draft(self, fake_kv_manager, make_batch): + def test_no_op_when_is_draft(self, fake_kv_manager, make_batch): fake_kv_manager.is_draft = True ctx = [ - _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, - est_reuse=768) + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, est_reuse=768) for i in range(4) ] batch = make_batch(ctx, []) fake_kv_manager.prepare_resources(batch) - fake_kv_manager.impl.count_reusable_blocks.assert_not_called() + fake_kv_manager.impl.analyze_prefix_reuse.assert_not_called() assert {r.py_request_id for r in batch.context_requests} == {0, 1, 2, 3} class TestResetContextRequestsFiltered: - """Test G — reset_context_requests reclassifies the FILTERED list.""" - - def test_g_reset_with_chunk_type_mutation( - self, fake_kv_manager, make_batch - ): - """A request may flip is_last_context_chunk inside add_sequence_batch - (e.g., reuse covers most of the prompt, leaving only the last chunk). - Verify that after prepare_resources, the request lands in - context_requests_last_chunk via the reset_context_requests call.""" - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + """``reset_context_requests`` reclassifies the post-skip request list.""" + + def test_reset_with_chunk_type_mutation(self, fake_kv_manager, make_batch): + """A request may flip ``is_last_context_chunk`` inside ``add_sequence_batch`` + (when reuse covers most of the prompt, leaving only the last chunk). + After ``prepare_resources``, the request must land in + ``context_requests_last_chunk`` via the ``reset_context_requests`` call.""" + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) # Start with a non-last chunk that gets flipped during add_sequence_batch. - fc = _make_first_chunk_req(rid=0, prompt_len=1024, chunk_size=256, - est_reuse=0) + fc = _make_first_chunk_req(rid=0, prompt_len=1024, chunk_size=256, est_reuse=0) fc.is_last_context_chunk = False # initial classification def _flip_to_last(infos, reqs): @@ -319,26 +310,20 @@ def _flip_to_last(infos, reqs): class TestKvConnectorElifBranch: - """Test H — kv_connector returns False routes to elif branch.""" - - def test_h_connector_should_not_add_charges_estimated_reuse( - self, fake_kv_manager, make_batch - ): - """When kv_connector_manager.should_add_sequence(req) returns False - for a first-chunk req, resource_manager.py:741-744 charges the - budget using req.estimated_reusable_tokens (NOT actual reuse, since - we skip the count_reusable_blocks call for these). The req is also - NOT added to add_sequence_batch.""" + """When the connector rejects a first-chunk req, the elif branch charges + the budget using the scheduler's estimate (without re-probing) and the + request is not handed to ``add_sequence_batch``.""" + + def test_connector_should_not_add_charges_estimated_reuse(self, fake_kv_manager, make_batch): connector = Mock() # req0: rejected by connector → elif branch # req1, req2: accepted by connector → if branch connector.should_add_sequence.side_effect = [False, True, True] fake_kv_manager.kv_connector_manager = connector - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) ctx = [ - _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, - est_reuse=448) + _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, est_reuse=448) for i in range(3) ] batch = make_batch(ctx, []) @@ -349,7 +334,7 @@ def test_h_connector_should_not_add_charges_estimated_reuse( # Initial budget = 512 - 0(gen) - 0(non-first) = 512 # req0 elif: helper(448, 256, 512): P+chunk=704 >= 512 → max(1, 64) = 64 # budget -= 64 → remaining 448 - # req1 if: count_reusable=0 → helper(0, 256, 512)=256 ≤ 448 → admit, budget 192 + # req1 if: actual_reuse=0 → helper(0, 256, 512)=256 ≤ 448 → admit, budget 192 # req2 if: helper=256 > 192 → SKIP surviving = [r.py_request_id for r in batch.context_requests] assert surviving == [0, 1], f"expected [0, 1], got {surviving}" @@ -359,21 +344,19 @@ def test_h_connector_should_not_add_charges_estimated_reuse( class TestKvConnectorCallbacks: - """Tests I, J — connector update_state_after_alloc and build_scheduler_output.""" + """Connector callbacks fire only for admitted requests, and the post-loop + ``build_scheduler_output`` sees the filtered batch.""" - def test_i_update_state_after_alloc_only_for_admitted( - self, fake_kv_manager, make_batch - ): + def test_update_state_after_alloc_only_for_admitted(self, fake_kv_manager, make_batch): connector = Mock() connector.should_add_sequence.return_value = True fake_kv_manager.kv_connector_manager = connector fake_kv_manager.get_cache_indices = Mock(return_value=[10, 20]) # 3 ctx reqs: budget allows 2 to be admitted, 1 skipped. - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) ctx = [ - _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, - est_reuse=768) + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, est_reuse=768) for i in range(3) ] batch = make_batch(ctx, []) @@ -381,24 +364,20 @@ def test_i_update_state_after_alloc_only_for_admitted( fake_kv_manager.prepare_resources(batch) # update_state_after_alloc should be called once per admitted req (2). - # Skipped req2 should NOT trigger the callback. + # The skipped req should NOT trigger the callback. admitted_ids = [ - call.args[0].py_request_id - for call in connector.update_state_after_alloc.call_args_list + call.args[0].py_request_id for call in connector.update_state_after_alloc.call_args_list ] assert sorted(admitted_ids) == [0, 1] - def test_j_build_scheduler_output_after_filter( - self, fake_kv_manager, make_batch - ): + def test_build_scheduler_output_after_filter(self, fake_kv_manager, make_batch): connector = Mock() connector.should_add_sequence.return_value = True fake_kv_manager.kv_connector_manager = connector - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) ctx = [ - _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, - est_reuse=768) + _make_first_chunk_req(rid=i, prompt_len=1024, chunk_size=256, est_reuse=768) for i in range(3) ] batch = make_batch(ctx, []) @@ -413,16 +392,15 @@ def test_j_build_scheduler_output_after_filter( class TestContextDraftTokens: - """Test K — ctx request with draft tokens triggers extra add_token calls.""" + """A context request with draft tokens triggers extra ``add_token`` calls.""" - def test_k_context_with_draft_tokens( - self, fake_kv_manager, make_batch - ): + def test_context_with_draft_tokens(self, fake_kv_manager, make_batch): # num_extra_kv_tokens=0, num_draft_tokens=3 → add_token called 3 times. - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) ctx = [ - _make_first_chunk_req(rid=0, prompt_len=512, chunk_size=256, - est_reuse=0, num_draft_tokens=3) + _make_first_chunk_req( + rid=0, prompt_len=512, chunk_size=256, est_reuse=0, num_draft_tokens=3 + ) ] batch = make_batch(ctx, []) @@ -431,37 +409,30 @@ def test_k_context_with_draft_tokens( assert fake_kv_manager.impl.add_token.call_count == 3 -class TestCountReusableBlocksCallShape: - """Test L — the budget guard calls count_reusable_blocks with the - expected argument shape. A signature drift would silently break - re-probing.""" +class TestAnalyzePrefixReuseCallShape: + """The budget guard calls ``analyze_prefix_reuse`` with the expected + argument shape — a signature drift would silently break re-probing.""" - def test_l_count_reusable_blocks_call_args( - self, fake_kv_manager, make_batch - ): + def test_analyze_prefix_reuse_call_args(self, fake_kv_manager, make_batch): sentinel_unique = [101, 102, 103, 104] - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 - req = _make_first_chunk_req(rid=0, prompt_len=1024, chunk_size=256, - est_reuse=0, - unique_tokens=sentinel_unique) + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) + req = _make_first_chunk_req( + rid=0, prompt_len=1024, chunk_size=256, est_reuse=0, unique_tokens=sentinel_unique + ) batch = make_batch([req], []) fake_kv_manager.prepare_resources(batch) - fake_kv_manager.impl.count_reusable_blocks.assert_called_once_with( - sentinel_unique, req, False) + fake_kv_manager.impl.analyze_prefix_reuse.assert_called_once_with(sentinel_unique, req) class TestNonFirstChunkFallthrough: - """Test M (catalogue) — non-first-chunk reqs fall through both if/elif - branches and are kept by reset_context_requests, but NOT added to - add_sequence_batch.""" - - def test_m_non_first_chunk_falls_through( - self, fake_kv_manager, make_batch - ): - nf = _make_non_first_chunk_req(rid=0, chunk_size=128, - current_pos=128, prompt_len=1024) + """Non-first-chunk reqs fall through both the if and elif branches and + are kept by ``reset_context_requests``, but are NOT added to + ``add_sequence_batch``.""" + + def test_non_first_chunk_falls_through(self, fake_kv_manager, make_batch): + nf = _make_non_first_chunk_req(rid=0, chunk_size=128, current_pos=128, prompt_len=1024) batch = make_batch([nf], []) fake_kv_manager.prepare_resources(batch) @@ -471,22 +442,26 @@ def test_m_non_first_chunk_falls_through( # =========================================================================== -# Tier 3: deterministic DYN-2868 regression +# Regression-net pair # =========================================================================== def _model_engine_total_tokens(scheduled_batch): - """Reimplement model_engine.py:2291-2297 + 2615-2616 token counting, - INDEPENDENT of _estimate_post_reuse_compute. The negative control - (Test N) patches the helper to 0; the verifier MUST not go through - the same code path or it would self-invalidate.""" + """Mirror of the actual ``_prepare_tp_inputs`` ``len(position_ids)`` + computation, INDEPENDENT of ``_estimate_post_reuse_compute``. The + negative-control test patches the helper out; the verifier MUST not + go through the same code path or it would self-invalidate. + + Per ctx req: ``len(prompt_tokens[begin:begin+chunk_size])`` = + ``min(chunk_size, len(get_tokens(0)) - context_current_position)``. + Per gen req: + - draft path: ``1 + len(py_draft_tokens)`` (extend path, NOT per beam). + - no-draft path: one position per beam. + """ total = 0 for req in scheduled_batch.context_requests: begin = req.context_current_position all_tokens = req.get_tokens(0) total += max(0, min(req.context_chunk_size, len(all_tokens) - begin)) for req in scheduled_batch.generation_requests: - # Draft path (model_engine.py:2399, 2459): 1 + len(py_draft_tokens), - # NOT multiplied by beam. - # No-draft path (model_engine.py:2615-2616): one position per beam. if len(req.py_draft_tokens) > 0: total += 1 + len(req.py_draft_tokens) else: @@ -494,19 +469,27 @@ def _model_engine_total_tokens(scheduled_batch): return total -class TestDYN2868Regression: - """Tests M, N — the regression-net pair.""" +class TestEvictionRaceRegression: + """Deterministic regression for the eviction-race overshoot the guard + prevents. + + ``test_invariant_passes_with_guard`` — synthesize a scheduler that + estimated full reuse and then full eviction (``analyze_prefix_reuse`` + returns 0). The guard must skip enough requests that the post-prepare + forward token total fits within ``max_num_tokens``. + + ``test_invariant_breaks_when_guard_disabled`` — negative control. With + the helper monkey-patched to ``0`` the skip predicate ``> remaining`` + is always false, so all four requests are admitted and the same + workload overshoots. The verifier uses ``_model_engine_total_tokens`` + instead of the patched helper, so the comparison is meaningful. + """ - def test_m_invariant_passes_with_guard( - self, fake_kv_manager, make_batch - ): - """The PR's purpose: the guard prevents post-prepare_resources - token count from exceeding max_num_tokens under run-14 conditions.""" - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 # eviction + def test_invariant_passes_with_guard(self, fake_kv_manager, make_batch): + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) # eviction ctx = [ - _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, - est_reuse=448) + _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, est_reuse=448) for i in range(4) ] batch = make_batch(ctx, []) @@ -524,35 +507,21 @@ def _set_pos(infos, reqs): total = _model_engine_total_tokens(batch) assert total <= fake_kv_manager.max_num_tokens, ( - f"DYN-2868 regression: post-prepare total {total} > " - f"max_num_tokens {fake_kv_manager.max_num_tokens}. Surviving: " + f"post-prepare forward total {total} > max_num_tokens " + f"{fake_kv_manager.max_num_tokens}; surviving: " f"{[r.py_request_id for r in batch.context_requests]}" ) - def test_n_invariant_breaks_when_guard_disabled( - self, fake_kv_manager, make_batch, monkeypatch - ): - """Negative control: with the guard disabled (helper patched → 0), - the same scenario admits ALL 4 reqs and overshoots. - - Trace with helper=0: - - Initial budget = 512 - - req_compute = 0 for all reqs - - line 725: `0 > 512` is False; no skips - - line 734: budget -= 0; remaining stays 512 - - All 4 admitted → total = 4 * 256 = 1024 > 512. - - Verifier uses _model_engine_total_tokens (NOT the patched helper), - so this is not self-invalidating.""" + def test_invariant_breaks_when_guard_disabled(self, fake_kv_manager, make_batch, monkeypatch): monkeypatch.setattr( - KVCacheManager, "_estimate_post_reuse_compute", + KVCacheManager, + "_estimate_post_reuse_compute", lambda self, *a, **k: 0, ) - fake_kv_manager.impl.count_reusable_blocks.return_value = 0 + fake_kv_manager.impl.analyze_prefix_reuse.return_value = _summary(0) ctx = [ - _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, - est_reuse=448) + _make_first_chunk_req(rid=i, prompt_len=512, chunk_size=256, est_reuse=448) for i in range(4) ] batch = make_batch(ctx, []) @@ -567,8 +536,8 @@ def _set_pos(infos, reqs): total = _model_engine_total_tokens(batch) assert total > fake_kv_manager.max_num_tokens, ( - f"Negative control: expected overshoot, got total={total} ≤ " + f"negative control: expected overshoot, got total={total} ≤ " f"max={fake_kv_manager.max_num_tokens}. The synthetic scenario " - f"no longer triggers DYN-2868 with the guard disabled — " + f"no longer triggers the bug condition with the guard disabled — " f"strengthen the workload (more requests / larger chunks)." )