[ROCm] Fix misrouting race-condition in multi-decode P/D disagg with mori-io (#51681)
Signed-off-by: Vincent Cave <vincent.cave@amd.com> Signed-off-by: Shiksha Patel <shikpate@amd.com> Co-authored-by: Shiksha Patel <shikpate@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Vincent Cavé committed
Sep 24, 2026 at 21:52 UTC
92e9c6bf13138da9caa42c815e242affee046c20
6 files changed
+343
-49
tests/v1/kv_connector/unit/test_moriio_kv_layout.py
+10
-21
@@ -4,8 +4,7 @@
4
import importlib.util
5
import math
6
import threading
7
-from collections import OrderedDict, defaultdict
8
-from queue import Queue
7
+from collections import OrderedDict
8
from types import SimpleNamespace
9
from typing import Any
10
@@ -19,6 +18,8 @@ from vllm.v1.kv_cache_interface import (
18
MLAAttentionSpec,
19
)
20
21
+from .utils import make_moriio_writer
22
+
23
aiter_available = importlib.util.find_spec("aiter") is not None
24
mori_available = importlib.util.find_spec("mori") is not None
25
@@ -90,18 +91,6 @@ def _remote_meta(num_blocks: int = 16) -> SimpleNamespace:
91
return SimpleNamespace(num_blocks=num_blocks)
92
93
93
-def _writer_with_fake_worker(fake_worker: Any) -> Any:
94
- writer = MoRIIOWriter.__new__(MoRIIOWriter)
95
- writer._worker_ref = lambda: fake_worker
96
- writer._write_task_q = Queue()
97
- writer._write_state_lock = threading.Lock()
98
- writer._scheduled_writes = defaultdict(int)
99
- writer._scheduled_layers = defaultdict(set)
100
- writer._sealed_writes = {}
101
- writer.ensure_worker_started = lambda: None
102
- return writer
103
-
104
-
94
def _wrapper_for_messages() -> Any:
95
wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper)
96
wrapper.lock = threading.Lock()
@@ -409,7 +398,7 @@ def test_write_scheduler_deduplicates_layers_and_seals_expected_count():
398
request_info = RemoteAllocInfo(block_ids=[4, 5])
399
wrapper = _wrapper_for_messages()
400
wrapper.done_remote_allocate_req_dict["xfer"] = request_info
412
- writer = _writer_with_fake_worker(SimpleNamespace(moriio_wrapper=wrapper))
401
+ writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper))
402
403
assert writer.schedule_write(_write_task("dense0"))
404
assert not writer.schedule_write(_write_task("dense0"))
@@ -459,12 +448,12 @@ def test_write_completion_notifies_once_after_all_sealed_writes_finish():
448
request_info = RemoteAllocInfo(block_ids=[4, 5], writes_expected=2)
449
request_info.transfer_statuses.extend(["status-a", "status-b"])
450
request_info.completion_request_id = "req"
462
- request_info.completion_remote_notify_port = 7000
451
+ # Final notify port is resolved in _execute_write_task; finalize reads it
452
+ # as-is, so set the already-offset value (base 7000 + tp_rank-2 offset).
453
+ request_info.completion_notify_port = 7002
454
request_info.completion_remote_ip = "127.0.0.1"
455
wrapper.done_remote_allocate_req_dict["xfer"] = request_info
465
- writer = _writer_with_fake_worker(
466
- SimpleNamespace(moriio_wrapper=wrapper, tp_rank=2)
467
- )
456
+ writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper, tp_rank=2))
457
writer._scheduled_writes["xfer"] = 2
458
writer._scheduled_layers["xfer"] = {"dense0", "indexer"}
459
writer._sealed_writes["xfer"] = 2
@@ -512,7 +501,7 @@ def test_moriio_wrapper_waits_scoped_statuses_without_global_drain():
501
def test_write_failure_marks_terminal_and_clears_scheduled_state():
502
wrapper = _wrapper_for_messages()
503
wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5])
515
- writer = _writer_with_fake_worker(SimpleNamespace(moriio_wrapper=wrapper))
504
+ writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper))
505
writer._scheduled_writes["xfer"] = 2
506
writer._scheduled_layers["xfer"] = {"dense0", "indexer"}
507
writer._sealed_writes["xfer"] = 2
@@ -530,7 +519,7 @@ def test_write_failure_marks_terminal_and_clears_scheduled_state():
519
def test_schedule_write_rejects_terminal_transfer_without_recreating_state():
520
wrapper = _wrapper_for_messages()
521
wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5])
533
- writer = _writer_with_fake_worker(SimpleNamespace(moriio_wrapper=wrapper))
522
+ writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper))
523
writer._scheduled_writes["xfer"] = 1
524
writer._scheduled_layers["xfer"] = {"dense0"}
525
writer._sealed_writes["xfer"] = 1
tests/v1/kv_connector/unit/test_moriio_write_done_routing.py
new
+271
@@ -0,0 +1,271 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for mori-io write_done routing (#51681)."""
4
+
5
+import threading
6
+from types import SimpleNamespace
7
+from typing import Any
8
+
9
+import pytest
10
+
11
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
12
+ RemoteAllocInfo,
13
+ ReqMeta,
14
+ WriteTask,
15
+ get_port_offset,
16
+)
17
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
18
+ MoRIIOConnectorWorker,
19
+)
20
+
21
+from .utils import make_moriio_writer
22
+
23
+pytestmark = pytest.mark.cpu_test
24
+
25
+
26
+def _fake_worker(
27
+ alloc: dict[str, Any],
28
+ *,
29
+ tp_rank: int = 0,
30
+ set_remote_dp_size_local: int | None = None,
31
+) -> Any:
32
+ wrapper = SimpleNamespace(
33
+ done_remote_allocate_req_dict=alloc,
34
+ lock=threading.Lock(),
35
+ done_req_ids=[],
36
+ sent_notifies=[],
37
+ )
38
+ wrapper.send_notify = lambda transfer_id, remote_ip, remote_port, message_type: (
39
+ wrapper.sent_notifies.append(
40
+ (transfer_id, remote_ip, remote_port, message_type)
41
+ )
42
+ )
43
+ wrapper.waiting_for_transfer_complete = lambda statuses: None
44
+ wrapper._mark_transfer_terminal_locked = lambda transfer_id: None
45
+ wrapper._is_transfer_terminal_locked = lambda transfer_id: False
46
+ worker = SimpleNamespace(moriio_wrapper=wrapper, tp_rank=tp_rank)
47
+ if set_remote_dp_size_local is not None:
48
+ worker.remote_dp_size_local = set_remote_dp_size_local
49
+ return worker
50
+
51
+
52
+def _req_meta(
53
+ *,
54
+ transfer_id: str = "tA",
55
+ remote_host: str = "10.0.0.1",
56
+ multi_pod_hosts: list[str] | None = None,
57
+ remote_dp_size: int = 16,
58
+ remote_dp_size_local: int = 8,
59
+) -> ReqMeta:
60
+ return ReqMeta(
61
+ transfer_id=transfer_id,
62
+ local_block_ids=[[1]],
63
+ remote_block_ids=[[2]],
64
+ remote_host=remote_host,
65
+ remote_port=1,
66
+ remote_handshake_port=2,
67
+ remote_notify_port=20000,
68
+ remote_engine_id="remote-engine",
69
+ tp_size=1,
70
+ remote_dp_size=remote_dp_size,
71
+ multi_pod_hosts=list(multi_pod_hosts or []),
72
+ remote_dp_size_local=remote_dp_size_local,
73
+ )
74
+
75
+
76
+def _write_task(
77
+ transfer_id: str,
78
+ remote_ip: str,
79
+ *,
80
+ multi_pod_hosts: list[str] | None = None,
81
+ remote_dp_size_local: int = 0,
82
+ request_id: str = "req",
83
+) -> WriteTask:
84
+ return WriteTask(
85
+ request_id=request_id,
86
+ transfer_id=transfer_id,
87
+ dst_engine_id="remote-engine",
88
+ local_block_ids=[1, 3],
89
+ remote_block_ids_hint=None,
90
+ layer_name="layer0",
91
+ event=None,
92
+ remote_notify_port=20000,
93
+ remote_ip=remote_ip,
94
+ multi_pod_hosts=multi_pod_hosts or [],
95
+ remote_dp_size_local=remote_dp_size_local,
96
+ )
97
+
98
+
99
+def test_execute_write_task_ip_is_per_task_under_overwrite():
100
+ alloc = {"tA": RemoteAllocInfo(block_ids=None, decode_dp_rank=0)}
101
+ worker = _fake_worker(alloc)
102
+ worker.multi_pod_hosts = ["10.0.0.2"]
103
+ worker.remote_dp_size_local = 1
104
+
105
+ writer = make_moriio_writer(worker)
106
+ task_a = _write_task(
107
+ "tA", "10.0.0.1", multi_pod_hosts=["10.0.0.1"], remote_dp_size_local=1
108
+ )
109
+ writer._execute_write_task(task_a)
110
+
111
+ assert alloc["tA"].completion_remote_ip == "10.0.0.1"
112
+ assert alloc["tA"].completion_remote_ip != "10.0.0.2"
113
+
114
+
115
+def test_deferred_task_ip_is_per_task_under_overwrite():
116
+ alloc = {"tA": RemoteAllocInfo(block_ids=None, decode_dp_rank=0)}
117
+ worker = _fake_worker(alloc)
118
+ worker.multi_pod_hosts = ["10.0.0.2"]
119
+ worker.remote_dp_size_local = 1
120
+
121
+ writer = make_moriio_writer(worker)
122
+ task_a = _write_task(
123
+ "tA", "10.0.0.1", multi_pod_hosts=["10.0.0.1"], remote_dp_size_local=1
124
+ )
125
+ writer._deferred_tasks.append(task_a)
126
+ writer._process_deferred_tasks()
127
+
128
+ assert alloc["tA"].completion_remote_ip == "10.0.0.1"
129
+ assert writer._deferred_tasks == []
130
+
131
+
132
+def test_two_requests_resolve_distinct_hosts():
133
+ alloc = {
134
+ "tA": RemoteAllocInfo(block_ids=None, decode_dp_rank=0),
135
+ "tB": RemoteAllocInfo(block_ids=None, decode_dp_rank=3),
136
+ }
137
+ worker = _fake_worker(alloc)
138
+ writer = make_moriio_writer(worker)
139
+
140
+ task_a = _write_task(
141
+ "tA", "10.0.0.1", multi_pod_hosts=["10.0.0.1"], remote_dp_size_local=1
142
+ )
143
+ task_b = _write_task(
144
+ "tB", "10.0.0.2", multi_pod_hosts=["10.0.0.2"], remote_dp_size_local=1
145
+ )
146
+ worker.multi_pod_hosts = ["10.0.0.2"]
147
+ writer._execute_write_task(task_a)
148
+ writer._execute_write_task(task_b)
149
+
150
+ assert alloc["tA"].completion_remote_ip == "10.0.0.1"
151
+ assert alloc["tB"].completion_remote_ip == "10.0.0.2"
152
+
153
+
154
+def test_finalize_notify_port_uses_pod_local_rank():
155
+ base = 20000
156
+ info = RemoteAllocInfo(block_ids=None, decode_dp_rank=9)
157
+ alloc = {"tA": info}
158
+ worker = _fake_worker(alloc, tp_rank=0)
159
+ writer = make_moriio_writer(worker)
160
+
161
+ task = _write_task(
162
+ "tA",
163
+ "10.0.0.1",
164
+ multi_pod_hosts=["10.0.0.1", "10.0.0.2"],
165
+ remote_dp_size_local=8,
166
+ )
167
+ writer._execute_write_task(task)
168
+ assert info.completion_remote_ip == "10.0.0.2"
169
+
170
+ info.writes_expected = 1
171
+ info.writes_done = 1
172
+ writer._finalize_if_complete("tA", info)
173
+
174
+ assert worker.moriio_wrapper.sent_notifies
175
+ _tid, ip, port, mtype = worker.moriio_wrapper.sent_notifies[-1]
176
+ assert mtype == "write_done"
177
+ assert ip == "10.0.0.2"
178
+ assert port == base + get_port_offset(9 % 8, 0)
179
+
180
+
181
+def test_finalize_port_ignores_stale_worker_dp_local():
182
+ base = 20000
183
+ info = RemoteAllocInfo(block_ids=None, decode_dp_rank=9)
184
+ alloc = {"tA": info}
185
+ worker = _fake_worker(alloc, tp_rank=0, set_remote_dp_size_local=1)
186
+ writer = make_moriio_writer(worker)
187
+
188
+ task = _write_task(
189
+ "tA",
190
+ "10.0.0.1",
191
+ multi_pod_hosts=["10.0.0.1", "10.0.0.2"],
192
+ remote_dp_size_local=8,
193
+ )
194
+ writer._execute_write_task(task)
195
+ assert info.completion_notify_port == base + get_port_offset(9 % 8, 0)
196
+
197
+ info.writes_expected = 1
198
+ info.writes_done = 1
199
+ writer._finalize_if_complete("tA", info)
200
+
201
+ _tid, _ip, port, mtype = worker.moriio_wrapper.sent_notifies[-1]
202
+ assert mtype == "write_done"
203
+ assert port == base + get_port_offset(9 % 8, 0)
204
+ assert port != base + get_port_offset(9 % 1, 0)
205
+
206
+
207
+def test_finalize_single_pod_port_uses_global_rank():
208
+ base = 20000
209
+ info = RemoteAllocInfo(block_ids=None, decode_dp_rank=3)
210
+ alloc = {"tA": info}
211
+ worker = _fake_worker(alloc, tp_rank=2)
212
+ writer = make_moriio_writer(worker)
213
+
214
+ task = _write_task("tA", "10.0.0.1", multi_pod_hosts=[], remote_dp_size_local=0)
215
+ writer._execute_write_task(task)
216
+ assert info.completion_remote_ip == "10.0.0.1"
217
+ assert info.completion_notify_port == base + get_port_offset(3, 2)
218
+
219
+ info.writes_expected = 1
220
+ info.writes_done = 1
221
+ writer._finalize_if_complete("tA", info)
222
+
223
+ _tid, _ip, port, _mtype = worker.moriio_wrapper.sent_notifies[-1]
224
+ assert port == base + get_port_offset(3, 2)
225
+
226
+
227
+def test_write_blocks_for_req_forwards_routing_to_schedule_write():
228
+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
229
+ captured: dict[str, Any] = {}
230
+
231
+ def _capture_schedule_write_blocks(**kwargs: Any) -> None:
232
+ captured.update(kwargs)
233
+
234
+ worker.schedule_write_blocks = _capture_schedule_write_blocks # type: ignore[method-assign]
235
+
236
+ meta = _req_meta(
237
+ remote_host="10.0.0.1",
238
+ multi_pod_hosts=["10.0.0.1", "10.0.0.2"],
239
+ remote_dp_size_local=8,
240
+ )
241
+ worker._write_blocks_for_req("req-a", meta, "layer0", kv_layer=None)
242
+
243
+ assert captured["multi_pod_hosts"] == ["10.0.0.1", "10.0.0.2"]
244
+ assert captured["remote_dp_size_local"] == 8
245
+ assert captured["remote_ip"] == "10.0.0.1"
246
+ assert captured["local_block_ids"] == [1]
247
+ assert captured["remote_block_ids"] == [2]
248
+ assert not hasattr(worker, "multi_pod_hosts")
249
+
250
+
251
+def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty():
252
+ worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
253
+ captured: dict[str, Any] = {}
254
+
255
+ def _capture_schedule_write_blocks(**kwargs: Any) -> None:
256
+ captured.update(kwargs)
257
+
258
+ worker.schedule_write_blocks = _capture_schedule_write_blocks # type: ignore[method-assign]
259
+
260
+ meta = _req_meta(
261
+ remote_host="10.0.0.9",
262
+ multi_pod_hosts=[],
263
+ remote_dp_size=4,
264
+ remote_dp_size_local=0,
265
+ )
266
+ worker._write_blocks_for_req("req-b", meta, "layer0", kv_layer=None)
267
+
268
+ assert captured["multi_pod_hosts"] == ["10.0.0.9"]
269
+ assert captured["remote_dp_size_local"] == 4
270
+ assert captured["local_block_ids"] == [1]
271
+ assert captured["remote_block_ids"] == [2]
tests/v1/kv_connector/unit/utils.py
+26
@@ -632,3 +632,29 @@ def make_nixl_push_scheduler(
632
sched.blocks_per_sw = []
633
634
return sched
635
+
636
+
637
+def make_moriio_writer(fake_worker: Any) -> Any:
638
+ """Build a MoRIIOWriter with internals stubbed for unit tests.
639
+
640
+ Bypasses ``__init__`` and wires only the write/finalize state the tests
641
+ touch, including the deferred-task fields used by the routing suite.
642
+ """
643
+ import threading
644
+ from queue import Queue
645
+
646
+ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
647
+ MoRIIOWriter,
648
+ )
649
+
650
+ writer = MoRIIOWriter.__new__(MoRIIOWriter)
651
+ writer._worker_ref = lambda: fake_worker
652
+ writer._write_task_q = Queue()
653
+ writer._write_state_lock = threading.Lock()
654
+ writer._scheduled_writes = defaultdict(int)
655
+ writer._scheduled_layers = defaultdict(set)
656
+ writer._sealed_writes = {}
657
+ writer._deferred_tasks = []
658
+ writer._defer_timeout = 60.0
659
+ writer.ensure_worker_started = lambda: None
660
+ return writer
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+3
-1
@@ -69,6 +69,8 @@ class WriteTask:
69
event: torch.cuda.Event
70
remote_notify_port: int
71
remote_ip: str
72
+ multi_pod_hosts: list[str] = field(default_factory=list)
73
+ remote_dp_size_local: int = 0
74
enqueue_time: float = field(default_factory=time.perf_counter)
75
retried: int = 0
76
@@ -96,7 +98,7 @@ class RemoteAllocInfo:
98
writes_expected: int | None = None
99
decode_dp_rank: int = 0
100
completion_request_id: str | None = None
99
- completion_remote_notify_port: int | None = None
101
+ completion_notify_port: int | None = None
102
completion_remote_ip: str | None = None
103
completion_notified: bool = False
104
transfer_statuses: list[Any] = field(default_factory=list)
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+19
-12
@@ -1776,6 +1776,8 @@ class MoRIIOConnectorWorker:
1776
kv_layer: torch.Tensor,
1777
remote_notify_port: int,
1778
remote_ip: str,
1779
+ multi_pod_hosts: list[str],
1780
+ remote_dp_size_local: int,
1781
) -> None:
1782
"""Schedule a block write operation.
1783
@@ -1789,6 +1791,8 @@ class MoRIIOConnectorWorker:
1791
kv_layer: KV cache tensor
1792
remote_notify_port: Port for completion notification
1793
remote_ip: IP address of remote node
1794
+ multi_pod_hosts: List of pod IPs for multi-pod Wide-EP
1795
+ remote_dp_size_local: Per-pod DP size for multi-pod
1796
1797
"""
1798
# synchronization to prevent dirty reads between
@@ -1810,6 +1814,8 @@ class MoRIIOConnectorWorker:
1814
event=event,
1815
remote_notify_port=remote_notify_port,
1816
remote_ip=remote_ip,
1817
+ multi_pod_hosts=multi_pod_hosts,
1818
+ remote_dp_size_local=remote_dp_size_local,
1819
)
1820
self._writer.schedule_write(task)
1821
@@ -3167,18 +3173,17 @@ class MoRIIOConnectorWorker:
3173
)
3174
3175
def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):
3170
- # Stash multi_pod_hosts + local DP size on the worker so
3171
- # MoRIIOEngine._finalize_if_complete (which sees only the WriteTask,
3172
- # not ReqMeta) can pick the per-rank pod IP for the completion notify.
3173
- # Last-writer-wins is safe: all requests share the same topology.
3174
- if meta.multi_pod_hosts:
3175
- self.multi_pod_hosts = list(meta.multi_pod_hosts)
3176
- else:
3177
- self.multi_pod_hosts = [meta.remote_host]
3178
- if meta.remote_dp_size_local:
3179
- self.remote_dp_size_local = int(meta.remote_dp_size_local)
3180
- else:
3181
- self.remote_dp_size_local = int(meta.remote_dp_size)
3176
+ # Compute per-request values to pass through the task (no shared state).
3177
+ # This avoids race conditions when concurrent requests target different
3178
+ # decode pods - each WriteTask carries its own routing info.
3179
+ hosts = (
3180
+ list(meta.multi_pod_hosts) if meta.multi_pod_hosts else [meta.remote_host]
3181
+ )
3182
+ dp_local = (
3183
+ int(meta.remote_dp_size_local)
3184
+ if meta.remote_dp_size_local
3185
+ else int(meta.remote_dp_size)
3186
+ )
3187
# WRITE does not support HMA, so unwrap blocks into flat lists.
3188
# remote_block_ids can itself be empty so need to be careful when unwrapping.
3189
local_block_ids = meta.local_block_ids[0]
@@ -3193,6 +3198,8 @@ class MoRIIOConnectorWorker:
3198
kv_layer=kv_layer,
3199
remote_notify_port=meta.remote_notify_port,
3200
remote_ip=meta.remote_host,
3201
+ multi_pod_hosts=hosts,
3202
+ remote_dp_size_local=dp_local,
3203
)
3204
3205
def merge_contiguous_blocks(
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+14
-15
@@ -314,18 +314,27 @@ class MoRIIOWriter:
314
request_info = self._get_remote_alloc_info(task.transfer_id)
315
with self._write_state_lock:
316
request_info.completion_request_id = task.request_id
317
- request_info.completion_remote_notify_port = task.remote_notify_port
317
# Wide-EP multi-pod: task.remote_ip addresses only the first pod,
318
# so resolve the per-rank host from multi_pod_hosts (falls back to
319
# task.remote_ip for single-pod or on any indexing miss).
320
_remote_ip = task.remote_ip
322
- _hosts = list(getattr(self.worker, "multi_pod_hosts", []) or [])
323
- _dp_local = int(getattr(self.worker, "remote_dp_size_local", 0) or 0)
321
+ _hosts = task.multi_pod_hosts
322
+ _dp_local = task.remote_dp_size_local
323
if _hosts and _dp_local > 0:
324
_pod_idx = int(request_info.decode_dp_rank) // _dp_local
325
if 0 <= _pod_idx < len(_hosts):
326
_remote_ip = _hosts[_pod_idx]
327
request_info.completion_remote_ip = _remote_ip
328
+ # Resolve the final notify port under the lock so finalize only
329
+ # reads it. The offset uses the per-pod local rank (% dp_local),
330
+ # since each pod binds notify sockets only for its local ranks;
331
+ # single-pod is bit-identical (modulus is a no-op).
332
+ _decode_dp_rank_for_port = int(request_info.decode_dp_rank)
333
+ if _dp_local > 0:
334
+ _decode_dp_rank_for_port = _decode_dp_rank_for_port % _dp_local
335
+ request_info.completion_notify_port = task.remote_notify_port + (
336
+ get_port_offset(_decode_dp_rank_for_port, self.worker.tp_rank)
337
+ )
338
if task.transfer_id in self._sealed_writes:
339
request_info.writes_expected = self._sealed_writes[task.transfer_id]
340
@@ -461,9 +470,9 @@ class MoRIIOWriter:
470
if request_info.completion_notified:
471
return
472
request_id = request_info.completion_request_id
464
- remote_notify_port = request_info.completion_remote_notify_port
473
+ remote_port = request_info.completion_notify_port
474
remote_ip = request_info.completion_remote_ip
466
- if request_id is None or remote_notify_port is None or remote_ip is None:
475
+ if request_id is None or remote_port is None or remote_ip is None:
476
return
477
transfer_statuses = list(request_info.transfer_statuses)
478
request_info.transfer_statuses.clear()
@@ -472,16 +481,6 @@ class MoRIIOWriter:
481
# Wait for this request's transfers to complete.
482
self.worker.moriio_wrapper.waiting_for_transfer_complete(transfer_statuses)
483
475
- # The notify port offset must use the per-pod local rank
476
- # (% dp_local), since each pod binds notify sockets only for its local
477
- # ranks. Single-pod is bit-identical (modulus is a no-op).
478
- _dp_local = int(getattr(self.worker, "remote_dp_size_local", 0) or 0)
479
- _decode_dp_rank_for_port = int(request_info.decode_dp_rank)
480
- if _dp_local > 0:
481
- _decode_dp_rank_for_port = _decode_dp_rank_for_port % _dp_local
482
- remote_port = remote_notify_port + get_port_offset(
483
- _decode_dp_rank_for_port, self.worker.tp_rank
484
- )
484
# Consider using RDMA immediate data in decode side
485
# to eliminate the need for this notification.
486
# Consider including the first gen token from prefill in the notification