diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 00037596ad4d..7999cd8031f0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -737,6 +737,11 @@ def __init__( # responses and flushing them at a synchronised point in the executor # loop avoids the mismatch. self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] + # Requests with a buffered terminal response are terminated only after + # the synchronized flush has published that response. This preserves + # queue-backed client delivery while retaining normal termination as + # the single owner of resource and result-queue cleanup. + self._pending_response_terminations: List[LlmRequest] = [] # Same buffer-then-synced-flush pattern as _pending_transfer_responses # above: _handle_responses and _append_iter_stats are reached from # per-rank-divergent gates, so their tp_allgather collectives are @@ -1166,7 +1171,10 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): (request.py_request_id, response)) if self.async_transfer_manager.end_transfer(request): self.active_requests.remove(request) - self._terminate_request(request) + if response: + self._pending_response_terminations.append(request) + else: + self._terminate_request(request) return if self.async_transfer_manager.end_transfer(request): if transfer_failed: @@ -1184,11 +1192,15 @@ def _flush_pending_transfer_responses(self): """ responses = self._pending_transfer_responses self._pending_transfer_responses = [] + requests_to_terminate = self._pending_response_terminations + self._pending_response_terminations = [] if responses or self.enable_attention_dp: # Even when this rank has no responses we must participate in the # collective when ADP is enabled so that the other rank's gather # can complete. self._enqueue_responses(responses) + for request in requests_to_terminate: + self._terminate_request(request) def _handle_kv_transfer_timeouts_synced(self): """ADP-safe drain of the KV-transfer-timeout consensus collective. @@ -4086,7 +4098,8 @@ def _handle_disagg_cache_errors_synced(self): self.is_shutdown = True self._handle_errors(error_msg, requests=None, - charge_budget=False) + charge_budget=False, + fatal_is_collective_aligned=True) return if not (self.enable_attention_dp and self.dist.world_size != 1): @@ -4181,6 +4194,13 @@ def _executor_loop(self): scheduled_batch, iter_stats = self._prepare_and_schedule_batch() if scheduled_batch is None: + # _handle_disagg_cache_errors_synced() can buffer a + # non-fatal response before scheduling observes shutdown. + # Drain it before leaving the loop so the client does not + # wait for its own timeout. Scheduling shutdown is + # model-parallel synchronized, so every ADP rank reaches + # this collective together. + self._flush_pending_transfer_responses() self._event_loop_completed = True break @@ -4192,6 +4212,11 @@ def _executor_loop(self): self.kv_cache_manager.revert_allocate_generation( req) self._finalize_adp_dummy_allocation(False) + # _check_benchmark_disagg_gate() makes this retry decision + # with a model-parallel all-gather. Flush before retrying so + # a response buffered at the top of this pass is not held + # until the benchmark fill gate opens. + self._flush_pending_transfer_responses() continue if not self._is_kv_manager_v2: @@ -4341,7 +4366,6 @@ def _executor_loop(self): self.kv_cache_manager.update_context_resources( scheduled_batch) self._send_kv_async(scheduled_batch.all_requests()) - self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -4373,6 +4397,14 @@ def _executor_loop(self): self._kv_connector_terminate_requests() + # This is the one ADP-synchronized response flush for a + # completed executor-loop pass. It is deliberately outside + # ``if can_queue``: non-fatal errors can be buffered on an + # idle pass, and every ADP rank must enter the tp_gather in + # the same order. Keep it after all per-pass error handling + # so a response is not needlessly delayed to the next pass. + self._flush_pending_transfer_responses() + if self.enable_iter_perf_stats and sample_state is not None: self._process_iter_stats( finished_requests, self.active_requests, @@ -4656,6 +4688,7 @@ def _executor_loop_overlap(self): scheduled_batch, iter_stats = self._prepare_and_schedule_batch() if scheduled_batch is None: + self._flush_pending_transfer_responses() self._event_loop_completed = True break @@ -4667,6 +4700,7 @@ def _executor_loop_overlap(self): self.kv_cache_manager.revert_allocate_generation( req) self._finalize_adp_dummy_allocation(False) + self._flush_pending_transfer_responses() continue if not self._is_kv_manager_v2: @@ -7150,7 +7184,8 @@ def _handle_errors(self, error_msg: Optional[str] = None, *, requests: Optional[List[LlmRequest]] = None, - charge_budget: bool = True) -> None: + charge_budget: bool = True, + fatal_is_collective_aligned: bool = False) -> None: """Fail requests and optionally initiate shutdown on fatal errors. When ``charge_budget`` is True (the default), classifies the error @@ -7186,6 +7221,14 @@ def _handle_errors(self, """ error_responses: Dict[int, LlmResponse] = {} error_msg = error_msg or "error" + multi_rank_adp = (self.enable_attention_dp + and self.dist.world_size != 1) + # ``fatal_is_collective_aligned`` is set only by the synchronized caller + # (_handle_disagg_cache_errors_synced, after its world allreduce), which + # guarantees every ADP rank enters the fatal path in the same collective + # order. Do NOT infer it from ``self._fatal_error is not None``: a + # rank-local setter would then route into a tp_gather while peers are + # elsewhere, recreating the desync this path exists to prevent. budget_fatal = (self._error_budget.consume(error_msg) if charge_budget else False) @@ -7240,14 +7283,17 @@ def _handle_errors(self, client_id=getattr(item.request, 'client_id', None)))) - adp_collective_required = (self.enable_attention_dp - and self.dist.world_size != 1) - if waiting_responses or adp_collective_required: - self._enqueue_responses(waiting_responses) + if not multi_rank_adp: if waiting_responses: + self._enqueue_responses(waiting_responses) logger.info( f"Drained {len(waiting_responses)} queued requests " "on fatal error") + elif fatal_is_collective_aligned: + # Synchronized fatal: every ADP rank enters this drain gather + # together, so issue it even when this rank has no queued + # responses, to stay peer-aligned. + self._enqueue_responses(waiting_responses) failed_requests = (list(self.active_requests) if requests is None else requests) @@ -7265,12 +7311,65 @@ def _handle_errors(self, request for request in self.active_requests if request not in requests ] - self._enqueue_responses(list(error_responses.items())) - for request in failed_requests: - self._terminate_request(request) + defer_termination = False + publish_immediately = False + if is_fatal: + if multi_rank_adp: + if fatal_is_collective_aligned: + # Synchronized fatal: all ranks agree and march through the + # aligned publish/terminate collectives together, so publish + # here instead of diverging. + publish_immediately = True + else: + # A novel rank-local fatal can be observed by one ADP rank + # before its peers reach their next collective. Do not issue + # tp_gather here: it would desynchronize the group exactly + # like a rank-local non-fatal response. Skip the gather and + # raise below; the distributed supervisor tears down peers. + logger.error( + "Skipping rank-local fatal response gather under ADP") + else: + publish_immediately = True + elif multi_rank_adp: + # Under attention DP, _enqueue_responses performs a tp_gather + # that every rank must enter in the same order. Non-fatal errors + # (e.g. a failed disagg KV transfer) are observed by a single + # rank, so enqueueing here would pair this rank's gather against + # a different collective on its peers — typically the per-step + # tp_allgather(batch_size) — corrupting both sides. Buffer the + # responses instead; every rank flushes the buffer together at + # _flush_pending_transfer_responses. + self._pending_transfer_responses.extend(error_responses.items()) + self._pending_response_terminations.extend(failed_requests) + defer_termination = True + else: + # Without multi-rank ADP there is no rank-divergent collective, so + # publish the error immediately. + publish_immediately = True + + if publish_immediately: + # A fatal executor exits before the next normal flush; a non-ADP + # executor has no rank-divergent collective. Both can publish the + # current errors plus any previously buffered terminal responses. + pending = self._pending_transfer_responses + self._pending_transfer_responses = [] + pending_terminations = self._pending_response_terminations + self._pending_response_terminations = [] + self._enqueue_responses(pending + list(error_responses.items())) + for request in pending_terminations: + self._terminate_request(request) + + if not defer_termination: + for request in failed_requests: + self._terminate_request(request) if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_request() + if multi_rank_adp and not fatal_is_collective_aligned: + # Only a novel rank-local fatal raises to force local teardown; + # a synchronized fatal has already published in lockstep and + # tears down every rank together via is_shutdown. + raise self._fatal_error def _terminate_request(self, request: LlmRequest): # Dummy requests don't participate in disagg KV cache transfers, @@ -7370,8 +7469,27 @@ def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): gather_responses = [] if responses_list is not None: for resp in responses_list: - if resp is not None: - gather_responses.extend(resp) + if resp is None: + continue + if not isinstance(resp, (list, tuple)): + # A non-list contribution means the TP collective + # was matched against a *different* collective on + # a peer rank (collectives pair by call order, not + # by type). A common mismatch partner is the + # per-step tp_allgather(batch_size), which makes + # the stray payload an int. The gathered data is + # corrupt on every rank, so fail fast instead of + # raising an opaque TypeError below or silently + # enqueueing garbage responses. + raise RuntimeError( + f"_enqueue_responses: TP collective desync — " + f"gathered a {type(resp).__name__} " + f"instead of a response list. A peer rank " + f"entered a different collective (e.g. " + f"tp_allgather(batch_size)); check for " + f"per-rank-divergent callers of " + f"_enqueue_responses.") + gather_responses.extend(resp) responses = gather_responses logger.debug( f'after gather, rank = {self.dist.rank}, responses = {responses}') diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 08a6f4701761..0294bcd0ecf4 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -362,6 +362,7 @@ def test_peer_buffer_poison_triggers_world_consistent_fatal_cleanup(monkeypatch) "Disagg KV cache transfer buffer is poisoned; process restart is required", requests=None, charge_budget=False, + fatal_is_collective_aligned=True, ) @@ -376,6 +377,8 @@ def test_preclassified_fatal_error_keeps_adp_response_collectives_aligned(): executor.executor_request_queue = Mock() executor.executor_request_queue.get_request_queue.return_value = raw_queue executor.active_requests = [] + executor._pending_transfer_responses = [] + executor._pending_response_terminations = [] executor.gather_all_responses = False executor.enable_attention_dp = True executor.dist = SimpleNamespace(rank=1, world_size=2) @@ -383,7 +386,11 @@ def test_preclassified_fatal_error_keeps_adp_response_collectives_aligned(): executor._terminate_request = Mock() PyExecutor._handle_errors( - executor, "poisoned transfer buffer", requests=None, charge_budget=False + executor, + "poisoned transfer buffer", + requests=None, + charge_budget=False, + fatal_is_collective_aligned=True, ) executor._error_budget.consume.assert_not_called() diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 161ac02eb995..a978efa3f6cc 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -30,7 +30,12 @@ SHUTDOWN_REQUEST_ID, RequestQueueItem, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig +from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequest, + LlmRequestState, + LlmResponse, + SamplingConfig, +) from tensorrt_llm._torch.pyexecutor.py_executor import ( ATTENTION_DP_DUMMY_REQUEST_ID, DisaggTransferAdmissionController, @@ -2474,6 +2479,203 @@ def test_handles_error_on_single_rank(self): assert len(stub.handle_errors_calls) == 1 +class TestPendingTransferResponseFlush: + def test_rank_local_fatal_error_does_not_issue_adp_response_gather(self): + """A lone fatal rank must fail locally rather than desynchronize TP.""" + executor = object.__new__(PyExecutor) + executor._error_budget = Mock() + executor._error_budget.consume.return_value = True + executor._error_budget.budget = 0.0 + executor._fatal_error = None + executor.is_shutdown = False + executor.enable_attention_dp = True + executor.dist = Mock(world_size=2) + executor.waiting_queue = [] + executor.executor_request_queue = Mock() + executor.executor_request_queue.get_request_queue.return_value.empty.return_value = True + executor.gather_all_responses = False + executor.active_requests = [] + executor._pending_transfer_responses = [] + executor._enqueue_responses = Mock() + executor._terminate_request = Mock() + + with pytest.raises(RuntimeError, match="Fatal error: local failure"): + PyExecutor._handle_errors(executor, "local failure") + + executor._enqueue_responses.assert_not_called() + executor.executor_request_queue.enqueue_shutdown_request.assert_called_once_with() + + def test_adp_flush_participates_with_an_empty_response_list(self): + """Ranks without an error still join the synchronized response gather.""" + executor = object.__new__(PyExecutor) + executor._pending_transfer_responses = [] + executor._pending_response_terminations = [] + executor.enable_attention_dp = True + executor._enqueue_responses = Mock() + + PyExecutor._flush_pending_transfer_responses(executor) + + executor._enqueue_responses.assert_called_once_with([]) + + def test_flush_delivers_and_clears_buffered_responses(self): + executor = object.__new__(PyExecutor) + responses = [(7, Mock())] + executor._pending_transfer_responses = responses + executor._pending_response_terminations = [] + executor.enable_attention_dp = False + executor._enqueue_responses = Mock() + + PyExecutor._flush_pending_transfer_responses(executor) + + executor._enqueue_responses.assert_called_once_with(responses) + assert executor._pending_transfer_responses == [] + + def test_rank_zero_keeps_result_queue_until_buffered_error_flushes(self): + """A queued client receives a rank-0 ADP error before cleanup.""" + executor = object.__new__(PyExecutor) + request_id = 7 + response = LlmResponse(request_id=request_id) + response.request_id = request_id + response.client_id = 42 + response.error_msg = "transfer failed" + result_queue = Mock() + executor._pending_transfer_responses = [(request_id, response)] + request = types.SimpleNamespace(py_request_id=request_id) + executor._pending_response_terminations = [request] + executor.enable_attention_dp = False + executor.gather_all_responses = False + executor.dist = Mock(rank=0) + executor.dist.mapping.tp_group = [0] + executor.responses = {} + executor.response_cv = threading.Condition() + executor.result_wait_queues = {request_id: result_queue} + executor._terminate_request = Mock( + side_effect=lambda _: executor.result_wait_queues.pop(request_id) + ) + + PyExecutor._flush_pending_transfer_responses(executor) + + result_queue.put_response.remote.assert_called_once_with(42, response) + assert request_id not in executor.result_wait_queues + executor._terminate_request.assert_called_once_with(request) + + @staticmethod + def _make_executor_loop_stub(): + executor = object.__new__(PyExecutor) + executor.device_id = 0 + profiler = MagicMock() + profiler.__enter__.return_value = Mock() + executor._profiler = Mock(return_value=profiler) + executor.hang_detector = MagicMock() + executor.enable_iter_perf_stats = False + executor._resource_governor_enabled = False + executor._is_kv_manager_v2 = False + executor.is_benchmark_disagg = False + executor._handle_disagg_cache_errors_synced = Mock() + executor._flush_pending_transfer_responses = Mock() + return executor + + @staticmethod + def _patch_executor_loop_cuda(monkeypatch): + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.set_device", + Mock(), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.cudart.cudaSetDevice", + Mock(), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.CUASSERT", + Mock(), + ) + + def test_flushes_before_clean_scheduler_shutdown(self, monkeypatch): + """A response buffered before scheduling must survive a clean exit.""" + executor = self._make_executor_loop_stub() + executor._prepare_and_schedule_batch = Mock(return_value=(None, None)) + self._patch_executor_loop_cuda(monkeypatch) + + PyExecutor._executor_loop(executor) + + executor._flush_pending_transfer_responses.assert_called_once_with() + + def test_flushes_before_benchmark_retry(self, monkeypatch): + """The synchronized benchmark retry path must not strand a response.""" + executor = self._make_executor_loop_stub() + scheduled_batch = types.SimpleNamespace(generation_requests=[]) + executor._prepare_and_schedule_batch = Mock( + side_effect=[(scheduled_batch, None), (None, None)] + ) + executor._check_benchmark_disagg_gate = Mock(return_value=(False, True)) + executor._finalize_adp_dummy_allocation = Mock() + self._patch_executor_loop_cuda(monkeypatch) + + PyExecutor._executor_loop(executor) + + # Once for the retry pass and once for the following clean exit. + assert executor._flush_pending_transfer_responses.call_count == 2 + + def test_idle_pass_has_one_flush(self, monkeypatch): + """An idle pass must not pay an additional response gather.""" + executor = self._make_executor_loop_stub() + scheduled_batch = types.SimpleNamespace( + encoder_requests=[], paused_requests=[], generation_requests=[] + ) + executor._prepare_and_schedule_batch = Mock( + side_effect=[(scheduled_batch, None), (None, None)] + ) + executor._check_benchmark_disagg_gate = Mock(return_value=(True, False)) + executor._terminate_requests = Mock() + executor._pause_requests = Mock() + executor._can_queue = Mock(return_value=(False, None)) + executor.kv_connector_manager = None + executor._revert_gen_alloc = Mock() + executor._finalize_adp_dummy_allocation = Mock() + executor._handle_kv_transfer_timeouts_synced = Mock() + executor.kv_cache_transceiver = None + executor._kv_connector_terminate_requests = Mock() + executor._flush_iter_stats_synced = Mock() + executor.iter_counter = 0 + self._patch_executor_loop_cuda(monkeypatch) + + PyExecutor._executor_loop(executor) + + # One completed idle pass plus the clean-exit drain, not two flushes + # during the idle pass itself. + assert executor._flush_pending_transfer_responses.call_count == 2 + + def test_overlap_flushes_before_clean_scheduler_shutdown(self, monkeypatch): + """The overlap loop must not drop a buffered response on clean exit.""" + executor = self._make_executor_loop_stub() + executor._can_pause_for_rebalance = Mock(return_value=False) + executor._wait_for_model_engine_input_copy = Mock() + executor._prepare_and_schedule_batch = Mock(return_value=(None, None)) + self._patch_executor_loop_cuda(monkeypatch) + + PyExecutor._executor_loop_overlap(executor) + + executor._flush_pending_transfer_responses.assert_called_once_with() + + def test_overlap_flushes_before_benchmark_retry(self, monkeypatch): + """The overlap retry path must not strand a buffered response.""" + executor = self._make_executor_loop_stub() + scheduled_batch = types.SimpleNamespace(generation_requests=[]) + executor._can_pause_for_rebalance = Mock(return_value=False) + executor._wait_for_model_engine_input_copy = Mock() + executor._prepare_and_schedule_batch = Mock( + side_effect=[(scheduled_batch, None), (None, None)] + ) + executor._check_benchmark_disagg_gate = Mock(return_value=(False, True)) + executor._finalize_adp_dummy_allocation = Mock() + self._patch_executor_loop_cuda(monkeypatch) + + PyExecutor._executor_loop_overlap(executor) + + # Once for the retry pass and once for the following clean exit. + assert executor._flush_pending_transfer_responses.call_count == 2 + + class TestOneModelMTPDraftTokenScheduling: """Regression tests for the one-model MTP over-scheduling bug (#16101).