diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index fc53e2975e14..19f704d60c2a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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, ) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 700c3ae68e9b..2d3e5957cb1c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/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) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 29b04404bbf1..9d3fb4dca7bb 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/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.