diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 228a321ec23f..2dd57b7f55b3 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -1494,13 +1494,23 @@ def _get_session(self, unique_rid: Optional[int]) -> Optional["RxSession"]: def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: self_ri = self._registrar.self_rank_info - assert task._params.ctx_request_id is not None, ( - f"ctx_request_id is None for task unique_rid={task._unique_rid}" - ) assert task._unique_rid is not None, "KVRecvTask unique_rid is None" + # Some requests arrive with ctx_request_id None while disagg_request_id + # is set; disagg_request_id is the receive-session key, so fall back to + # it instead of failing here (nvbugs/6482576). + sender_req_id = task._params.ctx_request_id + if sender_req_id is None: + sender_req_id = task._params.disagg_request_id + if sender_req_id is None: + # Not an assert: must survive python -O so a None id never reaches + # RecvReqInfo.sender_req_id / the wire. + raise ValueError( + "both ctx_request_id and disagg_request_id are None for task " + f"unique_rid={task._unique_rid}" + ) # Receiver's cached prefix is implicit in block_ids size; sender derives dst_start. return RecvReqInfo( - sender_req_id=task._params.ctx_request_id, + sender_req_id=sender_req_id, instance_name=self_ri.instance_name, instance_rank=self_ri.instance_rank, block_ids_per_layer_groups=task._kv_slice.block_ids_per_layer_groups, diff --git a/tests/unittest/disaggregated/test_request_id.py b/tests/unittest/disaggregated/test_request_id.py index 70ea8a36fcad..3b5deac0b1b2 100644 --- a/tests/unittest/disaggregated/test_request_id.py +++ b/tests/unittest/disaggregated/test_request_id.py @@ -1,10 +1,14 @@ -"""Tests for _get_request_id in ExecutorRequestQueue. +"""Tests for disaggregated request-id handling. -Demonstrates the known bug: disagg_request_id=0 is falsy and gets skipped. +Covers ExecutorRequestQueue._get_request_id and the Receiver's sender_req_id +fallback in the native KV transceiver (nvbugs/6482576). """ +from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + def _make_queue(max_batch_size=128): """Create an ExecutorRequestQueue with mocked Distributed.""" @@ -56,3 +60,48 @@ def test_get_request_id_zero_bug(): # BUG: should return 0, but returns auto-incremented id instead assert rid != 0, "If this fails, the bug has been fixed — update this test" assert rid == 128 # falls through to auto-increment + + +# --------------------------------------------------------------------------- # +# Receiver._build_recv_req_info: sender_req_id fallback (nvbugs/6482576) +# --------------------------------------------------------------------------- # +def _make_recv_task(ctx_request_id, disagg_request_id, unique_rid=123): + """Minimal stand-in for KVRecvTask with only the fields the method reads.""" + return SimpleNamespace( + _params=SimpleNamespace(ctx_request_id=ctx_request_id, disagg_request_id=disagg_request_id), + _unique_rid=unique_rid, + _kv_slice=SimpleNamespace(block_ids_per_layer_groups=[], mamba_state_index=None), + _aux_slot=None, + slice_id=0, + ) + + +def _build_recv_req_info(tfr, task): + """Call the unbound Receiver method against a mocked registrar.""" + recv_self = SimpleNamespace( + _registrar=SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="gen-0", instance_rank=0) + ) + ) + return tfr.Receiver._build_recv_req_info(recv_self, task) + + +def test_build_recv_req_info_prefers_ctx_request_id(): + """Normal disagg flow: ctx_request_id keys the sender's TxSession.""" + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + info = _build_recv_req_info(tfr, _make_recv_task(ctx_request_id=7, disagg_request_id=99)) + assert info.sender_req_id == 7 + + +def test_build_recv_req_info_falls_back_to_disagg_request_id(): + """nvbugs/6482576: fall back to disagg_request_id when ctx_request_id is None.""" + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + info = _build_recv_req_info(tfr, _make_recv_task(ctx_request_id=None, disagg_request_id=99)) + assert info.sender_req_id == 99 + + +def test_build_recv_req_info_both_ids_none_raises(): + """Raise when neither request id is available (survives python -O).""" + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + with pytest.raises(ValueError): + _build_recv_req_info(tfr, _make_recv_task(ctx_request_id=None, disagg_request_id=None))