diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 208f8fc58001..5c714af1eac8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2729,11 +2729,11 @@ def _executor_loop_pp(self): logger.debug(f"microbatch {microbatch_id} can be queued") if not self.pp_async_broadcast_sample_state: - # Drain pending relay isends before forward: rendezvous - # sends need this rank to keep entering MPI, and a forward - # blocked in a native call would starve the receiving rank. - for mb in range(self.num_micro_batches): - self.wait_on_pp_send_handles(self.send_handles, mb) + # Drain the relay isends whose recv is due before this + # forward: rendezvous sends need this rank to keep + # entering MPI, and a forward blocked in a native call + # would starve the receiving rank. + self._drain_relay_sends_before_forward(microbatch_id) self._add_inflight_ids(scheduled_batch) @@ -3400,6 +3400,34 @@ def wait_on_pp_send_handles(self, send_handles, microbatch_id): send_handles[microbatch_id].wait() send_handles[microbatch_id] = None + def _drain_relay_sends_before_forward(self, microbatch_id: int) -> None: + """Wait on the inline relay isends whose matching recv is posted this + iteration, so a forward blocked in a native call cannot starve it. + + Ranks other than the last relay each sample state in the iteration + they receive it, so every isend they still hold was matched last + iteration and the full sweep only clears finished handles. + + The last rank originates the ring, and the first rank posts the + matching recv only ``pp_size - 2`` iterations later, when its + executed-microbatch cursor reaches that slot. Waiting on such a slot + earlier blocks this rank on a recv the peer issues after its next + top-of-loop collectives (the disagg transfer consensus, for one), + which cannot complete without this rank: with pp >= 4 and a + rendezvous-size sample state that is a deadlock. Wait only on the slot + the first rank relays this iteration; later slots are drained on the + iteration their recv is due, ahead of that iteration's forward. + """ + if not self.dist.is_last_pp_rank: + for mb in range(self.num_micro_batches): + self.wait_on_pp_send_handles(self.send_handles, mb) + return + # Same offset the non-last ranks apply to pick their executed + # microbatch in stage 2 of this iteration. + due_microbatch_id = (microbatch_id + 1 - + self.dist.pp_size) % self.num_micro_batches + self.wait_on_pp_send_handles(self.send_handles, due_microbatch_id) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: """Handle dynamic draft length for the current batch. diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e7eab00bb195..ee6f6e489008 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -1,14 +1,8 @@ accuracy/test_disaggregated_serving.py::TestGLM52NVFP4::test_nvfp4_nixl[cache_mgr_v1] SKIP (https://nvbugs/6619883) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[GSM8K-gen_tp=1-ctx_pp=2] SKIP (https://nvbugs/6435097) -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[GSM8K-gen_tp=1-ctx_pp=4] SKIP (https://nvbugs/6428069) -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[GSM8K-gen_tp=2-ctx_pp=4] SKIP (https://nvbugs/6428069) -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=1-ctx_pp=4] SKIP (https://nvbugs/6428069) -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=2-ctx_pp=4] SKIP (https://nvbugs/6428069) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ngram SKIP (https://nvbugs/6245651) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[GSM8K-tp2pp1] SKIP (https://nvbugs/6644475) -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[GSM8K-tp2pp2] SKIP (https://nvbugs/6428069) accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp2pp1] SKIP (https://nvbugs/6611817) -accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_tp_pp_symmetric[MMLU-tp2pp2] SKIP (https://nvbugs/6428069) accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664) accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform] SKIP (https://nvbugs/6661863) accuracy/test_dwdp_aggregated.py::TestDwdpAggDeepSeekV3Lite::test_dwdp_agg_accuracy[mode_a_uniform_contention_opt] SKIP (https://nvbugs/6644489) @@ -60,8 +54,6 @@ cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90] SKIP (https: cpp/test_unit_tests.py::test_unit_tests[batch_manager-80] SKIP (https://nvbugs/6674826) disaggregated/test_auto_scaling.py::test_disagg_server_restart[etcd-round_robin] SKIP (https://nvbugs/6611817) disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6105768) -disaggregated/test_disaggregated.py::test_disaggregated_ctxpp4_genpp4[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6428069) -disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2pp2_gentp2pp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6428069) disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_cache_aware_balance[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_bf16_conditional[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_gen_only[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6162322) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 85a873978fb9..87137ada6e70 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -3329,3 +3329,40 @@ def test_real_count_excludes_dummies(self): gathered = executor.dist.tp_allgather.call_args[0][0] assert gathered[1] == 2, "scheduled count keeps counting the dummy" assert gathered[3] == 1, "real count must exclude the dummy" + + +def _make_pp_relay_executor(*, pp_size: int, is_last_pp_rank: bool) -> PyExecutor: + executor = PyExecutor.__new__(PyExecutor) + executor.dist = Mock(pp_size=pp_size, is_last_pp_rank=is_last_pp_rank) + executor.num_micro_batches = pp_size + executor.send_handles = [object() for _ in range(pp_size)] + executor.wait_on_pp_send_handles = Mock() + return executor + + +@pytest.mark.parametrize("pp_size", [3, 4, 5]) +def test_last_pp_rank_drains_only_the_relay_send_whose_recv_is_due(pp_size): + """The first rank relays slot ``(m + 1 - pp_size) % n`` in iteration ``m``; + the last rank must wait on exactly that slot before its forward. Waiting on + the others blocks it on recvs the peer posts only after its next + top-of-loop collectives, which need this rank (pp>=4 disagg deadlock).""" + executor = _make_pp_relay_executor(pp_size=pp_size, is_last_pp_rank=True) + + for microbatch_id in range(pp_size): + executor.wait_on_pp_send_handles.reset_mock() + + executor._drain_relay_sends_before_forward(microbatch_id) + + first_rank_relay_slot = (microbatch_id + 1 - pp_size) % pp_size + executor.wait_on_pp_send_handles.assert_called_once_with( + executor.send_handles, first_rank_relay_slot + ) + + +def test_non_last_pp_rank_drains_every_relay_send(): + executor = _make_pp_relay_executor(pp_size=4, is_last_pp_rank=False) + + executor._drain_relay_sends_before_forward(2) + + waited = sorted(call.args[1] for call in executor.wait_on_pp_send_handles.call_args_list) + assert waited == [0, 1, 2, 3]