[https://nvbugs/6428057][fix] Propagate host stop criteria across PP ranks

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -2600,10 +2600,16 @@ 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,
                 )
+            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)
@@ -2619,9 +2625,15 @@ def _ring_broadcast_sample_state(
                 py_result_diffs.append(diff)
                 request.py_result.reset_diff()
             self.wait_on_pp_send_handles(self.send_handles, microbatch_id)
+            # Device finish reasons are absent on the host-stop fast path, so
+            # every PP rank must select the same request-update branch.
             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,
                 )

tests/integration/test_lists/waives.txt

@@ -54,7 +54,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6426847)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[ep4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6402058)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6278337)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6428057)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6278337)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=True] SKIP (https://nvbugs/6278337)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/6388153)
@@ -70,7 +69,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpu
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=False-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6427411)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6427411)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=CUTLASS-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] SKIP (https://nvbugs/6388363)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-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/6428087)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTEDSL-mtp_nextn=0-tp2pp2-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=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)
@@ -422,7 +420,6 @@ triton_server/test_triton.py::test_python_bls_unit_tests[python-bls-unit-tests]
 unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py::test_prefill_varlen[varlen_hd512_overlap] SKIP (https://nvbugs/6426860)
 unittest/_torch/misc/test_share_tensor.py::TestShareTensor::test_share_tensor_different_dtypes SKIP (https://nvbugs/6418021)
 unittest/_torch/modules/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912)
-unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124)
 unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py -m "part0" SKIP (https://nvbugs/6372711)
 unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py -m "part1" SKIP (https://nvbugs/6426852)
 unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py::test_llm_partial_update_weights_nvfp4[auto-Qwen3/Qwen3-8B] SKIP (https://nvbugs/6372690)

tests/unittest/_torch/executor/test_py_executor.py

@@ -702,6 +702,74 @@ def test_pp_schedule_restores_propagated_gate_decision(self):
         assert wait_for_progress
 
 
+class TestSampleStateBroadcastPP:
+    @staticmethod
+    def _make_executor(*, is_last_pp_rank: bool, is_second_last_pp_rank: bool) -> PyExecutor:
+        executor = object.__new__(PyExecutor)
+        executor.dist = Mock(
+            is_last_pp_rank=is_last_pp_rank,
+            is_second_last_pp_rank=is_second_last_pp_rank,
+            prev_pp_rank=3,
+            next_pp_rank=1,
+        )
+        executor.executed_batch_response_queue = Mock()
+        executor.send_handles = {}
+        executor.wait_on_pp_send_handles = Mock()
+        return executor
+
+    @staticmethod
+    def _make_batch(*, use_host_stop_criteria: bool) -> types.SimpleNamespace:
+        py_result = Mock()
+        py_result.get_diff.return_value = {"tokens": [42]}
+        request = types.SimpleNamespace(py_result=py_result)
+        sample_state = types.SimpleNamespace(
+            requests=[request],
+            host=object(),
+            use_host_stop_criteria=use_host_stop_criteria,
+        )
+        batch = types.SimpleNamespace(microbatch_id=0, sample_state=sample_state)
+        return batch
+
+    def test_ring_preserves_host_stop_criteria(self) -> None:
+        source = self._make_executor(is_last_pp_rank=True, is_second_last_pp_rank=False)
+        source_batch = self._make_batch(use_host_stop_criteria=True)
+        PyExecutor._ring_broadcast_sample_state(source, source_batch)
+        source_payload = source.dist.isend_object.call_args.args[0]
+        assert source_payload[2] is True
+
+        relay = self._make_executor(is_last_pp_rank=False, is_second_last_pp_rank=False)
+        relay_batch = self._make_batch(use_host_stop_criteria=False)
+        relay.dist.recv_object.return_value = source_payload
+        PyExecutor._ring_broadcast_sample_state(relay, relay_batch)
+        relay_payload = relay.dist.isend_object.call_args.args[0]
+        assert relay_batch.sample_state.use_host_stop_criteria is True
+        assert relay_payload[2] is True
+
+        terminal = self._make_executor(is_last_pp_rank=False, is_second_last_pp_rank=True)
+        terminal_batch = self._make_batch(use_host_stop_criteria=False)
+        terminal.dist.recv_object.return_value = relay_payload
+        PyExecutor._ring_broadcast_sample_state(terminal, terminal_batch)
+        assert terminal_batch.sample_state.use_host_stop_criteria is True
+        terminal.dist.isend_object.assert_not_called()
+
+    def test_ring_defaults_for_other_sampler_states(self) -> None:
+        source = self._make_executor(is_last_pp_rank=True, is_second_last_pp_rank=False)
+        source_batch = self._make_batch(use_host_stop_criteria=False)
+        del source_batch.sample_state.use_host_stop_criteria
+        PyExecutor._ring_broadcast_sample_state(source, source_batch)
+        source_payload = source.dist.isend_object.call_args.args[0]
+        assert source_payload[2] is False
+
+        relay = self._make_executor(is_last_pp_rank=False, is_second_last_pp_rank=False)
+        relay_batch = self._make_batch(use_host_stop_criteria=False)
+        del relay_batch.sample_state.use_host_stop_criteria
+        relay.dist.recv_object.return_value = source_payload
+        PyExecutor._ring_broadcast_sample_state(relay, relay_batch)
+        relay_payload = relay.dist.isend_object.call_args.args[0]
+        assert not hasattr(relay_batch.sample_state, "use_host_stop_criteria")
+        assert relay_payload[2] is False
+
+
 class TestComputeScheduledTokens:
     """Tests for PyExecutor._compute_scheduled_tokens.