diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f69041fa7ab8..fb74cca97eea 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3002,10 +3002,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( + sample_state.host, py_result_diffs, use_host_stop_criteria = self.dist.recv_object( src=self.dist.prev_pp_rank, tag=tag, ) + # Propagate the last-rank fast-path decision so update_requests on + # this rank matches how state.host was populated on the last rank. + 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) @@ -3023,7 +3027,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 e1d3b85239b6..dd6e9d290ae3 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -66,7 +66,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backe accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/5945081) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6427411) -accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6428063) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6272673)