diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 3e9fc60e7e30..55a5336dedb3 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -2381,21 +2381,50 @@ def _detect_deadlock( "Configure kv_cache_config.host_cache_size, increase " "kv_cache_config.max_tokens, or lower max_batch_size." ) + cause = "The KV cache pool is likely exhausted." + too_long = self._contexts_over_token_budget(pending_ctx) + if too_long: + cause = ( + f"Context request(s) {too_long[:8]} need more than " + f"max_num_tokens={self.max_num_tokens} uncached tokens in one " + "iteration and chunked prefill is off, so they can never be " + "scheduled. Otherwise the KV cache pool is likely exhausted." + ) + remedy = ( + f"{remedy} For the requests that cannot fit, enable chunked prefill " + "(enable_chunked_prefill=True) or increase max_num_tokens." + ) # Logged before the raise so the diagnostic survives even if the # exception is caught or truncated upstream. logger.warning( f"[V2Scheduler] Declaring deadlock after {self._stalled_schedules} passes: " f"{self._summarize_stall(active_requests, pending_ctx, inflight_request_ids, preempted_ids)} " - f"{remedy}" + f"{cause} {remedy}" ) raise RuntimeError( f"V2 scheduler deadlock: {num_gen_candidates} generation and " f"{num_ctx_candidates} context request(s) active but none could " f"be scheduled, suspended or preempted in " - f"{self._stalled_schedules} consecutive attempts. The KV cache " - f"pool is likely exhausted. {remedy}" + f"{self._stalled_schedules} consecutive attempts. {cause} {remedy}" ) + def _contexts_over_token_budget(self, pending_ctx: RequestList) -> list[int]: + """IDs of pending context requests that no iteration can schedule. + + Without chunked prefill a context runs whole, so one whose uncached + prompt (net of reusable cached blocks) plus draft tokens exceeds + ``max_num_tokens`` never fits. Encoder requests are budgeted by their + encoder output instead and are not counted. + """ + if self.chunking_enabled or self.max_num_tokens is None: + return [] + return [ + req.py_request_id + for req in pending_ctx + if req.state_value == self._context_init_state_value + and self._uncached_tokens(req) + get_draft_token_length(req) > self.max_num_tokens + ] + def _summarize_stall( self, active_requests: RequestList, diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_scheduler.py index 1cc4847fe19b..eb60d6eaf5af 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kv_cache_v2_scheduler.py @@ -2504,6 +2504,53 @@ def test_raises_after_repeated_stalls_with_context_candidates(self): with pytest.raises(RuntimeError, match="V2 scheduler deadlock"): sched.schedule_request(reqs, set()) + @staticmethod + def _deadlock_message(sched, reqs) -> str: + sched._DEADLOCK_STALL_ITERS = 3 + for _ in range(2): + sched.schedule_request(reqs, set()) + with pytest.raises(RuntimeError, match="V2 scheduler deadlock") as excinfo: + sched.schedule_request(reqs, set()) + return str(excinfo.value) + + def test_names_a_context_that_cannot_fit_without_chunking(self): + """Without chunked prefill a context longer than max_num_tokens never fits, + so the error names it and adds chunked prefill to the pool remedy.""" + mgr = make_kv_cache_manager(has_cache_tier_below_gpu=False) + sched = make_scheduler(mgr, max_num_tokens=1000) + + message = self._deadlock_message( + sched, [make_ctx_request(7, 1200, is_first_context_chunk=False)] + ) + + assert "Context request(s) [7] need more than max_num_tokens=1000" in message + assert "enable_chunked_prefill=True" in message + assert "host_cache_size" in message # the pool remedy is kept + + def test_a_context_that_fits_after_prefix_reuse_is_not_named(self): + """1200 prompt tokens with 400 reusable leave 800 to compute, which fits.""" + mgr = make_kv_cache_manager( + resize_context_fn=lambda req, n: False, has_cache_tier_below_gpu=False + ) + mgr.probe_context_reuse.side_effect = lambda req: 400 + sched = make_scheduler(mgr, max_num_tokens=1000) + + message = self._deadlock_message(sched, [make_ctx_request(7, 1200)]) + + assert "The KV cache pool is likely exhausted." in message + assert "enable_chunked_prefill" not in message + + def test_only_unchunked_context_init_requests_are_counted(self): + sched = make_scheduler(make_kv_cache_manager(), max_num_tokens=1000) + too_long = make_ctx_request(0, 1200, is_first_context_chunk=False) + encoder = make_encoder_request(1, 1200) + + assert sched._contexts_over_token_budget([too_long, encoder]) == [0] + chunked = make_scheduler( + make_kv_cache_manager(), max_num_tokens=1000, ctx_chunk_config=(None, 64) + ) + assert chunked._contexts_over_token_budget([too_long]) == [] + def test_raises_after_repeated_stalls_with_generation_candidates(self): mgr = make_kv_cache_manager(try_allocate_generation_fn=lambda req: False, can_evict=True) sched = make_scheduler(mgr, max_num_tokens=100)