From 5c79ba6f72af7bdac7b4ae516942e39e37f1567a Mon Sep 17 00:00:00 2001
From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Date: Wed, 8 Jul 2026 15:34:25 -0700
Subject: [PATCH] [nvbugs/6426886][fix] Propagate use_host_stop_criteria across
 PP ranks

TestDeepSeekV3Lite::test_nvfp4_4gpus[...pp4-fp8kv=True-attention_dp=True-
cuda_graph=True-overlap_scheduler=True...] hung and MPI_Aborted on PP=4.
Just before the hang, ranks 0/1/2 raised
IndexError: list index out of range in finish_if_reason at
sampler.py:2775 while indexing an empty finish_reasons list.

The last PP rank runs the fast-greedy path with
use_host_stop_criteria=True, which suppresses population of
finish_reasons_host on the SampleStateTorch. _ring_broadcast_sample_state
only forwarded (sample_state.host, py_result_diffs) to non-last PP ranks,
so their sample_state.use_host_stop_criteria stayed at the constructor
default False. Non-last ranks then entered the process_draft_tokens branch
in update_requests and raised IndexError on the empty finish_reasons list.
The exception killed the executor loop, peer ranks stalled in MPI
collectives, HangDetector fired and MPI_Abort was triggered.

Include use_host_stop_criteria in the PP ring payload and restore it on
the receiver. Guarded with hasattr/getattr so sampler flavors without the
field (e.g. SampleStateTRTLLM) are unaffected.

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index fc53e2975e14..10615a0fba60 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -2600,10 +2600,15 @@ def _ring_broadcast_sample_state(
         if not self.dist.is_last_pp_rank:
             # Receive tokens from previous pp rank (w.r.t model forward direction)
             with nvtx_range("recv_sample_state"):
-                sample_state.host, py_result_diffs = self.dist.recv_object(
-                    src=self.dist.prev_pp_rank,
-                    tag=tag,
-                )
+                sample_state.host, py_result_diffs, use_host_stop_criteria = \
+                    self.dist.recv_object(src=self.dist.prev_pp_rank, tag=tag)
+
+            # The last PP rank's fast host-stop path leaves
+            # host.finish_reasons=None; without this flag, non-last ranks would
+            # take the finish_reasons-indexing branch in update_requests and
+            # raise IndexError on the empty list.
+            if hasattr(sample_state, "use_host_stop_criteria"):
+                sample_state.use_host_stop_criteria = use_host_stop_criteria
 
             for request, py_result_diff in zip(requests, py_result_diffs):
                 request.py_result.apply_diff(py_result_diff)
@@ -2621,7 +2626,8 @@ def _ring_broadcast_sample_state(
             self.wait_on_pp_send_handles(self.send_handles, microbatch_id)
             with nvtx_range("send_sample_state"):
                 self.send_handles[microbatch_id] = self.dist.isend_object(
-                    (sample_state.host, py_result_diffs),
+                    (sample_state.host, py_result_diffs,
+                     getattr(sample_state, "use_host_stop_criteria", False)),
                     dest=self.dist.next_pp_rank,
                     tag=tag,
                 )