@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0024f39a32

Fork 0

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 6, 17:00 UTC

[ROCm][P/D][MORI][BugFix] Add transfer_id for moriio_connector so moriio_connector to restore P/D functionality (#34907)

Signed-off-by: Randall Smith <Randall.Smith@amd.com>

rasmith committed Mar 15, 2026 at 21:36 UTC 0024f39a3224326a9f871919cf16a06c58edfdad

4 files changed +101 -33

examples/online_serving/disaggregated_serving/moriio_toy_proxy_server.py

@@ -14,6 +14,10 @@ import regex as re

14 import zmq

15 from quart import Quart, make_response, request

17 +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (

18 + MoRIIOConstants,

19 +)

20 +

21 logger = logging.getLogger(__name__)

22 logger.setLevel(logging.DEBUG)

23 prefill_instances: list[dict] = []

@@ -213,6 +217,8 @@ async def handle_request():

217

218 dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])

219

220 + transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"

221 +

222 req_data_to_prefill = copy.deepcopy(req_data)

223 req_data_to_prefill["kv_transfer_params"] = {}

224 req_data["kv_transfer_params"] = {}

@@ -222,6 +228,7 @@ async def handle_request():

228 req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = (

229 decode_instance_endpoint["tp_size"]

230 )

231 + req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id

232

233 send_prefill_task = asyncio.create_task(

234 send_request_to_prefill(

@@ -267,6 +274,7 @@ async def handle_request():

274

275 if selected_prefill_dp_rank is not None:

276 req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank

277 + req_data["kv_transfer_params"]["transfer_id"] = transfer_id

278

279 decode_request_task = asyncio.create_task(

280 start_decode_request(

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py

+16 -15

@@ -39,11 +39,13 @@ logger = init_logger(__name__)

39 Transfer = tuple[int, float]

40 EngineId = str

41 ReqId = str

42 +TransferId = str

45 @dataclass

46 class WriteTask:

46 - request_id: str

47 + request_id: ReqId

48 + transfer_id: TransferId

49 dst_engine_id: str

50 local_block_ids: list[int]

51 remote_block_ids_hint: list[int] | None

@@ -59,7 +61,8 @@ class WriteTask:

61 class LayerTransferPlan:

62 """Plan for transferring a single layer."""

62 - request_id: str

64 + request_id: ReqId

65 + transfer_id: TransferId

66 layer_name: str

67 sess_idx: int

68 transfer_local_offsets: list[int]

@@ -234,6 +237,7 @@ class MoRIIOConstants:

237 POP_DONE_RECV = b"pop_done_recv"

238 OVER = b"OVER"

239 COMPLETION_PREFIX = "cmpl"

240 + TRANSFER_PREFIX = "tx"

241

242 PING_INTERVAL = 5

243 MAX_PING_RETRIES = 100

@@ -247,6 +251,7 @@ class MoRIIOConstants:

251 class ReqMeta:

252 """Metadata for a single request."""

253

254 + transfer_id: TransferId

255 local_block_ids: list[int]

256 remote_block_ids: list[int]

257 remote_host: str

@@ -263,21 +268,15 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):

268 self.reqs_to_recv: dict[ReqId, ReqMeta] = {}

269 self.reqs_to_save: dict[ReqId, ReqMeta] = {}

270 self.reqs_to_send: dict[ReqId, float] = {}

271 + self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}

272

273 def __repr__(self):

268 - return_str = ""

269 - for req_id, req_meta in self.reqs_to_recv.items():

270 - return_str += (

271 - f"{req_id = },{req_meta.local_block_ids = },"

272 - f"{req_meta.remote_host = },{req_meta.remote_port = }"

273 - f"{req_meta.remote_engine_id = },{req_meta.tp_size = }"

274 - )

275 - return_str = f"MoRIIOConnectorMetadata:reqs_to_recv:{return_str},"

276 -

277 - for req_id, expiry in self.reqs_to_send.items():

278 - return_str += f"{req_id = },{expiry = }"

279 - return_str = f"MoRIIOConnectorMetadata:reqs_to_send:{return_str},"

280 - return return_str

274 + return (

275 + f"MoRIIOConnectorMetadata: reqs_to_recv={self.reqs_to_recv}, "

276 + f"reqs_to_save={self.reqs_to_save}, "

277 + f"reqs_to_send={self.reqs_to_send}, "

278 + f"transfer_id_to_request_id={self.transfer_id_to_request_id}"

279 + )

280

281 def add_new_req(

282 self,

@@ -286,7 +285,9 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):

285 kv_transfer_params: dict[str, Any],

286 write_mode=False,

287 ):

288 + transfer_id = kv_transfer_params["transfer_id"]

289 _req = ReqMeta(

290 + transfer_id=transfer_id,

291 local_block_ids=local_block_ids,

292 remote_block_ids=kv_transfer_params["remote_block_ids"],

293 remote_engine_id=kv_transfer_params["remote_engine_id"],

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

+56 -4

@@ -32,6 +32,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (

32 MoRIIOMode,

33 ReqId,

34 ReqMeta,

35 + TransferId,

36 WriteTask,

37 get_moriio_mode,

38 get_port_offset,

@@ -277,6 +278,30 @@ class MoRIIOConnectorScheduler:

278 # Reqs to send and their expiration time

279 self._reqs_need_send: dict[ReqId, float] = {}

280 self.paths: dict[str, zmq.Socket] = {}

281 + self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}

282 + self.request_id_to_transfer_id: dict[ReqId, TransferId] = {}

283 +

284 + def map_request_id(self, request_id: ReqId, transfer_id: TransferId):

285 + self.transfer_id_to_request_id[transfer_id] = request_id

286 + self.request_id_to_transfer_id[request_id] = transfer_id

287 +

288 + def unmap_request_id(self, request_id: ReqId):

289 + if request_id in self.request_id_to_transfer_id:

290 + transfer_id = self.request_id_to_transfer_id[request_id]

291 + del self.request_id_to_transfer_id[request_id]

292 + if transfer_id in self.transfer_id_to_request_id:

293 + del self.transfer_id_to_request_id[transfer_id]

294 + else:

295 + logger.warning(

296 + "transfer id not in transfer_id_to_request_id lookup"

297 + "table. there is likely a bug!"

298 + )

299 + else:

300 + logger.warning(

301 + "Could not find %s in transfer_id_to_request_id"

302 + "lookup table. This could lead to a possible hang.",

303 + request_id,

304 + )

305

306 def get_num_new_matched_tokens(

307 self,

@@ -309,7 +334,12 @@ class MoRIIOConnectorScheduler:

334 return len(token_ids) - 1 - num_computed_tokens, False

335

336 def send_notify_block(

312 - self, req_id: str, block_notify_list: list[int], host=None, port=None

337 + self,

338 + req_id: ReqId,

339 + transfer_id: TransferId,

340 + block_notify_list: list[int],

341 + host=None,

342 + port=None,

343 ):

344 path = make_zmq_path("tcp", host, port)

345 if path not in self.paths:

@@ -321,6 +351,7 @@ class MoRIIOConnectorScheduler:

351

352 data = {

353 "req_id": req_id,

354 + "transfer_id": transfer_id,

355 "block_notify_list": block_notify_list or [],

356 "decode_rank": self.dp_rank,

357 "type": "remote_blocks",

@@ -338,6 +369,9 @@ class MoRIIOConnectorScheduler:

369 params = request.kv_transfer_params

370 if not params:

371 return

372 + transfer_id = params["transfer_id"]

373 + request_id = request.request_id

374 + self.map_request_id(request_id, transfer_id)

375 if params.get("do_remote_decode"):

376 local_block_ids = blocks.get_block_ids()[0]

377 self._reqs_need_save[request.request_id] = (request, local_block_ids)

@@ -386,6 +420,7 @@ class MoRIIOConnectorScheduler:

420

421 self.send_notify_block(

422 req_id=request.request_id,

423 + transfer_id=request.kv_transfer_params["transfer_id"],

424 block_notify_list=blocks.get_block_ids()[0],

425 host=params.get("remote_host"),

426 port=target_port,

@@ -400,6 +435,7 @@ class MoRIIOConnectorScheduler:

435 scheduler_output: SchedulerOutput,

436 ) -> KVConnectorMetadata:

437 meta = MoRIIOConnectorMetadata()

438 + meta.transfer_id_to_request_id = self.transfer_id_to_request_id

439

440 if self.mode == MoRIIOMode.WRITE:

441 # when async_load_kv finished,

@@ -506,6 +542,9 @@ class MoRIIOConnectorScheduler:

542 should be freed now or will be sent asynchronously and freed later.

543 """

544

545 + request_id = request.request_id

546 + self.unmap_request_id(request_id)

547 +

548 params = request.kv_transfer_params

549 logger.debug(

550 "MoriioConnector request_finished, request_status=%s, "

@@ -728,6 +767,7 @@ class MoRIIOConnectorWorker:

767 self.cache_config.cache_dtype,

768 use_mla=self.use_mla,

769 )

770 + self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}

771

772 # TODO: consider the integration of flashinfer or other backends.

773 self.backend_name = backend.get_name()

@@ -735,7 +775,8 @@ class MoRIIOConnectorWorker:

775

776 def schedule_write_blocks(

777 self,

738 - request_id: str,

778 + request_id: ReqId,

779 + transfer_id: TransferId,

780 dst_engine_id: str,

781 local_block_ids: list[int],

782 remote_block_ids: list[int] | None,

@@ -748,6 +789,7 @@ class MoRIIOConnectorWorker:

789

790 Args:

791 request_id: Unique identifier for the request

792 + transfer_id: Unique identifier for the transfer

793 dst_engine_id: Destination engine ID

794 local_block_ids: Local block IDs to transfer

795 remote_block_ids: Hint for remote block IDs

@@ -768,6 +810,7 @@ class MoRIIOConnectorWorker:

810

811 task = WriteTask(

812 request_id=request_id,

813 + transfer_id=transfer_id,

814 dst_engine_id=dst_engine_id,

815 local_block_ids=local_block_ids,

816 remote_block_ids_hint=remote_block_ids,

@@ -1010,7 +1053,7 @@ class MoRIIOConnectorWorker:

1053 return {remote_agent_name}

1054

1055 def _background_moriio_handshake(

1013 - self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta

1056 + self, req_id: ReqId, remote_engine_id: EngineId, meta: ReqMeta

1057 ):

1058 # Do MoRIIO handshake in background and add to _ready_requests when done.

1059 fut = None

@@ -1189,6 +1232,13 @@ class MoRIIOConnectorWorker:

1232 else:

1233 done_recving = self._pop_done_transfers()

1234

1235 + done_recving = {

1236 + self.transfer_id_to_request_id[id]

1237 + for id in filter(

1238 + lambda id: id in self.transfer_id_to_request_id, done_recving

1239 + )

1240 + }

1241 +

1242 return done_sending, done_recving

1243

1244 def _pop_done_transfers(self) -> set[str]:

@@ -1269,6 +1319,7 @@ class MoRIIOConnectorWorker:

1319 Start loading by triggering non-blocking moriio_xfer.

1320 We check for these trnxs to complete in each step().

1321 """

1322 + self.transfer_id_to_request_id = metadata.transfer_id_to_request_id

1323 if self.is_producer:

1324 self.moriio_wrapper.async_wait_reqid()

1325 return

@@ -1332,9 +1383,10 @@ class MoRIIOConnectorWorker:

1383 remote_notify_port=meta.remote_notify_port,

1384 )

1385

1335 - def _write_blocks_for_req(self, req_id: str, meta: ReqMeta, layer_name, kv_layer):

1386 + def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):

1387 self.schedule_write_blocks(

1388 request_id=req_id,

1389 + transfer_id=meta.transfer_id,

1390 dst_engine_id=meta.remote_engine_id,

1391 local_block_ids=meta.local_block_ids,

1392 remote_block_ids=meta.remote_block_ids,

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py

+21 -14

@@ -29,6 +29,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (

29 MoRIIOError,

30 RemoteAllocInfo,

31 TransferError,

32 + TransferId,

33 WriteTask,

34 get_port_offset,

35 get_role,

@@ -162,14 +163,14 @@ class MoRIIOWriter:

163 True if remote blocks are ready

164 """

165 return (

165 - task.request_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict

166 + task.transfer_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict

167 )

168

168 - def _get_remote_alloc_info(self, request_id: str) -> RemoteAllocInfo:

169 + def _get_remote_alloc_info(self, transfer_id: str) -> RemoteAllocInfo:

170 """Get remote allocation info for a request.

171

172 Args:

172 - request_id: The request ID

173 + transfer_id:TransferId The request ID

174

175 Returns:

176 Remote allocation information

@@ -178,10 +179,10 @@ class MoRIIOWriter:

179 KeyError: If allocation info is missing

180 """

181 try:

181 - return self.worker.moriio_wrapper.done_remote_allocate_req_dict[request_id]

182 + return self.worker.moriio_wrapper.done_remote_allocate_req_dict[transfer_id]

183 except KeyError as e:

184 raise KeyError(

184 - f"Remote allocation info missing for request {request_id}"

185 + f"Remote allocation info missing for transfer {transfer_id}"

186 ) from e

187

188 def _execute_write_task(self, task: WriteTask) -> None:

@@ -192,10 +193,14 @@ class MoRIIOWriter:

193

194 """

195 # Get remote allocation info

195 - request_info = self._get_remote_alloc_info(task.request_id)

196 + request_info = self._get_remote_alloc_info(task.transfer_id)

197

198 if request_info.block_ids is None:

198 - logger.debug("Request %s remote block IDs not ready", task.request_id)

199 + logger.debug(

200 + "Request remote block IDs not ready:request_id = %s, transfer_id = %s",

201 + task.request_id,

202 + task.transfer_id,

203 + )

204 return

205

206 # Wait for CUDA event

@@ -257,6 +262,7 @@ class MoRIIOWriter:

262

263 return LayerTransferPlan(

264 request_id=task.request_id,

265 + transfer_id=task.transfer_id,

266 layer_name=task.layer_name,

267 sess_idx=sess_idx,

268 transfer_local_offsets=local_off,

@@ -312,17 +318,18 @@ class MoRIIOWriter:

318

319 # Send completion notification

320 self.worker.moriio_wrapper.send_notify(

315 - task.request_id, task.remote_ip, remote_port

321 + task.transfer_id, task.remote_ip, remote_port

322 )

323 # mark request as done, then we can free the blocks

324 with self.worker.moriio_wrapper.lock:

325 self.worker.moriio_wrapper.done_req_ids.append(task.request_id)

326 del self.worker.moriio_wrapper.done_remote_allocate_req_dict[

321 - task.request_id

327 + task.transfer_id

328 ]

329 logger.debug(

324 - "Completed transfer for request %s, notified port %d",

330 + "Completed transfer for (request, transfer) %s, %s, notified port %d",

331 task.request_id,

332 + task.transfer_id,

333 remote_port,

334 )

335

@@ -355,7 +362,7 @@ class MoRIIOWrapper:

362 self.notify_port: int | None = None

363 self.lock = threading.Lock()

364 self.done_req_ids: list[str] = []

358 - self.done_remote_allocate_req_dict: dict[str, RemoteAllocInfo] = {}

365 + self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {}

366 self.done_write_cache_req_ids: list[str] = []

367 self.notify_thread: threading.Thread | None = None

368 self.sessions: list[IOEngine.Session] = []

@@ -525,7 +532,7 @@ class MoRIIOWrapper:

532

533 try:

534 msg_str = msg.decode("UTF-8")

528 - if msg_str.startswith(MoRIIOConstants.COMPLETION_PREFIX):

535 + if msg_str.startswith(MoRIIOConstants.TRANSFER_PREFIX):

536 self._handle_completion_message(msg_str)

537 handled = True

538 except UnicodeDecodeError:

@@ -535,7 +542,7 @@ class MoRIIOWrapper:

542

543 def _handle_structured_message(self, data: dict):

544 assert get_role() == ROLE.PRODUCER, "Only prefill can get block messages"

538 - req_id = data["req_id"]

545 + transfer_id = data["transfer_id"]

546 block_notify_list = data.get("block_notify_list", [])

547 decode_dp_rank = data.get("decode_rank", 0)

548 assert len(block_notify_list) > 0, (

@@ -543,7 +550,7 @@ class MoRIIOWrapper:

550 )

551

552 with self.lock:

546 - self.done_remote_allocate_req_dict[req_id] = RemoteAllocInfo(

553 + self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(

554 block_ids=block_notify_list, decode_dp_rank=decode_dp_rank

555 )

556