From 5f0f52d15097dbd27d04cb50f6a3f3053dac74a7 Mon Sep 17 00:00:00 2001 From: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Date: Fri, 24 Jul 2026 01:46:35 -0700 Subject: [PATCH 1/2] [None][fix] Surface engine event-loop crash instead of a thread-restart cascade start_thread() restarted the already-finished await-response thread, raising "threads can only be started once" and masking the real cause: the engine event loop crashed (AwaitResponseHelper broadcasts _event_loop_error and returns, so the thread exits). The misleading RuntimeError then cascaded into a peer MPI-collective hang across ranks. Detect the finished thread and re-raise the stashed engine event-loop error instead. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/executor/worker.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) 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() From 6d1195c8c6b9576f032dd8819e7389a4b4705bd9 Mon Sep 17 00:00:00 2001 From: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> Date: Fri, 24 Jul 2026 01:48:47 -0700 Subject: [PATCH 2/2] [None][fix] Don't throttle async disagg gen KV transfers by max_tokens_in_buffer The DisaggTransferAdmissionController derives its in-flight transfer budget from max_tokens_in_buffer (a gather-buffer sizing knob), which for an 8k / high- concurrency config is only ~2 concurrent transfers per rank. That bound is meant for the SYNCHRONOUS transfer path (limit blocking transfers per executor iteration); on the ASYNC path transfers stay in flight and don't block, so applying it there serializes the ctx->gen handoff and backs requests up past kv_transfer_timeout_ms, failing them with "Disagg KV cache transfer error". Bypass the budget on the async path (already capacity-scheduled, so bounded by max_batch_size, matching v1/CPP behavior). Also harden _pad_attention_dp_dummy_request: a transient overshoot where KV-transfer-error requests briefly linger in active_requests (len > expected_num_active_requests) tripped a hard assert that crashed the gen loop on all ADP ranks. Warn and continue; the padding decision keys on the schedulable count. Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 21 ++++++- .../_torch/executor/test_py_executor.py | 55 ++++++++++++++++++- 2 files changed, 73 insertions(+), 3 deletions(-) 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/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)