[TRTLLM-15264][fix] Fail only the affected requests on disagg peer-layout mismatch

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)

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)

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)