[https://nvbugs/6482576][fix] Fall back to disagg_request_id in Python NIXL decode receiver

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,

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