diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 3c0c519e6562..21632f9358e2 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -176,6 +176,16 @@ class MessageType: CANCEL_SESSION = b"CANCEL_SESSION" +class PeerIncompatibleError(ValueError): + """A context peer failed the KV/recurrent-state layout compatibility check. + + Subclasses ValueError so existing ``except ValueError`` handlers still + match, while letting dispatch_task catch peer incompatibility narrowly and + fail only the requests targeting that peer instead of crashing the + executor loop. + """ + + class TaskStatus(Enum): INIT = "INIT" TRANSFERRING = "TRANSFERRING" @@ -1586,6 +1596,11 @@ def __init__( self._bounce = bounce self._dealers = {} self._sender_ep_instance_map = {} + # info_endpoint -> diagnostic message for peers that failed the + # compatibility check. Requests targeting such a peer fail fast + # without re-validating. Executor-thread-only access, no lock (same + # discipline as _sender_ep_instance_map). + self._incompatible_peers: dict[str, str] = {} self._messenger = ZMQMessenger(mode="ROUTER") self._sessions = {} # unique_rid -> RxSession @@ -1711,7 +1726,26 @@ def dispatch_task(self, task: KVRecvTask) -> None: ) receiver_req = self._build_recv_req_info(task) sender_dp_rank = params.ctx_dp_rank - peer_infos: RankInfo = self._get_sender_info(params) + try: + peer_infos: RankInfo = self._get_sender_info(params) + except PeerIncompatibleError as e: + # Fail only this request via the normal transfer-error path + # (task ERROR -> session ERROR -> WaitResult.FAILED -> + # DISAGG_TRANS_ERROR), so one incompatible context peer does not + # take down the whole generation worker. The async path's + # cross-rank consensus unions failed rids, so the request fails + # globally once any rank fails here — ranks whose page table has + # no recurrent layers (e.g. a hybrid-model PP stage) pass + # validation and may still run a transfer that is then + # discarded. The task is already in the session's _kv_tasks, no + # bounce reservation exists yet, and session._sender_endpoints + # is still empty, so no cleanup is needed here. + logger.error( + "dispatch_task: context peer incompatible, failing request " + f"unique_rid={task._unique_rid}: {e}" + ) + task.fail(e) + return if sender_dp_rank is not None: # Normal path: ctx_dp_rank is known, send to overlapping ranks. @@ -1817,6 +1851,10 @@ def _get_or_connect_dealer(self, endpoint: Optional[str]): def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: info_endpoint = self._extract_info_endpoint(params) + if info_endpoint in self._incompatible_peers: + # Known-incompatible peer: fail fast without another + # REQUEST_INSTANCE_INFO round-trip or re-validation. + raise PeerIncompatibleError(self._incompatible_peers[info_endpoint]) if self._should_register_peer(params): logger.info(f"Registering peer in first request to endpoint '{info_endpoint}'") messenger = ZMQMessenger(mode="DEALER", endpoint=info_endpoint) @@ -1830,15 +1868,28 @@ def _get_sender_info(self, params: DisaggregatedParams) -> RankInfo: # Recurrent-state (Mamba/KDA) layout gate on the receiver side. # The sender-side check (PeerRegistrar.register) runs in the # sender's listener thread, where exceptions are only logged, so - # reject here — before REGISTER_RANK_INFO is even sent — to fail - # the first gen request loudly instead of hanging on a transfer - # the sender will never serve. - MambaPolicy.validate_peer_compatible( - self._registrar.self_rank_info, - sender_info, - self._registrar.self_extractor.page_table, - sender_info.page_table, - ) + # reject here — before REGISTER_RANK_INFO is even sent, so no + # dealers are connected and no partial registration happens for + # the bad peer. The failure is converted to PeerIncompatibleError + # (handled in dispatch_task) so only requests targeting this peer + # fail, and cached so later requests fail fast. + try: + MambaPolicy.validate_peer_compatible( + self._registrar.self_rank_info, + sender_info, + self._registrar.self_extractor.page_table, + sender_info.page_table, + ) + except ValueError as e: + msg = ( + f"context peer at '{info_endpoint}' is incompatible: {e} " + "(cached: all further requests to this context peer fail " + "fast; restart this generation worker to re-validate, " + "e.g. after redeploying a compatible server on the same " + "endpoint)" + ) + self._incompatible_peers[info_endpoint] = msg + raise PeerIncompatibleError(msg) from e for endpoint in sender_info.sender_endpoints: dealer = self._get_or_connect_dealer(endpoint) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 398504421c9d..7d2e577d6b63 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -51,7 +51,7 @@ accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-dp4-cutl accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-ep4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[v2_kv_cache-tp4-cutlass-auto] SKIP (https://nvbugs/5596343) accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[cutlass-auto] SKIP (https://nvbugs/5596343) -accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 SKIP (https://nvbugs/6596064) +accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 SKIP (https://nvbugs/6525011) accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 SKIP (https://nvbugs/6596064) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[pp4-attn_backend=TRTLLM-torch_compile=False] SKIP (https://nvbugs/6490043) accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16_4gpus[tp2pp2-attn_backend=FLASHINFER-torch_compile=True] SKIP (https://nvbugs/6422337) diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 8abd9147dc72..7b3401d2331f 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -1496,5 +1496,112 @@ def test_session_has_transferring_tasks_false(): gen_transfer_worker.shutdown() +@pytest.mark.timeout(120) +def test_incompatible_peer_fails_only_affected_requests(): + """An incompatible context peer must fail only the requests targeting it. + + When MambaPolicy.validate_peer_compatible rejects a peer during first + registration, the receiver must (a) not raise out of receive() — the + session fails through the normal transfer-error path (WaitResult.FAILED) + with the diagnostic preserved; (b) cache the incompatibility so later + requests to the same endpoint fail fast without re-validating; and + (c) keep serving transfers from compatible peers through the same + Receiver. + """ + tensorrt_llm.logger.set_level("info") + # setup_good provides the gen worker under test and a compatible ctx + # peer; setup_bad provides a second, independent ctx instance whose + # endpoint gets poisoned (the incompatibility cache is permanent per + # endpoint, so the bad peer cannot be reused for the healthy transfer). + setup_good = create_transfer_worker_setup( + ctx_tp=1, ctx_pp=1, ctx_enable_dp=False, gen_tp=1, gen_pp=1, gen_enable_dp=False + ) + setup_bad = create_transfer_worker_setup( + ctx_tp=1, ctx_pp=1, ctx_enable_dp=False, gen_tp=1, gen_pp=1, gen_enable_dp=False + ) + gen_tw = setup_good["gen_transfer_workers"][0] + receiver = gen_tw._receiver + bad_endpoint = setup_bad["ctx_info_endpoint"] + + sampling_params = SamplingParams(temperature=0) + sc = tensorrt_llm.bindings.SamplingConfig(sampling_params._get_sampling_config()) + + def make_gen_request(request_id): + rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF + req = LlmRequest( + request_id=request_id, + max_new_tokens=1, + input_tokens=list(range(16)), + sampling_config=sc, + is_streaming=False, + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + ) + req.py_disaggregated_params = DisaggregatedParams( + ctx_request_id=request_id, + ctx_dp_rank=0, + ctx_info_endpoint=bad_endpoint, + disagg_request_id=rid, + ) + return req + + # dispatch_task fails at peer validation, before any block transfer, so + # empty per-layer-group block lists suffice (no KV sequence needed). + page_table = gen_tw._rank_info.page_table + empty_slice = KVSlice( + is_last_slice=True, + block_ids_per_layer_groups=[np.array([], dtype=np.int64) for _ in page_table.layer_groups], + ) + + validate_calls = [] + + def raiser(*args, **kwargs): + validate_calls.append(1) + raise ValueError("synthetic recurrent-state mismatch") + + try: + with pytest.MonkeyPatch.context() as mp: + mp.setattr(transfer_mod.MambaPolicy, "validate_peer_compatible", staticmethod(raiser)) + + # (a) First request to the incompatible peer: receive() must not + # raise; the session fails with the diagnostic preserved. + rx1 = gen_tw.create_rx_session(make_gen_request(500)) + rx1.receive(empty_slice) + assert rx1.wait_complete(blocking=True) == WaitResult.FAILED + # The non-blocking polling path must report the same terminal + # failure for an errored task (no None / spurious success). + assert rx1.wait_complete(blocking=False) == WaitResult.FAILED + assert rx1.has_failed() + exc = rx1._kv_tasks[0]._exception + assert exc is not None + assert "synthetic recurrent-state mismatch" in str(exc) + assert len(validate_calls) == 1 + rx1.close() + + # (b) Second request to the same endpoint: fails fast from the + # cache — no re-validation (call count unchanged) and no + # registration of the bad peer. + rx2 = gen_tw.create_rx_session(make_gen_request(501)) + rx2.receive(empty_slice) + assert rx2.wait_complete(blocking=True) == WaitResult.FAILED + assert rx2.wait_complete(blocking=False) == WaitResult.FAILED + assert rx2.has_failed() + assert "synthetic recurrent-state mismatch" in str(rx2._kv_tasks[0]._exception) + assert len(validate_calls) == 1 + assert bad_endpoint in receiver._incompatible_peers + assert bad_endpoint not in receiver._sender_ep_instance_map + rx2.close() + + # (c) The same gen Receiver still completes a real transfer from a + # compatible ctx peer (real validate_peer_compatible restored). + add_and_verify_request(setup_good, 0, 1, setup_good["request_len"], send_first=True) + assert setup_good["ctx_info_endpoint"] in receiver._sender_ep_instance_map + finally: + for s in (setup_good, setup_bad): + for worker in s["ctx_transfer_workers"]: + worker.shutdown() + for worker in s["gen_transfer_workers"]: + worker.shutdown() + + if __name__ == "__main__": test_transfer_worker_v1(1, 1, False, 1, 1, False, False)