diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 73184a0d046d..bcf82c176e84 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3335,10 +3335,14 @@ 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) + + # Propagate use_host_stop_criteria: the last rank's fast host-stop + # path leaves host.finish_reasons=None, so non-last ranks must + # know to skip the finish_reasons indexing in update_requests. + 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) @@ -3356,7 +3360,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, ) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 8acb9bc37202..ddcfc421ca61 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -399,7 +399,6 @@ unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741) unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5605741) unittest/llmapi/test_additional_model_outputs.py -m "gpu2" SKIP (https://nvbugs/6428091) unittest/llmapi/test_llm.py::test_generate_with_detokenization_stop_words_streaming[/scratch.trt_llm_data/llm-models/gemma/gemma-3-1b-it] SKIP (https://nvbugs/6566772) -unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2" SKIP (https://nvbugs/6428092) unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp2[False-False-True] SKIP (https://nvbugs/6432826) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[None] SKIP (https://nvbugs/6162504) unittest/llmapi/test_llm_pytorch.py::test_gqa_nemo_lora[cuda_graph_config0] SKIP (https://nvbugs/6162504)