[#21826][fix] Name the context that can never fit in the V2 deadlock error

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,

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)