From 35ecb1bbb0b1a58d2af2808df30d39acc88fc969 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Fri, 22 Aug 2025 07:45:51 +0000 Subject: [PATCH 1/2] Refactored the logits communication. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 75 +++++++++++-------- 1 file changed, 43 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 453434d9d6b9..d536ce4c6b84 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -846,38 +846,8 @@ def _executor_loop_pp(self): self._handle_canceled_requests() - # If logits were requested last PP rank has to send to first PP rank (who sends responses) the - # logits of the requests that have finished. - # NOTE: If the rank processing the logits ever becomes the same as - # the rank sending the responses, this code can be removed. - finished_reqs = [ - r for r in previous_batch.sample_state. - scheduled_requests.all_requests() - if r.state == LlmRequestState.GENERATION_COMPLETE - and (r.py_return_context_logits - or r.py_return_generation_logits) - ] - if self.dist.is_first_pp_rank and len(finished_reqs): - finished_reqs_py_results = [ - r.py_result for r in finished_reqs - ] - finished_reqs_py_results = self.dist.recv_object( - src=self.dist.prev_pp_rank, - tag=prev_microbatch_id, - ) - for req, py_result in zip(finished_reqs, - finished_reqs_py_results): - req.py_result = py_result - - elif self.dist.is_last_pp_rank and len(finished_reqs): - if self.send_handles[ - prev_microbatch_id] is not None: - self.send_handles[prev_microbatch_id].wait() - self.send_handles[ - prev_microbatch_id] = self.dist.isend_object( - [r.py_result for r in finished_reqs], - dest=self.dist.next_pp_rank, - tag=prev_microbatch_id) + # Handle logits communication between pipeline parallel ranks + self._handle_logits_communication(previous_batch, prev_microbatch_id) finished_requests = self._handle_responses() previous_scheduled_batch = previous_batch.sample_state.scheduled_requests @@ -1728,6 +1698,47 @@ def _terminate_ctx_finished_requests(self): self._terminate_request(request) self.ctx_in_transmission_requests.remove(request) + def _handle_logits_communication(self, previous_batch, prev_microbatch_id): + """Handle logits communication between pipeline parallel ranks. + + If logits were requested, the last PP rank sends to the first PP rank (who sends responses) + the logits of the requests that have finished. + + Args: + previous_batch: The previous batch state + prev_microbatch_id: The microbatch ID for the previous batch + """ + # NOTE: If the rank processing the logits ever becomes the same as + # the rank sending the responses, this code can be removed. + finished_reqs = [ + r for r in previous_batch.sample_state. + scheduled_requests.all_requests() + if r.state == LlmRequestState.GENERATION_COMPLETE + and (r.py_return_context_logits + or r.py_return_generation_logits) + ] + if self.dist.is_first_pp_rank and len(finished_reqs): + finished_reqs_py_results = [ + r.py_result for r in finished_reqs + ] + finished_reqs_py_results = self.dist.recv_object( + src=self.dist.prev_pp_rank, + tag=prev_microbatch_id, + ) + for req, py_result in zip(finished_reqs, + finished_reqs_py_results): + req.py_result = py_result + + elif self.dist.is_last_pp_rank and len(finished_reqs): + if self.send_handles[ + prev_microbatch_id] is not None: + self.send_handles[prev_microbatch_id].wait() + self.send_handles[ + prev_microbatch_id] = self.dist.isend_object( + [r.py_result for r in finished_reqs], + dest=self.dist.next_pp_rank, + tag=prev_microbatch_id) + def _await_any_response(self, timeout: Optional[float] = None ) -> List[LlmResponse]: From 6f9df1fa8cd36e60bd5e699451a9453ed3358049 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Fri, 22 Aug 2025 08:30:27 +0000 Subject: [PATCH 2/2] Formatting. Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 38 ++++++++----------- 1 file changed, 16 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d536ce4c6b84..4ddce877db6b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -846,8 +846,8 @@ def _executor_loop_pp(self): self._handle_canceled_requests() - # Handle logits communication between pipeline parallel ranks - self._handle_logits_communication(previous_batch, prev_microbatch_id) + self._handle_logits_communication( + previous_batch, prev_microbatch_id) finished_requests = self._handle_responses() previous_scheduled_batch = previous_batch.sample_state.scheduled_requests @@ -1700,10 +1700,10 @@ def _terminate_ctx_finished_requests(self): def _handle_logits_communication(self, previous_batch, prev_microbatch_id): """Handle logits communication between pipeline parallel ranks. - - If logits were requested, the last PP rank sends to the first PP rank (who sends responses) + + If logits were requested, the last PP rank sends to the first PP rank (who sends responses) the logits of the requests that have finished. - + Args: previous_batch: The previous batch state prev_microbatch_id: The microbatch ID for the previous batch @@ -1711,33 +1711,27 @@ def _handle_logits_communication(self, previous_batch, prev_microbatch_id): # NOTE: If the rank processing the logits ever becomes the same as # the rank sending the responses, this code can be removed. finished_reqs = [ - r for r in previous_batch.sample_state. - scheduled_requests.all_requests() - if r.state == LlmRequestState.GENERATION_COMPLETE - and (r.py_return_context_logits - or r.py_return_generation_logits) + r for r in + previous_batch.sample_state.scheduled_requests.all_requests() + if r.state == LlmRequestState.GENERATION_COMPLETE and ( + r.py_return_context_logits or r.py_return_generation_logits) ] if self.dist.is_first_pp_rank and len(finished_reqs): - finished_reqs_py_results = [ - r.py_result for r in finished_reqs - ] + finished_reqs_py_results = [r.py_result for r in finished_reqs] finished_reqs_py_results = self.dist.recv_object( src=self.dist.prev_pp_rank, tag=prev_microbatch_id, ) - for req, py_result in zip(finished_reqs, - finished_reqs_py_results): + for req, py_result in zip(finished_reqs, finished_reqs_py_results): req.py_result = py_result elif self.dist.is_last_pp_rank and len(finished_reqs): - if self.send_handles[ - prev_microbatch_id] is not None: + if self.send_handles[prev_microbatch_id] is not None: self.send_handles[prev_microbatch_id].wait() - self.send_handles[ - prev_microbatch_id] = self.dist.isend_object( - [r.py_result for r in finished_reqs], - dest=self.dist.next_pp_rank, - tag=prev_microbatch_id) + self.send_handles[prev_microbatch_id] = self.dist.isend_object( + [r.py_result for r in finished_reqs], + dest=self.dist.next_pp_rank, + tag=prev_microbatch_id) def _await_any_response(self, timeout: Optional[float] = None