From ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab Mon Sep 17 00:00:00 2001 From: limeward <32970461+edwinlim0919@users.noreply.github.com> Date: Sun, 26 Jul 2026 19:10:00 -0700 Subject: [PATCH] [Core][KV-transfer] MoRIIO: heterogeneous TP<->DP prefill/decode read routing (#46116) Signed-off-by: Edwin Lim --- .../unit/test_moriio_routing_fairness.py | 210 ++++++++++++++ .../kv_connector/v1/moriio/moriio_common.py | 5 + .../v1/moriio/moriio_connector.py | 259 +++++++++++++++++- 3 files changed, 467 insertions(+), 7 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_routing_fairness.py diff --git a/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py new file mode 100644 index 00000000000..e38ec3f7e9b --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_routing_fairness.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hardware-fair KV-read routing across MoRIIO heterogeneous P/D configs. + +Dependency-light (no GPU / ROCm / mori): builds bare ``MoRIIOConnectorWorker`` +instances via ``object.__new__`` and drives the REAL read-source routing +decision (``_resolve_read_source`` / ``_next_flex_tp_rank``) directly -- no +routing logic is re-implemented here, so a future change to those functions is +what these tests exercise. + +Scope -- the CONNECTOR (decode-side) read routing. The connector never selects +the prefill instance; the proxy does and hands the connector a ``remote_host`` + +``remote_dp_rank`` per request. The connector's only fairness lever is WHICH +prefill (dp, tp) rank each read targets, and the RFC's deployments collapse to +three real connector behaviours: + + symmetric TP (TP prefill + TP decode) -> decode tp_k reads prefill tp_k + flexible (TP prefill + DP decode) -> round-robin over prefill tp0..N-1 + owner DP (DP prefill, any decode) -> read the owner rank, tp0 + +These assume the proxy delivers a FAIR request stream (each prefill instance + +owner dp-rank an equal share) and verify the connector never re-introduces a +bottleneck. That proxy contract is checked separately against the real +``flat_interleaved_dp_route`` in the toy-proxy tests (PR #46115). + +A node runs one of two modes (8 GPUs each): + * TP8 -> (dp_size, tp_size) = (1, 8); MLA latent KV REPLICATED on all ranks. + * DP8EP -> (dp_size, tp_size) = (8, 1); KV PARTITIONED, one owner rank/request. +""" + +from collections import Counter +from dataclasses import dataclass +from unittest.mock import MagicMock + +import pytest + +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + ReqMeta, + get_port_offset, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorWorker, +) + +MODE_DIMS = {"TP8": (1, 8), "DP8EP": (8, 1)} # (dp_size, tp_size), 8 GPUs/node + + +@dataclass(frozen=True) +class PDConfig: + name: str + p_mode: str + d_mode: str + + @property + def p_dp(self) -> int: + return MODE_DIMS[self.p_mode][0] + + @property + def p_tp(self) -> int: + return MODE_DIMS[self.p_mode][1] + + @property + def d_dp(self) -> int: + return MODE_DIMS[self.d_mode][0] + + @property + def d_tp(self) -> int: + return MODE_DIMS[self.d_mode][1] + + @property + def n_prefill_gpus(self) -> int: + return self.p_dp * self.p_tp + + +CONFIGS = [ + PDConfig("1P_TP8:1D_TP8", "TP8", "TP8"), + PDConfig("2P_TP8:1D_DP8EP", "TP8", "DP8EP"), + PDConfig("2P_TP8:2D_TP8", "TP8", "TP8"), + PDConfig("2P_DP8EP:3D_DP8EP", "DP8EP", "DP8EP"), + PDConfig("2P_DP8EP:4D_TP8", "DP8EP", "TP8"), +] +CONFIG_IDS = [c.name for c in CONFIGS] + + +def make_decode_worker(*, world_size: int, tp_rank: int, dp_rank: int): + w = object.__new__(MoRIIOConnectorWorker) + w.world_size = world_size + w.tp_rank = tp_rank + w.dp_rank = dp_rank + w.use_mla = True + return w + + +def make_meta(*, p_tp: int, p_dp: int, remote_dp_rank: int, host: str = "phost0"): + return ReqMeta( + transfer_id="t", + local_block_ids=[1], + remote_block_ids=[2], + remote_host=host, + remote_port=1234, + remote_handshake_port=6301, + remote_notify_port=61005, + remote_engine_id=f"{host}:6301", + tp_size=p_tp, + remote_dp_size=p_dp, + remote_dp_rank=remote_dp_rank, + ) + + +def build_decode_workers(cfg: PDConfig) -> list: + """Decode workers that issue reads for one prefill instance. A TP decode + instance reads from every tp rank; a DP decode instance from each dp-rank + worker. Reused across requests so per-worker round-robin state advances.""" + if cfg.d_tp > 1: + return [ + make_decode_worker(world_size=cfg.d_tp, tp_rank=r, dp_rank=0) + for r in range(cfg.d_tp) + ] + return [ + make_decode_worker(world_size=1, tp_rank=0, dp_rank=d) for d in range(cfg.d_dp) + ] + + +def prefill_target_multiset(cfg: PDConfig, rounds: int) -> Counter: + """Drive the REAL _resolve_read_source over a fair input stream; tally which + prefill GPU each read targets. The only modelled assumption is the proxy + contract -- each owner dp-rank delivered equally (``for owner_dp in + range(p_dp)``); no proxy algorithm is reproduced.""" + workers = build_decode_workers(cfg) + hits: Counter = Counter() + for _ in range(rounds): + for owner_dp in range(cfg.p_dp): + for worker in workers: + meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=owner_dp) + chosen_tp, _flexible = worker._resolve_read_source(meta) + target = get_port_offset(owner_dp, chosen_tp, cfg.p_tp) + assert 0 <= target < cfg.n_prefill_gpus + hits[target] += 1 + return hits + + +@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS) +def test_kv_read_load_is_hardware_fair(cfg: PDConfig) -> None: + # rounds a multiple of p_tp so the round-robin closes on an exact cycle. + hits = prefill_target_multiset(cfg, rounds=16) + gpus = set(range(cfg.n_prefill_gpus)) + assert set(hits) == gpus, f"unused prefill GPUs: {sorted(gpus - set(hits))}" + assert max(hits.values()) == min(hits.values()), dict(sorted(hits.items())) + + +@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS) +def test_flexible_gate_fires_only_for_tp_prefill_dp_decode(cfg: PDConfig) -> None: + flags = set() + for worker in build_decode_workers(cfg): + meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=0) + _chosen, flexible = worker._resolve_read_source(meta) + flags.add(flexible) + is_mirror = cfg.d_tp == 1 and cfg.p_dp == 1 and cfg.p_tp > 1 + assert flags == {is_mirror} + + +def test_symmetric_tp_is_a_bijection() -> None: + targets = [] + for tp_rank in range(8): + worker = make_decode_worker(world_size=8, tp_rank=tp_rank, dp_rank=0) + chosen_tp, flexible = worker._resolve_read_source( + make_meta(p_tp=8, p_dp=1, remote_dp_rank=0) + ) + assert not flexible + targets.append(chosen_tp) + assert sorted(targets) == list(range(8)) + + +def test_owner_dp_read_is_faithful_and_covers_every_rank() -> None: + worker = make_decode_worker(world_size=8, tp_rank=3, dp_rank=0) + targets = [] + for owner_dp in range(8): + chosen_tp, flexible = worker._resolve_read_source( + make_meta(p_tp=1, p_dp=8, remote_dp_rank=owner_dp) + ) + assert not flexible + assert chosen_tp == 0 # p_tp == 1, only tp0 exists + targets.append(get_port_offset(owner_dp, chosen_tp, 1)) + assert sorted(targets) == list(range(8)) + + +def test_flexible_round_robin_is_deterministic_uniform_and_staggered() -> None: + w0 = make_decode_worker(world_size=1, tp_rank=0, dp_rank=0) + seq = [w0._next_flex_tp_rank(8) for _ in range(64)] + assert seq[:8] == list(range(8)) + assert Counter(seq) == Counter({t: 8 for t in range(8)}) + + first_pick = [ + make_decode_worker(world_size=1, tp_rank=0, dp_rank=d)._next_flex_tp_rank(8) + for d in range(8) + ] + assert sorted(first_pick) == list(range(8)) + + +def test_read_blocks_for_req_threads_chosen_tp() -> None: + # The resolved (chosen_tp, flexible) must reach _read_blocks, which keys the + # session AND the notify port off that single value -- so a read and its + # completion notify address the same prefill rank. + worker = make_decode_worker(world_size=1, tp_rank=0, dp_rank=3) + worker._read_blocks = MagicMock() + worker._read_blocks_for_req("r", make_meta(p_tp=8, p_dp=1, remote_dp_rank=0)) + kw = worker._read_blocks.call_args.kwargs + assert kw["flexible"] is True + assert kw["chosen_tp"] == 3 # first flexible pick = dp_rank seed + assert get_port_offset(0, kw["chosen_tp"], 8) == 3 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 15585123e5c..45df75eeb3c 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 @@ -422,6 +422,10 @@ class ReqMeta: remote_engine_id: str tp_size: int remote_dp_size: int + # Prefill DP rank that owns this request's KV (forwarded by the proxy). The + # read must target this rank's memory registration; the default 0 preserves + # the symmetric single-DP behaviour. + remote_dp_rank: int = 0 class MoRIIOConnectorMetadata(KVConnectorMetadata): @@ -475,6 +479,7 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata): remote_notify_port=int(remote_notify_port), tp_size=kv_transfer_params.get("tp_size", 1), remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), + remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0), ) if write_mode: self.reqs_to_save[request_id] = _req 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 ea41adf437c..c478ced2211 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 @@ -1014,6 +1014,8 @@ class MoRIIOConnectorWorker: self._handshake_futures: dict[EngineId, Future[set[str]]] = {} # Protects _handshake_futures and _remote_agents. self._handshake_lock = threading.RLock() + # Remote engines already covered by the eager pre-forward handshake. + self._eager_handshaked_engines: set[EngineId] = set() self.block_size = vllm_config.cache_config.block_size self.model_config = vllm_config.model_config @@ -1261,8 +1263,16 @@ class MoRIIOConnectorWorker: remote_tp_size: int, expected_engine_id: str, remote_dp_rank: int = 0, + remote_tp_rank: int | None = None, ) -> set[str]: - """Do a MoRIIO handshake with a remote instance.""" + """Do a MoRIIO handshake with a remote instance. + + remote_tp_rank: explicit remote TP index to dial. Flexible-read callers + pass the chosen prefill TP rank so the handshake, the (dp, tp) session + key and the notify port all address the SAME rank. None falls back to the + local-rank mapping _remote_tp_rank -- byte-identical for callers not yet + TP-aware. + """ start_time = time.perf_counter() @@ -1270,9 +1280,12 @@ class MoRIIOConnectorWorker: # a hack to keep us moving. We will switch when moving to etcd # or where we have a single ZMQ socket in the scheduler. - port_offset = get_port_offset( - remote_dp_rank, self._remote_tp_rank(remote_tp_size) + dial_tp_rank = ( + self._remote_tp_rank(remote_tp_size) + if remote_tp_rank is None + else int(remote_tp_rank) ) + port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size) path = make_zmq_path("tcp", host, port + port_offset) logger.debug("handshake Querying metadata on path: %s", path) @@ -1729,6 +1742,149 @@ class MoRIIOConnectorWorker: def get_engine_name_with_dp(self, engine_name, dp_rank): return f"{engine_name}_dp{dp_rank}" + def get_engine_name_with_dp_tp(self, engine_name, dp_rank, tp_rank): + # Per-(dp, tp) session key. The flexible mirror read keys sessions per + # (dp, tp) so one decode worker can hold a session to EACH prefill TP + # rank and spread reads across them; other configs keep the DP-only key. + return f"{engine_name}_dp{dp_rank}_tp{tp_rank}" + + def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> None: + """Handshake EVERY remote prefill DP rank BEFORE the decode forward pass, + identically across all local TP workers. + + Why this exists (the deadlock it prevents): with heterogeneous DP prefill + a decode TP worker reads KV from whichever prefill DP rank owns the + request, so across requests every worker must reach several prefill DP + ranks. The decode forward issues per-layer TP collectives (e.g. an + all-gather) that all local TP workers must enter together. If the + handshakes are left to fire lazily on the read path, the workers diverge: + a worker whose target rank is already cached races ahead into the forward + collective while a peer is still blocked in a handshake recv(). The first + worker then waits inside the collective for the stuck peer -> 600s NCCL + timeout / hang. This was observed directly with mixed TP<->DP configs. + + Fix: complete ALL prefill-DP-rank handshakes for every referenced remote + engine HERE, before any read enters the forward, so no worker is still + handshaking once its peers reach a collective. Fires ONCE per remote + engine (first contact), gated by _eager_handshaked_engines. The engine + set comes from scheduler-built metadata (identical on every TP worker), + so all workers run the same handshakes in the same order and reach the + all-reduce barrier below together. + + Failure handling: handshake exceptions are caught, never raised before + the collective (raising early would hang the peers still waiting for it). + Every worker reaches the all-reduce(MIN) vote; if ANY worker failed, ALL + raise the same error AFTER the collective, so the step fails fast and + uniformly in ~seconds instead of one rank hanging the forward for 600s. + """ + import torch.distributed as dist + + # Distinct remote engines referenced this step, in metadata (== + # scheduler) order so every TP worker iterates engines identically. + engines: dict[str, ReqMeta] = {} + for _req_id, meta in metadata.reqs_to_recv.items(): + remote_engine_id = ( + str(meta.remote_host) + ":" + str(meta.remote_handshake_port) + ) + engines.setdefault(remote_engine_id, meta) + + for remote_engine_id, meta in engines.items(): + if remote_engine_id in self._eager_handshaked_engines: + continue + + remote_dp_size = int(meta.remote_dp_size) + port = int(meta.remote_handshake_port) + tp_size = int(meta.tp_size) + + # Flexible mirror (TP prefill + MLA, world_size==1 decode): the read + # round-robins over prefill TP ranks, so pre-warm a session to EVERY + # (dp, tp) rank. Other configs pre-warm per DP rank (tp resolved by + # the fixed local-rank mapping) -- byte-identical to before. The + # mirror's decode is DP+EP, whose forward all-to-all is the collective + # that the eager barrier keeps everyone in step for. + flexible = ( + self.world_size == 1 + and self.use_mla + and remote_dp_size == 1 + and tp_size > 1 + ) + # (engine_id, dp_rank, tp_rank_or_None); tp_rank is None on the legacy + # path so _moriio_handshake falls back to its _remote_tp_rank mapping. + targets: list[tuple[Any, int, int | None]] + if flexible: + targets = [ + (self.get_engine_name_with_dp_tp(remote_engine_id, dp, tp), dp, tp) + for dp in range(remote_dp_size) + for tp in range(max(1, tp_size)) + ] + else: + targets = [ + (self.get_engine_name_with_dp(remote_engine_id, dp), dp, None) + for dp in range(remote_dp_size) + ] + + # Submit handshakes for every not-yet-known target UNDER the lock; do + # NOT hold it across the join or the collective (a stalled recv must + # not block another thread's lock acquisition). Gate on BOTH + # _remote_agents AND layer metadata: a rank with an agent entry but no + # layer metadata is half-handshaked and would KeyError at read time. + futures: list[tuple[str, Future[set[str]]]] = [] + with self._handshake_lock: + for eid, cur_dp_rank, cur_tp_rank in targets: + if ( + eid in self._remote_agents + and eid in self.layer_name_to_remote_kv_cache_metadata + ): + continue + fut = self._handshake_initiation_executor.submit( + self._moriio_handshake, + meta.remote_host, + port, + tp_size, + eid, + cur_dp_rank, + cur_tp_rank, + ) + futures.append((eid, fut)) + + # Join outside the lock. Bounded handshake errors are recorded here + # and reported after the all-reduce. + all_ok = True + results: dict[str, set[str]] = {} + for eid, fut in futures: + try: + results[eid] = fut.result() + except Exception: + logger.exception("Eager MoRIIO handshake failed for %s", eid) + all_ok = False + + with self._handshake_lock: + for eid, agents in results.items(): + self._remote_agents[eid] = agents + + logger.info( + "Eager MoRIIO handshake: engine=%s dp_size=%d new_ranks=%d " + "ok=%s tp_rank=%d", + remote_engine_id, + remote_dp_size, + len(futures), + all_ok, + self.tp_rank, + ) + # CPU all-reduce = TP-uniform success vote AND lockstep barrier: it + # blocks until every TP worker arrives, gives them the same verdict, + # and stays off the model compute stream. + vote = torch.tensor([1 if all_ok else 0], device="cpu", dtype=torch.int32) + dist.all_reduce(vote, group=self.tp_group.cpu_group, op=dist.ReduceOp.MIN) + if int(vote.item()) == 0: + raise HandshakeError( + f"Eager MoRIIO handshake failed for {remote_engine_id} on " + "at least one TP rank; failing this step fast to avoid a " + "TP collective hang" + ) + + self._eager_handshaked_engines.add(remote_engine_id) + def start_load_kv(self, metadata: MoRIIOConnectorMetadata): """ Start loading by triggering non-blocking moriio_xfer. @@ -1750,6 +1906,12 @@ class MoRIIOConnectorWorker: if self.mode == MoRIIOMode.WRITE: return + # Handshake every referenced remote prefill rank up front, before any + # read enters the forward pass. A lazy per-rank handshake on the read + # path lets TP workers diverge into a forward collective while a peer is + # still blocked handshaking -> NCCL hang (see below). + self._eager_handshake_all_dp_ranks(metadata) + wait_handshake_readd_req = False remote_engine_id = None @@ -1758,8 +1920,15 @@ class MoRIIOConnectorWorker: str(meta.remote_host) + ":" + str(meta.remote_handshake_port) ) meta.remote_engine_id = remote_engine_id + # The eager handshake above already covered every referenced engine + # (and keys the mirror per (dp, tp), which the DP-only dp0 probe below + # would miss). Only fall back to the lazy background handshake for an + # engine it did not cover. dp0_remote_engine_id = self.get_engine_name_with_dp(remote_engine_id, 0) - if dp0_remote_engine_id not in self._remote_agents: + if ( + remote_engine_id not in self._eager_handshaked_engines + and dp0_remote_engine_id not in self._remote_agents + ): # Initiate handshake with remote engine to exchange metadata. with self._handshake_lock: if remote_engine_id not in self._remote_agents: @@ -1809,12 +1978,53 @@ class MoRIIOConnectorWorker: self.save_kv_layer(metadata, layer_name, kv_layer, None) self._writer.seal_pending_transfers() + def _next_flex_tp_rank(self, remote_tp_size: int) -> int: + """Deterministic round-robin over prefill tp0..N-1 for the flexible read. + + Round-robin (not random): exactly uniform and testable, with the same + prefill-NIC balancing. Seeded from this decode rank's dp_rank so + concurrent decode DP ranks are phase-staggered -- at a given read index + distinct decode ranks target distinct prefill TP ranks. + """ + rr = getattr(self, "_flex_tp_rr", None) + if rr is None: + rr = int(getattr(self, "dp_rank", 0) or 0) + self._flex_tp_rr = rr + 1 + return rr % remote_tp_size + + def _resolve_read_source(self, meta: ReqMeta) -> tuple[int, bool]: + """Resolve (chosen_tp, flexible) for reading this request's KV. + + Flexible mirror (decode world_size==1 + MLA + pure-TP prefill): MLA + replicates the latent KV across the prefill TP ranks, so any is a valid + source; round-robin across them to spread RDMA/NIC load. Otherwise the + source TP rank is fixed by the local-rank mapping (_remote_tp_rank) -- + forward DP8EP->TP8 -> tp0; symmetric TP -> tp_rank -- byte-identical to + prior behaviour. chosen_tp is the single value threaded into the (dp, tp) + session key, the handshake dial and the notify port, so all three address + the SAME prefill rank (drift -> read one rank but notify another -> the + read rank's prefill buffer is never freed). + """ + remote_tp_size = int(meta.tp_size) + flexible = ( + self.world_size == 1 + and self.use_mla + and int(meta.remote_dp_size) == 1 + and remote_tp_size > 1 + ) + if flexible: + chosen_tp = self._next_flex_tp_rank(remote_tp_size) + else: + chosen_tp = self._remote_tp_rank(remote_tp_size) + return chosen_tp, flexible + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): logger.debug( "Remote agent %s available, calling _read_blocks for req %s", meta.remote_engine_id, req_id, ) + chosen_tp, flexible = self._resolve_read_source(meta) self._read_blocks( request_id=req_id, transfer_id=meta.transfer_id, @@ -1824,6 +2034,9 @@ class MoRIIOConnectorWorker: remote_host=meta.remote_host, remote_notify_port=meta.remote_notify_port, remote_tp_size=meta.tp_size, + remote_dp_rank=meta.remote_dp_rank, + chosen_tp=chosen_tp, + flexible=flexible, ) def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer): @@ -1976,12 +2189,37 @@ class MoRIIOConnectorWorker: remote_host: str, remote_notify_port: int, remote_tp_size: int, + remote_dp_rank: int = 0, + chosen_tp: int | None = None, + flexible: bool = False, ) -> None: if self.mode == MoRIIOMode.WRITE: return - dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0) - sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id) + # Read from the prefill rank that actually computed this request's KV + # (forwarded by the proxy). Hardcoding DP0 reads from a different rank's + # memory registration; per-rank num_blocks differ, so high block ids can + # overrun the wrong rank's region. + # + # eff_tp = the remote TP rank this read targets. The flexible mirror + # reads from a round-robin-chosen prefill TP rank and keys the session + # per (dp, tp); other configs use the fixed local-rank mapping (eff_tp == + # _remote_tp_rank), byte-identical to before. This key MUST match the one + # the eager handshake stored the session under. + eff_tp = ( + int(chosen_tp) + if chosen_tp is not None + else self._remote_tp_rank(remote_tp_size) + ) + if flexible: + remote_dp_engine_id = self.get_engine_name_with_dp_tp( + dst_engine_id, int(remote_dp_rank), eff_tp + ) + else: + remote_dp_engine_id = self.get_engine_name_with_dp( + dst_engine_id, int(remote_dp_rank) + ) + sessions, remote_moriio_meta = self._get_built_session(remote_dp_engine_id) # SQ-full backpressure deadline, shared across this request's layers. _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout @@ -2030,6 +2268,13 @@ class MoRIIOConnectorWorker: self._recving_transfers[request_id].append(transfer_status) self._recving_transfers_callback_addr[request_id] = ( remote_host, - str(remote_notify_port + self._remote_tp_rank(remote_tp_size)), + str( + remote_notify_port + + get_port_offset( + int(remote_dp_rank), + eff_tp, + remote_tp_size, + ) + ), transfer_id, )