tensorrt_llm/_torch/pyexecutor/py_executor.py (5)

849-851: Optional: add an NVTX range around the helper for profiling symmetry.

Stage 2 uses explicit ranges; adding one here helps correlate Stage 3 work in traces.

Apply this diff:

-                        # Handle logits communication between pipeline parallel ranks
-                        self._handle_logits_communication(previous_batch, prev_microbatch_id)
+                        # Handle logits communication between pipeline parallel ranks
+                        with torch.cuda.nvtx.range("_handle_logits_communication"):
+                            self._handle_logits_communication(previous_batch, prev_microbatch_id)

---

1701-1710: Add type hints and tighten the docstring (Google-style) to match repo guidelines.

Clarifies the API and eliminates trailing whitespace in the docstring stanza.

Apply this diff:

-    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
-        """
+    def _handle_logits_communication(
+        self,
+        previous_batch: BatchStatePP,
+        prev_microbatch_id: int,
+    ) -> None:
+        """Handle logits communication between pipeline-parallel ranks.
+
+        If any request in ``previous_batch`` finished and requested logits, the last PP rank
+        sends its ``py_result`` to the first PP rank (which emits responses).
+
+        Args:
+            previous_batch (BatchStatePP): The previous microbatch state.
+            prev_microbatch_id (int): Microbatch identifier corresponding to ``previous_batch``.
+        """

---

1703-1706: Fix ruff W291/W293 (trailing whitespace/blank line with spaces).

There’s trailing whitespace and whitespace-only blank lines in the docstring stanza. The diff in my previous comment removes these.

---

1721-1727: Remove unused assignment before recv to avoid confusion.

finished_reqs_py_results = [r.py_result ...] is immediately overwritten by recv_object(). Drop it to prevent misleading “assigned-but-overwritten” patterns.

Apply this diff:

-        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(
+        if self.dist.is_first_pp_rank and len(finished_reqs):
+            finished_reqs_py_results = self.dist.recv_object(
                 src=self.dist.prev_pp_rank,
                 tag=prev_microbatch_id,
             )

---

1728-1741: Make the transfer order-agnostic (key by request_id) and assert size parity.

Zipping relies on identical ordering of finished_reqs across the first/last PP ranks. If ordering diverges (e.g., due to local state timing, cancellations), results could mis-attach silently. Sending a mapping keyed by py_request_id removes the ordering dependency and allows sanity checks.

Apply this diff:

-            for req, py_result in zip(finished_reqs,
-                                      finished_reqs_py_results):
-                req.py_result = py_result
+            # Order-agnostic attach by request id with sanity check
+            assert isinstance(finished_reqs_py_results, dict), "Expected dict payload keyed by py_request_id"
+            if len(finished_reqs_py_results) != len(finished_reqs):
+                logger.warning(
+                    "Logits payload size mismatch (recv=%d, local finished=%d) for microbatch %d",
+                    len(finished_reqs_py_results), len(finished_reqs), prev_microbatch_id
+                )
+            for req in finished_reqs:
+                try:
+                    req.py_result = finished_reqs_py_results[req.py_request_id]
+                except KeyError:
+                    logger.error(
+                        "Missing logits for request_id=%d in microbatch %d",
+                        req.py_request_id, prev_microbatch_id
+                    )
+                    # Leave as-is; downstream may handle missing py_result.
 
-        elif self.dist.is_last_pp_rank and len(finished_reqs):
+        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],
+                    {r.py_request_id: r.py_result for r in finished_reqs},
                     dest=self.dist.next_pp_rank,
                     tag=prev_microbatch_id)

Follow-up: please verify (in a multi-GPU PP job) that the first/last rank see the same set of finished-and-logits-requesting req_ids per microbatch. If you’d prefer to keep list semantics, at least add an assert len(finished_reqs_py_results) == len(finished_reqs) after the recv to catch divergence early.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

849-851: Nice extraction — call site reads much cleaner now.

Moving the logits PP comms into a dedicated helper improves readability and localizes the concern. No behavioral red flags at the call site.