[None][fix] Bypass the sync transfer-admission budget on the async disagg path

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3417,6 +3417,14 @@ def _apply_disagg_transfer_admission(
         if self._is_disagg_gen_only_no_context_benchmark():
             return fitting_disagg_gen_init_requests, False
 
+        # Async transfers stay in flight and don't block the iteration, so the
+        # per-iteration budget must not gate them. It is sized from
+        # max_tokens_in_buffer (a buffer knob, not a concurrency cap), so gating
+        # the async path throttles transfers and times out the ctx->gen handoff
+        # at high concurrency. Keep the budget only for the synchronous path.
+        if self._uses_async_disagg_gen_transfer():
+            return fitting_disagg_gen_init_requests, False
+
         controller = self._get_disagg_transfer_admission_controller()
         if not (getattr(self, "kv_cache_transceiver", None)
                 and controller.enabled() and fitting_disagg_gen_init_requests):
@@ -5757,7 +5765,18 @@ def _pad_attention_dp_dummy_request(self):
         if not self.enable_attention_dp:
             return
 
-        assert self.expected_num_active_requests >= len(self.active_requests)
+        # Disagg KV-transfer-error requests can transiently linger in
+        # active_requests (len > the ADP-consensus expected count) before cleanup
+        # drains them. It self-corrects and the padding below keys on the
+        # schedulable count, so warn rather than assert -- a hard assert here
+        # would crash the gen loop on every ADP rank.
+        if self.expected_num_active_requests < len(self.active_requests):
+            logger.warning_once(
+                "expected_num_active_requests "
+                f"({self.expected_num_active_requests}) < active_requests "
+                f"({len(self.active_requests)}); transient disagg-error "
+                "overshoot, continuing",
+                key="adp_dummy_active_overshoot")
         num_active_request = self._count_schedulable_active_requests()
 
         if self._should_skip_dummy_for_benchmark_disagg(num_active_request):

tensorrt_llm/executor/worker.py

@@ -80,8 +80,20 @@ def __init__(
             name="await_response_thread")
 
     def start_thread(self, thread: ManagedThread):
-        if self.engine.can_enqueue_requests() and not thread.is_alive():
-            thread.start()
+        if not self.engine.can_enqueue_requests():
+            return
+        if thread.is_alive():
+            return
+        if thread.ident is not None:
+            # The thread already ran and exited; it can't be restarted, and it
+            # only exits when the engine event loop crashed (stashed in
+            # _event_loop_error). Restarting masks that crash and cascades into a
+            # peer MPI-collective hang, so surface the real error instead.
+            err = getattr(self.engine, "_event_loop_error", None)
+            if err is not None:
+                raise err
+            return
+        thread.start()
 
     def await_response_task(self) -> bool:
         return self._await_response_helper()

tests/unittest/_torch/executor/test_py_executor.py

@@ -508,7 +508,10 @@ def test_uses_global_cp_prompt_length_for_transfer_cost(self):
         assert result.admitted_requests == [request]
         assert result.admitted_transfer_blocks == 3
 
-    def test_apply_reverts_deferred_v2_allocations(self):
+    def test_apply_reverts_deferred_v2_allocations(self, monkeypatch):
+        # The transfer budget only gates the synchronous path now, so pin sync
+        # mode to still exercise the deferral + revert logic.
+        monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1")
         executor = object.__new__(PyExecutor)
         executor.kv_cache_transceiver = Mock()
         executor._is_kv_manager_v2 = True
@@ -540,7 +543,9 @@ def test_apply_missing_controller_preserves_candidates(self):
         assert admitted == [candidate]
         assert not wait_for_progress
 
-    def test_apply_missing_v2_flag_defaults_to_non_v2(self):
+    def test_apply_missing_v2_flag_defaults_to_non_v2(self, monkeypatch):
+        # Budget gating is a synchronous-path behavior; pin sync mode.
+        monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "1")
         executor = object.__new__(PyExecutor)
         executor.kv_cache_transceiver = Mock()
         executor._revert_ctx_alloc = Mock()
@@ -604,6 +609,36 @@ def test_gen_only_no_context_bypasses_transfer_budget(self, monkeypatch):
         assert not wait_for_progress
         executor._revert_ctx_alloc.assert_not_called()
 
+    def test_async_mode_bypasses_transfer_budget(self):
+        # On the async path the max_tokens_in_buffer-derived budget must not gate
+        # transfers (gating it throttled high-concurrency disagg and timed out
+        # the ctx->gen handoff). The class fixture clears the mode env, so this
+        # is the default async path: even with the budget exhausted by an active
+        # transfer, every candidate is admitted with nothing deferred.
+        executor = object.__new__(PyExecutor)
+        executor.kv_cache_transceiver = Mock()
+        executor._is_kv_manager_v2 = True
+        executor._revert_ctx_alloc = Mock()
+        executor.active_requests = [
+            _make_disagg_transfer_request(1, 32, in_progress=True)
+        ]
+        executor._disagg_transfer_admission_controller = DisaggTransferAdmissionController(
+            max_tokens_in_buffer=32, tokens_per_block=32
+        )
+        candidates = [
+            _make_disagg_transfer_request(2, 32),
+            _make_disagg_transfer_request(3, 32),
+        ]
+
+        admitted, wait_for_progress = PyExecutor._apply_disagg_transfer_admission(
+            executor, candidates
+        )
+
+        # Sync mode would admit 0 (budget exhausted); async bypasses entirely.
+        assert admitted == candidates
+        assert not wait_for_progress
+        executor._revert_ctx_alloc.assert_not_called()
+
 
 @pytest.mark.usefixtures("_clear_disagg_transfer_mode_env")
 class TestDisaggTransferIdleProgress:
@@ -1463,6 +1498,22 @@ def test_pad_dummy_added_when_only_to_complete_requests_disagg():
     assert len(stub.active_requests) == 2
 
 
+def test_pad_dummy_tolerates_active_request_overshoot():
+    # A transient overshoot (len(active_requests) > expected_num_active_requests,
+    # when disagg transfer-error requests linger a tick before cleanup) used to
+    # trip a hard assert that crashed the gen loop on every ADP rank. It must now
+    # warn and continue instead of raising.
+    stub = _StubADPExecutor()
+    stub.active_requests = [
+        _make_adp_request(_STATE_GENERATION_IN_PROGRESS),
+        _make_adp_request(_STATE_GENERATION_IN_PROGRESS),
+    ]
+    stub.expected_num_active_requests = 1  # < len(active_requests) == 2
+
+    # Must not raise AssertionError (the pre-fix behavior on overshoot).
+    _run_pad(stub)
+
+
 def test_pad_dummy_added_when_only_wait_scheduler_requests_disagg():
     # Gen-first mode on the context server: DISAGG_CONTEXT_WAIT_SCHEDULER
     # sits BELOW the scheduler's window [CONTEXT_INIT, GENERATION_TO_COMPLETE)