[Core][KV-transfer] MoRIIO: heterogeneous TP<->DP prefill/decode read routing (#46116)
Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io>
limeward committed
Jul 26, 2026 at 19:10 UTC
ffc4f08c8ee130d4ea6347c1bf31ffd4f8af28ab
3 files changed
+467
-7
tests/v1/kv_connector/unit/test_moriio_routing_fairness.py
new
+210
@@ -0,0 +1,210 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Hardware-fair KV-read routing across MoRIIO heterogeneous P/D configs.
4
+
5
+Dependency-light (no GPU / ROCm / mori): builds bare ``MoRIIOConnectorWorker``
6
+instances via ``object.__new__`` and drives the REAL read-source routing
7
+decision (``_resolve_read_source`` / ``_next_flex_tp_rank``) directly -- no
8
+routing logic is re-implemented here, so a future change to those functions is
9
+what these tests exercise.
10
+
11
+Scope -- the CONNECTOR (decode-side) read routing. The connector never selects
12
+the prefill instance; the proxy does and hands the connector a ``remote_host`` +
13
+``remote_dp_rank`` per request. The connector's only fairness lever is WHICH
14
+prefill (dp, tp) rank each read targets, and the RFC's deployments collapse to
15
+three real connector behaviours:
16
+
17
+ symmetric TP (TP prefill + TP decode) -> decode tp_k reads prefill tp_k
18
+ flexible (TP prefill + DP decode) -> round-robin over prefill tp0..N-1
19
+ owner DP (DP prefill, any decode) -> read the owner rank, tp0
20
+
21
+These assume the proxy delivers a FAIR request stream (each prefill instance +
22
+owner dp-rank an equal share) and verify the connector never re-introduces a
23
+bottleneck. That proxy contract is checked separately against the real
24
+``flat_interleaved_dp_route`` in the toy-proxy tests (PR #46115).
25
+
26
+A node runs one of two modes (8 GPUs each):
27
+ * TP8 -> (dp_size, tp_size) = (1, 8); MLA latent KV REPLICATED on all ranks.
28
+ * DP8EP -> (dp_size, tp_size) = (8, 1); KV PARTITIONED, one owner rank/request.
29
+"""
30
+
31
+from collections import Counter
32
+from dataclasses import dataclass
33
+from unittest.mock import MagicMock
34
+
35
+import pytest
36
+
37
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
38
+ ReqMeta,
39
+ get_port_offset,
40
+)
41
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
42
+ MoRIIOConnectorWorker,
43
+)
44
+
45
+MODE_DIMS = {"TP8": (1, 8), "DP8EP": (8, 1)} # (dp_size, tp_size), 8 GPUs/node
46
+
47
+
48
+@dataclass(frozen=True)
49
+class PDConfig:
50
+ name: str
51
+ p_mode: str
52
+ d_mode: str
53
+
54
+ @property
55
+ def p_dp(self) -> int:
56
+ return MODE_DIMS[self.p_mode][0]
57
+
58
+ @property
59
+ def p_tp(self) -> int:
60
+ return MODE_DIMS[self.p_mode][1]
61
+
62
+ @property
63
+ def d_dp(self) -> int:
64
+ return MODE_DIMS[self.d_mode][0]
65
+
66
+ @property
67
+ def d_tp(self) -> int:
68
+ return MODE_DIMS[self.d_mode][1]
69
+
70
+ @property
71
+ def n_prefill_gpus(self) -> int:
72
+ return self.p_dp * self.p_tp
73
+
74
+
75
+CONFIGS = [
76
+ PDConfig("1P_TP8:1D_TP8", "TP8", "TP8"),
77
+ PDConfig("2P_TP8:1D_DP8EP", "TP8", "DP8EP"),
78
+ PDConfig("2P_TP8:2D_TP8", "TP8", "TP8"),
79
+ PDConfig("2P_DP8EP:3D_DP8EP", "DP8EP", "DP8EP"),
80
+ PDConfig("2P_DP8EP:4D_TP8", "DP8EP", "TP8"),
81
+]
82
+CONFIG_IDS = [c.name for c in CONFIGS]
83
+
84
+
85
+def make_decode_worker(*, world_size: int, tp_rank: int, dp_rank: int):
86
+ w = object.__new__(MoRIIOConnectorWorker)
87
+ w.world_size = world_size
88
+ w.tp_rank = tp_rank
89
+ w.dp_rank = dp_rank
90
+ w.use_mla = True
91
+ return w
92
+
93
+
94
+def make_meta(*, p_tp: int, p_dp: int, remote_dp_rank: int, host: str = "phost0"):
95
+ return ReqMeta(
96
+ transfer_id="t",
97
+ local_block_ids=[1],
98
+ remote_block_ids=[2],
99
+ remote_host=host,
100
+ remote_port=1234,
101
+ remote_handshake_port=6301,
102
+ remote_notify_port=61005,
103
+ remote_engine_id=f"{host}:6301",
104
+ tp_size=p_tp,
105
+ remote_dp_size=p_dp,
106
+ remote_dp_rank=remote_dp_rank,
107
+ )
108
+
109
+
110
+def build_decode_workers(cfg: PDConfig) -> list:
111
+ """Decode workers that issue reads for one prefill instance. A TP decode
112
+ instance reads from every tp rank; a DP decode instance from each dp-rank
113
+ worker. Reused across requests so per-worker round-robin state advances."""
114
+ if cfg.d_tp > 1:
115
+ return [
116
+ make_decode_worker(world_size=cfg.d_tp, tp_rank=r, dp_rank=0)
117
+ for r in range(cfg.d_tp)
118
+ ]
119
+ return [
120
+ make_decode_worker(world_size=1, tp_rank=0, dp_rank=d) for d in range(cfg.d_dp)
121
+ ]
122
+
123
+
124
+def prefill_target_multiset(cfg: PDConfig, rounds: int) -> Counter:
125
+ """Drive the REAL _resolve_read_source over a fair input stream; tally which
126
+ prefill GPU each read targets. The only modelled assumption is the proxy
127
+ contract -- each owner dp-rank delivered equally (``for owner_dp in
128
+ range(p_dp)``); no proxy algorithm is reproduced."""
129
+ workers = build_decode_workers(cfg)
130
+ hits: Counter = Counter()
131
+ for _ in range(rounds):
132
+ for owner_dp in range(cfg.p_dp):
133
+ for worker in workers:
134
+ meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=owner_dp)
135
+ chosen_tp, _flexible = worker._resolve_read_source(meta)
136
+ target = get_port_offset(owner_dp, chosen_tp, cfg.p_tp)
137
+ assert 0 <= target < cfg.n_prefill_gpus
138
+ hits[target] += 1
139
+ return hits
140
+
141
+
142
+@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS)
143
+def test_kv_read_load_is_hardware_fair(cfg: PDConfig) -> None:
144
+ # rounds a multiple of p_tp so the round-robin closes on an exact cycle.
145
+ hits = prefill_target_multiset(cfg, rounds=16)
146
+ gpus = set(range(cfg.n_prefill_gpus))
147
+ assert set(hits) == gpus, f"unused prefill GPUs: {sorted(gpus - set(hits))}"
148
+ assert max(hits.values()) == min(hits.values()), dict(sorted(hits.items()))
149
+
150
+
151
+@pytest.mark.parametrize("cfg", CONFIGS, ids=CONFIG_IDS)
152
+def test_flexible_gate_fires_only_for_tp_prefill_dp_decode(cfg: PDConfig) -> None:
153
+ flags = set()
154
+ for worker in build_decode_workers(cfg):
155
+ meta = make_meta(p_tp=cfg.p_tp, p_dp=cfg.p_dp, remote_dp_rank=0)
156
+ _chosen, flexible = worker._resolve_read_source(meta)
157
+ flags.add(flexible)
158
+ is_mirror = cfg.d_tp == 1 and cfg.p_dp == 1 and cfg.p_tp > 1
159
+ assert flags == {is_mirror}
160
+
161
+
162
+def test_symmetric_tp_is_a_bijection() -> None:
163
+ targets = []
164
+ for tp_rank in range(8):
165
+ worker = make_decode_worker(world_size=8, tp_rank=tp_rank, dp_rank=0)
166
+ chosen_tp, flexible = worker._resolve_read_source(
167
+ make_meta(p_tp=8, p_dp=1, remote_dp_rank=0)
168
+ )
169
+ assert not flexible
170
+ targets.append(chosen_tp)
171
+ assert sorted(targets) == list(range(8))
172
+
173
+
174
+def test_owner_dp_read_is_faithful_and_covers_every_rank() -> None:
175
+ worker = make_decode_worker(world_size=8, tp_rank=3, dp_rank=0)
176
+ targets = []
177
+ for owner_dp in range(8):
178
+ chosen_tp, flexible = worker._resolve_read_source(
179
+ make_meta(p_tp=1, p_dp=8, remote_dp_rank=owner_dp)
180
+ )
181
+ assert not flexible
182
+ assert chosen_tp == 0 # p_tp == 1, only tp0 exists
183
+ targets.append(get_port_offset(owner_dp, chosen_tp, 1))
184
+ assert sorted(targets) == list(range(8))
185
+
186
+
187
+def test_flexible_round_robin_is_deterministic_uniform_and_staggered() -> None:
188
+ w0 = make_decode_worker(world_size=1, tp_rank=0, dp_rank=0)
189
+ seq = [w0._next_flex_tp_rank(8) for _ in range(64)]
190
+ assert seq[:8] == list(range(8))
191
+ assert Counter(seq) == Counter({t: 8 for t in range(8)})
192
+
193
+ first_pick = [
194
+ make_decode_worker(world_size=1, tp_rank=0, dp_rank=d)._next_flex_tp_rank(8)
195
+ for d in range(8)
196
+ ]
197
+ assert sorted(first_pick) == list(range(8))
198
+
199
+
200
+def test_read_blocks_for_req_threads_chosen_tp() -> None:
201
+ # The resolved (chosen_tp, flexible) must reach _read_blocks, which keys the
202
+ # session AND the notify port off that single value -- so a read and its
203
+ # completion notify address the same prefill rank.
204
+ worker = make_decode_worker(world_size=1, tp_rank=0, dp_rank=3)
205
+ worker._read_blocks = MagicMock()
206
+ worker._read_blocks_for_req("r", make_meta(p_tp=8, p_dp=1, remote_dp_rank=0))
207
+ kw = worker._read_blocks.call_args.kwargs
208
+ assert kw["flexible"] is True
209
+ assert kw["chosen_tp"] == 3 # first flexible pick = dp_rank seed
210
+ assert get_port_offset(0, kw["chosen_tp"], 8) == 3
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+5
@@ -422,6 +422,10 @@ class ReqMeta:
422
remote_engine_id: str
423
tp_size: int
424
remote_dp_size: int
425
+ # Prefill DP rank that owns this request's KV (forwarded by the proxy). The
426
+ # read must target this rank's memory registration; the default 0 preserves
427
+ # the symmetric single-DP behaviour.
428
+ remote_dp_rank: int = 0
429
430
431
class MoRIIOConnectorMetadata(KVConnectorMetadata):
@@ -475,6 +479,7 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
479
remote_notify_port=int(remote_notify_port),
480
tp_size=kv_transfer_params.get("tp_size", 1),
481
remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
482
+ remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0),
483
)
484
if write_mode:
485
self.reqs_to_save[request_id] = _req
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+252
-7
@@ -1014,6 +1014,8 @@ class MoRIIOConnectorWorker:
1014
self._handshake_futures: dict[EngineId, Future[set[str]]] = {}
1015
# Protects _handshake_futures and _remote_agents.
1016
self._handshake_lock = threading.RLock()
1017
+ # Remote engines already covered by the eager pre-forward handshake.
1018
+ self._eager_handshaked_engines: set[EngineId] = set()
1019
1020
self.block_size = vllm_config.cache_config.block_size
1021
self.model_config = vllm_config.model_config
@@ -1261,8 +1263,16 @@ class MoRIIOConnectorWorker:
1263
remote_tp_size: int,
1264
expected_engine_id: str,
1265
remote_dp_rank: int = 0,
1266
+ remote_tp_rank: int | None = None,
1267
) -> set[str]:
1265
- """Do a MoRIIO handshake with a remote instance."""
1268
+ """Do a MoRIIO handshake with a remote instance.
1269
+
1270
+ remote_tp_rank: explicit remote TP index to dial. Flexible-read callers
1271
+ pass the chosen prefill TP rank so the handshake, the (dp, tp) session
1272
+ key and the notify port all address the SAME rank. None falls back to the
1273
+ local-rank mapping _remote_tp_rank -- byte-identical for callers not yet
1274
+ TP-aware.
1275
+ """
1276
1277
start_time = time.perf_counter()
1278
@@ -1270,9 +1280,12 @@ class MoRIIOConnectorWorker:
1280
# a hack to keep us moving. We will switch when moving to etcd
1281
# or where we have a single ZMQ socket in the scheduler.
1282
1273
- port_offset = get_port_offset(
1274
- remote_dp_rank, self._remote_tp_rank(remote_tp_size)
1283
+ dial_tp_rank = (
1284
+ self._remote_tp_rank(remote_tp_size)
1285
+ if remote_tp_rank is None
1286
+ else int(remote_tp_rank)
1287
)
1288
+ port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size)
1289
path = make_zmq_path("tcp", host, port + port_offset)
1290
logger.debug("handshake Querying metadata on path: %s", path)
1291
@@ -1729,6 +1742,149 @@ class MoRIIOConnectorWorker:
1742
def get_engine_name_with_dp(self, engine_name, dp_rank):
1743
return f"{engine_name}_dp{dp_rank}"
1744
1745
+ def get_engine_name_with_dp_tp(self, engine_name, dp_rank, tp_rank):
1746
+ # Per-(dp, tp) session key. The flexible mirror read keys sessions per
1747
+ # (dp, tp) so one decode worker can hold a session to EACH prefill TP
1748
+ # rank and spread reads across them; other configs keep the DP-only key.
1749
+ return f"{engine_name}_dp{dp_rank}_tp{tp_rank}"
1750
+
1751
+ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> None:
1752
+ """Handshake EVERY remote prefill DP rank BEFORE the decode forward pass,
1753
+ identically across all local TP workers.
1754
+
1755
+ Why this exists (the deadlock it prevents): with heterogeneous DP prefill
1756
+ a decode TP worker reads KV from whichever prefill DP rank owns the
1757
+ request, so across requests every worker must reach several prefill DP
1758
+ ranks. The decode forward issues per-layer TP collectives (e.g. an
1759
+ all-gather) that all local TP workers must enter together. If the
1760
+ handshakes are left to fire lazily on the read path, the workers diverge:
1761
+ a worker whose target rank is already cached races ahead into the forward
1762
+ collective while a peer is still blocked in a handshake recv(). The first
1763
+ worker then waits inside the collective for the stuck peer -> 600s NCCL
1764
+ timeout / hang. This was observed directly with mixed TP<->DP configs.
1765
+
1766
+ Fix: complete ALL prefill-DP-rank handshakes for every referenced remote
1767
+ engine HERE, before any read enters the forward, so no worker is still
1768
+ handshaking once its peers reach a collective. Fires ONCE per remote
1769
+ engine (first contact), gated by _eager_handshaked_engines. The engine
1770
+ set comes from scheduler-built metadata (identical on every TP worker),
1771
+ so all workers run the same handshakes in the same order and reach the
1772
+ all-reduce barrier below together.
1773
+
1774
+ Failure handling: handshake exceptions are caught, never raised before
1775
+ the collective (raising early would hang the peers still waiting for it).
1776
+ Every worker reaches the all-reduce(MIN) vote; if ANY worker failed, ALL
1777
+ raise the same error AFTER the collective, so the step fails fast and
1778
+ uniformly in ~seconds instead of one rank hanging the forward for 600s.
1779
+ """
1780
+ import torch.distributed as dist
1781
+
1782
+ # Distinct remote engines referenced this step, in metadata (==
1783
+ # scheduler) order so every TP worker iterates engines identically.
1784
+ engines: dict[str, ReqMeta] = {}
1785
+ for _req_id, meta in metadata.reqs_to_recv.items():
1786
+ remote_engine_id = (
1787
+ str(meta.remote_host) + ":" + str(meta.remote_handshake_port)
1788
+ )
1789
+ engines.setdefault(remote_engine_id, meta)
1790
+
1791
+ for remote_engine_id, meta in engines.items():
1792
+ if remote_engine_id in self._eager_handshaked_engines:
1793
+ continue
1794
+
1795
+ remote_dp_size = int(meta.remote_dp_size)
1796
+ port = int(meta.remote_handshake_port)
1797
+ tp_size = int(meta.tp_size)
1798
+
1799
+ # Flexible mirror (TP prefill + MLA, world_size==1 decode): the read
1800
+ # round-robins over prefill TP ranks, so pre-warm a session to EVERY
1801
+ # (dp, tp) rank. Other configs pre-warm per DP rank (tp resolved by
1802
+ # the fixed local-rank mapping) -- byte-identical to before. The
1803
+ # mirror's decode is DP+EP, whose forward all-to-all is the collective
1804
+ # that the eager barrier keeps everyone in step for.
1805
+ flexible = (
1806
+ self.world_size == 1
1807
+ and self.use_mla
1808
+ and remote_dp_size == 1
1809
+ and tp_size > 1
1810
+ )
1811
+ # (engine_id, dp_rank, tp_rank_or_None); tp_rank is None on the legacy
1812
+ # path so _moriio_handshake falls back to its _remote_tp_rank mapping.
1813
+ targets: list[tuple[Any, int, int | None]]
1814
+ if flexible:
1815
+ targets = [
1816
+ (self.get_engine_name_with_dp_tp(remote_engine_id, dp, tp), dp, tp)
1817
+ for dp in range(remote_dp_size)
1818
+ for tp in range(max(1, tp_size))
1819
+ ]
1820
+ else:
1821
+ targets = [
1822
+ (self.get_engine_name_with_dp(remote_engine_id, dp), dp, None)
1823
+ for dp in range(remote_dp_size)
1824
+ ]
1825
+
1826
+ # Submit handshakes for every not-yet-known target UNDER the lock; do
1827
+ # NOT hold it across the join or the collective (a stalled recv must
1828
+ # not block another thread's lock acquisition). Gate on BOTH
1829
+ # _remote_agents AND layer metadata: a rank with an agent entry but no
1830
+ # layer metadata is half-handshaked and would KeyError at read time.
1831
+ futures: list[tuple[str, Future[set[str]]]] = []
1832
+ with self._handshake_lock:
1833
+ for eid, cur_dp_rank, cur_tp_rank in targets:
1834
+ if (
1835
+ eid in self._remote_agents
1836
+ and eid in self.layer_name_to_remote_kv_cache_metadata
1837
+ ):
1838
+ continue
1839
+ fut = self._handshake_initiation_executor.submit(
1840
+ self._moriio_handshake,
1841
+ meta.remote_host,
1842
+ port,
1843
+ tp_size,
1844
+ eid,
1845
+ cur_dp_rank,
1846
+ cur_tp_rank,
1847
+ )
1848
+ futures.append((eid, fut))
1849
+
1850
+ # Join outside the lock. Bounded handshake errors are recorded here
1851
+ # and reported after the all-reduce.
1852
+ all_ok = True
1853
+ results: dict[str, set[str]] = {}
1854
+ for eid, fut in futures:
1855
+ try:
1856
+ results[eid] = fut.result()
1857
+ except Exception:
1858
+ logger.exception("Eager MoRIIO handshake failed for %s", eid)
1859
+ all_ok = False
1860
+
1861
+ with self._handshake_lock:
1862
+ for eid, agents in results.items():
1863
+ self._remote_agents[eid] = agents
1864
+
1865
+ logger.info(
1866
+ "Eager MoRIIO handshake: engine=%s dp_size=%d new_ranks=%d "
1867
+ "ok=%s tp_rank=%d",
1868
+ remote_engine_id,
1869
+ remote_dp_size,
1870
+ len(futures),
1871
+ all_ok,
1872
+ self.tp_rank,
1873
+ )
1874
+ # CPU all-reduce = TP-uniform success vote AND lockstep barrier: it
1875
+ # blocks until every TP worker arrives, gives them the same verdict,
1876
+ # and stays off the model compute stream.
1877
+ vote = torch.tensor([1 if all_ok else 0], device="cpu", dtype=torch.int32)
1878
+ dist.all_reduce(vote, group=self.tp_group.cpu_group, op=dist.ReduceOp.MIN)
1879
+ if int(vote.item()) == 0:
1880
+ raise HandshakeError(
1881
+ f"Eager MoRIIO handshake failed for {remote_engine_id} on "
1882
+ "at least one TP rank; failing this step fast to avoid a "
1883
+ "TP collective hang"
1884
+ )
1885
+
1886
+ self._eager_handshaked_engines.add(remote_engine_id)
1887
+
1888
def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
1889
"""
1890
Start loading by triggering non-blocking moriio_xfer.
@@ -1750,6 +1906,12 @@ class MoRIIOConnectorWorker:
1906
if self.mode == MoRIIOMode.WRITE:
1907
return
1908
1909
+ # Handshake every referenced remote prefill rank up front, before any
1910
+ # read enters the forward pass. A lazy per-rank handshake on the read
1911
+ # path lets TP workers diverge into a forward collective while a peer is
1912
+ # still blocked handshaking -> NCCL hang (see below).
1913
+ self._eager_handshake_all_dp_ranks(metadata)
1914
+
1915
wait_handshake_readd_req = False
1916
remote_engine_id = None
1917
@@ -1758,8 +1920,15 @@ class MoRIIOConnectorWorker:
1920
str(meta.remote_host) + ":" + str(meta.remote_handshake_port)
1921
)
1922
meta.remote_engine_id = remote_engine_id
1923
+ # The eager handshake above already covered every referenced engine
1924
+ # (and keys the mirror per (dp, tp), which the DP-only dp0 probe below
1925
+ # would miss). Only fall back to the lazy background handshake for an
1926
+ # engine it did not cover.
1927
dp0_remote_engine_id = self.get_engine_name_with_dp(remote_engine_id, 0)
1762
- if dp0_remote_engine_id not in self._remote_agents:
1928
+ if (
1929
+ remote_engine_id not in self._eager_handshaked_engines
1930
+ and dp0_remote_engine_id not in self._remote_agents
1931
+ ):
1932
# Initiate handshake with remote engine to exchange metadata.
1933
with self._handshake_lock:
1934
if remote_engine_id not in self._remote_agents:
@@ -1809,12 +1978,53 @@ class MoRIIOConnectorWorker:
1978
self.save_kv_layer(metadata, layer_name, kv_layer, None)
1979
self._writer.seal_pending_transfers()
1980
1981
+ def _next_flex_tp_rank(self, remote_tp_size: int) -> int:
1982
+ """Deterministic round-robin over prefill tp0..N-1 for the flexible read.
1983
+
1984
+ Round-robin (not random): exactly uniform and testable, with the same
1985
+ prefill-NIC balancing. Seeded from this decode rank's dp_rank so
1986
+ concurrent decode DP ranks are phase-staggered -- at a given read index
1987
+ distinct decode ranks target distinct prefill TP ranks.
1988
+ """
1989
+ rr = getattr(self, "_flex_tp_rr", None)
1990
+ if rr is None:
1991
+ rr = int(getattr(self, "dp_rank", 0) or 0)
1992
+ self._flex_tp_rr = rr + 1
1993
+ return rr % remote_tp_size
1994
+
1995
+ def _resolve_read_source(self, meta: ReqMeta) -> tuple[int, bool]:
1996
+ """Resolve (chosen_tp, flexible) for reading this request's KV.
1997
+
1998
+ Flexible mirror (decode world_size==1 + MLA + pure-TP prefill): MLA
1999
+ replicates the latent KV across the prefill TP ranks, so any is a valid
2000
+ source; round-robin across them to spread RDMA/NIC load. Otherwise the
2001
+ source TP rank is fixed by the local-rank mapping (_remote_tp_rank) --
2002
+ forward DP8EP->TP8 -> tp0; symmetric TP -> tp_rank -- byte-identical to
2003
+ prior behaviour. chosen_tp is the single value threaded into the (dp, tp)
2004
+ session key, the handshake dial and the notify port, so all three address
2005
+ the SAME prefill rank (drift -> read one rank but notify another -> the
2006
+ read rank's prefill buffer is never freed).
2007
+ """
2008
+ remote_tp_size = int(meta.tp_size)
2009
+ flexible = (
2010
+ self.world_size == 1
2011
+ and self.use_mla
2012
+ and int(meta.remote_dp_size) == 1
2013
+ and remote_tp_size > 1
2014
+ )
2015
+ if flexible:
2016
+ chosen_tp = self._next_flex_tp_rank(remote_tp_size)
2017
+ else:
2018
+ chosen_tp = self._remote_tp_rank(remote_tp_size)
2019
+ return chosen_tp, flexible
2020
+
2021
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
2022
logger.debug(
2023
"Remote agent %s available, calling _read_blocks for req %s",
2024
meta.remote_engine_id,
2025
req_id,
2026
)
2027
+ chosen_tp, flexible = self._resolve_read_source(meta)
2028
self._read_blocks(
2029
request_id=req_id,
2030
transfer_id=meta.transfer_id,
@@ -1824,6 +2034,9 @@ class MoRIIOConnectorWorker:
2034
remote_host=meta.remote_host,
2035
remote_notify_port=meta.remote_notify_port,
2036
remote_tp_size=meta.tp_size,
2037
+ remote_dp_rank=meta.remote_dp_rank,
2038
+ chosen_tp=chosen_tp,
2039
+ flexible=flexible,
2040
)
2041
2042
def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):
@@ -1976,12 +2189,37 @@ class MoRIIOConnectorWorker:
2189
remote_host: str,
2190
remote_notify_port: int,
2191
remote_tp_size: int,
2192
+ remote_dp_rank: int = 0,
2193
+ chosen_tp: int | None = None,
2194
+ flexible: bool = False,
2195
) -> None:
2196
if self.mode == MoRIIOMode.WRITE:
2197
return
2198
1983
- dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0)
1984
- sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id)
2199
+ # Read from the prefill rank that actually computed this request's KV
2200
+ # (forwarded by the proxy). Hardcoding DP0 reads from a different rank's
2201
+ # memory registration; per-rank num_blocks differ, so high block ids can
2202
+ # overrun the wrong rank's region.
2203
+ #
2204
+ # eff_tp = the remote TP rank this read targets. The flexible mirror
2205
+ # reads from a round-robin-chosen prefill TP rank and keys the session
2206
+ # per (dp, tp); other configs use the fixed local-rank mapping (eff_tp ==
2207
+ # _remote_tp_rank), byte-identical to before. This key MUST match the one
2208
+ # the eager handshake stored the session under.
2209
+ eff_tp = (
2210
+ int(chosen_tp)
2211
+ if chosen_tp is not None
2212
+ else self._remote_tp_rank(remote_tp_size)
2213
+ )
2214
+ if flexible:
2215
+ remote_dp_engine_id = self.get_engine_name_with_dp_tp(
2216
+ dst_engine_id, int(remote_dp_rank), eff_tp
2217
+ )
2218
+ else:
2219
+ remote_dp_engine_id = self.get_engine_name_with_dp(
2220
+ dst_engine_id, int(remote_dp_rank)
2221
+ )
2222
+ sessions, remote_moriio_meta = self._get_built_session(remote_dp_engine_id)
2223
2224
# SQ-full backpressure deadline, shared across this request's layers.
2225
_sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout
@@ -2030,6 +2268,13 @@ class MoRIIOConnectorWorker:
2268
self._recving_transfers[request_id].append(transfer_status)
2269
self._recving_transfers_callback_addr[request_id] = (
2270
remote_host,
2033
- str(remote_notify_port + self._remote_tp_rank(remote_tp_size)),
2271
+ str(
2272
+ remote_notify_port
2273
+ + get_port_offset(
2274
+ int(remote_dp_rank),
2275
+ eff_tp,
2276
+ remote_tp_size,
2277
+ )
2278
+ ),
2279
transfer_id,
2280
)