[https://nvbugs/5977180][fix] size KV cache and admit requests per beam width

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:

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."""

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.

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).

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."""