From fcde6091139344ad6c0dd42194ce8c41f5c9bb2f Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 10 Aug 2026 16:11:27 -0700 Subject: [PATCH 1/3] [TRTLLM-15264][fix] Fail only the affected requests on disagg peer-layout mismatch Receiver._get_sender_info raised MambaPolicy.validate_peer_compatible's ValueError on the executor thread with nothing catching it, so one incompatible context instance took down the whole generation worker, including traffic from healthy peers. Route the failure through the existing transfer-error path instead: _get_sender_info converts the validation ValueError into a typed PeerIncompatibleError (still a ValueError subclass) that dispatch_task catches to fail just that request's KVRecvTask. The session then reports WaitResult.FAILED and the request ends in DISAGG_TRANS_ERROR, the same mechanism used for remote agent failures and cancellation, while the worker keeps serving other peers. The check itself is unchanged and still runs before REGISTER_RANK_INFO, so no dealers are connected and no partial registration happens for the rejected peer. The incompatibility is cached per info_endpoint so later requests to the same peer fail fast without another REQUEST_INSTANCE_INFO round-trip or re-validation; the cached diagnostic notes that a generation-worker restart is needed to re-validate (e.g. after redeploying a compatible server on the same endpoint). Add a unit test covering: receive() does not raise and the session fails with the diagnostic preserved; the second request to the same endpoint fails fast without re-validation; and the same Receiver still completes a real transfer from a compatible peer. Signed-off-by: Brian Nguyen --- .../_torch/disaggregation/native/transfer.py | 71 ++++++++++-- .../disaggregated/test_kv_transfer.py | 103 ++++++++++++++++++ 2 files changed, 164 insertions(+), 10 deletions(-) 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/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 8abd9147dc72..310fd8a9dd8f 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -1496,5 +1496,108 @@ 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 + 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.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) From 37ae58154b4db291a731a33db2d4ae20657577aa Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 11 Aug 2026 11:34:38 -0700 Subject: [PATCH 2/3] test waives: skip accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4, Initial failure: `accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 (pre-existing failure, tracking bug pending) Signed-off-by: Brian Nguyen --- tests/integration/test_lists/waives.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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) From 6cde49514a6696cda641d56f836ba0b528989fe8 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Tue, 11 Aug 2026 17:02:19 -0700 Subject: [PATCH 3/3] [TRTLLM-15264][test] Cover non-blocking wait_complete failure path test_incompatible_peer_fails_only_affected_requests only checked the blocking wait_complete path. Add wait_complete(blocking=False) == WaitResult.FAILED assertions after each failed receive() to cover the non-blocking polling path returning the terminal failure for an errored task (no None / spurious success). Signed-off-by: Brian Nguyen --- tests/unittest/disaggregated/test_kv_transfer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 310fd8a9dd8f..7b3401d2331f 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -1567,6 +1567,9 @@ def raiser(*args, **kwargs): 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 @@ -1580,6 +1583,7 @@ def raiser(*args, **kwargs): 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