diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 87f74f21466f..6a5377fac19d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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, *, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 197f4280d4a9..4decd379edea 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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() diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 3716f2397635..0f47bbc39559 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/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: diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index 9dca6497110b..7f321cf631b7 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/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")