[TRTLLM-11946][feat] Disaggregated gen-first serving with ADP

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)

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,

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()

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": {

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"])

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]

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]

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)