[None][fix] PyExecutor: non-fatal KV-cache block-pool exhaustion guard

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -72,7 +72,8 @@
                             derive_attention_dp_per_rank_request_cap,
                             get_from_waiting_queue, merge_requests)
 from .resource_manager import (KVCacheManagerV2, ResourceManager,
-                               ResourceManagerType, request_context)
+                               ResourceManagerType,
+                               kv_exhaustion_nonfatal_enabled, request_context)
 from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState,
                       SampleStateTensors, TRTLLMSampler)
 from .scheduler import (RequestScheduler, ScheduledRequests,
@@ -438,6 +439,10 @@ def __init__(
             and self.kv_cache_manager.enable_partial_reuse)
 
         self.max_input_len = max_input_len
+        # Requests dropped by the KV-cache exhaustion guard
+        # (TRTLLM_KV_EXHAUSTION_NONFATAL=1), failed one iteration later so no
+        # in-flight overlap forward still references their blocks.
+        self._pending_kv_rejected: List[LlmRequest] = []
         # _executor_loop private data
         self.max_num_active_requests = model_engine.get_max_num_sequences()
         # nvbug-6133201: under attention DP, tighten the per-rank request
@@ -2766,6 +2771,10 @@ def _executor_loop(self):
                 if self._is_kv_manager_v2 and self._can_pause_for_rebalance():
                     self._maybe_rebalance_kv_pools()
 
+                # Fail any requests the KV-cache exhaustion guard dropped last
+                # iteration (now safe: their in-flight forward has been
+                # consumed). Before scheduling so they are not re-picked.
+                self._drain_kv_rejected_requests()
                 scheduled_batch, iter_stats = self._prepare_and_schedule_batch()
                 self._handle_control_request()
 
@@ -2797,6 +2806,7 @@ def _executor_loop(self):
                     self._handle_dynamic_draft_len(scheduled_batch)
 
                     self.resource_manager.prepare_resources(scheduled_batch)
+                    self._collect_kv_rejected_requests(scheduled_batch)
 
                 if self.kv_connector_manager:
                     self.kv_connector_manager.handle_metadata()
@@ -3131,6 +3141,10 @@ def _executor_loop_overlap(self):
                 if self._is_kv_manager_v2 and self._can_pause_for_rebalance():
                     self._maybe_rebalance_kv_pools()
 
+                # Fail any requests the KV-cache exhaustion guard dropped last
+                # iteration (now safe: their in-flight forward has been
+                # consumed). Before scheduling so they are not re-picked.
+                self._drain_kv_rejected_requests()
                 scheduled_batch, iter_stats = self._prepare_and_schedule_batch()
                 self._handle_control_request()
 
@@ -3173,6 +3187,7 @@ def _executor_loop_overlap(self):
                     self._handle_dynamic_draft_len(scheduled_batch)
 
                     self.resource_manager.prepare_resources(scheduled_batch)
+                    self._collect_kv_rejected_requests(scheduled_batch)
 
                 if self.kv_connector_manager:
                     self.kv_connector_manager.handle_metadata()
@@ -4474,6 +4489,40 @@ def _update_requests(self,
             logger.error(f"Encountered an error in sampling: {error_msg}")
             self._handle_errors(error_msg)
 
+    def _collect_kv_rejected_requests(
+            self, scheduled_batch: ScheduledRequests) -> None:
+        """Move requests dropped by the KV-cache exhaustion guard out of the
+        just-prepared batch into the pending queue.  They are failed at the top
+        of the next iteration (see _drain_kv_rejected_requests) rather than now,
+        so any in-flight overlap forward that referenced their blocks has been
+        consumed before the blocks are freed."""
+        rejected = scheduled_batch.kv_cache_rejected_requests
+        if rejected:
+            self._pending_kv_rejected.extend(rejected)
+            scheduled_batch.kv_cache_rejected_requests = []
+
+    def _drain_kv_rejected_requests(self) -> None:
+        """Fail requests dropped by the KV-cache exhaustion guard, reusing the
+        request-scoped error path (charge_budget=False) so a saturated block
+        pool degrades to a graceful per-request rejection instead of the C++
+        allocation assert killing the executor thread.
+
+        Only runs when the guard is enabled (TRTLLM_KV_EXHAUSTION_NONFATAL=1) --
+        a process-wide, rank-symmetric switch.  When enabled it is entered
+        unconditionally each iteration, even with nothing to reject: under
+        attention DP the response gather inside _handle_errors must be entered
+        by every rank in lockstep (an empty `requests` list is a paired no-op).
+        """
+        if not kv_exhaustion_nonfatal_enabled():
+            return
+        rejected = self._pending_kv_rejected
+        self._pending_kv_rejected = []
+        self._handle_errors(
+            "KV cache block pool exhausted under load; request rejected to keep "
+            "the engine running (TRTLLM_KV_EXHAUSTION_NONFATAL=1).",
+            requests=rejected,
+            charge_budget=False)
+
     def _handle_errors(self,
                        error_msg: Optional[str] = None,
                        *,

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -80,6 +80,91 @@
     int]]  # window_size -> (blocks_in_primary_pool, blocks_in_secondary_pool)
 
 
+# ---------------------------------------------------------------------------
+# Non-fatal KV-cache block-pool exhaustion guard
+# ---------------------------------------------------------------------------
+# Under MAX_UTILIZATION the capacity scheduler deliberately over-subscribes the
+# KV-cache block pool, relying on eviction/pause to recover.  When the pool is
+# truly saturated the C++ block manager raises a fatal assertion from inside
+# allocation (WindowBlockManager::allocateBlock / LRUEvictionPolicy::getFreeBlock)
+# which propagates out of `prepare_resources`, escapes the PyExecutor event loop
+# and kills the worker thread on every rank: the engine stops stepping but the
+# process stays up, so health/metrics go silent until the pod is killed.
+#
+# With TRTLLM_KV_EXHAUSTION_NONFATAL=1 we instead reject the request(s) that
+# could not be allocated and keep stepping -- the same opt-in, log-and-continue
+# philosophy as TRTLLM_NAN_GUARD_NONFATAL.  This preserves MAX_UTILIZATION's high
+# concurrency for the common (short-request) traffic while degrading the rare
+# over-subscription event from a pod-wide outage to a single graceful rejection.
+_KV_CACHE_EXHAUSTION_SIGNATURES = (
+    "No free blocks left",  # WindowBlockManager::allocateBlock (kvCacheManager.cpp)
+    "No free block found",  # LRUEvictionPolicy::getFreeBlock (evictionPolicy.cpp)
+    "Can't allocate new blocks",  # WindowBlockManager::allocateBlock / addToken
+)
+
+
+def is_kv_cache_exhaustion_error(exc: BaseException) -> bool:
+    """Whether ``exc`` is a KV-cache block-pool exhaustion raised by the C++
+    block manager and surfaced through nanobind as a generic exception.
+
+    Matched by message signature because the C++ assertions (TLLM_THROW /
+    TLLM_CHECK_WITH_INFO) do not map to a dedicated Python exception type.
+    Unrelated failures (CUDA OOM, request validation) must NOT match, so they
+    keep their existing fatal behavior.
+    """
+    msg = str(exc)
+    return any(sig in msg for sig in _KV_CACHE_EXHAUSTION_SIGNATURES)
+
+
+def kv_exhaustion_nonfatal_enabled() -> bool:
+    """True when TRTLLM_KV_EXHAUSTION_NONFATAL=1 opts into non-fatal handling of
+    KV-cache block-pool exhaustion (default off -> fatal, as before)."""
+    return os.environ.get("TRTLLM_KV_EXHAUSTION_NONFATAL", "0") == "1"
+
+
+def _record_kv_exhausted_requests(scheduled_batch: ScheduledRequests,
+                                  requests: List[LlmRequest],
+                                  exc: BaseException) -> None:
+    """Remove ``requests`` (which could not be allocated because the block pool
+    is saturated) from ``scheduled_batch`` and queue them for graceful rejection
+    by the executor.
+
+    Blocks are intentionally NOT freed here: with overlap scheduling the
+    offending request may still be referenced by the previous batch's in-flight
+    forward, so the executor frees + fails them only once no kernel can
+    reference their blocks (see PyExecutor._drain_kv_rejected_requests).  The
+    failing allocation threw on a pre-allocation capacity check, so no partial
+    block was left behind and the manager stays consistent for the rest of the
+    step.
+    """
+    if not requests:
+        return
+    reject_ids = {req.py_request_id for req in requests}
+    scheduled_batch.context_requests_chunking = [
+        r for r in scheduled_batch.context_requests_chunking
+        if r.py_request_id not in reject_ids
+    ]
+    scheduled_batch.context_requests_last_chunk = [
+        r for r in scheduled_batch.context_requests_last_chunk
+        if r.py_request_id not in reject_ids
+    ]
+    scheduled_batch.generation_requests = [
+        r for r in scheduled_batch.generation_requests
+        if r.py_request_id not in reject_ids
+    ]
+    already = {
+        req.py_request_id
+        for req in scheduled_batch.kv_cache_rejected_requests
+    }
+    scheduled_batch.kv_cache_rejected_requests.extend(
+        req for req in requests if req.py_request_id not in already)
+    logger.error(
+        "KV cache block pool exhausted (%s); rejecting %d request(s) %s to keep "
+        "the engine running (TRTLLM_KV_EXHAUSTION_NONFATAL=1).",
+        str(exc).strip(), len(requests),
+        [req.py_request_id for req in requests])
+
+
 @dataclass
 class PoolConfiguration:
     """Configuration of a single KV pool.
@@ -1028,21 +1113,39 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                         batch_ctx_requests.append(req)
                         self._active_sequences.add(req.py_request_id)
 
-            if batch_request_infos:
-                self.impl.add_sequence_batch(batch_request_infos,
-                                             batch_llm_requests)
-                for req in batch_ctx_requests:
-                    for _ in range(self.num_extra_kv_tokens):
-                        self.impl.add_token(req.py_request_id)
-                    for _ in range(get_draft_token_length(req)):
-                        self.impl.add_token(req.py_request_id)
+            # TRTLLM_KV_EXHAUSTION_NONFATAL=1: reject the request(s) that hit a
+            # saturated block pool instead of letting the C++ allocation assert
+            # kill the executor thread (see is_kv_cache_exhaustion_error).
+            nonfatal = kv_exhaustion_nonfatal_enabled()
 
-                    if self.kv_connector_manager is not None:
-                        block_ids = self.get_cache_indices(req)
-                        self.kv_connector_manager.update_state_after_alloc(
-                            req, block_ids)
-
-            for req in scheduled_batch.generation_requests:
+            if batch_request_infos:
+                try:
+                    self.impl.add_sequence_batch(batch_request_infos,
+                                                 batch_llm_requests)
+                    for req in batch_ctx_requests:
+                        for _ in range(self.num_extra_kv_tokens):
+                            self.impl.add_token(req.py_request_id)
+                        for _ in range(get_draft_token_length(req)):
+                            self.impl.add_token(req.py_request_id)
+
+                        if self.kv_connector_manager is not None:
+                            block_ids = self.get_cache_indices(req)
+                            self.kv_connector_manager.update_state_after_alloc(
+                                req, block_ids)
+                except Exception as e:
+                    if not (nonfatal and is_kv_cache_exhaustion_error(e)):
+                        raise
+                    # Onboarding this step's new context requests overflowed the
+                    # pool.  Onboarding only touches first-chunk context requests
+                    # (no in-flight forward references them yet), so reject the
+                    # whole new-admission batch and let the running generation
+                    # requests below proceed.
+                    _record_kv_exhausted_requests(scheduled_batch,
+                                                   list(batch_ctx_requests), e)
+
+            # Iterate a snapshot: _record_kv_exhausted_requests mutates
+            # scheduled_batch.generation_requests on rejection.
+            for req in list(scheduled_batch.generation_requests):
                 if self.mapping.has_cp_helix():
                     # Distribute the decode blocks across CP ranks in a round-robin manner.
                     decode_block_id = (req.py_decoding_iter -
@@ -1054,10 +1157,20 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                         req.py_helix_is_inactive_rank = True
                         # Skip allocating KV cache at decode for inactive helix ranks.
                         continue
-                draft_len = get_draft_token_length(req)
-                self.impl.add_token(req.py_request_id)
-                for _ in range(max(draft_len, self._kv_reserve_draft_tokens)):
+                try:
+                    draft_len = get_draft_token_length(req)
                     self.impl.add_token(req.py_request_id)
+                    for _ in range(max(draft_len,
+                                       self._kv_reserve_draft_tokens)):
+                        self.impl.add_token(req.py_request_id)
+                except Exception as e:
+                    if not (nonfatal and is_kv_cache_exhaustion_error(e)):
+                        raise
+                    # A running request cannot grow into the saturated pool.
+                    # Reject just this one; the executor frees it once no
+                    # in-flight forward references its blocks.
+                    _record_kv_exhausted_requests(scheduled_batch, [req], e)
+                    continue
 
             # prefill and generation kernels wait for scheduled offload/onboard/partial copy work before launching
             self.impl.refresh_blocks()

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py

@@ -78,12 +78,17 @@ class ScheduledRequests:
     """Requests that are in the generation phase."""
     paused_requests: RequestList
     """Requests that are paused."""
+    kv_cache_rejected_requests: RequestList
+    """Requests dropped this step because the KV-cache block pool was exhausted
+    (only populated when TRTLLM_KV_EXHAUSTION_NONFATAL=1).  Drained and failed
+    by the executor; see PyExecutor._drain_kv_rejected_requests."""
 
     def __init__(self):
         self.context_requests_chunking: RequestList = []
         self.context_requests_last_chunk: RequestList = []
         self.generation_requests: RequestList = []
         self.paused_requests: RequestList = []
+        self.kv_cache_rejected_requests: RequestList = []
 
     @property
     def is_generation_only(self) -> bool:

tests/unittest/_torch/executor/test_resource_manager.py

@@ -35,6 +35,32 @@
 sys.path.append(str(root_dir / "tests" / "integration"))
 
 
+def test_is_kv_cache_exhaustion_error_classification():
+    """The non-fatal KV-exhaustion guard must recognize the fatal block-pool
+    asserts thrown by the C++ block manager (kvCacheManager.cpp /
+    evictionPolicy.cpp) and surfaced through nanobind as generic exceptions,
+    while leaving unrelated errors alone so they still propagate fatally."""
+    from tensorrt_llm._torch.pyexecutor.resource_manager import \
+        is_kv_cache_exhaustion_error
+
+    # The three real assert signatures (kvCacheManager.cpp:2480 / :2810,
+    # evictionPolicy.cpp:160).
+    assert is_kv_cache_exhaustion_error(
+        RuntimeError("Can't allocate new blocks for window size 262152. "
+                     "No free blocks left."))
+    assert is_kv_cache_exhaustion_error(
+        RuntimeError("No free block found. This shouldn't happen!"))
+    assert is_kv_cache_exhaustion_error(
+        RuntimeError("Can't allocate new blocks. No free blocks left."))
+
+    # Unrelated errors must NOT be swallowed by the guard: a genuine CUDA OOM
+    # or a request-validation error should still surface as before.
+    assert not is_kv_cache_exhaustion_error(
+        ValueError("Token ID out of range"))
+    assert not is_kv_cache_exhaustion_error(
+        RuntimeError("CUDA out of memory"))
+
+
 class TestResourceManager(unittest.TestCase):
     CPP_RESOURCES_DIR = os.path.join(str(root_dir), "cpp", "tests", "resources")
     CPP_DATA_DIR = os.path.join(CPP_RESOURCES_DIR, "data")