llmd+vllm+mori-ep(inter node wide-ep)+mori-io(write) for 2p2d with dp=ep=16 tp=1 (#45043)
Signed-off-by: Shiksha Patel <shikpate@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
shikamd123 committed
Aug 3, 2026 at 17:13 UTC
f42761204fb508bc205418c77b80bd50d0d11774
8 files changed
+711
-121
examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+23
-2
@@ -1,5 +1,20 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+#
4
+# Minimal single-node reference proxy for MoRI-IO prefill/decode disaggregation.
5
+# It demonstrates the DP-rank pinning contract the connector expects: pick one
6
+# prefill DP rank per request (flat_interleaved_dp_route) and pin BOTH legs to it
7
+# via the ``X-data-parallel-rank`` header plus ``kv_transfer_params``
8
+# (``remote_dp_rank`` / ``remote_dp_size`` / ``remote_tp_size``).
9
+#
10
+# Scope: single-node only. It intentionally does NOT supply the cross-pod
11
+# Wide-EP (2P2D) peer info -- ``remote_hosts`` / ``multi_pod_hosts``,
12
+# ``data_parallel_size_local``, ``is_request_leader``. In real multi-pod
13
+# deployments that peer info is produced by a production routing sidecar, e.g.:
14
+# * llm-d-router (https://github.com/llm-d/llm-d-router)
15
+# * vLLM's own vllm-router (examples in the vllm-project org)
16
+# * or any NIXL-shaped router that sets the same kv_transfer_params contract.
17
+# Use this file as a protocol reference, not a production router.
18
import argparse
19
import asyncio
20
import copy
@@ -181,7 +196,7 @@ async def send_request_to_prefill(
196
raise RuntimeError(error_message)
197
198
184
-async def start_decode_request(endpoint, req_data, request_id):
199
+async def start_decode_request(endpoint, req_data, request_id, selected_dp_rank=None):
200
session = aiohttp.ClientSession(
201
timeout=aiohttp.ClientTimeout(total=6 * 6000 * 6000)
202
)
@@ -189,6 +204,10 @@ async def start_decode_request(endpoint, req_data, request_id):
204
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
205
"X-Request-Id": request_id,
206
}
207
+ # Pin the decode leg to the same DP rank as its prefill leg so both land on
208
+ # the same rank (vLLM reads this header directly; no engine-side routing).
209
+ if selected_dp_rank is not None:
210
+ headers["X-data-parallel-rank"] = str(selected_dp_rank)
211
response = await session.post(url=endpoint, json=req_data, headers=headers)
212
return session, response
213
@@ -361,7 +380,9 @@ async def handle_request(api: str, request: Request):
380
381
decode_request_url = decode_instance_endpoint["request_address"] + api
382
decode_request_task = asyncio.create_task(
364
- start_decode_request(decode_request_url, req_data, request_id)
383
+ start_decode_request(
384
+ decode_request_url, req_data, request_id, selected_prefill_dp_rank
385
+ )
386
)
387
388
session, decode_response = await decode_request_task
tests/v1/kv_connector/unit/test_moriio_connector.py
+6
@@ -508,9 +508,15 @@ def test_send_transfer_release_sends_structured_release_message():
508
scheduler._send_transfer_release("xfer-7", "127.0.0.1", 7000)
509
510
payload = sock.send.call_args.args[0]
511
+ # WRITE-mode release advertises the consumer (decode) TP size so the prefill
512
+ # side counts the right number of ACKs via get_moriio_expected_ack_count,
513
+ # mirroring the READ-mode release in _pop_done_transfers (see
514
+ # test_read_completion_sends_structured_release_with_consumer_tp_size). The
515
+ # fixture's tp_size is 2, so consumer_tp_size == 2.
516
assert msgspec.msgpack.decode(payload) == {
517
"type": "release",
518
"transfer_id": "xfer-7",
519
+ "consumer_tp_size": 2,
520
}
521
522
tests/v1/kv_connector/unit/test_moriio_routing.py
new
+92
@@ -0,0 +1,92 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for MoRIIO router-authoritative DP-rank routing.
4
+
5
+These lock in the redesign that addresses the PR #45043 review feedback
6
+(njhill): the connector must NOT self-derive a DP rank, and the prefill leg
7
+must *return* the rank it actually ran on via ``request_finished`` so the
8
+decode->prefill notify is routed by pure propagation (READ / serial WRITE),
9
+not by an independent hash on each side.
10
+
11
+Following test_moriio_unmap.py, we bind the unbound method to a lightweight
12
+stand-in rather than constructing a full scheduler.
13
+"""
14
+
15
+from types import SimpleNamespace
16
+
17
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
18
+ MoRIIOConnectorScheduler,
19
+)
20
+from vllm.v1.request import RequestStatus
21
+
22
+_request_finished = MoRIIOConnectorScheduler.request_finished
23
+
24
+
25
+def _producer(global_dp_rank: int, dp_size: int, dp_size_local: int):
26
+ """Minimal prefill-side (producer) scheduler stand-in for request_finished."""
27
+ return SimpleNamespace(
28
+ is_producer=True,
29
+ _global_dp_rank=global_dp_rank,
30
+ engine_id="engine-abc",
31
+ host_ip="10.0.0.1",
32
+ handshake_port=6300,
33
+ side_notify_port=7300,
34
+ vllm_config=SimpleNamespace(
35
+ parallel_config=SimpleNamespace(
36
+ data_parallel_size=dp_size,
37
+ data_parallel_size_local=dp_size_local,
38
+ tensor_parallel_size=1,
39
+ )
40
+ ),
41
+ )
42
+
43
+
44
+def _req(
45
+ request_id: str, params: dict | None, status=RequestStatus.FINISHED_LENGTH_CAPPED
46
+):
47
+ return SimpleNamespace(
48
+ request_id=request_id,
49
+ kv_transfer_params=params,
50
+ status=status,
51
+ )
52
+
53
+
54
+def test_request_finished_returns_prefill_rank_for_propagation():
55
+ # 2P2D: global rank 11 lives on pod 1 (ranks 8..15), local size 8.
56
+ s = _producer(global_dp_rank=11, dp_size=16, dp_size_local=8)
57
+ req = _req("r1", {"do_remote_decode": True, "transfer_id": "t1"})
58
+
59
+ delay_free, params = _request_finished(s, req, block_ids=[])
60
+
61
+ # The prefill leg returns the rank IT actually ran on, marked authoritative,
62
+ # so the router forwards it to decode and the notify propagates (no hashing).
63
+ assert params is not None
64
+ assert params["remote_dp_rank"] == 11
65
+ assert params["remote_dp_rank_override"] is True
66
+ assert params["remote_dp_size"] == 16
67
+ assert params["remote_dp_size_local"] == 8
68
+ assert params["transfer_id"] == "t1"
69
+ # Empty block_ids => nothing deferred, blocks freed immediately.
70
+ assert delay_free is False
71
+
72
+
73
+def test_request_finished_rank_tracks_the_actual_rank():
74
+ # A different rank must be echoed verbatim -- the value is the producer's
75
+ # own rank, never a recomputation of the request id.
76
+ s = _producer(global_dp_rank=4, dp_size=8, dp_size_local=8)
77
+ req = _req("r2", {"do_remote_decode": True, "transfer_id": "t2"})
78
+
79
+ _, params = _request_finished(s, req, block_ids=[])
80
+
81
+ assert params["remote_dp_rank"] == 4
82
+ assert params["remote_dp_size_local"] == 8
83
+
84
+
85
+def test_request_finished_no_kv_params_is_noop():
86
+ s = _producer(global_dp_rank=0, dp_size=16, dp_size_local=8)
87
+ req = _req("r3", None)
88
+
89
+ delay_free, params = _request_finished(s, req, block_ids=[])
90
+
91
+ assert delay_free is False
92
+ assert params is None
tests/v1/kv_connector/unit/test_moriio_tp_ack.py
+2
@@ -296,6 +296,8 @@ def test_read_completion_sends_structured_release_with_consumer_tp_size():
296
worker._recving_transfers_callback_addr = {
297
"req": ("127.0.0.1", "7000", "tx-release")
298
}
299
+ # Transfer-timeout reaping state consulted by _pop_done_transfers.
300
+ worker._recving_transfers_start = {}
301
302
assert worker._pop_done_transfers() == {"tx-release"}
303
assert worker.moriio_wrapper.sent == [
tests/v1/kv_connector/unit/test_moriio_unmap.py
new
+68
@@ -0,0 +1,68 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for MoRIIOConnectorScheduler.{map,unmap}_request_id.
4
+
5
+vLLM's input_processor mutates a request_id after the connector has stored it
6
+(e.g. ``cmpl-abc-0`` -> ``cmpl-abc-0-956053a4``), so an exact-match unmap can
7
+miss. These tests cover the stable-``transfer_id`` fallback that handles that
8
+case (and the clean exact-match and miss paths).
9
+
10
+The map/unmap methods touch only the two id<->id dicts, so we bind them to a
11
+lightweight stand-in rather than constructing a full scheduler.
12
+"""
13
+
14
+from types import SimpleNamespace
15
+
16
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
17
+ MoRIIOConnectorScheduler,
18
+)
19
+
20
+_map = MoRIIOConnectorScheduler.map_request_id
21
+_unmap = MoRIIOConnectorScheduler.unmap_request_id
22
+
23
+
24
+def _sched():
25
+ return SimpleNamespace(
26
+ transfer_id_to_request_id={},
27
+ request_id_to_transfer_id={},
28
+ )
29
+
30
+
31
+def test_exact_match_unmap_clears_both_tables():
32
+ s = _sched()
33
+ _map(s, "rid-1", "tid-1")
34
+ _unmap(s, "rid-1")
35
+ assert s.request_id_to_transfer_id == {}
36
+ assert s.transfer_id_to_request_id == {}
37
+
38
+
39
+def test_unmap_via_transfer_id_after_request_id_mutation():
40
+ s = _sched()
41
+ # Connector stores the ORIGINAL request_id at map time.
42
+ _map(s, "cmpl-abc-0", "tid-42")
43
+ # input_processor later mutates the rid; exact match would miss, so the
44
+ # stable transfer_id is used to find and clear the original entry.
45
+ _unmap(s, "cmpl-abc-0-956053a4", transfer_id="tid-42")
46
+ assert s.request_id_to_transfer_id == {}
47
+ assert s.transfer_id_to_request_id == {}
48
+
49
+
50
+def test_unmap_miss_is_noop_and_does_not_raise():
51
+ s = _sched()
52
+ _map(s, "rid-1", "tid-1")
53
+ # Unknown rid and unknown/None transfer_id: must not raise, must not touch
54
+ # the existing (unrelated) mapping.
55
+ _unmap(s, "rid-unknown", transfer_id=None)
56
+ _unmap(s, "rid-unknown", transfer_id="tid-unknown")
57
+ assert s.request_id_to_transfer_id == {"rid-1": "tid-1"}
58
+ assert s.transfer_id_to_request_id == {"tid-1": "rid-1"}
59
+
60
+
61
+def test_exact_match_preferred_over_transfer_id_fallback():
62
+ s = _sched()
63
+ _map(s, "rid-1", "tid-1")
64
+ # Exact rid is present, so it is removed directly even if a transfer_id is
65
+ # also supplied; no stale entries left behind.
66
+ _unmap(s, "rid-1", transfer_id="tid-1")
67
+ assert s.request_id_to_transfer_id == {}
68
+ assert s.transfer_id_to_request_id == {}
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+110
-21
@@ -191,6 +191,25 @@ def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int:
191
return (dp_rank) * tp_size + tp_rank
192
193
194
+def fold_local_rank(global_dp_rank: int, dp_size_local: int) -> int:
195
+ """Fold a global DP rank into its pod-local rank [0, dp_size_local).
196
+
197
+ ``dp_size_local == 0`` is the external-DP sentinel (local size unknown):
198
+ return the rank unchanged, since a global DP rank is always < the global
199
+ DP size so no folding is needed and the modulo is skipped (never divides
200
+ by zero).
201
+ """
202
+ return global_dp_rank % dp_size_local if dp_size_local else global_dp_rank
203
+
204
+
205
+def pod_index(global_dp_rank: int, dp_size_local: int) -> int:
206
+ """Pod index (0-based) a global DP rank lives on for Wide-EP multi-pod.
207
+
208
+ ``dp_size_local == 0`` (external-DP sentinel) collapses to a single pod.
209
+ """
210
+ return global_dp_rank // dp_size_local if dp_size_local else 0
211
+
212
+
213
def resolve_host_ip(extra_config: dict) -> str:
214
"""The IP this MoRIIO process advertises for KV transfer.
215
@@ -280,9 +299,11 @@ class MoRIIOConfig:
299
kv_transfer_config = vllm_config.kv_transfer_config
300
extra_config = kv_transfer_config.kv_connector_extra_config
301
tp_rank = get_tensor_model_parallel_rank()
283
- dp_rank = vllm_config.parallel_config.data_parallel_rank
302
+ # Fold the global data_parallel_rank back to [0, dp_size_local) for
303
+ # per-node port allocation (handles the external-DP sentinel).
304
+ pc = vllm_config.parallel_config
305
+ dp_rank = fold_local_rank(pc.data_parallel_rank, pc.data_parallel_size_local)
306
base_notify_port = int(extra_config["notify_port"])
285
- dp_size = vllm_config.parallel_config.data_parallel_size
307
tp_size = get_tensor_model_parallel_world_size()
308
port_offset = get_port_offset(dp_rank, tp_rank)
309
backend = str(extra_config.get("backend", "rdma")).lower()
@@ -304,15 +325,18 @@ class MoRIIOConfig:
325
return cls(
326
local_ip=resolve_host_ip(extra_config),
327
local_kv_port=get_open_port(),
307
- proxy_ip=extra_config["proxy_ip"],
328
+ proxy_ip=extra_config.get("proxy_ip", ""),
329
local_ping_port=get_open_port(),
309
- proxy_ping_port=int(extra_config["proxy_ping_port"]),
330
+ proxy_ping_port=int(extra_config.get("proxy_ping_port", 0)),
331
http_port=int(extra_config["http_port"]),
332
handshake_port=int(extra_config["handshake_port"]),
333
notify_port=base_notify_port + port_offset,
334
tp_rank=tp_rank,
335
dp_rank=dp_rank,
315
- dp_size=dp_size,
336
+ # Advertise the GLOBAL DP world size to the routing proxy so it can
337
+ # address every rank across pods; dp_size_local above is only for
338
+ # per-pod port folding, not for routing.
339
+ dp_size=pc.data_parallel_size,
340
tp_size=tp_size,
341
read_mode=get_moriio_mode(kv_transfer_config) == MoRIIOMode.READ,
342
qp_per_transfer=int(extra_config.get("qp_per_transfer", 1)),
@@ -391,20 +415,26 @@ def parse_moriio_zmq_address(
415
return host, handshake_port, notify_port
416
417
394
-def get_peer_zmq_from_request_id(request_id: str, is_producer: bool) -> str:
418
+def get_peer_zmq_from_request_id(request_id: str, is_producer: bool) -> str | None:
419
"""Extract the *peer's* zmq_address from the vLLM router request_id.
420
421
The producer (prefill) needs the decode's address; the consumer (decode)
422
needs the prefill's address.
423
+
424
+ Returns ``None`` when the request_id does not encode peer info. The
425
+ llm-d routing sidecar (``llm-d-inference-scheduler``) does not embed
426
+ addresses in ``request_id``; instead it passes ``remote_host``,
427
+ ``remote_handshake_port`` and ``remote_notify_port`` explicitly in
428
+ ``kv_transfer_params``. Callers must handle the ``None`` return by
429
+ falling back to those fields. See ``add_new_req`` for the canonical
430
+ fallback path.
431
"""
432
if is_producer:
433
m = _DECODE_ZMQ_RE.search(request_id)
434
else:
435
m = _PREFILL_ZMQ_RE.search(request_id)
436
if m is None:
405
- raise ValueError(
406
- f"Cannot parse peer zmq_address from request_id: {request_id!r}"
407
- )
437
+ return None
438
return m.group(1)
439
440
@@ -426,6 +456,10 @@ class ReqMeta:
456
# read must target this rank's memory registration; the default 0 preserves
457
# the symmetric single-DP behaviour.
458
remote_dp_rank: int = 0
459
+ # Multi-pod: list of remote pod IPs indexed by pod_idx.
460
+ multi_pod_hosts: list[str] = field(default_factory=list)
461
+ # Per-pod DP size; 0 means fallback to remote_dp_size.
462
+ remote_dp_size_local: int = 0
463
464
465
class MoRIIOConnectorMetadata(KVConnectorMetadata):
@@ -450,23 +484,76 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
484
kv_transfer_params: dict[str, Any],
485
write_mode=False,
486
):
487
+ """Ingest a peer's ``kv_transfer_params`` into a typed ``ReqMeta``.
488
+
489
+ This is the single place peer info enters the connector. The
490
+ ``kv_transfer_params`` contract (as produced by the llm-d sidecar /
491
+ vLLM router, or echoed by the prefill leg's ``request_finished``):
492
+
493
+ Required (always):
494
+ * ``transfer_id`` -- stable id shared by both legs.
495
+ * ``remote_engine_id`` -- peer engine id for the handshake table.
496
+ * ``remote_block_ids`` -- peer block ids (may be [] in WRITE mode,
497
+ where decode allocates its own blocks).
498
+
499
+ Peer address -- ONE of the following two must resolve:
500
+ * embedded in ``request_id`` (vLLM-router PD id form), OR
501
+ * explicit ``remote_host`` + ``remote_handshake_port`` +
502
+ ``remote_notify_port`` (llm-d sidecar / returnable path).
503
+ If neither resolves we raise -- there is no safe default host/port.
504
+
505
+ Optional (defaulted):
506
+ * ``tp_size`` (default 1) -- peer TP size.
507
+ * ``remote_dp_size`` (default 1) -- peer GLOBAL DP size.
508
+ * ``remote_dp_size_local`` (default = ``remote_dp_size``) -- per-pod
509
+ DP size for Wide-EP multi-pod port/host
510
+ folding; 0/absent means single-pod.
511
+ * ``remote_hosts`` (default [remote_host]) -- per-pod IP list
512
+ indexed by ``pod_idx = rank // dp_local``.
513
+
514
+ Routing keys consumed elsewhere (NOT here): ``remote_dp_rank`` /
515
+ ``remote_dp_rank_override`` gate the decode->prefill notify target in
516
+ MoRIIOConnectorScheduler; they are router-authoritative and never
517
+ self-derived (see that class's request routing contract).
518
+ """
519
transfer_id = kv_transfer_params["transfer_id"]
520
455
- remote_host = kv_transfer_params.get("remote_host")
456
- remote_handshake_port = kv_transfer_params.get("remote_handshake_port")
457
- remote_notify_port = kv_transfer_params.get("remote_notify_port")
458
- if (
459
- remote_host is None
460
- or remote_handshake_port is None
461
- or remote_notify_port is None
462
- ):
463
- # Parse host/ports from the request_id. The router embeds both
464
- # zmq_addresses in PD request IDs, but WRITE decode requests may carry
465
- # a plain request ID and get the remote address via kv_transfer_params.
466
- peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode)
521
+ # Try request_id embedded address first, fallback to explicit params.
522
+ peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode)
523
+ if peer_zmq is not None:
524
remote_host, remote_handshake_port, remote_notify_port = (
525
parse_moriio_zmq_address(peer_zmq)
526
)
527
+ else:
528
+ try:
529
+ remote_host = kv_transfer_params["remote_host"]
530
+ if not remote_host:
531
+ raise ValueError(
532
+ f"request_id {request_id!r} does not embed a peer "
533
+ f"zmq_address and kv_transfer_params['remote_host'] is "
534
+ f"empty; cannot route MoRI-IO transfer"
535
+ )
536
+ remote_handshake_port = int(kv_transfer_params["remote_handshake_port"])
537
+ remote_notify_port = int(kv_transfer_params["remote_notify_port"])
538
+ except (KeyError, TypeError, ValueError) as e:
539
+ raise ValueError(
540
+ f"request_id {request_id!r} does not embed a peer "
541
+ f"zmq_address and kv_transfer_params is missing one or "
542
+ f"more sidecar-fallback keys (need remote_host, "
543
+ f"remote_handshake_port, remote_notify_port): {e}"
544
+ ) from e
545
+
546
+ # Multi-pod: use multi_pod_hosts list or fallback to single host.
547
+ _pod_hosts = kv_transfer_params.get("remote_hosts") or [remote_host]
548
+ if not isinstance(_pod_hosts, list):
549
+ _pod_hosts = [_pod_hosts]
550
+ _pod_hosts = [str(h) for h in _pod_hosts]
551
+ _remote_dp_size_local = int(
552
+ kv_transfer_params.get(
553
+ "remote_dp_size_local",
554
+ kv_transfer_params.get("remote_dp_size", 1),
555
+ )
556
+ )
557
558
_req = ReqMeta(
559
transfer_id=transfer_id,
@@ -488,6 +575,8 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
575
),
576
remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
577
remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0),
578
+ multi_pod_hosts=_pod_hosts,
579
+ remote_dp_size_local=_remote_dp_size_local,
580
)
581
if write_mode:
582
self.reqs_to_save[request_id] = _req
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+391
-96
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import logging
4
import math
5
+import os
6
import queue
7
import threading
8
import time
@@ -36,11 +37,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
37
ReqMeta,
38
TransferId,
39
WriteTask,
40
+ fold_local_rank,
41
get_moriio_mode,
42
get_peer_zmq_from_request_id,
43
get_port_offset,
44
get_role,
45
parse_moriio_zmq_address,
46
+ pod_index,
47
resolve_host_ip,
48
set_role,
49
zmq_ctx,
@@ -373,13 +376,29 @@ class MoRIIOConnectorScheduler:
376
"notify_port"
377
]
378
self.tp_size = self.vllm_config.parallel_config.tensor_parallel_size
376
- self.dp_rank = self.vllm_config.parallel_config.data_parallel_rank
379
+ # Local DP rank (0..dp_local-1) for port arithmetic; fold_local_rank
380
+ # handles the external-DP sentinel (dp_local==0) so the arithmetic
381
+ # can't ZeroDivisionError at init.
382
+ _pc = self.vllm_config.parallel_config
383
+ self.dp_rank = fold_local_rank(
384
+ _pc.data_parallel_rank, _pc.data_parallel_size_local
385
+ )
386
+ # Only first-pod ranks originate notify to avoid duplicates. Guarding
387
+ # via pod_index keeps _is_kv_master True for the external-DP sentinel
388
+ # (single pod) instead of being silently stuck False.
389
+ self._is_kv_master = (
390
+ pod_index(_pc.data_parallel_rank, _pc.data_parallel_size_local) == 0
391
+ )
392
+ # Global DP rank for pinned request ownership check.
393
+ self._global_dp_rank = self.vllm_config.parallel_config.data_parallel_rank
394
self.is_producer = self.kv_transfer_config.kv_role == "kv_producer"
395
# Requests that need to start recv/send.
396
# New requests are added by update_state_after_alloc in
397
# the scheduler. Used to make metadata passed to Worker.
398
self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
399
self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
400
+ # Snapshot of kv_transfer_params for chunked prefill recovery.
401
+ self._req_kv_params: dict[ReqId, dict] = {}
402
403
# For chunked prefill, we perform layer-wise access within the final chunk.
404
# TODO: Perform transfer at end chunk.
@@ -396,12 +415,15 @@ class MoRIIOConnectorScheduler:
415
# finished_sending before the deadline, we inject them into
416
# connector_output.finished_sending so the scheduler frees the blocks to avoid
417
# hanging indefinitely waiting for a free notification that never comes.
399
- self._deferred_send_deadlines: dict[ReqId, float] = {}
418
+ # Value: (deadline, transfer_id) for unmapping after mutation.
419
+ self._deferred_send_deadlines: dict[ReqId, tuple[float, TransferId | None]] = {}
420
self._defer_timeout = float(
421
self.kv_transfer_config.kv_connector_extra_config.get(
422
"defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT
423
)
424
)
425
+ # Buffer for early ACKs that arrive before request_finished.
426
+ self._pending_sent_acks: dict[ReqId, float] = {}
427
self.paths: dict[str, zmq.Socket] = {}
428
self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
429
self.request_id_to_transfer_id: dict[ReqId, TransferId] = {}
@@ -410,23 +432,36 @@ class MoRIIOConnectorScheduler:
432
self.transfer_id_to_request_id[transfer_id] = request_id
433
self.request_id_to_transfer_id[request_id] = transfer_id
434
413
- def unmap_request_id(self, request_id: ReqId):
435
+ def unmap_request_id(
436
+ self, request_id: ReqId, transfer_id: TransferId | None = None
437
+ ):
438
+ """Unmap request_id/transfer_id. Uses transfer_id for lookup if
439
+ exact request_id match fails (handles input_processor mutation)."""
440
if request_id in self.request_id_to_transfer_id:
415
- transfer_id = self.request_id_to_transfer_id[request_id]
441
+ tid = self.request_id_to_transfer_id[request_id]
442
del self.request_id_to_transfer_id[request_id]
417
- if transfer_id in self.transfer_id_to_request_id:
418
- del self.transfer_id_to_request_id[transfer_id]
419
- else:
420
- logger.warning(
421
- "transfer id not in transfer_id_to_request_id lookup"
422
- "table. there is likely a bug!"
443
+ if tid in self.transfer_id_to_request_id:
444
+ del self.transfer_id_to_request_id[tid]
445
+ return
446
+
447
+ # Fallback: use transfer_id to find original request_id.
448
+ if transfer_id is not None and transfer_id in self.transfer_id_to_request_id:
449
+ original_rid = self.transfer_id_to_request_id[transfer_id]
450
+ if original_rid != request_id:
451
+ logger.debug(
452
+ "MoRI-IO unmap via transfer_id: %r -> %r", request_id, original_rid
453
)
424
- else:
425
- logger.warning(
426
- "Could not find %s in transfer_id_to_request_id"
427
- "lookup table. This could lead to a possible hang.",
428
- request_id,
429
- )
454
+ if original_rid in self.request_id_to_transfer_id:
455
+ del self.request_id_to_transfer_id[original_rid]
456
+ del self.transfer_id_to_request_id[transfer_id]
457
+ return
458
+
459
+ logger.warning(
460
+ "MoRI-IO unmap MISS: rid=%r transfer_id=%r table_size=%d",
461
+ request_id,
462
+ transfer_id,
463
+ len(self.request_id_to_transfer_id),
464
+ )
465
466
def get_num_new_matched_tokens(
467
self,
@@ -478,7 +513,12 @@ class MoRIIOConnectorScheduler:
513
"req_id": req_id,
514
"transfer_id": transfer_id,
515
"block_notify_list": block_notify_list or [],
481
- "decode_rank": self.dp_rank,
516
+ # GLOBAL decode dp rank: producer derives the per-pod notify offset
517
+ # (% dp_local), owning pod index (// dp_local), and write-target
518
+ # from it. Sending the LOCAL rank made child-pod consumers look
519
+ # like master rank 0 and hang in WAITING_FOR_REMOTE_KVS. Single-pod:
520
+ # local == global.
521
+ "decode_rank": self._global_dp_rank,
522
"type": "remote_blocks",
523
}
524
serialized_data = msgpack.dumps(data)
@@ -493,8 +533,18 @@ class MoRIIOConnectorScheduler:
533
)
534
self.paths[path] = sock
535
536
+ # Advertise the consumer (decode) TP size so prefill counts the right
537
+ # number of ACKs (get_moriio_expected_ack_count; see the upstream
538
+ # READ-mode release in _pop_done_transfers). Resolves to 1 for our
539
+ # homogeneous TP=1 configs; only matters under heterogeneous-TP fan-in.
540
self.paths[path].send(
497
- msgpack.dumps({"type": "release", "transfer_id": transfer_id})
541
+ msgpack.dumps(
542
+ {
543
+ "type": "release",
544
+ "transfer_id": transfer_id,
545
+ "consumer_tp_size": self.tp_size,
546
+ }
547
+ )
548
)
549
550
def _release_write_prefill_blocks(self, request_id: ReqId, params: dict[str, Any]):
@@ -513,6 +563,8 @@ class MoRIIOConnectorScheduler:
563
if remote_host is None or remote_notify_port is None:
564
try:
565
peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False)
566
+ if peer_zmq is None:
567
+ raise ValueError("no peer zmq address for request")
568
remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)
569
except ValueError:
570
logger.warning(
@@ -534,15 +586,54 @@ class MoRIIOConnectorScheduler:
586
num_external_tokens: int,
587
connector_worker: "MoRIIOConnectorWorker | None" = None,
588
):
589
+ """Scheduler-side post-allocation hook (decode leg in WRITE mode).
590
+
591
+ In WRITE mode this fires the decode->prefill "blocks ready" notify
592
+ that lets the producer RDMA-Write its KV into the freshly allocated
593
+ decode blocks.
594
+
595
+ DP-rank routing contract (router-authoritative)
596
+ ------------------------------------------------
597
+ Both legs of a disagg pair must agree on a single prefill DP rank,
598
+ otherwise the notify lands on a rank that never handshook and the
599
+ request hangs until ``VLLM_MORIIO_DEFERRED_TIMEOUT_S``.
600
+
601
+ The routing authority is the external router/sidecar, NOT the
602
+ connector. The llm-d routing sidecar computes ``H = pickDPRank(uuid,
603
+ dp_size)`` (see ``dp_rank.go``) once and pins BOTH legs to ``H`` two
604
+ ways: the ``X-data-parallel-rank`` dispatch header (which vLLM engine
605
+ the leg lands on) and ``kv_transfer_params.remote_dp_rank=H`` (which
606
+ prefill rank this notify targets), with ``remote_dp_rank_override=True``.
607
+
608
+ The connector consumes ``remote_dp_rank`` verbatim and does NOT
609
+ self-derive a rank: an independent hash here could disagree with the
610
+ router's dispatch pin and misroute the notify. For the returnable paths
611
+ (READ / serial WRITE) the prefill leg also echoes back the rank it ran
612
+ on via ``request_finished`` so routing is pure propagation.
613
+
614
+ The owning rank is the only one that originates the notify
615
+ (``remote_dp_rank_override`` -> global-rank match; otherwise the
616
+ ``_is_kv_master`` anti-duplicate gate). See the inline comments
617
+ below for the exactly-once-across-pods reasoning.
618
+ """
619
params = request.kv_transfer_params
620
if not params:
621
return
540
- transfer_id = params["transfer_id"]
622
+ # LLM-D sidecar compat: the nixlv2 routing sidecar emits NIXL-shaped
623
+ # kv_transfer_params without MoRI-IO's transfer_id. Synthesize one
624
+ # deterministically from request_id so producer and consumer agree
625
+ # without a sidecar wire-protocol change.
626
+ transfer_id = params.get("transfer_id") or f"sidecar-{request.request_id}"
627
+ params.setdefault("transfer_id", transfer_id)
628
request_id = request.request_id
629
self.map_request_id(request_id, transfer_id)
630
if params.get("do_remote_decode"):
631
local_block_ids = blocks.get_block_ids()[0]
632
self._reqs_need_save[request.request_id] = (request, local_block_ids)
633
+ # Snapshot params now so chunked-prefill build_connector_meta
634
+ # can recover them on the final chunk even if the live
635
+ # request.kv_transfer_params has been mutated/cleared.
636
+ self._req_kv_params[request.request_id] = dict(params)
637
638
if params is not None and params.get("do_remote_prefill"):
639
if self.mode == MoRIIOMode.READ:
@@ -568,6 +659,9 @@ class MoRIIOConnectorScheduler:
659
request,
660
local_block_ids,
661
)
662
+ # Snapshot params for chunked prefill consumption
663
+ # in build_connector_meta (see comment above).
664
+ self._req_kv_params[request.request_id] = dict(params)
665
else:
666
logger.warning(
667
"Got invalid KVTransferParams: %s. This "
@@ -582,37 +676,113 @@ class MoRIIOConnectorScheduler:
676
)
677
678
remote_dp_rank = request.kv_transfer_params.get("remote_dp_rank", 0)
585
- remote_host = request.kv_transfer_params.get("remote_host")
586
- remote_notify_port = request.kv_transfer_params.get(
587
- "remote_notify_port"
588
- )
589
- if remote_host is None or remote_notify_port is None:
679
+
680
+ # Effective DP fan-out for the per-pod port/host math below
681
+ # (see _remote_dp_rank_for_port / _pod_idx). Capped to the
682
+ # per-pod local size when the router advertises it (Wide-EP).
683
+ _dp_size = int(request.kv_transfer_params.get("remote_dp_size", 1) or 1)
684
+ try:
685
+ _dp_local = int(
686
+ request.kv_transfer_params.get("remote_dp_size_local", 0) or 0
687
+ )
688
+ if _dp_local > 0:
689
+ _dp_size = min(_dp_size, _dp_local)
690
+ except (TypeError, ValueError):
691
+ _dp_local = 0
692
+
693
+ # Rank routing is ROUTER-AUTHORITATIVE: honor the router-pinned
694
+ # remote_dp_rank (matched to the X-data-parallel-rank dispatch
695
+ # pin). Self-deriving a rank could disagree and notify a rank
696
+ # that never served the request. If unpinned in a multi-rank
697
+ # deployment, warn instead of guessing.
698
+ if (
699
+ _dp_size > 1
700
+ and "remote_dp_rank" not in request.kv_transfer_params
701
+ and "is_request_leader" not in request.kv_transfer_params
702
+ ):
703
+ logger.warning(
704
+ "MoRI-IO decode notify: remote_dp_size=%d but the router "
705
+ "did not pin remote_dp_rank for request %s; defaulting to "
706
+ "rank 0. The router must set remote_dp_rank (and the "
707
+ "X-data-parallel-rank dispatch header) so the prefill and "
708
+ "decode legs agree on the same rank.",
709
+ _dp_size,
710
+ request.request_id,
711
+ )
712
+
713
+ # Only the owning rank originates notify (exactly-once).
714
+ # Priority: is_request_leader > remote_dp_rank_override > _is_kv_master
715
+ _leader_flag = request.kv_transfer_params.get("is_request_leader")
716
+ if _leader_flag is not None:
717
+ _should_notify = bool(_leader_flag)
718
+ elif request.kv_transfer_params.get("remote_dp_rank_override"):
719
+ _should_notify = self._global_dp_rank == remote_dp_rank
720
+ else:
721
+ _should_notify = self._is_kv_master
722
+ if _should_notify:
723
peer_zmq = get_peer_zmq_from_request_id(
724
request.request_id, is_producer=False
725
)
593
- remote_host, _, remote_notify_port = parse_moriio_zmq_address(
594
- peer_zmq
726
+ if peer_zmq is not None:
727
+ remote_host, _, remote_notify_port = parse_moriio_zmq_address(
728
+ peer_zmq
729
+ )
730
+ else:
731
+ # Sidecar fallback: use explicit params fields.
732
+ params = request.kv_transfer_params or {}
733
+ remote_host = params.get("remote_host") or ""
734
+ try:
735
+ remote_notify_port = int(
736
+ params.get("remote_notify_port") or 0
737
+ )
738
+ except (TypeError, ValueError):
739
+ remote_notify_port = 0
740
+ if not remote_host or not remote_notify_port:
741
+ raise ValueError(
742
+ f"request {request.request_id!r}: "
743
+ f"request_id has no embedded peer "
744
+ f"zmq_address and kv_transfer_params is "
745
+ f"missing remote_host / remote_notify_port "
746
+ f"(got remote_host={remote_host!r}, "
747
+ f"remote_notify_port={remote_notify_port!r})"
748
+ )
749
+
750
+ # num_external_tokens == 0: nothing to push, so don't tell
751
+ # the producer to write into these blocks.
752
+ block_notify_list = (
753
+ blocks.get_block_ids()[0] if num_external_tokens > 0 else []
754
)
596
- remote_notify_port = int(remote_notify_port)
755
598
- # num_external_tokens == 0: nothing to push, so don't tell the
599
- # producer to write into these blocks.
600
- block_notify_list = (
601
- blocks.get_block_ids()[0] if num_external_tokens > 0 else []
602
- )
603
-
604
- for tp_index in range(self.tp_size):
605
- target_port = remote_notify_port + get_port_offset(
606
- remote_dp_rank, tp_index
756
+ # Wide-EP multi-pod: a pod binds notify sockets only for
757
+ # its LOCAL ranks, so the port offset must use the per-pod
758
+ # local rank (% dp_local), not the global rank. Single-pod
759
+ # is bit-identical (modulus is a no-op).
760
+ _remote_dp_rank_for_port = fold_local_rank(
761
+ remote_dp_rank, _dp_local
762
)
763
+ # The target rank may live on a child pod at a different IP,
764
+ # so resolve the per-pod host (pod_idx = global // dp_local).
765
+ # Otherwise a notify for child ranks lands on the master
766
+ # pod and the request hangs in WAITING_FOR_REMOTE_KVS.
767
+ _notify_host = remote_host
768
+ _kvp = request.kv_transfer_params or {}
769
+ _remote_hosts = _kvp.get("remote_hosts") or []
770
+ if _dp_local > 0 and _remote_hosts:
771
+ _pod_idx = pod_index(remote_dp_rank, _dp_local)
772
+ if 0 <= _pod_idx < len(_remote_hosts):
773
+ _notify_host = _remote_hosts[_pod_idx]
774
+ for tp_index in range(self.tp_size):
775
+ target_port = remote_notify_port + get_port_offset(
776
+ _remote_dp_rank_for_port, tp_index
777
+ )
778
609
- self.send_notify_block(
610
- req_id=request.request_id,
611
- transfer_id=request.kv_transfer_params["transfer_id"],
612
- block_notify_list=block_notify_list,
613
- host=remote_host,
614
- port=target_port,
615
- )
779
+ self.send_notify_block(
780
+ req_id=request.request_id,
781
+ transfer_id=request.kv_transfer_params["transfer_id"],
782
+ block_notify_list=block_notify_list,
783
+ host=_notify_host,
784
+ port=target_port,
785
+ )
786
787
# Only trigger 1 KV transfer per request.
788
@@ -636,6 +806,12 @@ class MoRIIOConnectorScheduler:
806
if new_block_ids is not None:
807
block_ids = new_block_ids[0]
808
# TODO : hybrid attn, etc
809
+ # A non-disagg request (no kv_transfer_params, e.g. smoke
810
+ # test) is never registered in _reqs_need_pending_save;
811
+ # indexing it unconditionally would KeyError and crash the
812
+ # EngineCore. Skip it silently.
813
+ if req_id not in self._reqs_need_pending_save:
814
+ continue
815
req, existing_blocks = self._reqs_need_pending_save[req_id]
816
updated_blocks = list(existing_blocks) + (block_ids)
817
self._reqs_need_pending_save[req_id] = (req, updated_blocks)
@@ -643,25 +819,30 @@ class MoRIIOConnectorScheduler:
819
len(self._reqs_need_pending_save[req_id][1]) * self.block_size
820
>= req.num_prompt_tokens
821
):
822
+ # Final chunk: live kv_transfer_params may be cleared,
823
+ # so prefer the snapshot from update_state_after_alloc.
824
+ kv_params = self._req_kv_params.pop(
825
+ req_id, req.kv_transfer_params or {}
826
+ )
827
meta.add_new_req(
828
request_id=req_id,
829
local_block_ids=self._reqs_need_pending_save[req_id][1],
649
- kv_transfer_params=req.kv_transfer_params or {},
830
+ kv_transfer_params=kv_params,
831
write_mode=True,
832
)
833
del self._reqs_need_pending_save[req_id]
834
835
# Loop through scheduled reqs and convert to ReqMeta.
836
for req_id, (req, block_ids) in self._reqs_need_recv.items():
656
- assert req.kv_transfer_params is not None
837
+ kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})
838
meta.add_new_req(
839
request_id=req_id,
840
local_block_ids=block_ids,
660
- kv_transfer_params=req.kv_transfer_params,
841
+ kv_transfer_params=kv_params,
842
)
843
844
for req_id, (req, block_ids) in self._reqs_need_save.items():
664
- assert req.kv_transfer_params is not None
845
+ kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})
846
if req.num_prompt_tokens > len(block_ids) * self.block_size:
847
# not last chunk prefill
848
self._reqs_need_pending_save[req_id] = (req, block_ids)
@@ -669,13 +850,22 @@ class MoRIIOConnectorScheduler:
850
meta.add_new_req(
851
request_id=req_id,
852
local_block_ids=block_ids,
672
- kv_transfer_params=req.kv_transfer_params,
853
+ kv_transfer_params=kv_params,
854
write_mode=True,
855
)
856
# Clear the list once workers start the transfers
857
858
meta.reqs_to_send = self._reqs_need_send
859
860
+ # Reclaim snapshot cache entries that completed this step. Keep
861
+ # entries that are still pending (chunked prefill not yet at the
862
+ # final chunk) — those will be popped above on the final chunk.
863
+ for req_id in self._reqs_need_recv:
864
+ self._req_kv_params.pop(req_id, None)
865
+ for req_id in self._reqs_need_save:
866
+ if req_id not in self._reqs_need_pending_save:
867
+ self._req_kv_params.pop(req_id, None)
868
+
869
self._reqs_need_recv.clear()
870
self._reqs_need_save.clear()
871
self._reqs_need_send = {}
@@ -702,14 +892,14 @@ class MoRIIOConnectorScheduler:
892
"""
893
894
request_id = request.request_id
895
+ params = request.kv_transfer_params
896
# Consumer: can unmap transfer_id<->request_id immediately since done_recving
897
# has fired at this point (i.e. KV has been transferred)
898
# Producer: must keep the mapping until we get notification that blocks can
899
# be freed, which may be several scheduler steps later.
900
if not self.is_producer:
710
- self.unmap_request_id(request_id)
711
-
712
- params = request.kv_transfer_params
901
+ transfer_id = params.get("transfer_id") if params else None
902
+ self.unmap_request_id(request_id, transfer_id=transfer_id)
903
logger.debug(
904
"MoriioConnector request_finished, request_status=%s, "
905
"kv_transfer_params=%s",
@@ -754,11 +944,17 @@ class MoRIIOConnectorScheduler:
944
+ MoRIIOConstants.VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT
945
)
946
self._deferred_send_deadlines[request.request_id] = (
757
- time.monotonic() + self._defer_timeout
947
+ time.monotonic() + self._defer_timeout,
948
+ params.get("transfer_id") if params else None,
949
)
950
951
# Return KV transfer params forwarded verbatim to the decode instance by
761
- # the router.
952
+ # the router. remote_dp_rank is the rank THIS prefill leg actually ran
953
+ # on: on the returnable paths (READ / serial WRITE) the router forwards
954
+ # it to the decode leg, so the decode->prefill notify is routed by pure
955
+ # propagation of the producer's real rank -- no hashing, no reliance on
956
+ # the router independently pinning the same rank. remote_dp_rank_override
957
+ # makes the decode side honor it via the global-rank match gate.
958
return delay_free_blocks, dict(
959
do_remote_prefill=True,
960
do_remote_decode=False,
@@ -767,61 +963,96 @@ class MoRIIOConnectorScheduler:
963
remote_host=self.host_ip,
964
remote_handshake_port=self.handshake_port,
965
remote_notify_port=self.side_notify_port,
966
+ remote_dp_rank=self._global_dp_rank,
967
+ remote_dp_rank_override=True,
968
remote_dp_size=self.vllm_config.parallel_config.data_parallel_size,
969
+ remote_dp_size_local=(
970
+ self.vllm_config.parallel_config.data_parallel_size_local
971
+ ),
972
tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
973
transfer_id=params["transfer_id"],
974
)
975
976
def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
776
- """Free KV blocks from sends that never received a completion signal.
777
-
778
- Called every scheduler step. When a send is deferred (request_finished
779
- returns True), blocks remain allocated until the worker reports
780
- finished_sending. If that notification is lost (e.g. ibv_post_send
781
- failure), blocks leak permanently. This method injects timed-out
782
- entries into connector_output.finished_sending so the scheduler
783
- frees them via the normal path.
977
+ """Reconcile worker finished_sending ACKs with the producer lifecycle.
978
+
979
+ Called every scheduler step (when there is KV connector output),
980
+ BEFORE the scheduler's finished_sending free loop. We rewrite
981
+ ``connector_output.finished_sending`` in place so it contains
982
+ exactly the producer requests that are safe for the scheduler to
983
+ free this step, letting the scheduler keep the plain upstream
984
+ ``assert req_id in self.requests`` + ``_free_blocks`` path:
985
+
986
+ * An ACK whose request is in ``_deferred_send_deadlines``
987
+ (request_finished ran with delay_free_blocks=True, so the
988
+ scheduler is holding its blocks) is surfaced -> freed now.
989
+ * An ACK that arrives BEFORE its request finished is parked in
990
+ ``_pending_sent_acks`` and released on a later step once the
991
+ request enters ``_deferred_send_deadlines``.
992
+ * A deferred send whose ACK never arrives is reaped after
993
+ ``_defer_timeout`` and surfaced, so leaked blocks are freed.
994
+ * A parked ACK that never matches a deferral before its own
995
+ deadline is a stale duplicate (e.g. a real ACK landing after the
996
+ send was already reaped) and is dropped.
997
+
998
+ Consumers never populate finished_sending (they report
999
+ finished_recving), and they unmap in request_finished, so this is a
1000
+ no-op for them.
1001
"""
785
- # Producer: unmap transfer_id<->request_id for sends that are now (async)
786
- # reported as completed. This unmapping have to be deferred until now
787
- # so get_finished can use it in any scheduler step.
788
- # Consumer: unmapping already done in request_finished
789
- if self.is_producer and connector_output.finished_sending:
790
- for req_id in connector_output.finished_sending:
791
- self.unmap_request_id(req_id)
792
-
793
- if not self._deferred_send_deadlines:
1002
+ if not self.is_producer:
1003
return
1004
796
- # Remove entries the worker already reported as finished_sending, these will be
797
- # freed anyways.
798
- for req_id in connector_output.finished_sending or ():
799
- self._deferred_send_deadlines.pop(req_id, None)
800
-
1005
+ incoming = set(connector_output.finished_sending or ())
1006
now = time.monotonic()
802
- expired_reqs = [
1007
+
1008
+ # Surface ACKs whose request is already finished (blocks held for
1009
+ # delayed free); park the rest until their request finishes.
1010
+ safe: set[ReqId] = set()
1011
+ for req_id in incoming:
1012
+ if req_id in self._deferred_send_deadlines:
1013
+ safe.add(req_id)
1014
+ else:
1015
+ self._pending_sent_acks.setdefault(req_id, now + self._defer_timeout)
1016
+
1017
+ # Release previously parked ACKs whose request has since finished.
1018
+ for req_id in self._pending_sent_acks:
1019
+ if req_id in self._deferred_send_deadlines:
1020
+ safe.add(req_id)
1021
+
1022
+ # Reap deferred sends whose ACK never arrived (avoid leaking blocks).
1023
+ expired = [
1024
req_id
804
- for req_id, deadline in self._deferred_send_deadlines.items()
1025
+ for req_id, (deadline, _) in self._deferred_send_deadlines.items()
1026
if now >= deadline
1027
]
807
- if not expired_reqs:
808
- return
1028
+ if expired:
1029
+ safe.update(expired)
1030
+ logger.warning(
1031
+ "Reaped %d deferred sends with no finished_sending "
1032
+ "notification after %.0fs. This indicates lost async KV "
1033
+ "completion notifications from the KV connector.",
1034
+ len(expired),
1035
+ self._defer_timeout,
1036
+ )
1037
810
- if connector_output.finished_sending is None:
811
- connector_output.finished_sending = set()
812
- # Register the expired requests as finished so the scheduler frees their blocks.
813
- for req_id in expired_reqs:
814
- connector_output.finished_sending.add(req_id)
815
- del self._deferred_send_deadlines[req_id]
816
- if self.is_producer:
817
- self.unmap_request_id(req_id)
818
- logger.warning(
819
- "Reaped %d deferred sends with no finished_sending notification "
820
- "after %.0fs. This indicates lost async KV completion "
821
- "notifications from the KV connector.",
822
- len(expired_reqs),
823
- self._defer_timeout,
824
- )
1038
+ # Finalize the requests we are surfacing: drop their deferral/park
1039
+ # bookkeeping and unmap their transfer ids.
1040
+ for req_id in safe:
1041
+ deferred_info = self._deferred_send_deadlines.pop(req_id, None)
1042
+ transfer_id = deferred_info[1] if deferred_info else None
1043
+ self._pending_sent_acks.pop(req_id, None)
1044
+ self.unmap_request_id(req_id, transfer_id=transfer_id)
1045
+
1046
+ # Drop stale parked ACKs that never matched a deferral in time.
1047
+ stale = [
1048
+ req_id
1049
+ for req_id, deadline in self._pending_sent_acks.items()
1050
+ if now >= deadline
1051
+ ]
1052
+ for req_id in stale:
1053
+ self._pending_sent_acks.pop(req_id, None)
1054
+
1055
+ connector_output.finished_sending = safe or None
1056
1057
1058
class MoRIIOConnectorWorker:
@@ -998,6 +1229,10 @@ class MoRIIOConnectorWorker:
1229
self._recving_transfers: defaultdict[ReqId, list] = defaultdict(list)
1230
# Values are (remote_host, remote_notify_port, transfer_id).
1231
self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {}
1232
+ # Monotonic-clock start times for each in-flight recv transfer.
1233
+ # Used by _pop_done_transfers to abort transfers whose RDMA
1234
+ # completion is lost, instead of hanging forever.
1235
+ self._recving_transfers_start: dict[str, float] = {}
1236
1237
# Track the expiration time of requests that are waiting to be sent.
1238
self._reqs_to_send: dict[ReqId, float] = {}
@@ -1368,6 +1603,11 @@ class MoRIIOConnectorWorker:
1603
port = int(meta.remote_handshake_port)
1604
tp_size = int(meta.tp_size)
1605
remote_dp_size = int(meta.remote_dp_size)
1606
+ # Wide-EP multi-pod: remote DP ranks span pods at different IPs
1607
+ # (ranks per pod = dp_local), so resolve the host per cur_dp_rank
1608
+ # below instead of using a single host for all ranks.
1609
+ pod_hosts = list(meta.multi_pod_hosts) if meta.multi_pod_hosts else [host]
1610
+ remote_dp_size_local = int(meta.remote_dp_size_local) or remote_dp_size
1611
1612
def request_ready(_f: Future[Any], entry=(req_id, meta)):
1613
logger.info("MoRIIO handshake done for request %s", req_id)
@@ -1381,8 +1621,21 @@ class MoRIIOConnectorWorker:
1621
1622
for cur_dp_rank in range(remote_dp_size):
1623
dp_engine_id = self.get_engine_name_with_dp(remote_engine_id, cur_dp_rank)
1624
+ _pod_idx = pod_index(cur_dp_rank, remote_dp_size_local)
1625
+ if _pod_idx >= len(pod_hosts):
1626
+ _pod_idx = 0
1627
+ _per_rank_host = pod_hosts[_pod_idx]
1628
+ # The handshake port offset must use the per-pod local rank, since
1629
+ # each pod binds sockets only for its local ranks; dp_engine_id
1630
+ # keeps the global rank for uniqueness. Single-pod is bit-identical.
1631
+ _per_rank_local_dp = fold_local_rank(cur_dp_rank, remote_dp_size_local)
1632
future = self._handshake_initiation_executor.submit(
1385
- self._moriio_handshake, host, port, tp_size, dp_engine_id, cur_dp_rank
1633
+ self._moriio_handshake,
1634
+ _per_rank_host,
1635
+ port,
1636
+ tp_size,
1637
+ dp_engine_id,
1638
+ _per_rank_local_dp,
1639
)
1640
fut_list.append(future)
1641
@@ -1463,7 +1716,19 @@ class MoRIIOConnectorWorker:
1716
block_shape = first_kv_cache.shape[-block_rank:]
1717
self.num_blocks = first_geometry.num_blocks
1718
self.slot_size_bytes = first_geometry.slot_size_bytes
1466
- assert first_geometry.block_size == self.block_size
1719
+ if first_geometry.block_size != self.block_size:
1720
+ # DeepSeek-V3 / MLA backends (e.g. FlashMLA) override the
1721
+ # configured block_size at attention-layer creation time, so
1722
+ # the KV cache tensor is laid out with a different (usually
1723
+ # larger) block_size than vllm_config.cache_config.block_size.
1724
+ # Trust the actual tensor shape and reconcile.
1725
+ logger.info(
1726
+ "KV cache block_size=%d differs from config block_size=%d; "
1727
+ "using actual tensor shape (attention backend override).",
1728
+ first_geometry.block_size,
1729
+ self.block_size,
1730
+ )
1731
+ self.block_size = first_geometry.block_size
1732
# TODO(tms): self.block_len needs to be per-layer for sliding window,
1733
# hybrid attn, etc
1734
# block size in bytes
@@ -1634,6 +1899,7 @@ class MoRIIOConnectorWorker:
1899
1900
def _pop_done_transfers(self) -> set[str]:
1901
done_req_ids: set[str] = set()
1902
+ _xfer_timeout = int(os.environ.get("VLLM_MORIIO_TRANSFER_TIMEOUT_S", "120"))
1903
with self.moriio_wrapper.lock:
1904
to_remove = []
1905
for req_id, status_list in self._recving_transfers.items():
@@ -1675,9 +1941,25 @@ class MoRIIOConnectorWorker:
1941
to_remove.append(req_id)
1942
# Do NOT add to done_req_ids: decode KV cache is incomplete.
1943
# The request will expire via the normal request timeout.
1944
+ elif req_id in self._recving_transfers_start:
1945
+ # Abort still-in-flight transfers that exceed the
1946
+ # configured deadline. Otherwise a lost RDMA
1947
+ # completion would leave the decode worker hung
1948
+ # indefinitely on this request.
1949
+ _age = time.monotonic() - self._recving_transfers_start[req_id]
1950
+ if _age > _xfer_timeout:
1951
+ logger.error(
1952
+ "RDMA read TIMED OUT for req %s after %.1fs "
1953
+ "(VLLM_MORIIO_TRANSFER_TIMEOUT_S=%d)",
1954
+ req_id,
1955
+ _age,
1956
+ _xfer_timeout,
1957
+ )
1958
+ to_remove.append(req_id)
1959
for req_id in to_remove:
1960
del self._recving_transfers[req_id]
1961
del self._recving_transfers_callback_addr[req_id]
1962
+ self._recving_transfers_start.pop(req_id, None)
1963
1964
return done_req_ids
1965
@@ -2043,6 +2325,18 @@ class MoRIIOConnectorWorker:
2325
)
2326
2327
def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):
2328
+ # Stash multi_pod_hosts + local DP size on the worker so
2329
+ # MoRIIOEngine._finalize_if_complete (which sees only the WriteTask,
2330
+ # not ReqMeta) can pick the per-rank pod IP for the completion notify.
2331
+ # Last-writer-wins is safe: all requests share the same topology.
2332
+ if meta.multi_pod_hosts:
2333
+ self.multi_pod_hosts = list(meta.multi_pod_hosts)
2334
+ else:
2335
+ self.multi_pod_hosts = [meta.remote_host]
2336
+ if meta.remote_dp_size_local:
2337
+ self.remote_dp_size_local = int(meta.remote_dp_size_local)
2338
+ else:
2339
+ self.remote_dp_size_local = int(meta.remote_dp_size)
2340
self.schedule_write_blocks(
2341
request_id=req_id,
2342
transfer_id=meta.transfer_id,
@@ -2271,6 +2565,7 @@ class MoRIIOConnectorWorker:
2565
_backoff = min(_backoff * 2, 0.05)
2566
with self.moriio_wrapper.lock:
2567
self._recving_transfers[request_id].append(transfer_status)
2568
+ self._recving_transfers_start.setdefault(request_id, time.monotonic())
2569
self._recving_transfers_callback_addr[request_id] = (
2570
remote_host,
2571
str(
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py
+19
-2
@@ -311,7 +311,17 @@ class MoRIIOWriter:
311
with self._write_state_lock:
312
request_info.completion_request_id = task.request_id
313
request_info.completion_remote_notify_port = task.remote_notify_port
314
- request_info.completion_remote_ip = task.remote_ip
314
+ # Wide-EP multi-pod: task.remote_ip addresses only the first pod,
315
+ # so resolve the per-rank host from multi_pod_hosts (falls back to
316
+ # task.remote_ip for single-pod or on any indexing miss).
317
+ _remote_ip = task.remote_ip
318
+ _hosts = list(getattr(self.worker, "multi_pod_hosts", []) or [])
319
+ _dp_local = int(getattr(self.worker, "remote_dp_size_local", 0) or 0)
320
+ if _hosts and _dp_local > 0:
321
+ _pod_idx = int(request_info.decode_dp_rank) // _dp_local
322
+ if 0 <= _pod_idx < len(_hosts):
323
+ _remote_ip = _hosts[_pod_idx]
324
+ request_info.completion_remote_ip = _remote_ip
325
if task.transfer_id in self._sealed_writes:
326
request_info.writes_expected = self._sealed_writes[task.transfer_id]
327
@@ -454,8 +464,15 @@ class MoRIIOWriter:
464
# Wait for this request's transfers to complete.
465
self.worker.moriio_wrapper.waiting_for_transfer_complete(transfer_statuses)
466
467
+ # The notify port offset must use the per-pod local rank
468
+ # (% dp_local), since each pod binds notify sockets only for its local
469
+ # ranks. Single-pod is bit-identical (modulus is a no-op).
470
+ _dp_local = int(getattr(self.worker, "remote_dp_size_local", 0) or 0)
471
+ _decode_dp_rank_for_port = int(request_info.decode_dp_rank)
472
+ if _dp_local > 0:
473
+ _decode_dp_rank_for_port = _decode_dp_rank_for_port % _dp_local
474
remote_port = remote_notify_port + get_port_offset(
458
- request_info.decode_dp_rank, self.worker.tp_rank
475
+ _decode_dp_rank_for_port, self.worker.tp_rank
476
)
477
# Consider using RDMA immediate data in decode side
478
# to eliminate the need for this notification.