@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 552eb81918

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 2, 13:30 UTC

[Bugfix][ROCm] Resolve MoRI connector hangs at high concurrency (#40344)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>

Simon Danielsson committed May 28, 2026 at 16:30 UTC 552eb819184c90fe37efa6394ef20864f869bea8

3 files changed +270 -32

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

+25

@@ -211,6 +211,8 @@ class MoRIIOConfig:

211 dp_rank: int

212 dp_size: int

213 tp_size: int

214 + transfer_timeout: float

215 + defer_timeout: float

216 read_mode: bool = False

217 qp_per_transfer: int = 1

218 post_batch_size: int = -1

@@ -231,6 +233,10 @@ class MoRIIOConfig:

233 # read_mode -> If true, run the connector in READ mode (consumer

234 # pulls KV from producer) instead of the default

235 # WRITE mode.

236 + # transfer_timeout -> Timeout for waiting_for_transfer_complete before

237 + # raising TransferError (sec).

238 + # defer_timeout -> Timeout before a deferred send with no finished_sending

239 + # notification is reaped and its blocks force-freed (sec).

240

241 # Knobs for RDMA transfers, ignored if on xgmi backend

242 # qp_per_transfer -> Number of RDMA Queue Pairs per KV transfer.

@@ -260,6 +266,15 @@ class MoRIIOConfig:

266 "must be one of 'rdma' or 'xgmi'."

267 )

268

269 + transfer_timeout = float(

270 + extra_config.get(

271 + "transfer_timeout", MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT

272 + )

273 + )

274 + defer_timeout = float(

275 + extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT)

276 + )

277 +

278 return cls(

279 local_ip=get_ip(),

280 local_kv_port=get_open_port(),

@@ -278,6 +293,8 @@ class MoRIIOConfig:

293 post_batch_size=int(extra_config.get("post_batch_size", -1)),

294 num_workers=int(extra_config.get("num_workers", 1)),

295 backend=backend,

296 + transfer_timeout=transfer_timeout,

297 + defer_timeout=defer_timeout,

298 )

299

300

@@ -298,6 +315,14 @@ class MoRIIOConstants:

315

316 VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT = 3600

317

318 + # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError.

319 + # Overridable via kv_connector_extra_config["transfer_timeout"].

320 + DEFAULT_TRANSFER_TIMEOUT = 30.0

321 + # Timeout (seconds) before a deferred send with no finished_sending

322 + # notification is reaped and its blocks force-freed.

323 + # Overridable via kv_connector_extra_config["defer_timeout"].

324 + DEFAULT_DEFER_TIMEOUT = 60.0

325 +

326

327 # The router embeds both zmq_addresses in the request_id (similar to P2pNcclConnector):

328 # "___prefill_addr_{zmq}___decode_addr_{zmq}_{32-hex-uuid}"

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

+165 -17

@@ -60,6 +60,7 @@ from vllm.utils.network_utils import (

60 )

61 from vllm.v1.attention.selector import get_attn_backend

62 from vllm.v1.core.sched.output import SchedulerOutput

63 +from vllm.v1.outputs import KVConnectorOutput

64 from vllm.v1.request import RequestStatus

66 if TYPE_CHECKING:

@@ -171,6 +172,10 @@ class MoRIIOConnector(KVConnectorBase_V1):

172 assert self.connector_scheduler is not None

173 return self.connector_scheduler.request_finished(request, block_ids)

174

175 + def update_connector_output(self, connector_output: KVConnectorOutput) -> None:

176 + assert self.connector_scheduler is not None

177 + self.connector_scheduler.update_connector_output(connector_output)

178 +

179 ############################################################

180 # Worker Side Methods

181 ############################################################

@@ -279,6 +284,17 @@ class MoRIIOConnectorScheduler:

284 set_role(ROLE.CONSUMER)

285 # Reqs to send and their expiration time

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

287 + # Deadlines for requests whose block freeing was deferred.

288 + # Survives across scheduler steps. If the worker never reports

289 + # finished_sending before the deadline, we inject them into

290 + # connector_output.finished_sending so the scheduler frees the blocks to avoid

291 + # hanging indefinitely waiting for a free notification that never comes.

292 + self._deferred_send_deadlines: dict[ReqId, float] = {}

293 + self._defer_timeout = float(

294 + self.kv_transfer_config.kv_connector_extra_config.get(

295 + "defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT

296 + )

297 + )

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

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

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

@@ -408,9 +424,7 @@ class MoRIIOConnectorScheduler:

424 )

425

426 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

427 + # WRITE mode, decode side: notify P that blocks are ready

428 assert request.kv_transfer_params is not None, (

429 "kv_transfer_params should not be None"

430 )

@@ -418,7 +432,7 @@ class MoRIIOConnectorScheduler:

432 remote_dp_rank = request.kv_transfer_params.get("remote_dp_rank", 0)

433

434 peer_zmq = get_peer_zmq_from_request_id(

421 - request.request_id, is_producer=True

435 + request.request_id, is_producer=False

436 )

437 remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)

438

@@ -552,7 +566,12 @@ class MoRIIOConnectorScheduler:

566 """

567

568 request_id = request.request_id

555 - self.unmap_request_id(request_id)

569 + # Consumer: can unmap transfer_id<->request_id immediately since done_recving

570 + # has fired at this point (i.e. KV has been transferred)

571 + # Producer: must keep the mapping until we get notification that blocks can

572 + # be freed, which may be several scheduler steps later.

573 + if not self.is_producer:

574 + self.unmap_request_id(request_id)

575

576 params = request.kv_transfer_params

577 logger.debug(

@@ -592,6 +611,9 @@ class MoRIIOConnectorScheduler:

611 time.perf_counter()

612 + MoRIIOConstants.VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT

613 )

614 + self._deferred_send_deadlines[request.request_id] = (

615 + time.monotonic() + self._defer_timeout

616 + )

617

618 # Return KV transfer params forwarded verbatim to the decode instance by

619 # the router.

@@ -604,6 +626,57 @@ class MoRIIOConnectorScheduler:

626 transfer_id=params["transfer_id"],

627 )

628

629 + def update_connector_output(self, connector_output: KVConnectorOutput) -> None:

630 + """Free KV blocks from sends that never received a completion signal.

631 +

632 + Called every scheduler step. When a send is deferred (request_finished

633 + returns True), blocks remain allocated until the worker reports

634 + finished_sending. If that notification is lost (e.g. ibv_post_send

635 + failure), blocks leak permanently. This method injects timed-out

636 + entries into connector_output.finished_sending so the scheduler

637 + frees them via the normal path.

638 + """

639 + # Producer: unmap transfer_id<->request_id for sends that are now (async)

640 + # reported as completed. This unmapping have to be deferred until now

641 + # so get_finished can use it in any scheduler step.

642 + # Consumer: unmapping already done in request_finished

643 + if self.is_producer and connector_output.finished_sending:

644 + for req_id in connector_output.finished_sending:

645 + self.unmap_request_id(req_id)

646 +

647 + if not self._deferred_send_deadlines:

648 + return

649 +

650 + # Remove entries the worker already reported as finished_sending, these will be

651 + # freed anyways.

652 + for req_id in connector_output.finished_sending or ():

653 + self._deferred_send_deadlines.pop(req_id, None)

654 +

655 + now = time.monotonic()

656 + expired_reqs = [

657 + req_id

658 + for req_id, deadline in self._deferred_send_deadlines.items()

659 + if now >= deadline

660 + ]

661 + if not expired_reqs:

662 + return

663 +

664 + if connector_output.finished_sending is None:

665 + connector_output.finished_sending = set()

666 + # Register the expired requests as finished so the scheduler frees their blocks.

667 + for req_id in expired_reqs:

668 + connector_output.finished_sending.add(req_id)

669 + del self._deferred_send_deadlines[req_id]

670 + if self.is_producer:

671 + self.unmap_request_id(req_id)

672 + logger.warning(

673 + "Reaped %d deferred sends with no finished_sending notification "

674 + "after %.0fs. This indicates lost async KV completion "

675 + "notifications from the KV connector.",

676 + len(expired_reqs),

677 + self._defer_timeout,

678 + )

679 +

680

681 class MoRIIOConnectorWorker:

682 """Implementation of Worker side methods"""

@@ -664,6 +737,9 @@ class MoRIIOConnectorWorker:

737 self._handle_request_thread = None

738 self._ping_thread = None

739 self._writer = MoRIIOWriter(self)

740 + # Completions that arrived before transfer_id_to_request_id was populated.

741 + # Retried each step until the mapping is established.

742 + self._unmatched_write_completions: set[str] = set()

743

744 role = "producer" if self.is_producer else "consumer"

745 engine_suffix = (

@@ -696,7 +772,11 @@ class MoRIIOConnectorWorker:

772 )

773

774 # Agent.

699 - self.moriio_wrapper = MoRIIOWrapper(tp_rank=self.tp_rank, dp_rank=self.dp_rank)

775 + self.moriio_wrapper = MoRIIOWrapper(

776 + tp_rank=self.tp_rank,

777 + dp_rank=self.dp_rank,

778 + transfer_timeout=self.moriio_config.transfer_timeout,

779 + )

780 self.moriio_wrapper.set_moriio_engine(self.moriio_engine)

781 backend = (

782 BackendType.XGMI

@@ -757,7 +837,8 @@ class MoRIIOConnectorWorker:

837 self.dst_num_blocks: dict[EngineId, int] = {}

838 # In progress transfers.

839 self._recving_transfers: defaultdict[ReqId, list] = defaultdict(list)

760 - self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str]] = {}

840 + # Values are (remote_host, remote_notify_port, transfer_id).

841 + self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {}

842

843 # Track the expiration time of requests that are waiting to be sent.

844 self._reqs_to_send: dict[ReqId, float] = {}

@@ -1257,13 +1338,30 @@ class MoRIIOConnectorWorker:

1338 done_sending, done_recving = set(), set()

1339

1340 if self.is_producer:

1260 - done_sending = self.moriio_wrapper.pop_finished_req_ids()

1261 -

1341 + # pop_finished_req_ids returns transfer_ids (the ZMQ payload sent

1342 + # by decode via send_notify); map back to req_ids for the scheduler.

1343 + finished_transfer_ids = self.moriio_wrapper.pop_finished_req_ids()

1344 + done_sending = {

1345 + self.transfer_id_to_request_id[xfer_id]

1346 + for xfer_id in finished_transfer_ids

1347 + if xfer_id in self.transfer_id_to_request_id

1348 + }

1349 else:

1350 if self.mode == MoRIIOMode.WRITE:

1264 - done_recving = self.moriio_wrapper.pop_finished_write_req_ids()

1351 + fresh = self.moriio_wrapper.pop_finished_write_req_ids()

1352 + # Accumulate with any completions that arrived before their

1353 + # transfer_id was registered in transfer_id_to_request_id.

1354 + self._unmatched_write_completions |= fresh

1355 + done_recving = self._unmatched_write_completions

1356 else:

1266 - done_recving = self._pop_done_transfers()

1357 + # READ mode: the scheduler treats KV loads as synchronous

1358 + # (load_kv_async=False), so requests go directly to RUNNING

1359 + # instead of WAITING_FOR_REMOTE_KVS. We still call

1360 + # _pop_done_transfers() to send the notify to the prefill

1361 + # side and clean up internal state, but we must NOT report

1362 + # these as done_recving because the scheduler doesn't

1363 + # expect a finished_recving signal for RUNNING requests.

1364 + self._pop_done_transfers()

1365

1366 done_recving = {

1367 self.transfer_id_to_request_id[id]

@@ -1271,6 +1369,14 @@ class MoRIIOConnectorWorker:

1369 lambda id: id in self.transfer_id_to_request_id, done_recving

1370 )

1371 }

1372 + if self.mode == MoRIIOMode.WRITE and not self.is_producer:

1373 + # Remove the ones we successfully matched; leave unmatched for retry.

1374 + matched_xfer_ids = {

1375 + id

1376 + for id in self._unmatched_write_completions

1377 + if id in self.transfer_id_to_request_id

1378 + }

1379 + self._unmatched_write_completions -= matched_xfer_ids

1380

1381 return done_sending, done_recving

1382

@@ -1279,15 +1385,32 @@ class MoRIIOConnectorWorker:

1385 with self.moriio_wrapper.lock:

1386 to_remove = []

1387 for req_id, status_list in self._recving_transfers.items():

1282 - if status_list[-1].Succeeded():

1283 - done_req_ids.add(req_id)

1284 -

1285 - self.moriio_wrapper.send_notify(

1388 + last = status_list[-1]

1389 + if last.Succeeded():

1390 + host, port, xfer_id = self._recving_transfers_callback_addr[req_id]

1391 + done_req_ids.add(xfer_id)

1392 + self.moriio_wrapper.send_notify(xfer_id, host, port)

1393 + to_remove.append(req_id)

1394 + elif last.Failed():

1395 + logger.error(

1396 + "RDMA transfer failed for request %s: %s (code=%s). "

1397 + "Notifying prefill to free blocks; request will be "

1398 + "aborted by timeout.",

1399 req_id,

1287 - self._recving_transfers_callback_addr[req_id][0],

1288 - self._recving_transfers_callback_addr[req_id][1],

1400 + last.Message(),

1401 + last.Code(),

1402 )

1403 + host, port, xfer_id = self._recving_transfers_callback_addr[req_id]

1404 + try:

1405 + self.moriio_wrapper.send_notify(xfer_id, host, port)

1406 + except Exception:

1407 + logger.exception(

1408 + "Failed to send error notification for request %s",

1409 + req_id,

1410 + )

1411 to_remove.append(req_id)

1412 + # Do NOT add to done_req_ids: decode KV cache is incomplete.

1413 + # The request will expire via the normal request timeout.

1414 for req_id in to_remove:

1415 del self._recving_transfers[req_id]

1416 del self._recving_transfers_callback_addr[req_id]

@@ -1328,11 +1451,22 @@ class MoRIIOConnectorWorker:

1451 continue

1452 self._write_blocks_for_req(req_id, meta, layer_name, kv_layer)

1453

1454 + if remote_engine_id is None:

1455 + return

1456 + _deadline = time.monotonic() + self.moriio_config.transfer_timeout

1457 while True:

1458 if (

1459 self._ready_requests.empty()

1460 and remote_engine_id not in self.write_ready_flags

1461 ):

1462 + if time.monotonic() > _deadline:

1463 + logger.warning(

1464 + "Timed out waiting for write_ready_flags[%s]; "

1465 + "adjust with kv_connector_extra_config.transfer_timeout",

1466 + remote_engine_id,

1467 + )

1468 + break

1469 + time.sleep(0.001)

1470 continue

1471 elif not self._ready_requests.empty() and (

1472 remote_engine_id in self.write_ready_flags

@@ -1383,12 +1517,23 @@ class MoRIIOConnectorWorker:

1517 self._read_blocks_for_req(req_id, meta)

1518 # Start transfers for requests whose handshakes have now finished.

1519

1520 + if remote_engine_id is None and not wait_handshake_readd_req:

1521 + return

1522 + _deadline = time.monotonic() + self.moriio_config.transfer_timeout

1523 while True:

1524 if (

1525 self._ready_requests.empty()

1526 and remote_engine_id not in self.load_ready_flag

1527 and wait_handshake_readd_req

1528 ):

1529 + if time.monotonic() > _deadline:

1530 + logger.warning(

1531 + "Timed out waiting for load_ready_flag[%s]; "

1532 + "adjust with kv_connector_extra_config.transfer_timeout",

1533 + remote_engine_id,

1534 + )

1535 + break

1536 + time.sleep(0.001)

1537 continue

1538 elif (

1539 not self._ready_requests.empty()

@@ -1409,6 +1554,7 @@ class MoRIIOConnectorWorker:

1554 )

1555 self._read_blocks(

1556 request_id=req_id,

1557 + transfer_id=meta.transfer_id,

1558 dst_engine_id=meta.remote_engine_id,

1559 local_block_ids=meta.local_block_ids,

1560 remote_block_ids=meta.remote_block_ids,

@@ -1566,6 +1712,7 @@ class MoRIIOConnectorWorker:

1712 remote_block_ids: list[int],

1713 dst_engine_id: str,

1714 request_id: str,

1715 + transfer_id: str,

1716 remote_host: str,

1717 remote_notify_port: int,

1718 ) -> None:

@@ -1593,4 +1740,5 @@ class MoRIIOConnectorWorker:

1740 self._recving_transfers_callback_addr[request_id] = (

1741 remote_host,

1742 str(remote_notify_port + self.tp_rank),

1743 + transfer_id,

1744 )

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

+80 -15

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import threading

4 +import time

5 from typing import TYPE_CHECKING, Any

6 from weakref import ref as weakref_ref

@@ -76,6 +77,7 @@ class MoRIIOWriter:

77 self._write_worker_started = False

78 self._write_worker_lock = threading.Lock()

79 self._deferred_tasks: list[WriteTask] = []

80 + self._defer_timeout = worker.moriio_config.defer_timeout

82 @property

83 def worker(self) -> "MoRIIOConnectorWorker":

@@ -137,23 +139,55 @@ class MoRIIOWriter:

139 continue

140

141 # Execute the task

140 -

141 - self._execute_write_task(task)

142 + try:

143 + self._execute_write_task(task)

144 + except Exception:

145 + logger.exception(

146 + "Write task failed for request %s, marking done",

147 + task.request_id,

148 + )

149 + self._mark_request_done(task.transfer_id)

150

151 def _process_deferred_tasks(self) -> None:

152 """Process tasks that were previously deferred."""

153 if not self._deferred_tasks:

154 return

155

156 + defer_timeout = self._defer_timeout

157 + now = time.perf_counter()

158 still_deferred: list[WriteTask] = []

159 +

160 for task in self._deferred_tasks:

161 + if now - task.enqueue_time > defer_timeout:

162 + logger.error(

163 + "Deferred write task for request %s expired after %.1fs "

164 + "(remote blocks never arrived), marking done",

165 + task.request_id,

166 + now - task.enqueue_time,

167 + )

168 + self._mark_request_done(task.transfer_id)

169 + continue

170 if self._is_remote_ready(task):

151 - self._execute_write_task(task)

171 + try:

172 + self._execute_write_task(task)

173 + except Exception:

174 + logger.exception(

175 + "Deferred write task failed for request %s, marking done",

176 + task.request_id,

177 + )

178 + self._mark_request_done(task.transfer_id)

179 else:

180 still_deferred.append(task)

181

182 self._deferred_tasks = still_deferred

183

184 + def _mark_request_done(self, transfer_id: str) -> None:

185 + """Mark a request done so its blocks are freed, even on transfer failure."""

186 + wrapper = self.worker.moriio_wrapper

187 + with wrapper.lock:

188 + wrapper.done_req_ids.append(transfer_id)

189 + wrapper.done_remote_allocate_req_dict.pop(transfer_id, None)

190 +

191 def _is_remote_ready(self, task: WriteTask) -> bool:

192 """Check if remote blocks are allocated for this task.

193

@@ -323,7 +357,7 @@ class MoRIIOWriter:

357 )

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

359 with self.worker.moriio_wrapper.lock:

326 - self.worker.moriio_wrapper.done_req_ids.append(task.request_id)

360 + self.worker.moriio_wrapper.done_req_ids.append(task.transfer_id)

361 del self.worker.moriio_wrapper.done_remote_allocate_req_dict[

362 task.transfer_id

363 ]

@@ -351,6 +385,7 @@ class MoRIIOWrapper:

385 moriio_engine: "IOEngine | None" = None,

386 tp_rank: int = 0,

387 dp_rank: int = 0,

388 + transfer_timeout: float = MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT,

389 ):

390 self.tp_rank = tp_rank

391 self.dp_rank = dp_rank

@@ -365,6 +400,7 @@ class MoRIIOWrapper:

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

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

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

403 + self._transfer_timeout = transfer_timeout

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

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

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

@@ -399,6 +435,14 @@ class MoRIIOWrapper:

435 post_batch_size,

436 num_workers,

437 PollCqMode.POLLING,

438 + # vLLM uses ZMQ for completion signaling

439 + # and never calls PopInboundTransferStatus.

440 + # With notifications enabled, ibv_post_send

441 + # ENOMEM at high concurrency permanently

442 + # poisons TransferStatus before the data WR

443 + # completes, hanging requests in

444 + # WAITING_FOR_REMOTE_KVS indefinitely.

445 + enable_notification=False,

446 )

447 self.moriio_engine.create_backend(backend_type, rdma_cfg)

448

@@ -483,22 +527,43 @@ class MoRIIOWrapper:

527 if not self.transfer_status:

528 return

529

486 - transfers_to_wait = []

530 with self.lock:

531 transfers_to_wait = self.transfer_status[:]

532 self.transfer_status.clear()

533

491 - for status in transfers_to_wait:

492 - try:

493 - status.Wait()

494 - if not status.Succeeded():

495 - logger.error(

496 - "Transfer failed: %s, Code: %s", status.Message(), status.Code()

534 + timeout = self._transfer_timeout

535 + deadline = time.monotonic() + timeout

536 + remaining = list(transfers_to_wait)

537 + errors: list[str] = []

538 +

539 + while remaining:

540 + timed_out = time.monotonic() > deadline

541 + still_waiting = []

542 + for status in remaining:

543 + if status.Succeeded():

544 + continue

545 + if status.Failed():

546 + errors.append(

547 + f"RDMA transfer failed: {status.Message()} "

548 + f"(code={status.Code()})"

549 )

498 - raise TransferError("MoRIIO transfer failed!")

499 - except Exception as e:

500 - logger.error("Transfer %s failed: %s", status, e)

501 - raise

550 + continue

551 + if timed_out:

552 + errors.append(

553 + f"RDMA transfer timed out after {timeout:.0f}s. "

554 + f"Check NIC queue depth or reduce concurrency; "

555 + f"adjust with kv_connector_extra_config.transfer_timeout."

556 + )

557 + else:

558 + still_waiting.append(status)

559 + remaining = still_waiting

560 + if remaining:

561 + time.sleep(0.001)

562 + if errors:

563 + raise TransferError(

564 + f"{len(errors)}/{len(transfers_to_wait)} transfers failed:\n"

565 + + "\n".join(errors)

566 + )

567

568 def async_wait_reqid(self):

569 assert self.notify_port is not None, "Notify port cannot be None"