diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 95c92473efa5..5835d677b6f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1010,10 +1010,13 @@ def _get_token_num_for_estimation(self) -> int: ) num_cache_blocks *= num_pool_groups - # Multiply by beam width, to prevent rescaling of the max_seq_len caused by the influence of beam width during the preparation for kv_cache_estimation - max_num_tokens_for_estimation = ( - num_cache_blocks * self._tokens_per_block * - self._dummy_reqs[0].sampling_config.beam_width) + # Dummy context requests use the configured maximum beam width. Scale + # their block budget by the same value so the temporary KV cache used + # during warm-up can accommodate those requests. + num_cache_blocks *= self._max_beam_width + + max_num_tokens_for_estimation = (num_cache_blocks * + self._tokens_per_block) # V2 capacity is controlled by max_gpu_total_bytes; max_tokens only # describes the dummy workload needed for estimation. if self._is_kv_cache_manager_v2: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 19ab0cffc873..a83494c4181c 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -588,6 +588,8 @@ def __init__( self.force_terminate_ctx_for_partial_reuse = ( self.enable_disagg_partial_reuse_store and self.dist.pp_size == 1) + self._warn_if_kv_block_budget_unchecked() + self.max_input_len = max_input_len # _executor_loop private data self.max_num_active_requests = model_engine.get_max_num_sequences() @@ -5459,6 +5461,51 @@ def _validate_token_id_range(self, request: LlmRequest) -> None: self.model_engine.model.lm_head.num_embeddings): raise ValueError("Token ID out of range") + def _warn_if_kv_block_budget_unchecked(self) -> None: + """Warn when beam search runs against a pool no admission check covers. + + Managers that keep their own cost model -- sparse/compressed KV, + variable- or uniform-sliding-window, linear attention, mamba hybrids -- + opt out of ``get_request_kv_block_budget``, so + ``_validate_request_budget`` cannot reject a request whose per-beam + demand exceeds the pool. Such a request is admitted and then cannot + complete: under GUARANTEED_NO_EVICT it waits for capacity that will + never exist rather than failing. Leave a trace so that is diagnosable. + + Logged once because KV cache estimation builds a throwaway executor + before the real one (see create_py_executor), which would otherwise + report the same pools twice per startup. + """ + if self.max_beam_width <= 1: + return + unchecked = self.resource_manager.get_unchecked_kv_block_budget_pools() + if not unchecked: + return + logger.warning_once( + f"[PyExecutor] max_beam_width={self.max_beam_width} with " + f"{', '.join(pool.value for pool in unchecked)}: these pools " + "estimate their own KV cache demand, so requests are admitted " + "without a per-beam block feasibility check. A request needing " + "more blocks than a pool holds waits unscheduled instead of " + "failing; lower max_beam_width or max_seq_len if generation makes " + "no progress.", + key="kv_block_budget_unchecked") + + def _validate_request_budget(self, request: LlmRequest) -> None: + # Compare worst-case, beam-aware demand against every supported KV + # pool. Reject requests that can never fit so they fail cleanly instead + # of stalling the scheduler forever. + for resource_type, required_blocks, primary_capacity in ( + self.resource_manager.get_request_kv_block_budgets(request)): + if required_blocks > primary_capacity: + raise ValueError( + f"{resource_type.value} requires {required_blocks} KV cache " + f"blocks to complete the request, which exceeds its " + f"GPU-primary capacity of {primary_capacity} blocks " + f"(prompt_len={request.orig_prompt_len}, " + f"max_new_tokens={request.max_new_tokens}, " + f"beam_width={request.py_beam_width}).") + def _validate_request(self, request: LlmRequest): # Validate beam width sampling_config = request.sampling_config @@ -5521,6 +5568,9 @@ def _validate_request(self, request: LlmRequest): # Perform sampler-specific validation self.sampler.validate_request(request) + # Check if request has enough budget + self._validate_request_budget(request) + def _fetch_and_enqueue_requests(self, waiting_queue: WaitingQueue, total_num_active_requests: int) -> None: """Fetch requests from request_queue and enqueue to waiting_queue.""" diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 7e4fed2a0eb2..2baa082e8092 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -151,6 +151,29 @@ def get_max_resource_count(self) -> int: def get_needed_resource_to_completion(self, request: LlmRequest) -> int: raise NotImplementedError + def get_request_kv_block_budget( + self, request: LlmRequest) -> Optional[Tuple[int, int]]: + """Return required and capacity KV block counts for admission. + + ``None`` indicates that this resource manager does not support the + request-level KV block feasibility check. Opting in is deliberate: + a manager must only return a pair once it is known that the request + cannot possibly be served when ``required > capacity``, since the + executor rejects such requests outright. + """ + return None + + def kv_block_budget_applies(self) -> bool: + """Whether ``get_request_kv_block_budget`` reports a budget at all. + + Request-independent, so startup can ask before any request exists -- + see ``PyExecutor._warn_if_kv_block_budget_unchecked``, which reports + the pools that beam search runs against unchecked. A manager that + opts into the check must override this too, or it will be reported as + unchecked while in fact rejecting requests. + """ + return False + def add_dummy_requests(self, request_ids: List[int]): pass @@ -583,7 +606,6 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], blocks_per_window=blocks_per_window, tokens_per_block=tokens_per_block, max_seq_len=self.max_seq_len, - max_beam_width=max_beam_width, ) # Rewrite each pool's window_size to match the post-clamp window. @@ -789,6 +811,66 @@ def get_num_tokens(self, request: LlmRequest) -> int: # LlmRequest.get_num_tokens is out of sync with GenerationRequest when overlap scheduler is enabled. return self.impl.get_token_count(request.py_request_id) + def _num_blocks_to_completion(self, prompt_len: int, max_new_tokens: int, + beam_width: int) -> int: + # Beam-aware block estimate mirroring the C++ shared/unshared cost model + # (getNeededBlocksOneStep in kvCacheManager.cpp): the full prompt blocks + # are shared across all beams, while the partial-last-prompt block and + # the generated tokens are private to each beam. + shared_context_blocks = prompt_len // self.tokens_per_block + per_beam_tokens = prompt_len % self.tokens_per_block + max_new_tokens + per_beam_blocks = self.get_num_kv_blocks(per_beam_tokens) * beam_width + return shared_context_blocks + per_beam_blocks + + def kv_block_budget_applies(self) -> bool: + # Restricted to the exact base manager: the estimate below assumes every + # token of the sequence occupies a block of this pool. Subclasses break + # that assumption in ways that would make it wrong in either direction + # -- sparse/compressed KV managers (RocketKVCacheManager, + # DSACacheManager) retain less than the full sequence and would be + # over-estimated into false rejections, while the mamba hybrids split + # capacity across an extra state cache this single required/capacity + # pair cannot express. A subclass opts into the dense estimate by + # overriding this method once its own cost model has been checked + # against it, or supplies a cost model of its own by overriding both + # this and `get_request_kv_block_budget`. + if type(self) is not KVCacheManager: + return False + if self.kv_cache_type == CacheTypeCpp.CROSS: + return True + # Multi-window and linear-attention managers have per-pool capacity + # semantics that cannot be represented by one required/capacity pair. + # Uniform sliding-window managers also require window-aware accounting. + return not (self.kv_cache_type != CacheTypeCpp.SELF or self.is_vswa + or self.is_linear_attention + or any(window < self.max_seq_len + for window in self.max_attention_window_vec)) + + def get_request_kv_block_budget( + self, request: LlmRequest) -> Optional[Tuple[int, int]]: + """Required and available KV blocks under the dense full-attention cost + model, or ``None`` if this manager is not covered by it.""" + if not self.kv_block_budget_applies(): + return None + + if self.kv_cache_type == CacheTypeCpp.CROSS: + if request.encoder_output_len is None: + logger.warning( + f"Encoder output length is not set for request {request.request_id}" + ) + return None + required_blocks = self.get_num_kv_blocks(request.encoder_output_len) + return required_blocks, self.blocks_in_primary_pool + + extra_tokens = (self.num_extra_kv_tokens + + self._kv_reserve_draft_tokens) + required_blocks = self._num_blocks_to_completion( + request.orig_prompt_len, request.max_new_tokens + extra_tokens, + request.py_beam_width) + # Use GPU-primary capacity only. Secondary/offloaded blocks cannot make + # an otherwise impossible request schedulable on the GPU. + return required_blocks, self.blocks_in_primary_pool + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: # TODO: the C++ implementation of this method can be used, but the # Python and C++ schedulers currently do not agree on what "needed @@ -797,12 +879,12 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: # the Python scheduler needs to be fixed. # # return self.impl.get_remaining_blocks_to_completion(request) - context_token_count = request.orig_prompt_len - num_context_blocks = context_token_count // self.tokens_per_block - remaining_tokens = context_token_count + request.max_new_tokens - num_context_blocks * self.tokens_per_block - need_blocks = num_context_blocks + math.ceil( - remaining_tokens / self.tokens_per_block) - return need_blocks + # + # This intentionally uses beam_width=1 to preserve the historical + # beam-unaware value the Python scheduler expects. + return self._num_blocks_to_completion(request.orig_prompt_len, + request.max_new_tokens, + beam_width=1) @staticmethod def _has_mm_bidirectional_block(req: LlmRequest) -> bool: @@ -1560,20 +1642,6 @@ def calculate_max_num_blocks(self, return blocks_in_primary_pool, blocks_in_secondary_pool - def get_max_atten_window_upper_bound(self, blocks_in_primary_pool, - tokens_per_block, max_beam_width, - max_seq_len: Optional[int]): - token_capacity = blocks_in_primary_pool * tokens_per_block - max_blocks_per_seq = math.floor(token_capacity / - (max_beam_width * tokens_per_block)) - assert max_blocks_per_seq > 0, "Impossible to fit in any sequence in kvCache" - - max_atten_window_upper_bound = max_blocks_per_seq * tokens_per_block - if max_seq_len is not None and max_seq_len > max_atten_window_upper_bound and max_beam_width > 1: - max_atten_window_upper_bound -= tokens_per_block - assert max_atten_window_upper_bound > 0, "Impossible to fit in any sequence in kvCache" - return max_atten_window_upper_bound - def _resolve_window_size( self, window_size: Optional[int], @@ -2400,7 +2468,6 @@ def _validate_and_adjust_attention_windows( blocks_per_window: BlocksPerWindow, tokens_per_block: int, max_seq_len: int, - max_beam_width: int, ) -> Tuple[BlocksPerWindow, int, List[int], Dict[int, int]]: """ Validate and adjust attention windows against their upper bounds if needed. @@ -2426,11 +2493,8 @@ def _validate_and_adjust_attention_windows( _) in blocks_per_window.items(): if window_size < 0: continue - upper_bound = self.get_max_atten_window_upper_bound( - blocks_in_primary_pool=blocks_in_primary_pool, - tokens_per_block=tokens_per_block, - max_beam_width=max_beam_width, - max_seq_len=max_seq_len) + upper_bound = blocks_in_primary_pool * tokens_per_block + assert upper_bound > 0, "Impossible to fit in any sequence in kvCache" if window_size > upper_bound: logger.warning( f"Attention window size {window_size} exceeds upper bound {upper_bound} " @@ -2852,6 +2916,13 @@ def free_resources(self, request: "LlmRequest") -> None: class ResourceManager: + # KV pools a request's block demand is charged against. + _KV_BUDGET_RESOURCE_TYPES = ( + ResourceManagerType.KV_CACHE_MANAGER, + ResourceManagerType.DRAFT_KV_CACHE_MANAGER, + ResourceManagerType.CROSS_KV_CACHE_MANAGER, + ) + def __init__(self, resource_managers: dict[ResourceManagerType, BaseResourceManager]): self.resource_managers = OrderedDict(resource_managers) @@ -2867,6 +2938,36 @@ def get_resource_manager( self, type: ResourceManagerType) -> Optional[BaseResourceManager]: return self.resource_managers.get(type) + def get_request_kv_block_budgets( + self, + request: LlmRequest) -> List[Tuple[ResourceManagerType, int, int]]: + budgets = [] + for resource_type in self._KV_BUDGET_RESOURCE_TYPES: + kv_cache_manager = self.get_resource_manager(resource_type) + if kv_cache_manager is None: + continue + budget = kv_cache_manager.get_request_kv_block_budget(request) + if budget is not None: + required_blocks, primary_capacity = budget + budgets.append( + (resource_type, required_blocks, primary_capacity)) + else: + logger.warning_once( + f"Resource manager {resource_type} does not provide a per-request KV block budget", + key=f"missing_kv_block_budget:{resource_type.value}", + ) + return budgets + + def get_unchecked_kv_block_budget_pools(self) -> List[ResourceManagerType]: + """KV pools present but excluded from the admission check, i.e. those + whose ``get_request_kv_block_budget`` returns ``None`` for every + request.""" + return [ + resource_type for resource_type in self._KV_BUDGET_RESOURCE_TYPES + if (manager := self.get_resource_manager(resource_type)) is not None + and not manager.kv_block_budget_applies() + ] + @nvtx_range("maybe_fit_token_budget") def maybe_fit_token_budget(self, scheduled_batch: ScheduledRequests): """Apply the post-allocation token-budget trim (#13318) to the batch. diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5b39eea23efa..2488cec777ed 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -32,6 +32,7 @@ l0_a10: - unittest/_torch/executor/test_kv_cache_compression_manager.py - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py + - unittest/_torch/executor/test_resource_manager.py - unittest/_torch/modules/moe/test_communication_factory.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no # test list either). diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 9f0a4b9ac0df..901092cc04a0 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -139,12 +139,14 @@ def _make_creator( sliding_window=None, use_sliding_window=None, max_attention_window=None, + max_beam_width=1, ): """Build a minimal KvCacheCreator (bypasses __init__) wired up for _get_token_num_for_estimation only.""" c = object.__new__(KvCacheCreator) c._tokens_per_block = tokens_per_block + c._max_beam_width = max_beam_width c._net_max_seq_len = 2048 c._speculative_config = None c._dummy_reqs = dummy_reqs @@ -240,6 +242,23 @@ def test_without_adp_all_blocks_counted(): assert c._get_token_num_for_estimation() == n_reqs * 3 * tpb +def test_max_beam_width_scales_estimation_blocks(): + """Warm-up capacity uses the same configured beam width as dummy requests.""" + tpb = 64 + beam_width = 4 + c = _make_creator( + tpb, + [_make_mock_request(128, beam_width=beam_width)], + enable_attention_dp=False, + tp_size=1, + max_beam_width=beam_width, + ) + # Beam search falls back to V1; exercise the V1 memory-cap path. + c._is_kv_cache_manager_v2 = False + + assert c._get_token_num_for_estimation() == 3 * beam_width * tpb + + @pytest.mark.parametrize("tp_size", [2, 4, 8]) def test_adp_various_tp_sizes(tp_size): """ADP division must hold for several representative tp_size values.""" diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index 09aa31109f25..34db8b0bf8e3 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -17,9 +17,12 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.peft.lora.config import LoraConfig +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ( - KVCacheManager, PeftCacheManager, _merge_kv_cache_pool_pointers, + KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType, + _merge_kv_cache_pool_pointers, _warn_if_unsupported_v1_kv_cache_event_hash_algo) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import LayerType @@ -1240,5 +1243,306 @@ def test_genuine_vswa_requires_window_size(self): self.assertEqual(mgr._resolve_window_size(8192), 8192) +class TestRequestBudget(unittest.TestCase): + """Unit tests for the beam-aware KV block budget estimation and the + request admission budget check.""" + + TOKENS_PER_BLOCK = 8 + + class _OptedInKVCacheManager(KVCacheManager): + """A subclass that has checked its cost model against the dense one.""" + + def kv_block_budget_applies(self): + return True + + @classmethod + def _make_kv_cache_manager(cls, manager_cls=KVCacheManager): + # Build a bare KVCacheManager without touching the GPU (mirrors + # TestResolveWindowSize) so we can exercise the pure block math. + mgr = manager_cls.__new__(manager_cls) + mgr.tokens_per_block = cls.TOKENS_PER_BLOCK + mgr.kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF + mgr.is_vswa = False + mgr.is_linear_attention = False + mgr.num_extra_kv_tokens = 0 + mgr.max_total_draft_tokens = 0 + mgr._kv_reserve_draft_tokens = 0 + mgr.max_seq_len = 4096 + mgr.max_attention_window_vec = [mgr.max_seq_len] + mgr.blocks_in_primary_pool = 32 + mgr.blocks_in_secondary_pool = 100 + return mgr + + @staticmethod + def _make_request(prompt_len, + max_new_tokens, + beam_width, + encoder_output_len=None): + sampling_params = SamplingParams(n=beam_width, + best_of=beam_width, + use_beam_search=beam_width > 1) + return LlmRequest( + request_id=1, + max_new_tokens=max_new_tokens, + input_tokens=list(range(prompt_len)), + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + encoder_output_len=encoder_output_len, + ) + + def test_num_blocks_beam_one_matches_needed_resource(self): + """For beam_width=1 the beam-aware estimate must equal the historical + beam-unaware get_needed_resource_to_completion value.""" + mgr = self._make_kv_cache_manager() + for prompt_len, max_new_tokens in [(20, 5), (16, 1), (1, 100), (8, 8)]: + req = self._make_request(prompt_len, max_new_tokens, beam_width=1) + self.assertEqual( + mgr._num_blocks_to_completion(req.orig_prompt_len, + req.max_new_tokens, + req.py_beam_width), + mgr.get_needed_resource_to_completion(req)) + + def test_num_blocks_shares_prompt_across_beams(self): + """Only the partial-last-prompt block and generated tokens scale with + beam width; full prompt blocks are shared.""" + mgr = self._make_kv_cache_manager() + # prompt_len=20, tpb=8 -> 2 full shared blocks, 4 leftover prompt tokens. + # max_new_tokens=5 -> per-beam tokens = 4 + 5 = 9 -> ceil(9/8) = 2 blocks. + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=4) + # shared(2) + per_beam(2) * beam(4) = 10 + self.assertEqual( + mgr._num_blocks_to_completion(req.orig_prompt_len, + req.max_new_tokens, + req.py_beam_width), 10) + + # Beam width only multiplies the per-beam portion, not the shared prompt. + req1 = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=1) + self.assertEqual( + mgr._num_blocks_to_completion(req1.orig_prompt_len, + req1.max_new_tokens, + req1.py_beam_width), 4) + + def test_num_blocks_block_aligned_prompt(self): + """A block-aligned prompt contributes no per-beam partial block.""" + mgr = self._make_kv_cache_manager() + # prompt_len=16 (2 full blocks, 0 leftover), max_new_tokens=1. + # per-beam tokens = 0 + 1 = 1 -> ceil(1/8) = 1 block per beam. + req = self._make_request(prompt_len=16, max_new_tokens=1, beam_width=3) + # shared(2) + per_beam(1) * beam(3) = 5 + self.assertEqual( + mgr._num_blocks_to_completion(req.orig_prompt_len, + req.max_new_tokens, + req.py_beam_width), 5) + + def test_container_returns_v1_primary_pool_budget(self): + mgr = self._make_kv_cache_manager() + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: mgr}) + + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=4) + self.assertEqual( + resource_manager.get_request_kv_block_budgets(req), + [(ResourceManagerType.KV_CACHE_MANAGER, 10, + mgr.blocks_in_primary_pool)], + ) + + def test_container_no_kv_cache_manager_is_explicit_noop(self): + resource_manager = ResourceManager({}) + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=4) + self.assertEqual(resource_manager.get_request_kv_block_budgets(req), []) + + def test_v2_manager_does_not_run_v1_budget_check(self): + mgr = KVCacheManagerV2.__new__(KVCacheManagerV2) + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: mgr}) + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=1) + + self.assertEqual(resource_manager.get_request_kv_block_budgets(req), []) + + def test_subclass_manager_opts_out_by_default(self): + """Sparse/compressed KV and mamba hybrid subclasses hold a different + amount of KV per token, so they must not inherit the dense estimate.""" + + class _SubclassedKVCacheManager(KVCacheManager): + pass + + mgr = self._make_kv_cache_manager(_SubclassedKVCacheManager) + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: mgr}) + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=4) + + self.assertEqual(resource_manager.get_request_kv_block_budgets(req), []) + + def test_subclass_manager_can_opt_in(self): + """A subclass whose cost model matches opts in by declaring the budget + applies, which routes it through the dense estimate.""" + + mgr = self._make_kv_cache_manager(self._OptedInKVCacheManager) + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: mgr}) + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=4) + + self.assertEqual( + resource_manager.get_request_kv_block_budgets(req), + [(ResourceManagerType.KV_CACHE_MANAGER, 10, + mgr.blocks_in_primary_pool)], + ) + + def test_windowed_v1_manager_does_not_use_single_pool_budget(self): + mgr = self._make_kv_cache_manager() + mgr.max_attention_window_vec = [128] + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: mgr}) + req = self._make_request(prompt_len=200, max_new_tokens=5, beam_width=4) + + self.assertEqual(resource_manager.get_request_kv_block_budgets(req), []) + + def test_validate_request_budget_rejects_oversized_request(self): + """_validate_request_budget raises for a request that can never fit and + passes for one that fits.""" + mgr = self._make_kv_cache_manager() + mgr.blocks_in_primary_pool = 8 + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: mgr}) + executor = PyExecutor.__new__(PyExecutor) + executor.resource_manager = resource_manager + + # shared(25) + per_beam(1) * beam(4) exceeds primary capacity. + req = self._make_request(prompt_len=200, max_new_tokens=5, beam_width=4) + with self.assertRaisesRegex( + ValueError, "KV_CACHE_MANAGER requires 29 KV cache blocks.*" + "GPU-primary capacity of 8"): + executor._validate_request_budget(req) + + # shared(2) + per_beam(2) * beam(2) = 6, so it fits. + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=2) + executor._validate_request_budget(req) + + def test_validate_request_budget_checks_draft_manager(self): + mgr = self._make_kv_cache_manager() + mgr.blocks_in_primary_pool = 8 + resource_manager = ResourceManager( + {ResourceManagerType.DRAFT_KV_CACHE_MANAGER: mgr}) + executor = PyExecutor.__new__(PyExecutor) + executor.resource_manager = resource_manager + + oversized = self._make_request(prompt_len=200, + max_new_tokens=5, + beam_width=4) + with self.assertRaisesRegex( + ValueError, + "DRAFT_KV_CACHE_MANAGER requires 29 KV cache blocks"): + executor._validate_request_budget(oversized) + + fitting = self._make_request(prompt_len=20, + max_new_tokens=5, + beam_width=2) + executor._validate_request_budget(fitting) + + def test_validate_request_budget_checks_cross_manager(self): + mgr = self._make_kv_cache_manager() + mgr.kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.CROSS + mgr.blocks_in_primary_pool = 2 + resource_manager = ResourceManager( + {ResourceManagerType.CROSS_KV_CACHE_MANAGER: mgr}) + executor = PyExecutor.__new__(PyExecutor) + executor.resource_manager = resource_manager + + oversized = self._make_request(prompt_len=20, + max_new_tokens=5, + beam_width=1, + encoder_output_len=17) + with self.assertRaisesRegex( + ValueError, + "CROSS_KV_CACHE_MANAGER requires 3 KV cache blocks"): + executor._validate_request_budget(oversized) + + fitting = self._make_request(prompt_len=20, + max_new_tokens=5, + beam_width=1, + encoder_output_len=16) + executor._validate_request_budget(fitting) + + def test_unchecked_pools_reported_exactly_when_budget_is_skipped(self): + """get_unchecked_kv_block_budget_pools must agree with which managers + actually produce a budget -- the startup warning names pools from it, + so a disagreement would either warn about a checked pool or stay + silent about an unchecked one.""" + + class _SubclassedKVCacheManager(KVCacheManager): + pass + + windowed = self._make_kv_cache_manager() + windowed.max_attention_window_vec = [128] + cases = [ + ("dense base manager", self._make_kv_cache_manager(), True), + ("subclass opting in", + self._make_kv_cache_manager(self._OptedInKVCacheManager), True), + ("subclass opting out", + self._make_kv_cache_manager(_SubclassedKVCacheManager), False), + ("sliding window", windowed, False), + ("V2 manager", KVCacheManagerV2.__new__(KVCacheManagerV2), False), + ] + for name, mgr, is_checked in cases: + with self.subTest(manager=name): + self._assert_unchecked_matches_budget(mgr, is_checked) + + def _assert_unchecked_matches_budget(self, mgr, is_checked): + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: mgr}) + req = self._make_request(prompt_len=20, max_new_tokens=5, beam_width=4) + + unchecked = resource_manager.get_unchecked_kv_block_budget_pools() + has_budget = bool(resource_manager.get_request_kv_block_budgets(req)) + self.assertEqual(has_budget, is_checked) + self.assertEqual( + unchecked, + [] if is_checked else [ResourceManagerType.KV_CACHE_MANAGER]) + + def test_startup_warns_when_beam_search_runs_unchecked(self): + """Beam search against a pool that opts out of the budget check gets a + warning naming it, since such a request stalls rather than failing.""" + + class _SubclassedKVCacheManager(KVCacheManager): + pass + + executor = PyExecutor.__new__(PyExecutor) + executor.resource_manager = ResourceManager({ + ResourceManagerType.KV_CACHE_MANAGER: + self._make_kv_cache_manager(_SubclassedKVCacheManager) + }) + + executor.max_beam_width = 4 + with patch("tensorrt_llm._torch.pyexecutor.py_executor.logger" + ) as mock_logger: + executor._warn_if_kv_block_budget_unchecked() + mock_logger.warning_once.assert_called_once() + message = mock_logger.warning_once.call_args.args[0] + self.assertIn("max_beam_width=4", message) + self.assertIn("KV_CACHE_MANAGER", message) + + # Greedy and single-beam requests cannot outgrow the pool by beam + # width, so the warning would be noise. + executor.max_beam_width = 1 + with patch("tensorrt_llm._torch.pyexecutor.py_executor.logger" + ) as mock_logger: + executor._warn_if_kv_block_budget_unchecked() + mock_logger.warning_once.assert_not_called() + + def test_startup_silent_when_every_pool_is_checked(self): + executor = PyExecutor.__new__(PyExecutor) + executor.resource_manager = ResourceManager({ + ResourceManagerType.KV_CACHE_MANAGER: + self._make_kv_cache_manager() + }) + executor.max_beam_width = 4 + + with patch("tensorrt_llm._torch.pyexecutor.py_executor.logger" + ) as mock_logger: + executor._warn_if_kv_block_budget_unchecked() + mock_logger.warning_once.assert_not_called() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index b615a5bf3b5c..af33d59bbb9c 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -1869,12 +1869,14 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], # executor. A test that mirrored the predicate would keep passing if the # production check were deleted. # Everything _validate_request touches besides the beam checks runs after - # them and needs a live engine/sampler, so stub those two out; the beam - # width and beam_width_array branches are reached with the real code. + # them and needs a live engine/sampler/KV cache manager, so stub those out; + # the beam width and beam_width_array branches are reached with the real + # code. executor = types.SimpleNamespace( max_beam_width=request.py_beam_width, _validate_token_id_range=lambda _request: None, sampler=types.SimpleNamespace(validate_request=lambda _request: None), + _validate_request_budget=lambda _request: None, ) validate = functools.partial( PyExecutor._validate_request,