A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 18:45 UTC
[Fix][MoRI] Align MoRI-IO message format with P2pNcclConnector and vllm-router (#39565)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: Matvei Pashkovskii <mpashkov@amd.com>
Simon Danielsson committed
Apr 23, 2026 at 01:06 UTC
ac58e2a1704ba18db3c18748cee2fb6c874496d6
3 files changed
+203
-91
examples/online_serving/disaggregated_serving/moriio_toy_proxy_server.py
+91
-66
@@ -10,7 +10,6 @@ import uuid
10
11
import aiohttp
12
import msgpack
13
-import regex as re
13
import zmq
14
from quart import Quart, Request, make_response, request
15
@@ -25,32 +24,10 @@ decode_instances: list[dict] = []
24
request_nums = 0
25
app = Quart(__name__)
26
28
-IP_PORT_PATTERN = re.compile(r"//(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}):(\d+)")
29
-
27
28
TRANSFER_TYPE = None
29
30
34
-def _append_whole_dict_unique(target_list, data_dict):
35
- new_filtered = {k: v for k, v in data_dict.items() if k != "index"}
36
- for existed in target_list:
37
- existed_filtered = {k: v for k, v in existed.items() if k != "index"}
38
- if existed_filtered == new_filtered:
39
- return False
40
- print("!!APPEND!!", data_dict)
41
- target_list.append(data_dict)
42
- transfer_mode = data_dict.get("transfer_mode", "unknown")
43
- global TRANSFER_TYPE
44
-
45
- if TRANSFER_TYPE is None:
46
- TRANSFER_TYPE = transfer_mode
47
- logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
48
- elif transfer_mode != TRANSFER_TYPE:
49
- raise ValueError(f"mismatched transfer mode {TRANSFER_TYPE} vs {transfer_mode}")
50
-
51
- return True
52
-
53
-
31
_list_lock = threading.RLock()
32
33
@@ -68,23 +45,81 @@ def _listen_for_register(hostname, port):
45
if router_socket in socks:
46
remote_addr, msg = router_socket.recv_multipart()
47
data = msgpack.loads(msg)
71
- if data["type"] == "HELLO":
48
+ if data.get("type") == "HELLO":
49
pass
73
- elif (
74
- data["type"] == "register"
75
- and data["role"] == "P"
76
- and data["request_address"] not in prefill_instances
77
- ):
50
+ elif data.get("type") in ("P", "D"):
51
+ role = data["type"]
52
+ required_keys = {
53
+ "http_address",
54
+ "zmq_address",
55
+ "dp_size",
56
+ "tp_size",
57
+ "transfer_mode",
58
+ }
59
+ missing = required_keys - data.keys()
60
+ if missing:
61
+ logger.error(
62
+ "Registration message missing required keys %s; skipping",
63
+ missing,
64
+ )
65
+ continue
66
+ # Derive request_address from http_address
67
+ # api path suffix is appended at request time
68
+ instance = {
69
+ "role": role,
70
+ "request_address": f"http://{data['http_address']}/v1",
71
+ "http_address": data["http_address"],
72
+ "zmq_address": data["zmq_address"],
73
+ "dp_size": data["dp_size"],
74
+ "tp_size": data["tp_size"],
75
+ "transfer_mode": data["transfer_mode"],
76
+ }
77
+ # zmq_address format: "host:IP,handshake:PORT,notify:PORT"
78
+ # Stored verbatim; embedded into the request_id by handle_request.
79
+
80
+ global TRANSFER_TYPE
81
+ transfer_mode = instance["transfer_mode"]
82
+ target_list = prefill_instances if role == "P" else decode_instances
83
with _list_lock:
79
- _append_whole_dict_unique(prefill_instances, data)
80
-
81
- elif (
82
- data["type"] == "register"
83
- and data["role"] == "D"
84
- and data["request_address"] not in decode_instances
85
- ):
86
- with _list_lock:
87
- _append_whole_dict_unique(decode_instances, data)
84
+ if TRANSFER_TYPE is None:
85
+ TRANSFER_TYPE = transfer_mode
86
+ logger.info("SET TRANSFER TYPE TO %s", TRANSFER_TYPE)
87
+ elif transfer_mode != TRANSFER_TYPE:
88
+ logger.error(
89
+ "Mismatched transfer mode: expected %s, got %s;"
90
+ " skipping registration of %s",
91
+ TRANSFER_TYPE,
92
+ transfer_mode,
93
+ data["http_address"],
94
+ )
95
+ continue
96
+ existing_idx = next(
97
+ (
98
+ idx
99
+ for idx, i in enumerate(target_list)
100
+ if i.get("http_address") == data["http_address"]
101
+ ),
102
+ None,
103
+ )
104
+ if existing_idx is not None:
105
+ target_list[existing_idx] = instance
106
+ logger.info(
107
+ "Updated existing %s instance: %s",
108
+ "Prefill" if role == "P" else "Decode",
109
+ instance,
110
+ )
111
+ else:
112
+ target_list.append(instance)
113
+ logger.info(
114
+ "Registered %s instance: %s",
115
+ "Prefill" if role == "P" else "Decode",
116
+ instance,
117
+ )
118
+ else:
119
+ logger.warning(
120
+ "Received message with unrecognized type %r; ignoring",
121
+ data.get("type"),
122
+ )
123
124
125
def start_service_discovery(hostname, port):
@@ -101,7 +136,7 @@ def start_service_discovery(hostname, port):
136
137
138
async def send_request_to_prefill(
104
- endpoint, req_data, request_id, d_endpoint, dip, dport, selected_prefill_dp_rank
139
+ endpoint, req_data, request_id, selected_prefill_dp_rank
140
):
141
req_data_copy = req_data
142
@@ -109,12 +144,8 @@ async def send_request_to_prefill(
144
{
145
"do_remote_decode": True,
146
"do_remote_prefill": False,
112
- "remote_handshake_port": d_endpoint["handshake_port"],
113
- "remote_notify_port": d_endpoint["notify_port"],
147
"remote_engine_id": None,
148
"remote_block_ids": None,
116
- "remote_host": dip,
117
- "remote_port": dport,
149
}
150
)
151
req_data_copy["stream"] = False
@@ -197,14 +228,7 @@ async def handle_request(api: str, request: Request):
228
global request_nums
229
request_nums += 1
230
200
- def extract_ip_port_fast(url):
201
- match = IP_PORT_PATTERN.search(url)
202
- if not match:
203
- raise ValueError(f"Invalid URL format: {url}")
204
- return match.groups()
205
-
231
req_data = await request.get_json()
207
- request_id = str(uuid.uuid4())
232
233
prefill_instance_endpoint = None
234
decode_instance_endpoint = None
@@ -230,7 +254,14 @@ async def handle_request(api: str, request: Request):
254
prefill_instance_endpoint["dp_size"],
255
)
256
233
- dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
257
+ # Embed both zmq_addresses in the request_id so the connector can parse
258
+ # the peer's host/ports from it, similar to P2P-NCCL
259
+ uid = str(uuid.uuid4()).replace("-", "")
260
+ request_id = (
261
+ f"___prefill_addr_{prefill_instance_endpoint['zmq_address']}"
262
+ f"___decode_addr_{decode_instance_endpoint['zmq_address']}"
263
+ f"_{uid}"
264
+ )
265
266
transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"
267
@@ -251,35 +282,30 @@ async def handle_request(api: str, request: Request):
282
prefill_request_url,
283
req_data_to_prefill,
284
request_id,
254
- decode_instance_endpoint,
255
- dip,
256
- dport,
285
selected_prefill_dp_rank,
286
)
287
)
260
- ip, port = extract_ip_port_fast(prefill_request_url)
288
289
req_data["max_tokens"] -= 1
290
291
req_data["kv_transfer_params"] = {
292
"do_remote_decode": False,
293
"do_remote_prefill": True,
267
- "remote_handshake_port": prefill_instance_endpoint["handshake_port"],
268
- "remote_notify_port": prefill_instance_endpoint["notify_port"],
294
"remote_engine_id": None,
295
"remote_block_ids": None,
271
- "remote_host": ip,
272
- "remote_port": port,
296
+ "transfer_id": transfer_id,
297
}
298
if TRANSFER_TYPE == "READ":
299
# In read mode, prefill and decode are executed serially.
300
prefill_response = await send_prefill_task
277
- req_data["kv_transfer_params"]["remote_engine_id"] = prefill_response[
278
- "kv_transfer_params"
279
- ]["remote_engine_id"]
280
- req_data["kv_transfer_params"]["remote_block_ids"] = prefill_response[
281
- "kv_transfer_params"
282
- ]["remote_block_ids"]
301
+ prefill_kv = prefill_response["kv_transfer_params"]
302
+ req_data["kv_transfer_params"]["remote_engine_id"] = prefill_kv[
303
+ "remote_engine_id"
304
+ ]
305
+ req_data["kv_transfer_params"]["remote_block_ids"] = prefill_kv[
306
+ "remote_block_ids"
307
+ ]
308
+ req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"]
309
310
req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[
311
"dp_size"
@@ -290,7 +316,6 @@ async def handle_request(api: str, request: Request):
316
317
if selected_prefill_dp_rank is not None:
318
req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
293
- req_data["kv_transfer_params"]["transfer_id"] = transfer_id
319
320
decode_request_url = decode_instance_endpoint["request_address"] + api
321
decode_request_task = asyncio.create_task(
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+72
-5
@@ -8,6 +8,7 @@ from dataclasses import dataclass
8
from typing import TYPE_CHECKING, Any
9
10
import msgspec
11
+import regex as re
12
import torch
13
import zmq
14
@@ -239,7 +240,7 @@ class MoRIIOConstants:
240
COMPLETION_PREFIX = "cmpl"
241
TRANSFER_PREFIX = "tx"
242
242
- PING_INTERVAL = 5
243
+ PING_INTERVAL = 3
244
MAX_PING_RETRIES = 100
245
DEFAULT_HANDSHAKE_PORT = "6301"
246
DEFAULT_NOTIFY_PORT = "61005"
@@ -247,6 +248,64 @@ class MoRIIOConstants:
248
VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT = 3600
249
250
251
+# The router embeds both zmq_addresses in the request_id (similar to P2pNcclConnector):
252
+# "___prefill_addr_{zmq}___decode_addr_{zmq}_{32-hex-uuid}"
253
+# MoRIIO zmq_address format: "host:IP,handshake:PORT,notify:PORT"
254
+#
255
+# This lets each connector side parse the peer's connection info without
256
+# requiring the router to pass it explicitly in kv_transfer_params.
257
+_PREFILL_ZMQ_RE = re.compile(r"___prefill_addr_(.+?)___decode_addr_")
258
+# vLLM wraps the router's X-Request-Id as "cmpl-<id>-<seq>-<hex>" so there may
259
+# be a trailing "-<seq>-<hex>" suffix after the 32-char UUID. Allow it.
260
+_DECODE_ZMQ_RE = re.compile(r"___decode_addr_(.+)_[0-9a-f]{32}(?:-.*)?$")
261
+
262
+
263
+def parse_moriio_zmq_address(
264
+ zmq_address: str,
265
+) -> tuple[str, int, int]:
266
+ """Parse the MoRI-IO zmq address into its components.
267
+
268
+ Parses ``"host:IP,handshake:PORT,notify:PORT"`` into
269
+ (host, handshake_port, notify_port).
270
+
271
+ Each key-value pair is split on the *first* colon so that IPv6 addresses
272
+ (e.g. ``host:::1``) are handled correctly. Raises ``ValueError`` if any
273
+ of ``host``, ``handshake``, or ``notify`` keys are absent or if the port
274
+ values are non-numeric.
275
+ """
276
+ parts: dict[str, str] = {}
277
+ for segment in zmq_address.split(","):
278
+ key, _, val = segment.partition(":")
279
+ parts[key.strip()] = val.strip()
280
+ try:
281
+ host = parts["host"]
282
+ handshake_port = int(parts["handshake"])
283
+ notify_port = int(parts["notify"])
284
+ except (KeyError, ValueError) as e:
285
+ raise ValueError(
286
+ f"Malformed zmq_address {zmq_address!r}: expected "
287
+ f"'host:IP,handshake:PORT,notify:PORT' format"
288
+ ) from e
289
+ return host, handshake_port, notify_port
290
+
291
+
292
+def get_peer_zmq_from_request_id(request_id: str, is_producer: bool) -> str:
293
+ """Extract the *peer's* zmq_address from the vLLM router request_id.
294
+
295
+ The producer (prefill) needs the decode's address; the consumer (decode)
296
+ needs the prefill's address.
297
+ """
298
+ if is_producer:
299
+ m = _DECODE_ZMQ_RE.search(request_id)
300
+ else:
301
+ m = _PREFILL_ZMQ_RE.search(request_id)
302
+ if m is None:
303
+ raise ValueError(
304
+ f"Cannot parse peer zmq_address from request_id: {request_id!r}"
305
+ )
306
+ return m.group(1)
307
+
308
+
309
@dataclass
310
class ReqMeta:
311
"""Metadata for a single request."""
@@ -286,15 +345,23 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
345
write_mode=False,
346
):
347
transfer_id = kv_transfer_params["transfer_id"]
348
+
349
+ # Parse host/ports from the request_id. The router embeds both zmq_addresses
350
+ # in the request_id
351
+ peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode)
352
+ remote_host, remote_handshake_port, remote_notify_port = (
353
+ parse_moriio_zmq_address(peer_zmq)
354
+ )
355
+
356
_req = ReqMeta(
357
transfer_id=transfer_id,
358
local_block_ids=local_block_ids,
359
remote_block_ids=kv_transfer_params["remote_block_ids"],
360
remote_engine_id=kv_transfer_params["remote_engine_id"],
294
- remote_host=kv_transfer_params["remote_host"],
295
- remote_port=kv_transfer_params["remote_port"],
296
- remote_handshake_port=kv_transfer_params["remote_handshake_port"],
297
- remote_notify_port=kv_transfer_params["remote_notify_port"],
361
+ remote_host=remote_host,
362
+ remote_port=remote_handshake_port,
363
+ remote_handshake_port=remote_handshake_port,
364
+ remote_notify_port=remote_notify_port,
365
tp_size=kv_transfer_params.get("tp_size", 1),
366
remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
367
)
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+40
-20
@@ -35,8 +35,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
35
TransferId,
36
WriteTask,
37
get_moriio_mode,
38
+ get_peer_zmq_from_request_id,
39
get_port_offset,
40
get_role,
41
+ parse_moriio_zmq_address,
42
set_role,
43
zmq_ctx,
44
)
@@ -379,13 +381,12 @@ class MoRIIOConnectorScheduler:
381
if params is not None and params.get("do_remote_prefill"):
382
if self.mode == MoRIIOMode.READ:
383
if remote_block_ids := params.get("remote_block_ids"):
382
- if all(
383
- p in params
384
- for p in ("remote_engine_id", "remote_host", "remote_port")
385
- ):
386
- # If remote_blocks and num_external_tokens = 0, we
384
+ # remote_engine_id is returned by the prefill's request_finished.
385
+ # host/ports come from the request_id (parsed in add_new_req).
386
+ if "remote_engine_id" in params:
387
+ # If remote_blocks and num_external_tokens = 0, we have
388
# a full prefix cache hit on the D worker. We need to call
388
- # send_notif in _read_blocks to free the memory on the P.
389
+ # send_notify in _read_blocks to free the memory on the P.
390
391
# Get unhashed blocks to pull from remote.
392
local_block_ids = blocks.get_block_ids()[0]
@@ -407,22 +408,30 @@ class MoRIIOConnectorScheduler:
408
)
409
410
else:
411
+ # WRITE mode: prefill scheduler notifies the decode side that
412
+ # blocks are ready. Parse the decode's host/notify_port from
413
+ # the request_id
414
assert request.kv_transfer_params is not None, (
415
"kv_transfer_params should not be None"
416
)
417
418
remote_dp_rank = request.kv_transfer_params.get("remote_dp_rank", 0)
419
420
+ peer_zmq = get_peer_zmq_from_request_id(
421
+ request.request_id, is_producer=True
422
+ )
423
+ remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)
424
+
425
for tp_index in range(self.tp_size):
417
- target_port = request.kv_transfer_params[
418
- "remote_notify_port"
419
- ] + get_port_offset(remote_dp_rank, tp_index)
426
+ target_port = remote_notify_port + get_port_offset(
427
+ remote_dp_rank, tp_index
428
+ )
429
430
self.send_notify_block(
431
req_id=request.request_id,
432
transfer_id=request.kv_transfer_params["transfer_id"],
433
block_notify_list=blocks.get_block_ids()[0],
425
- host=params.get("remote_host"),
434
+ host=remote_host,
435
port=target_port,
436
)
437
@@ -584,15 +593,15 @@ class MoRIIOConnectorScheduler:
593
+ MoRIIOConstants.VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT
594
)
595
587
- # If we execute in P-D serial mode, no notification port is needed.
596
+ # Return KV transfer params forwarded verbatim to the decode instance by
597
+ # the router.
598
return delay_free_blocks, dict(
599
do_remote_prefill=True,
600
do_remote_decode=False,
601
remote_block_ids=computed_block_ids,
602
remote_engine_id=self.engine_id,
593
- remote_host=self.host_ip,
594
- remote_port=self.handshake_port,
603
tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
604
+ transfer_id=params["transfer_id"],
605
)
606
607
@@ -846,7 +855,15 @@ class MoRIIOConnectorWorker:
855
]
856
857
def _ping(self, zmq_context):
849
- http_request_address = f"http://{self.request_address}/v1"
858
+ # Use host:port format for http_address (compatible with official router)
859
+ http_address = f"{self.request_address}"
860
+ # Include host so the router embeds it in the request_id; the connector
861
+ # on the other side parses host/ports from there.
862
+ zmq_address = (
863
+ f"host:{self.local_ip},"
864
+ f"handshake:{self.handshake_port},"
865
+ f"notify:{self.notify_port}"
866
+ )
867
role = "P" if self.is_producer else "D"
868
869
retry_count = 0
@@ -857,14 +874,17 @@ class MoRIIOConnectorWorker:
874
while True:
875
try:
876
data = {
860
- "type": "register",
861
- "role": role,
862
- "index": str(index),
863
- "request_address": http_request_address,
864
- "handshake_port": self.handshake_port,
865
- "notify_port": self.notify_port,
877
+ "type": role, # "P" or "D"
878
+ "http_address": http_address,
879
+ "zmq_address": zmq_address,
880
+ # dp_size/tp_size are not used by the official vLLM router
881
+ # (routing operates at the http_address level); they are
882
+ # consumed only by the toy proxy server.
883
"dp_size": self.moriio_config.dp_size,
884
"tp_size": self.moriio_config.tp_size,
885
+ # transfer_mode is included so the router can distinguish
886
+ # READ (prefill-then-decode, sequential) from WRITE (concurrent)
887
+ # scheduling.
888
"transfer_mode": self.mode.name,
889
}
890