diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index feb02319a3e3..ab5f52e9d367 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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): diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index f173aabc224a..a6202ef74eb3 100644 --- a/tensorrt_llm/executor/worker.py +++ b/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() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 101edffbe155..1f6c3bbda223 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/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)