diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 42acbda93146..363502c6b93b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -424,6 +424,12 @@ 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]] = [] + # 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 + # lifted to _handle_kv_transfer_timeouts_synced / _flush_iter_stats_synced. + self._pending_timed_out_requests: List[LlmRequest] = [] + self._pending_iter_stats_dict: Optional[Dict] = None # ADP dummy role for _pad_attention_dp_dummy_request. Default is gen; # updated from observed request types. self._adp_dummy_is_gen: bool = True @@ -707,6 +713,52 @@ def _flush_pending_transfer_responses(self): # can complete. self._enqueue_responses(responses) + def _handle_kv_transfer_timeouts_synced(self): + """ADP-safe drain of the KV-transfer-timeout consensus collective. + + Lifted out of _handle_responses, which is reached from per-rank- + divergent gates (_process_previous_batch, _handle_executed_batch). + Non-ADP runs handle timeouts inline; the buffer is empty here. + """ + if not (self.enable_attention_dp and self.dist.world_size != 1): + return + timed_out = self._pending_timed_out_requests + self._pending_timed_out_requests = [] + any_timed_out = any(self.dist.tp_allgather(bool(timed_out))) + if any_timed_out: + self._handle_errors(error_msg="Request timed out (KV transfer)", + requests=timed_out) + + def _flush_iter_stats_synced(self): + """ADP-safe drain of the TLLM_METRICS_ALL_RANKS dict-gather collective. + + Lifted out of _append_iter_stats (same divergent-gate problem as + _handle_kv_transfer_timeouts_synced). Under ADP every rank sends + either the local dict or ``None`` every iter so the gather stays in + lockstep with peer collectives. When the gate doesn't hold, + _append_iter_stats takes its legacy path and never populates the + buffer. + """ + tp_size = self.dist.tp_size + gather_all_ranks = os.environ.get("TLLM_METRICS_ALL_RANKS", "0") == "1" + if not (gather_all_ranks and self.enable_iter_perf_stats and tp_size > 1 + and self.enable_attention_dp): + return + local_dict = self._pending_iter_stats_dict + self._pending_iter_stats_dict = None + gathered = self.dist.tp_allgather(local_dict) + if self.dist.tp_rank != 0: + return + rank_dicts = [d for d in gathered if d is not None] + if not rank_dicts: + return + with self.stats_lock: + for d in rank_dicts: + self.stats.append(("per_rank_dict", d)) + cap = self.max_stats_len * tp_size + if len(self.stats) > cap: + del self.stats[:len(self.stats) - cap] + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1432,23 +1484,9 @@ def _append_iter_stats(self, append_kv_cache_iteration_stats(local_dict, self._latest_kv_iter_stats) local_dict["rank"] = self.dist.tp_rank - - gathered = self.dist.tp_allgather(local_dict) - - if self.dist.tp_rank == 0: - with self.stats_lock: - # Wrap as ("per_rank_dict", dict) so the serializer can - # distinguish from the legacy (stats, req_stats, kv) tuple. - # Trim before appending so every rank's entry for the - # same iteration lands in the buffer atomically; evicting - # during the append loop would let partial iterations - # survive at the head of self.stats. - cap = self.max_stats_len * tp_size - overflow = max(0, len(self.stats) + len(gathered) - cap) - if overflow: - del self.stats[:overflow] - for d in gathered: - self.stats.append(("per_rank_dict", d)) + # Buffer for _flush_iter_stats_synced; in-place tp_allgather would + # desync (reached from per-rank-divergent gates). + self._pending_iter_stats_dict = local_dict return # Legacy path: rank-0-only (single-rank or iter stats disabled). @@ -2012,6 +2050,13 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): executed_batch.microbatch_id % self.dist.pp_size, ) + # Drain the synced-collective buffers outside the per-rank-divergent + # ``executed_batch is not None`` branch. handle_executed_batches + # invokes this helper the same number of times on every rank + # (broadcast count), so the flushes are rank-symmetric. + self._handle_kv_transfer_timeouts_synced() + self._flush_iter_stats_synced() + @nvtx_range("wait_on_pp_send_handles") def wait_on_pp_send_handles(self, send_handles, microbatch_id): if send_handles[microbatch_id] is not None: @@ -2567,6 +2612,10 @@ def _executor_loop(self): if self.enable_kv_cache_events: self._add_kv_cache_events() + # Drain timeout buffer outside ``if can_queue`` so the synced + # collective fires every iter regardless of future restructuring. + self._handle_kv_transfer_timeouts_synced() + if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() @@ -2582,6 +2631,9 @@ def _executor_loop(self): sample_state=sample_state, iter_stats=iter_stats, iter_start_time=iter_start_time)) + # Same lockstep guarantee for iter-stats; no-op when + # TLLM_METRICS_ALL_RANKS=0. + self._flush_iter_stats_synced() self.iter_counter += 1 @@ -2890,6 +2942,12 @@ def _executor_loop_overlap(self): else: self._enqueue_responses([]) + # Drain buffers from the (per-rank-divergent) + # _process_previous_batch above; rank-symmetric companion to + # the _enqueue_responses([]) call in the else branch. + self._handle_kv_transfer_timeouts_synced() + self._flush_iter_stats_synced() + # Call set_exclude_last_generation_logits after _process_previous_batch. # If set before, the response of a request may be incorrect, as it will # use the wrong indices for generation logits when streaming is enabled. @@ -4436,11 +4494,9 @@ def _handle_responses(self): for request in requests_to_terminate: self._terminate_request(request) if self.enable_attention_dp and self.dist.world_size != 1: - any_timed_out = any(self.dist.tp_allgather( - bool(timed_out_requests))) - if any_timed_out: - self._handle_errors(error_msg="Request timed out (KV transfer)", - requests=timed_out_requests) + # Buffer for _handle_kv_transfer_timeouts_synced; in-place + # tp_allgather would desync (reached from per-rank-divergent gates). + self._pending_timed_out_requests.extend(timed_out_requests) else: for req in timed_out_requests: self._handle_errors(