From 8f3e62bd545521989cafa79d9ab476755a1c0dfe Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 12:17:44 -0700 Subject: [PATCH 1/8] [ROCm][P/D][MoRIIO] Read-mode KV-release + best_of_n fixes On top of PR #41753 (XGMI backend selection), fixes for MoRIIO XGMI READ mode under best_of_n and high concurrency. Only touches moriio_connector.py. FIX(best_of-1many): transfer_id_to_request_id is now 1:many (list of sibling req_ids). map_request_id accumulates siblings; unmap drops one at a time and removes the key only when its last sibling is gone. The consumer WRITE-mode get_finished maps a single shared transfer_id to ALL its best_of siblings (was 1:1, leaving n-1 deferred forever on the decode). FIX(read-release): the decode notifies the prefill with the SHARED transfer_id (captured at read setup via worker-persistent _recving_transfer_id, since the scheduler unmaps at request_finished BEFORE the read completes). The prefill maps that transfer_id to its local sibling req_ids and releases them gated on vLLM's finished_req_ids (avoids the scheduler's \`assert req_id in self.requests\`). Signed-off-by: changning --- .../v1/moriio/moriio_connector.py | 97 ++++++++++++++++--- 1 file changed, 85 insertions(+), 12 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 804103275ae3..275092612657 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -181,7 +181,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: """Get the finished recving and sending requests.""" assert self.connector_worker is not None - return self.connector_worker.get_finished() + return self.connector_worker.get_finished(finished_req_ids) def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: assert self.connector_worker is not None @@ -284,15 +284,28 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} def map_request_id(self, request_id: ReqId, transfer_id: TransferId): - self.transfer_id_to_request_id[transfer_id] = request_id + # FIX(best_of-1many): n best_of siblings share one transfer_id; accumulate so + # the single completion can release all of them. + siblings = self.transfer_id_to_request_id.setdefault(transfer_id, []) + if request_id not in siblings: + siblings.append(request_id) + if len(siblings) > 1: + logger.debug( + "MoRIIO best_of fan-in: transfer_id=%s now has %d siblings", + transfer_id, len(siblings), + ) self.request_id_to_transfer_id[request_id] = transfer_id def unmap_request_id(self, request_id: ReqId): if request_id in self.request_id_to_transfer_id: transfer_id = self.request_id_to_transfer_id[request_id] del self.request_id_to_transfer_id[request_id] - if transfer_id in self.transfer_id_to_request_id: - del self.transfer_id_to_request_id[transfer_id] + siblings = self.transfer_id_to_request_id.get(transfer_id) + if siblings is not None: + if request_id in siblings: + siblings.remove(request_id) + if not siblings: + del self.transfer_id_to_request_id[transfer_id] else: logger.warning( "transfer id not in transfer_id_to_request_id lookup" @@ -750,6 +763,15 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): # In progress transfers. self._recving_transfers: defaultdict[ReqId, list] = defaultdict(list) self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str]] = {} + # FIX(read-release) state: consumer captures the shared tid at read setup so it + # can notify the producer with it on read-complete; producer accumulates held + # siblings and only frees them once vLLM marks them finished. + self._recving_transfer_id: dict[ReqId, str] = {} # consumer: req -> shared tid + self._read_sibs: dict[str, set[str]] = {} # producer: tid -> held siblings + self._read_done_tids: set[str] = set() # producer: tids done reading + self._read_finished_seen: set[str] = set() # producer: vLLM finished_req_ids + self._read_req_tid: dict[str, str] = {} # producer: sibling -> tid (pre-unmap) + self._read_seen_tid_sibcount: dict[str, int] = {} # producer: per-tid sibcount diff-track # Track the expiration time of requests that are waiting to be sent. self._reqs_to_send: dict[ReqId, float] = {} @@ -1239,7 +1261,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ready_event.wait() # Wait for listener ZMQ socket to be ready. self.moriio_wrapper.async_wait_reqid() - def get_finished(self) -> tuple[set[str], set[str]]: + def get_finished(self, finished_req_ids: set[str] | None = None) -> tuple[set[str], set[str]]: """ Get requests that are done sending or recving on this specific worker. The scheduler process (via the MultiprocExecutor) will use this output @@ -1250,6 +1272,29 @@ def get_finished(self) -> tuple[set[str], set[str]]: if self.is_producer: done_sending = self.moriio_wrapper.pop_finished_req_ids() + # FIX(read-release): release siblings of completed reads, gated on vLLM's + # is_finished signal so we don't trip the scheduler's assert on still-held reqs. + if self.mode == MoRIIOMode.READ: + self._read_done_tids.update(done_sending) + self._read_finished_seen |= (finished_req_ids or set()) + _m: set[str] = set() + for _t in list(self._read_done_tids): + if _t not in self._read_sibs: + continue + _sibs = self._read_sibs[_t] + _ready = _sibs & self._read_finished_seen + _m |= _ready + _sibs -= _ready + self._read_finished_seen -= _ready + if not _sibs: + logger.debug( + "MoRIIO read-release: transfer_id=%s fully drained", + _t, + ) + del self._read_sibs[_t] + self._read_done_tids.discard(_t) + self._read_seen_tid_sibcount.pop(_t, None) + done_sending = _m else: if self.mode == MoRIIOMode.WRITE: @@ -1257,12 +1302,18 @@ def get_finished(self) -> tuple[set[str], set[str]]: else: done_recving = self._pop_done_transfers() - done_recving = { - self.transfer_id_to_request_id[id] - for id in filter( - lambda id: id in self.transfer_id_to_request_id, done_recving - ) - } + # FIX(best_of-1many): WRITE -> one shared tid releases ALL its best_of siblings. + # READ -> consumer is synchronous (load_kv_async=False); the scheduler doesn't + # expect finished_recving here, and _pop_done_transfers's return was consumed + # internally above (triggered send_notify to the producer). + if not self.is_producer: + if self.mode == MoRIIOMode.WRITE: + _released: set[str] = set() + for _tid in done_recving: + _released.update(self.transfer_id_to_request_id.get(_tid, [])) + done_recving = _released + else: + done_recving = set() return done_sending, done_recving @@ -1274,8 +1325,15 @@ def _pop_done_transfers(self) -> set[str]: if status_list[-1].Succeeded(): done_req_ids.add(req_id) + # FIX(read-release): notify with the shared transfer_id; the producer + # may already have unmapped the local req_id by now. + _tid_for_notify = self._recving_transfer_id.get(req_id, req_id) + logger.debug( + "MoRIIO read-release notify: req=%s -> shared tid=%s", + req_id, _tid_for_notify, + ) self.moriio_wrapper.send_notify( - req_id, + _tid_for_notify, self._recving_transfers_callback_addr[req_id][0], self._recving_transfers_callback_addr[req_id][1], ) @@ -1283,6 +1341,7 @@ def _pop_done_transfers(self) -> set[str]: for req_id in to_remove: del self._recving_transfers[req_id] del self._recving_transfers_callback_addr[req_id] + self._recving_transfer_id.pop(req_id, None) return done_req_ids @@ -1346,6 +1405,17 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): """ self.transfer_id_to_request_id = metadata.transfer_id_to_request_id if self.is_producer: + for _t, _rs in metadata.transfer_id_to_request_id.items(): + _n = len(_rs) + if self._read_seen_tid_sibcount.get(_t) == _n: + continue # tid unchanged since last step -> nothing new to capture + for _r in _rs: + self._read_req_tid[_r] = _t + self._read_seen_tid_sibcount[_t] = _n + for _r in metadata.reqs_to_send: + _t = self._read_req_tid.pop(_r, None) + if _t is not None: + self._read_sibs.setdefault(_t, set()).add(_r) self.moriio_wrapper.async_wait_reqid() return if self.mode == MoRIIOMode.WRITE: @@ -1401,6 +1471,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): ) self._read_blocks( request_id=req_id, + transfer_id=meta.transfer_id, dst_engine_id=meta.remote_engine_id, local_block_ids=meta.local_block_ids, remote_block_ids=meta.remote_block_ids, @@ -1560,6 +1631,7 @@ def _read_blocks( request_id: str, remote_host: str, remote_notify_port: int, + transfer_id: str = "", ) -> None: if self.mode == MoRIIOMode.WRITE: return @@ -1582,6 +1654,7 @@ def _read_blocks( ) with self.moriio_wrapper.lock: self._recving_transfers[request_id].append(transfer_status) + self._recving_transfer_id[request_id] = transfer_id # FIX(read-release) self._recving_transfers_callback_addr[request_id] = ( remote_host, str(remote_notify_port + self.tp_rank), From c5f6bc626480a71ec3b47f17bb3ba7abcb42c338 Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 13:11:47 -0700 Subject: [PATCH 2/8] [Test][ROCm][P/D][MoRIIO] Unit tests for best_of_n / read-release fixes Signed-off-by: changning --- .../unit/test_moriio_connector.py | 159 ++++++++++++++++++ 1 file changed, 159 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index da78b62b9a03..6f92fa5ae76e 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -3,7 +3,9 @@ import importlib.util import os import subprocess +import threading import uuid +from collections import defaultdict from unittest.mock import MagicMock, patch import msgspec @@ -25,11 +27,13 @@ MoRIIOAgentMetadata, MoRIIOConnectorMetadata, MoRIIOConstants, + MoRIIOMode, zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( KVConnectorRole, MoRIIOConnector, + MoRIIOConnectorScheduler, MoRIIOConnectorWorker, ) from vllm.platforms import current_platform @@ -588,3 +592,158 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups): assert isinstance(metadata, MoRIIOAgentMetadata), ( "Decoded metadata is not MoRIIOAgentMetadata" ) + + +def _make_scheduler(role: str = "kv_producer") -> "MoRIIOConnectorScheduler": + """Construct a scheduler with the minimal kv_connector_extra_config required by __init__.""" + vllm_config = create_vllm_config(role=role) + vllm_config.kv_transfer_config.kv_connector_extra_config.update( + { + "proxy_ip": "127.0.0.1", + "proxy_ping_port": 36367, + "http_port": 8100, + "handshake_port": 6300, + "notify_port": 6100, + } + ) + return MoRIIOConnectorScheduler(vllm_config, "test_engine") + + +def _make_worker() -> MoRIIOConnectorWorker: + """Build a worker via __new__ so we can set just the attributes get_finished touches.""" + w = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + w.transfer_id_to_request_id = {} + w._read_sibs = {} + w._read_done_tids = set() + w._read_finished_seen = set() + w._read_req_tid = {} + w._read_seen_tid_sibcount = {} + w._recving_transfers = defaultdict(list) + w._recving_transfers_callback_addr = {} + w._recving_transfer_id = {} + return w + + +def test_best_of_n_map_accumulates_siblings(): + """FIX(best_of-1many) scheduler-side: map_request_id accumulates sibling req_ids under + a shared transfer_id (was a 1:1 dict that clobbered all but the last).""" + sched = _make_scheduler(role="kv_producer") + siblings = [f"{i}_cmpl-shared" for i in range(8)] + tid = "tx-shared" + + for s in siblings: + sched.map_request_id(s, tid) + + assert sched.transfer_id_to_request_id[tid] == siblings, ( + "all 8 siblings should be stored under the shared transfer_id" + ) + assert all(sched.request_id_to_transfer_id[s] == tid for s in siblings) + + # unmap one at a time: tid key persists until the last sibling is gone. + for s in siblings[:-1]: + sched.unmap_request_id(s) + assert tid in sched.transfer_id_to_request_id, ( + f"tid evicted too early after unmap({s})" + ) + assert s not in sched.transfer_id_to_request_id[tid] + sched.unmap_request_id(siblings[-1]) + assert tid not in sched.transfer_id_to_request_id, ( + "tid should be evicted when its last sibling is unmapped" + ) + + +def test_consumer_write_mode_fans_out_completion(): + """FIX(best_of-1many) consumer-side: get_finished in WRITE mode maps ONE shared + transfer_id to ALL its best_of siblings (was 1:1, leaving n-1 stuck Deferred).""" + worker = _make_worker() + worker.is_producer = False + worker.mode = MoRIIOMode.WRITE + worker.transfer_id_to_request_id = {"tx-shared": ["0_cmpl", "1_cmpl", "2_cmpl"]} + + fake_wrapper = MagicMock() + fake_wrapper.pop_finished_write_req_ids.return_value = {"tx-shared"} + worker.moriio_wrapper = fake_wrapper + + done_sending, done_recving = worker.get_finished() + assert done_sending == set() + assert done_recving == {"0_cmpl", "1_cmpl", "2_cmpl"}, ( + "one shared transfer_id must release ALL its best_of siblings" + ) + + +def test_consumer_read_mode_discards_done_recving(): + """READ-mode consumer is synchronous (load_kv_async=False) so reqs never enter + WAITING_FOR_REMOTE_KVS. _pop_done_transfers's return is consumed internally for + send_notify; passing it on would trip the scheduler's is_finished assert.""" + worker = _make_worker() + worker.is_producer = False + worker.mode = MoRIIOMode.READ + worker.moriio_wrapper = MagicMock() + # Pretend the internal pull-complete path returned a non-empty set. + worker._pop_done_transfers = MagicMock(return_value={"0_cmpl-running"}) + + done_sending, done_recving = worker.get_finished() + assert done_sending == set() + assert done_recving == set(), "READ-mode consumer must discard done_recving" + + +def test_producer_read_mode_finished_gated_release(): + """FIX(read-release): producer READ-mode get_finished releases siblings only after vLLM + marks them finished. Step 1: notify arrives, no finished -> no release. Step 2/3: each + finished sibling is released individually; tid is pruned only when the last is gone.""" + worker = _make_worker() + worker.is_producer = True + worker.mode = MoRIIOMode.READ + worker._read_sibs = {"tx-shared": {"a", "b"}} + worker._read_seen_tid_sibcount = {"tx-shared": 2} + + fake_wrapper = MagicMock() + worker.moriio_wrapper = fake_wrapper + + # Step 1: decode read-complete notify arrives (shared tid), no finished yet. + fake_wrapper.pop_finished_req_ids.return_value = {"tx-shared"} + done_sending, _ = worker.get_finished(finished_req_ids=set()) + assert done_sending == set(), "release must be gated on finished_req_ids" + assert worker._read_sibs == {"tx-shared": {"a", "b"}}, "no sibling freed yet" + assert "tx-shared" in worker._read_done_tids, "tid persisted for later release" + + # Step 2: vLLM marks 'a' finished -> only 'a' is freed; 'b' still held. + fake_wrapper.pop_finished_req_ids.return_value = set() + done_sending, _ = worker.get_finished(finished_req_ids={"a"}) + assert done_sending == {"a"} + assert worker._read_sibs == {"tx-shared": {"b"}} + assert "tx-shared" in worker._read_done_tids + + # Step 3: vLLM marks 'b' finished -> released; tid fully drained from all state. + done_sending, _ = worker.get_finished(finished_req_ids={"b"}) + assert done_sending == {"b"} + assert "tx-shared" not in worker._read_sibs + assert "tx-shared" not in worker._read_done_tids + assert "tx-shared" not in worker._read_seen_tid_sibcount + + +def test_pop_done_transfers_sends_shared_transfer_id(): + """FIX(read-release) decode-side: notify the prefill with the SHARED transfer_id + (captured at read setup), not the local sibling req_id with a -N-XXXX suffix the prefill + doesn't own. Bookkeeping cleaned up after.""" + worker = _make_worker() + fake_wrapper = MagicMock() + fake_wrapper.lock = threading.Lock() + worker.moriio_wrapper = fake_wrapper + + req_id = "0_cmpl-local-suffix-N-XXXX" + shared_tid = "tx-shared" + succeeded = MagicMock() + succeeded.Succeeded.return_value = True + worker._recving_transfers = defaultdict(list, {req_id: [succeeded]}) + worker._recving_transfers_callback_addr = {req_id: ("127.0.0.1", "6100")} + worker._recving_transfer_id = {req_id: shared_tid} + + done = worker._pop_done_transfers() + assert done == {req_id} + fake_wrapper.send_notify.assert_called_once_with(shared_tid, "127.0.0.1", "6100") + assert req_id not in worker._recving_transfers + assert req_id not in worker._recving_transfers_callback_addr + assert req_id not in worker._recving_transfer_id, ( + "_recving_transfer_id must be cleaned up alongside the other dicts" + ) From b01cc88deb19d48a90e6bbe6b32fbaceb465fe73 Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 13:49:59 -0700 Subject: [PATCH 3/8] [ROCm][P/D][MoRIIO] Address review: type hints + finished_seen filter Two follow-ups from automated review on the prior commit: Type hints (gemini high): transfer_id_to_request_id is now 1:many. Update the class-level type hints in MoRIIOConnectorScheduler.__init__, MoRIIOConnectorWorker.__init__, and MoRIIOConnectorMetadata.__init__ from dict[TransferId, ReqId] to dict[TransferId, list[ReqId]] for consistency with the new behavior and static analysis. _read_finished_seen leak (gemini high + codex P2): the prior commit unioned vLLM's engine-wide finished_req_ids into _read_finished_seen every step but only pruned IDs that matched active MoRIIO siblings. Non-MoRIIO IDs (and MoRIIO IDs whose request_finished returned delay_free_blocks=False) accumulated indefinitely. Filter at insertion against the union of held siblings (_read_sibs) and in-flight siblings (transfer_id_to_request_id), so the set stays bounded by tracked siblings only. Signed-off-by: changning --- .../kv_connector/v1/moriio/moriio_common.py | 2 +- .../kv_connector/v1/moriio/moriio_connector.py | 16 +++++++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 2733b2e0a878..4dcf1b340b67 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -335,7 +335,7 @@ def __init__(self): self.reqs_to_recv: dict[ReqId, ReqMeta] = {} self.reqs_to_save: dict[ReqId, ReqMeta] = {} self.reqs_to_send: dict[ReqId, float] = {} - self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + self.transfer_id_to_request_id: dict[TransferId, list[ReqId]] = {} def __repr__(self): return ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 275092612657..68e6eaf09269 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -280,7 +280,7 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): # Reqs to send and their expiration time self._reqs_need_send: dict[ReqId, float] = {} self.paths: dict[str, zmq.Socket] = {} - self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + self.transfer_id_to_request_id: dict[TransferId, list[ReqId]] = {} self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} def map_request_id(self, request_id: ReqId, transfer_id: TransferId): @@ -803,7 +803,7 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.cache_config.cache_dtype, use_mla=self.use_mla, ) - self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + self.transfer_id_to_request_id: dict[TransferId, list[ReqId]] = {} # TODO: consider the integration of flashinfer or other backends. self.backend_name = backend.get_name() @@ -1276,7 +1276,17 @@ def get_finished(self, finished_req_ids: set[str] | None = None) -> tuple[set[st # is_finished signal so we don't trip the scheduler's assert on still-held reqs. if self.mode == MoRIIOMode.READ: self._read_done_tids.update(done_sending) - self._read_finished_seen |= (finished_req_ids or set()) + # Filter finished_req_ids against currently-tracked siblings so non-MoRIIO + # req_ids don't accumulate in _read_finished_seen. Union covers both + # delay-freed siblings (now in _read_sibs) and in-flight ones (still in + # transfer_id_to_request_id). + if finished_req_ids: + tracked: set[str] = set() + for _rs in self._read_sibs.values(): + tracked.update(_rs) + for _rs in self.transfer_id_to_request_id.values(): + tracked.update(_rs) + self._read_finished_seen |= (finished_req_ids & tracked) _m: set[str] = set() for _t in list(self._read_done_tids): if _t not in self._read_sibs: From 1f211c98fe1fd7063892b1362c3fec75c846afbe Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 14:22:47 -0700 Subject: [PATCH 4/8] [ROCm][P/D][MoRIIO] Revert _read_finished_seen filter; prune phantom tids Signed-off-by: changning --- .../v1/moriio/moriio_connector.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 68e6eaf09269..96e09b478b22 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1276,20 +1276,18 @@ def get_finished(self, finished_req_ids: set[str] | None = None) -> tuple[set[st # is_finished signal so we don't trip the scheduler's assert on still-held reqs. if self.mode == MoRIIOMode.READ: self._read_done_tids.update(done_sending) - # Filter finished_req_ids against currently-tracked siblings so non-MoRIIO - # req_ids don't accumulate in _read_finished_seen. Union covers both - # delay-freed siblings (now in _read_sibs) and in-flight ones (still in - # transfer_id_to_request_id). - if finished_req_ids: - tracked: set[str] = set() - for _rs in self._read_sibs.values(): - tracked.update(_rs) - for _rs in self.transfer_id_to_request_id.values(): - tracked.update(_rs) - self._read_finished_seen |= (finished_req_ids & tracked) + self._read_finished_seen |= (finished_req_ids or set()) _m: set[str] = set() for _t in list(self._read_done_tids): if _t not in self._read_sibs: + # No siblings tracked. Either still in-flight on the scheduler + # (siblings haven't been delay-freed yet — wait), or phantom + # (siblings were inline-freed via request_finished's + # delay_free=False path and will never arrive). If the scheduler + # has also fully unmapped the tid, it's a phantom — drop now. + if _t not in self.transfer_id_to_request_id: + self._read_done_tids.discard(_t) + self._read_seen_tid_sibcount.pop(_t, None) continue _sibs = self._read_sibs[_t] _ready = _sibs & self._read_finished_seen From 0e3e011e8ea79064d564dfcf9b4b6b5d5bf57f2b Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 14:40:55 -0700 Subject: [PATCH 5/8] [ROCm][P/D][MoRIIO] Lifecycle filter for _read_finished_seen Signed-off-by: changning --- .../kv_connector/v1/moriio/moriio_connector.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 96e09b478b22..c5f65a18ffbd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1276,7 +1276,16 @@ def get_finished(self, finished_req_ids: set[str] | None = None) -> tuple[set[st # is_finished signal so we don't trip the scheduler's assert on still-held reqs. if self.mode == MoRIIOMode.READ: self._read_done_tids.update(done_sending) - self._read_finished_seen |= (finished_req_ids or set()) + # Filter finished_req_ids against siblings we've captured (in + # _read_req_tid, before scheduler unmap) or are currently holding + # (in _read_sibs). Non-MoRIIO IDs never enter — bounds growth. + if finished_req_ids: + _held: set[str] = set() + for _rs in self._read_sibs.values(): + _held.update(_rs) + self._read_finished_seen |= ( + finished_req_ids & (self._read_req_tid.keys() | _held) + ) _m: set[str] = set() for _t in list(self._read_done_tids): if _t not in self._read_sibs: From 2fbb5cd3cc21e399562a0d1ef2e39d63aeab1155 Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 14:56:35 -0700 Subject: [PATCH 6/8] [ROCm][P/D][MoRIIO] Fix CI: ruff line-length + mypy pop type Wrap long comment/docstring lines (E501) and switch _read_req_tid.pop(_r, None) -> if/pop to satisfy mypy (dict[str, str]). Signed-off-by: changning --- .../unit/test_moriio_connector.py | 22 ++++++----- .../v1/moriio/moriio_connector.py | 38 +++++++++++-------- 2 files changed, 34 insertions(+), 26 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 6f92fa5ae76e..39cb4e0df3de 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -595,7 +595,7 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups): def _make_scheduler(role: str = "kv_producer") -> "MoRIIOConnectorScheduler": - """Construct a scheduler with the minimal kv_connector_extra_config required by __init__.""" + """Scheduler with the minimal kv_connector_extra_config required by __init__.""" vllm_config = create_vllm_config(role=role) vllm_config.kv_transfer_config.kv_connector_extra_config.update( { @@ -610,7 +610,7 @@ def _make_scheduler(role: str = "kv_producer") -> "MoRIIOConnectorScheduler": def _make_worker() -> MoRIIOConnectorWorker: - """Build a worker via __new__ so we can set just the attributes get_finished touches.""" + """Worker via __new__ so we set only the attributes get_finished touches.""" w = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) w.transfer_id_to_request_id = {} w._read_sibs = {} @@ -625,8 +625,9 @@ def _make_worker() -> MoRIIOConnectorWorker: def test_best_of_n_map_accumulates_siblings(): - """FIX(best_of-1many) scheduler-side: map_request_id accumulates sibling req_ids under - a shared transfer_id (was a 1:1 dict that clobbered all but the last).""" + """FIX(best_of-1many) scheduler-side: map_request_id accumulates sibling + req_ids under a shared transfer_id (was 1:1 dict that clobbered all but the + last).""" sched = _make_scheduler(role="kv_producer") siblings = [f"{i}_cmpl-shared" for i in range(8)] tid = "tx-shared" @@ -688,9 +689,10 @@ def test_consumer_read_mode_discards_done_recving(): def test_producer_read_mode_finished_gated_release(): - """FIX(read-release): producer READ-mode get_finished releases siblings only after vLLM - marks them finished. Step 1: notify arrives, no finished -> no release. Step 2/3: each - finished sibling is released individually; tid is pruned only when the last is gone.""" + """FIX(read-release): producer READ-mode get_finished releases siblings only + after vLLM marks them finished. Step 1: notify arrives, no finished -> + no release. Step 2/3: each finished sibling is released individually; tid + is pruned only when the last is gone.""" worker = _make_worker() worker.is_producer = True worker.mode = MoRIIOMode.READ @@ -723,9 +725,9 @@ def test_producer_read_mode_finished_gated_release(): def test_pop_done_transfers_sends_shared_transfer_id(): - """FIX(read-release) decode-side: notify the prefill with the SHARED transfer_id - (captured at read setup), not the local sibling req_id with a -N-XXXX suffix the prefill - doesn't own. Bookkeeping cleaned up after.""" + """FIX(read-release) decode-side: notify the prefill with the SHARED + transfer_id (captured at read setup), not the local sibling req_id with a + -N-XXXX suffix the prefill doesn't own. Bookkeeping cleaned up after.""" worker = _make_worker() fake_wrapper = MagicMock() fake_wrapper.lock = threading.Lock() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index c5f65a18ffbd..cce97e3ec7cf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -292,7 +292,8 @@ def map_request_id(self, request_id: ReqId, transfer_id: TransferId): if len(siblings) > 1: logger.debug( "MoRIIO best_of fan-in: transfer_id=%s now has %d siblings", - transfer_id, len(siblings), + transfer_id, + len(siblings), ) self.request_id_to_transfer_id[request_id] = transfer_id @@ -766,12 +767,14 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): # FIX(read-release) state: consumer captures the shared tid at read setup so it # can notify the producer with it on read-complete; producer accumulates held # siblings and only frees them once vLLM marks them finished. - self._recving_transfer_id: dict[ReqId, str] = {} # consumer: req -> shared tid - self._read_sibs: dict[str, set[str]] = {} # producer: tid -> held siblings - self._read_done_tids: set[str] = set() # producer: tids done reading - self._read_finished_seen: set[str] = set() # producer: vLLM finished_req_ids - self._read_req_tid: dict[str, str] = {} # producer: sibling -> tid (pre-unmap) - self._read_seen_tid_sibcount: dict[str, int] = {} # producer: per-tid sibcount diff-track + self._recving_transfer_id: dict[ReqId, str] = {} # consumer: req -> shared tid + self._read_sibs: dict[str, set[str]] = {} # producer: tid -> held siblings + self._read_done_tids: set[str] = set() # producer: tids done reading + self._read_finished_seen: set[str] = set() # producer: vLLM finished_req_ids + self._read_req_tid: dict[str, str] = {} # producer: sibling -> tid (pre-unmap) + self._read_seen_tid_sibcount: dict[ + str, int + ] = {} # producer: per-tid sibcount diff-track # Track the expiration time of requests that are waiting to be sent. self._reqs_to_send: dict[ReqId, float] = {} @@ -1261,7 +1264,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ready_event.wait() # Wait for listener ZMQ socket to be ready. self.moriio_wrapper.async_wait_reqid() - def get_finished(self, finished_req_ids: set[str] | None = None) -> tuple[set[str], set[str]]: + def get_finished( + self, finished_req_ids: set[str] | None = None + ) -> tuple[set[str], set[str]]: """ Get requests that are done sending or recving on this specific worker. The scheduler process (via the MultiprocExecutor) will use this output @@ -1273,7 +1278,7 @@ def get_finished(self, finished_req_ids: set[str] | None = None) -> tuple[set[st if self.is_producer: done_sending = self.moriio_wrapper.pop_finished_req_ids() # FIX(read-release): release siblings of completed reads, gated on vLLM's - # is_finished signal so we don't trip the scheduler's assert on still-held reqs. + # is_finished signal so we don't assert on still-held reqs. if self.mode == MoRIIOMode.READ: self._read_done_tids.update(done_sending) # Filter finished_req_ids against siblings we've captured (in @@ -1283,8 +1288,8 @@ def get_finished(self, finished_req_ids: set[str] | None = None) -> tuple[set[st _held: set[str] = set() for _rs in self._read_sibs.values(): _held.update(_rs) - self._read_finished_seen |= ( - finished_req_ids & (self._read_req_tid.keys() | _held) + self._read_finished_seen |= finished_req_ids & ( + self._read_req_tid.keys() | _held ) _m: set[str] = set() for _t in list(self._read_done_tids): @@ -1342,12 +1347,13 @@ def _pop_done_transfers(self) -> set[str]: if status_list[-1].Succeeded(): done_req_ids.add(req_id) - # FIX(read-release): notify with the shared transfer_id; the producer - # may already have unmapped the local req_id by now. + # FIX(read-release): notify with the shared transfer_id; the + # producer may already have unmapped the local req_id by now. _tid_for_notify = self._recving_transfer_id.get(req_id, req_id) logger.debug( "MoRIIO read-release notify: req=%s -> shared tid=%s", - req_id, _tid_for_notify, + req_id, + _tid_for_notify, ) self.moriio_wrapper.send_notify( _tid_for_notify, @@ -1430,8 +1436,8 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): self._read_req_tid[_r] = _t self._read_seen_tid_sibcount[_t] = _n for _r in metadata.reqs_to_send: - _t = self._read_req_tid.pop(_r, None) - if _t is not None: + if _r in self._read_req_tid: + _t = self._read_req_tid.pop(_r) self._read_sibs.setdefault(_t, set()).add(_r) self.moriio_wrapper.async_wait_reqid() return From 81fabd7a0d6fe979842108992ecc16f0d21c17a9 Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 15:24:34 -0700 Subject: [PATCH 7/8] [ROCm][P/D][MoRIIO] Refactor read-release into named helpers Extract _read_release_step + _track_finished_siblings + _drop_phantom_tid + _drain_tid so get_finished reads as one call instead of a 40-line block. Signed-off-by: changning --- .../v1/moriio/moriio_connector.py | 88 +++++++++++-------- 1 file changed, 49 insertions(+), 39 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index cce97e3ec7cf..b6f3cd30dd88 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1277,46 +1277,8 @@ def get_finished( if self.is_producer: done_sending = self.moriio_wrapper.pop_finished_req_ids() - # FIX(read-release): release siblings of completed reads, gated on vLLM's - # is_finished signal so we don't assert on still-held reqs. if self.mode == MoRIIOMode.READ: - self._read_done_tids.update(done_sending) - # Filter finished_req_ids against siblings we've captured (in - # _read_req_tid, before scheduler unmap) or are currently holding - # (in _read_sibs). Non-MoRIIO IDs never enter — bounds growth. - if finished_req_ids: - _held: set[str] = set() - for _rs in self._read_sibs.values(): - _held.update(_rs) - self._read_finished_seen |= finished_req_ids & ( - self._read_req_tid.keys() | _held - ) - _m: set[str] = set() - for _t in list(self._read_done_tids): - if _t not in self._read_sibs: - # No siblings tracked. Either still in-flight on the scheduler - # (siblings haven't been delay-freed yet — wait), or phantom - # (siblings were inline-freed via request_finished's - # delay_free=False path and will never arrive). If the scheduler - # has also fully unmapped the tid, it's a phantom — drop now. - if _t not in self.transfer_id_to_request_id: - self._read_done_tids.discard(_t) - self._read_seen_tid_sibcount.pop(_t, None) - continue - _sibs = self._read_sibs[_t] - _ready = _sibs & self._read_finished_seen - _m |= _ready - _sibs -= _ready - self._read_finished_seen -= _ready - if not _sibs: - logger.debug( - "MoRIIO read-release: transfer_id=%s fully drained", - _t, - ) - del self._read_sibs[_t] - self._read_done_tids.discard(_t) - self._read_seen_tid_sibcount.pop(_t, None) - done_sending = _m + done_sending = self._read_release_step(done_sending, finished_req_ids) else: if self.mode == MoRIIOMode.WRITE: @@ -1339,6 +1301,54 @@ def get_finished( return done_sending, done_recving + def _read_release_step( + self, done_sending: set[str], finished_req_ids: set[str] | None + ) -> set[str]: + """FIX(read-release): release siblings of completed reads, gated on vLLM's + is_finished signal so we don't assert on still-held reqs.""" + self._read_done_tids.update(done_sending) + self._track_finished_siblings(finished_req_ids) + released: set[str] = set() + for tid in list(self._read_done_tids): + sibs = self._read_sibs.get(tid) + if sibs is None: + self._drop_phantom_tid(tid) + continue + ready = sibs & self._read_finished_seen + released |= ready + sibs -= ready + self._read_finished_seen -= ready + if not sibs: + self._drain_tid(tid) + return released + + def _track_finished_siblings(self, finished_req_ids: set[str] | None) -> None: + """Absorb vLLM's finished_req_ids, filtered against siblings we've captured + (in _read_req_tid, before scheduler unmap) or are currently holding (in + _read_sibs). Non-MoRIIO IDs never enter — bounds growth.""" + if not finished_req_ids: + return + held = {r for sibs in self._read_sibs.values() for r in sibs} + self._read_finished_seen |= finished_req_ids & ( + self._read_req_tid.keys() | held + ) + + def _drop_phantom_tid(self, tid: str) -> None: + """A tid in _read_done_tids with no siblings tracked is either still + in-flight on the scheduler or a phantom (siblings inline-freed via + request_finished's delay_free=False path). If the scheduler has also + unmapped the tid, no siblings will arrive — drop the bookkeeping.""" + if tid not in self.transfer_id_to_request_id: + self._read_done_tids.discard(tid) + self._read_seen_tid_sibcount.pop(tid, None) + + def _drain_tid(self, tid: str) -> None: + """All siblings of this tid released — clear all bookkeeping.""" + logger.debug("MoRIIO read-release: transfer_id=%s fully drained", tid) + del self._read_sibs[tid] + self._read_done_tids.discard(tid) + self._read_seen_tid_sibcount.pop(tid, None) + def _pop_done_transfers(self) -> set[str]: done_req_ids: set[str] = set() with self.moriio_wrapper.lock: From d6cb4f85632e06460820622fae1cc1a7ec37f527 Mon Sep 17 00:00:00 2001 From: changning Date: Sun, 24 May 2026 15:45:33 -0700 Subject: [PATCH 8/8] [Test][ROCm][P/D][MoRIIO] Use realistic literal IDs in best_of_n tests Replace synthetic strings ("tx-shared", "a"/"b", "0_cmpl") with literal tx-/cmpl- IDs matching the runtime format. Signed-off-by: changning --- .../unit/test_moriio_connector.py | 70 ++++++++++++------- 1 file changed, 44 insertions(+), 26 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 39cb4e0df3de..e6b848e3c19f 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -629,8 +629,17 @@ def test_best_of_n_map_accumulates_siblings(): req_ids under a shared transfer_id (was 1:1 dict that clobbered all but the last).""" sched = _make_scheduler(role="kv_producer") - siblings = [f"{i}_cmpl-shared" for i in range(8)] - tid = "tx-shared" + tid = "tx-d1f3abc456789012" + siblings = [ + "0_cmpl-fb4bc0e0a1234567", + "1_cmpl-fb4bc0e0a1234567", + "2_cmpl-fb4bc0e0a1234567", + "3_cmpl-fb4bc0e0a1234567", + "4_cmpl-fb4bc0e0a1234567", + "5_cmpl-fb4bc0e0a1234567", + "6_cmpl-fb4bc0e0a1234567", + "7_cmpl-fb4bc0e0a1234567", + ] for s in siblings: sched.map_request_id(s, tid) @@ -659,15 +668,21 @@ def test_consumer_write_mode_fans_out_completion(): worker = _make_worker() worker.is_producer = False worker.mode = MoRIIOMode.WRITE - worker.transfer_id_to_request_id = {"tx-shared": ["0_cmpl", "1_cmpl", "2_cmpl"]} + tid = "tx-35dba48109876543" + siblings = [ + "0_cmpl-35dba481b9876543", + "1_cmpl-35dba481b9876543", + "2_cmpl-35dba481b9876543", + ] + worker.transfer_id_to_request_id = {tid: siblings} fake_wrapper = MagicMock() - fake_wrapper.pop_finished_write_req_ids.return_value = {"tx-shared"} + fake_wrapper.pop_finished_write_req_ids.return_value = {tid} worker.moriio_wrapper = fake_wrapper done_sending, done_recving = worker.get_finished() assert done_sending == set() - assert done_recving == {"0_cmpl", "1_cmpl", "2_cmpl"}, ( + assert done_recving == set(siblings), ( "one shared transfer_id must release ALL its best_of siblings" ) @@ -681,7 +696,7 @@ def test_consumer_read_mode_discards_done_recving(): worker.mode = MoRIIOMode.READ worker.moriio_wrapper = MagicMock() # Pretend the internal pull-complete path returned a non-empty set. - worker._pop_done_transfers = MagicMock(return_value={"0_cmpl-running"}) + worker._pop_done_transfers = MagicMock(return_value={"0_cmpl-a1b2c3d4e5f60718"}) done_sending, done_recving = worker.get_finished() assert done_sending == set() @@ -696,45 +711,48 @@ def test_producer_read_mode_finished_gated_release(): worker = _make_worker() worker.is_producer = True worker.mode = MoRIIOMode.READ - worker._read_sibs = {"tx-shared": {"a", "b"}} - worker._read_seen_tid_sibcount = {"tx-shared": 2} + tid = "tx-7d3e9a2b54c1f680" + sib0 = "0_cmpl-7d3e9a2bc1f68054" + sib1 = "1_cmpl-7d3e9a2bc1f68054" + worker._read_sibs = {tid: {sib0, sib1}} + worker._read_seen_tid_sibcount = {tid: 2} fake_wrapper = MagicMock() worker.moriio_wrapper = fake_wrapper # Step 1: decode read-complete notify arrives (shared tid), no finished yet. - fake_wrapper.pop_finished_req_ids.return_value = {"tx-shared"} + fake_wrapper.pop_finished_req_ids.return_value = {tid} done_sending, _ = worker.get_finished(finished_req_ids=set()) assert done_sending == set(), "release must be gated on finished_req_ids" - assert worker._read_sibs == {"tx-shared": {"a", "b"}}, "no sibling freed yet" - assert "tx-shared" in worker._read_done_tids, "tid persisted for later release" + assert worker._read_sibs == {tid: {sib0, sib1}}, "no sibling freed yet" + assert tid in worker._read_done_tids, "tid persisted for later release" - # Step 2: vLLM marks 'a' finished -> only 'a' is freed; 'b' still held. + # Step 2: vLLM marks sib0 finished -> only sib0 is freed; sib1 still held. fake_wrapper.pop_finished_req_ids.return_value = set() - done_sending, _ = worker.get_finished(finished_req_ids={"a"}) - assert done_sending == {"a"} - assert worker._read_sibs == {"tx-shared": {"b"}} - assert "tx-shared" in worker._read_done_tids + done_sending, _ = worker.get_finished(finished_req_ids={sib0}) + assert done_sending == {sib0} + assert worker._read_sibs == {tid: {sib1}} + assert tid in worker._read_done_tids - # Step 3: vLLM marks 'b' finished -> released; tid fully drained from all state. - done_sending, _ = worker.get_finished(finished_req_ids={"b"}) - assert done_sending == {"b"} - assert "tx-shared" not in worker._read_sibs - assert "tx-shared" not in worker._read_done_tids - assert "tx-shared" not in worker._read_seen_tid_sibcount + # Step 3: vLLM marks sib1 finished -> released; tid fully drained from all state. + done_sending, _ = worker.get_finished(finished_req_ids={sib1}) + assert done_sending == {sib1} + assert tid not in worker._read_sibs + assert tid not in worker._read_done_tids + assert tid not in worker._read_seen_tid_sibcount def test_pop_done_transfers_sends_shared_transfer_id(): """FIX(read-release) decode-side: notify the prefill with the SHARED - transfer_id (captured at read setup), not the local sibling req_id with a - -N-XXXX suffix the prefill doesn't own. Bookkeeping cleaned up after.""" + transfer_id (captured at read setup), not the local sibling req_id (which has + the per-sibling prefix the prefill doesn't own). Bookkeeping cleaned up after.""" worker = _make_worker() fake_wrapper = MagicMock() fake_wrapper.lock = threading.Lock() worker.moriio_wrapper = fake_wrapper - req_id = "0_cmpl-local-suffix-N-XXXX" - shared_tid = "tx-shared" + shared_tid = "tx-aabbccdd11223344" + req_id = "0_cmpl-aabbccdd11223344" succeeded = MagicMock() succeeded.Succeeded.return_value = True worker._recving_transfers = defaultdict(list, {req_id: [succeeded]})