diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 4b53b4e574e1..d6162e4f0678 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -193,6 +193,12 @@ def __init__( class Sender(SenderBase): + # Time-to-live for orphaned RecvReqInfo entries (seconds). + # In gen-first ADP broadcast, non-assigned DP ranks accumulate + # RecvReqInfo that never gets consumed. Entries older than this + # are evicted during periodic sweeps. + _STALE_REQ_INFO_TTL_S = 120.0 + def __init__( self, peer_registrar: PeerRegistrar, @@ -202,6 +208,7 @@ def __init__( self._device_id = peer_registrar.self_rank_info.device_id self._agent = agent self._peer_requests: dict = {} + self._peer_requests_timestamps: dict[int, float] = {} # unique_rid -> insert time self._peer_requests_lock = threading.Lock() self._messenger = ZMQMessenger(mode="ROUTER") self._dealers = {} @@ -235,6 +242,7 @@ def _add_req_info(self, unique_rid: int, instance_rank: int, req_info: RecvReqIn with self._peer_requests_lock: if unique_rid not in self._peer_requests: self._peer_requests[unique_rid] = {} + self._peer_requests_timestamps[unique_rid] = time.monotonic() self._peer_requests[unique_rid][instance_rank] = req_info def _is_req_ready(self, unique_rid: int, expected_count: int) -> bool: @@ -258,6 +266,30 @@ def _get_first_req_info(self, unique_rid: Optional[int]) -> Optional[RecvReqInfo def _remove_req_info(self, unique_rid: int): with self._peer_requests_lock: self._peer_requests.pop(unique_rid, None) + self._peer_requests_timestamps.pop(unique_rid, None) + + def sweep_stale_req_infos(self): + """Evict RecvReqInfo entries that have no matching TxSession and exceed the TTL. + + Called opportunistically from the listener thread when a new REQUEST_DATA + arrives. With gen-first ADP broadcast, non-assigned DP ranks accumulate + entries that are never consumed; this sweep prevents unbounded growth. + """ + now = time.monotonic() + with self._peer_requests_lock: + stale_rids = [ + rid + for rid, ts in self._peer_requests_timestamps.items() + if now - ts > self._STALE_REQ_INFO_TTL_S + ] + if not stale_rids: + return + for rid in stale_rids: + with self._sessions_lock, self._peer_requests_lock: + if rid not in self._sessions and rid in self._peer_requests: + self._peer_requests.pop(rid, None) + self._peer_requests_timestamps.pop(rid, None) + logger.debug(f"Swept stale RecvReqInfo for rid={rid}") def setup_session(self, tx_session: "TxSession"): unique_rid = tx_session.disagg_request_id @@ -1040,17 +1072,41 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: def dispatch_task(self, task: KVRecvTask): params = task._params - logger.debug(f"Preparing async data transfer request for disagg_params={params}") + logger.debug( + f"Receiver.dispatch_task: unique_rid={task._unique_rid}, ctx_dp_rank={params.ctx_dp_rank}" + ) receiver_req = self._build_recv_req_info(task) sender_dp_rank = params.ctx_dp_rank - if sender_dp_rank is None: - raise ValueError( - f"ctx_dp_rank is None for request {task._unique_rid}; " - "disaggregated params may be missing context rank info" - ) peer_infos: RankInfo = self._get_sender_info(params) - peer_overlap = self._registrar.get_peer_overlap(peer_infos, sender_dp_rank) - task.expected_transfers = len(peer_overlap.ranks) + + if sender_dp_rank is not None: + # Normal path: ctx_dp_rank is known, send to overlapping ranks. + peer_overlap = self._registrar.get_peer_overlap(peer_infos, sender_dp_rank) + else: + # Gen-first with ADP: ctx_dp_rank unknown — broadcast REQUEST_DATA + # to ALL ctx sender ranks so every DP group receives it. + # get_peer_overlap returns ranks for one DP group (topology is + # symmetric), so use dp_rank=0 as representative. + dp_size = peer_infos.dp_size + dp0_overlap = self._registrar.get_peer_overlap(peer_infos, 0) + # Union of overlapping ranks across all DP groups for broadcast (deduplicated) + all_ranks_set: set[int] = set(dp0_overlap.ranks) + for dp in range(1, dp_size): + all_ranks_set.update(self._registrar.get_peer_overlap(peer_infos, dp).ranks) + all_ranks = list(all_ranks_set) + logger.debug( + f"Receiver.dispatch_task: ADP broadcast path, dp_size={dp_size}, " + f"all_ranks={all_ranks}" + ) + peer_overlap = type(dp0_overlap)(ranks=all_ranks) + + # In gen-first ADP broadcast, peer_overlap contains the union of all DP + # groups, but expected_transfers should reflect per-DP-group count since + # only one DP group will actually process the context request. + if sender_dp_rank is not None: + task.expected_transfers = len(peer_overlap.ranks) + else: + task.expected_transfers = len(dp0_overlap.ranks) session = self._get_session(task._unique_rid) if session is None: raise RuntimeError( @@ -1520,6 +1576,10 @@ def create_rx_session(self, request: LlmRequest) -> RxSession: def has_all_peer_req_infos_for_send(self, unique_rid: int) -> bool: return self._sender.has_all_peer_req_infos(unique_rid) + def sweep_stale_req_infos(self): + """Forward to Sender to evict orphaned RecvReqInfo from ADP broadcast.""" + self._sender.sweep_stale_req_infos() + def _setup_peer_infrastructure(self, kvm: KVCacheManager): self._rank_info_server = RankInfoServer(self._rank_info) if kvm.mapping.rank == 0 else None self._kv_extractor = KVRegionExtractorV1(kvm) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 32d21ced085f..878ab837ca3a 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -393,6 +393,10 @@ def check_context_transfer_status( del self._send_sessions[rid] self._close_failed_sessions(self._send_sessions, self._send_reqs, failed) + # Sweep orphaned RecvReqInfo entries from ADP broadcast on non-assigned + # DP ranks (entries that will never have a TxSession created for them). + self._transfer_worker.sweep_stale_req_infos() + return completed, failed def check_gen_transfer_status(self, at_least_request_num: Optional[int]): @@ -439,8 +443,16 @@ def get_disaggregated_params(self) -> Dict[str, Any]: # Keep this aligned with fields populated in respond_and_send_async(). # These values are server-level metadata used to seed generation-first # requests before context-phase response data arrives. + # + # With ADP (enable_attention_dp), ctx_dp_rank is not known at + # registration time because the context scheduler has not yet assigned + # the request to a DP rank. Return None so that the gen-side Receiver + # broadcasts REQUEST_DATA to all ctx DP ranks. The actual ctx_dp_rank + # is stamped into ContextPhaseParams by respond_and_send_async() after + # the prefill is scheduled. + ctx_dp_rank = None if self._mapping.enable_attention_dp else self._dp_rank return { - "ctx_dp_rank": self._dp_rank, + "ctx_dp_rank": ctx_dp_rank, "ctx_info_endpoint": [self._context_info_endpoint] if self._context_info_endpoint else None, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d1c22827ba34..0e4048110da3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -417,6 +417,14 @@ def __init__( self.expected_num_active_requests = 0 # TODO: Remove the condition on the PP size once disagg support from KVCache reuse # path is fixed. + # Buffer for responses generated inside _end_transfer_and_maybe_terminate. + # With ADP, _enqueue_responses does a tp_gather collective. When called + # from _send_kv_async the owning DP rank has a response but the other + # rank does not, causing a collective mismatch deadlock. Buffering the + # responses and flushing them at a synchronised point in the executor + # loop avoids the mismatch. + self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] + self.async_transfer_manager = AsyncTransferManager( self.resource_manager, should_store_blocks=self.enable_partial_reuse_for_disagg @@ -637,7 +645,15 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): response = request.create_response(False, self.dist.rank) if response: response.result.cached_tokens = request.cached_tokens - self._enqueue_responses([(request.py_request_id, response)]) + # Buffer the response instead of enqueueing immediately. + # With ADP, _enqueue_responses does a tp_gather collective. + # Calling it here would deadlock because only the owning DP + # rank reaches this point; the other DP rank never enters + # the matching collective. The buffer is flushed later at + # _flush_pending_transfer_responses where all ranks + # participate. + self._pending_transfer_responses.append( + (request.py_request_id, response)) if self.async_transfer_manager.end_transfer(request): self.active_requests.remove(request) self._terminate_request(request) @@ -650,6 +666,20 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): if not self.async_transfer_manager.should_store_blocks: self._terminate_request(request) + def _flush_pending_transfer_responses(self): + """Enqueue buffered transfer-completion responses. + + Must be called at a point where ALL DP ranks execute in lockstep so + that the tp_gather inside _enqueue_responses does not deadlock. + """ + responses = self._pending_transfer_responses + self._pending_transfer_responses = [] + if responses or self.enable_attention_dp: + # Even when this rank has no responses we must participate in the + # collective when ADP is enabled so that the other rank's gather + # can complete. + self._enqueue_responses(responses) + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1798,6 +1828,7 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): if self.kv_cache_transceiver: finished_ctx_reqs = scheduled_requests.context_requests_last_chunk self._send_kv_async(finished_ctx_reqs) + self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -2272,6 +2303,7 @@ def _executor_loop(self): self._update_requests(sample_state, self.resource_manager) self._send_kv_async(scheduled_batch.all_requests()) + self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -2519,6 +2551,11 @@ def _executor_loop_overlap(self): self._send_kv_async( self.previous_batch.scheduled_requests.all_requests()) + # Flush outside the conditional so that all DP ranks + # participate in the tp_gather collective even when + # should_process_previous_batch differs between ranks. + self._flush_pending_transfer_responses() + if self.drafter is not None and self.use_spec_decode and should_process_previous_batch: # Cleanup previous draft resources used in the draft model self.drafter.cleanup_previous_draft_resources() diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index ced0f85e858b..84deaeb223d2 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -183,8 +183,8 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): disagg_cluster = { "cluster_uri": cluster_uri, "cluster_name": "test_cluster", - "heartbeat_interval_sec": 1, - "inactive_timeout_sec": 2, + "heartbeat_interval_sec": 5, + "inactive_timeout_sec": 10, } # Auto-deduce minimal_instances from num_instances @@ -365,13 +365,9 @@ def multi_popen(server_configs, server_name="", enable_redirect_log=False): with ( MyThreadPoolExecutor(max_workers=max_workers) as thread_pool, temp_dir, - multi_popen(ctx_servers, "ctx", - enable_redirect_log=False) as ctx_processes, - multi_popen(gen_servers, "gen", enable_redirect_log=False) as - gen_processes, - multi_popen([(base_env, server_cmd)], - "disagg", - enable_redirect_log=False) as server_processes, + multi_popen(ctx_servers, "ctx") as ctx_processes, + multi_popen(gen_servers, "gen") as gen_processes, + multi_popen([(base_env, server_cmd)], "disagg") as server_processes, ): start_time = time.time() server_is_ready = False @@ -1845,7 +1841,9 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config, "ctx_tp1pp1", "ctx_tp1pp2", "ctx_tp2pp1", "ctx_tp2pp2", "ctx_tp1pp4" ], ) - def test_gen_first(self, ctx_tp_pp, gen_tp_pp): + @pytest.mark.parametrize("enable_attention_dp", [False, True], + ids=["noadp", "adp"]) + def test_gen_first(self, ctx_tp_pp, gen_tp_pp, enable_attention_dp): ctx_tp, ctx_pp = ctx_tp_pp gen_tp, gen_pp = gen_tp_pp total_gpus = ctx_tp * ctx_pp + gen_tp * gen_pp @@ -1860,6 +1858,7 @@ def test_gen_first(self, ctx_tp_pp, gen_tp_pp): ctx_server_config = { "tensor_parallel_size": ctx_tp, "pipeline_parallel_size": ctx_pp, + "enable_attention_dp": enable_attention_dp, "disable_overlap_scheduler": True, "cuda_graph_config": None, "cache_transceiver_config": { @@ -1871,6 +1870,7 @@ def test_gen_first(self, ctx_tp_pp, gen_tp_pp): gen_server_config = { "tensor_parallel_size": gen_tp, "pipeline_parallel_size": gen_pp, + "enable_attention_dp": enable_attention_dp, "disable_overlap_scheduler": True, "cuda_graph_config": None, "cache_transceiver_config": { diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 6205460962a7..0e671dfc2a2c 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -904,6 +904,7 @@ def test_disaggregated_overlap(disaggregated_test_root, llm_venv, cwd=llm_venv.get_working_directory()) +@pytest.mark.skip_less_device(8) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], indirect=True) @pytest.mark.parametrize("ctx_pp", [1, 4], ids=["ctx_pp1", "ctx_pp4"]) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index c818b9cdd3e9..f7439c5b8f72 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -95,7 +95,8 @@ accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp2tp1cp2] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_chunked_prefill -accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[ctx_tp1pp1-gen_tp1pp1] +accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[noadp-ctx_tp2pp1-gen_tp1pp1] +accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[adp-ctx_tp2pp1-gen_tp2pp1] accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_nixl_backend accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=False] accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index b48798797fae..5aa3f6138f41 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -131,7 +131,8 @@ l0_dgx_h100: - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[MMLU-gen_tp=2-ctx_pp=2] - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[GSM8K] - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_multi_instance[MMLU] - - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[ctx_tp1pp1-gen_tp1pp1] + - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[noadp-ctx_tp2pp1-gen_tp1pp1] + - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_gen_first[adp-ctx_tp2pp1-gen_tp2pp1] - disaggregated/test_auto_scaling.py::test_service_discovery[etcd-round_robin] - disaggregated/test_auto_scaling.py::test_worker_restart[etcd-load_balancing] - disaggregated/test_auto_scaling.py::test_worker_restart[etcd-round_robin] diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 3386e16c2917..5d017ba1549f 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -270,7 +270,8 @@ visual_gen/test_visual_gen_benchmark.py::test_online_benchmark[openai-videos] SK examples/test_visual_gen.py::test_vbench_dimension_score_wan SKIP (https://nvbugs/6050483) examples/test_visual_gen.py::test_vbench_dimension_score_wan22_a14b_fp8 SKIP (https://nvbugs/6050483) visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark SKIP (https://nvbugs/6050483) -disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6057459) +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/6050489) +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/6050489) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-v32-fp4_8k1k_con4096_ctx1_dep4_gen1_dep32_eplb256_mtp0_ccb-NIXL] SKIP (https://nvbugs/6085022) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-v32-fp4_32k4k_con2048_ctx1_dep4_gen1_dep32_eplb288_mtp1_ccb-NIXL] SKIP (https://nvbugs/6085022) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_32k4k_con2048_ctx1_dep4_gen1_dep32_eplb288_mtp1_ccb-NIXL] SKIP (https://nvbugs/6085022)