diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 89c314a66fe8..40f880f6deda 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -211,6 +211,8 @@ class MoRIIOConfig: dp_rank: int dp_size: int tp_size: int + transfer_timeout: float + defer_timeout: float read_mode: bool = False qp_per_transfer: int = 1 post_batch_size: int = -1 @@ -231,6 +233,10 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": # read_mode -> If true, run the connector in READ mode (consumer # pulls KV from producer) instead of the default # WRITE mode. + # transfer_timeout -> Timeout for waiting_for_transfer_complete before + # raising TransferError (sec). + # defer_timeout -> Timeout before a deferred send with no finished_sending + # notification is reaped and its blocks force-freed (sec). # Knobs for RDMA transfers, ignored if on xgmi backend # qp_per_transfer -> Number of RDMA Queue Pairs per KV transfer. @@ -260,6 +266,15 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": "must be one of 'rdma' or 'xgmi'." ) + transfer_timeout = float( + extra_config.get( + "transfer_timeout", MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT + ) + ) + defer_timeout = float( + extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT) + ) + return cls( local_ip=get_ip(), local_kv_port=get_open_port(), @@ -278,6 +293,8 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": post_batch_size=int(extra_config.get("post_batch_size", -1)), num_workers=int(extra_config.get("num_workers", 1)), backend=backend, + transfer_timeout=transfer_timeout, + defer_timeout=defer_timeout, ) @@ -298,6 +315,14 @@ class MoRIIOConstants: VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT = 3600 + # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError. + # Overridable via kv_connector_extra_config["transfer_timeout"]. + DEFAULT_TRANSFER_TIMEOUT = 30.0 + # Timeout (seconds) before a deferred send with no finished_sending + # notification is reaped and its blocks force-freed. + # Overridable via kv_connector_extra_config["defer_timeout"]. + DEFAULT_DEFER_TIMEOUT = 60.0 + # The router embeds both zmq_addresses in the request_id (similar to P2pNcclConnector): # "___prefill_addr_{zmq}___decode_addr_{zmq}_{32-hex-uuid}" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index dc7264c8f7cf..167eef6e1ca8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -60,6 +60,7 @@ ) from vllm.v1.attention.selector import get_attn_backend from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus if TYPE_CHECKING: @@ -171,6 +172,10 @@ def request_finished( assert self.connector_scheduler is not None return self.connector_scheduler.request_finished(request, block_ids) + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + ############################################################ # Worker Side Methods ############################################################ @@ -279,6 +284,17 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): set_role(ROLE.CONSUMER) # Reqs to send and their expiration time self._reqs_need_send: dict[ReqId, float] = {} + # Deadlines for requests whose block freeing was deferred. + # Survives across scheduler steps. If the worker never reports + # finished_sending before the deadline, we inject them into + # connector_output.finished_sending so the scheduler frees the blocks to avoid + # hanging indefinitely waiting for a free notification that never comes. + self._deferred_send_deadlines: dict[ReqId, float] = {} + self._defer_timeout = float( + self.kv_transfer_config.kv_connector_extra_config.get( + "defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT + ) + ) self.paths: dict[str, zmq.Socket] = {} self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} @@ -408,9 +424,7 @@ def update_state_after_alloc( ) else: - # WRITE mode: prefill scheduler notifies the decode side that - # blocks are ready. Parse the decode's host/notify_port from - # the request_id + # WRITE mode, decode side: notify P that blocks are ready assert request.kv_transfer_params is not None, ( "kv_transfer_params should not be None" ) @@ -418,7 +432,7 @@ def update_state_after_alloc( remote_dp_rank = request.kv_transfer_params.get("remote_dp_rank", 0) peer_zmq = get_peer_zmq_from_request_id( - request.request_id, is_producer=True + request.request_id, is_producer=False ) remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq) @@ -552,7 +566,12 @@ def request_finished( """ request_id = request.request_id - self.unmap_request_id(request_id) + # Consumer: can unmap transfer_id<->request_id immediately since done_recving + # has fired at this point (i.e. KV has been transferred) + # Producer: must keep the mapping until we get notification that blocks can + # be freed, which may be several scheduler steps later. + if not self.is_producer: + self.unmap_request_id(request_id) params = request.kv_transfer_params logger.debug( @@ -592,6 +611,9 @@ def request_finished( time.perf_counter() + MoRIIOConstants.VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT ) + self._deferred_send_deadlines[request.request_id] = ( + time.monotonic() + self._defer_timeout + ) # Return KV transfer params forwarded verbatim to the decode instance by # the router. @@ -604,6 +626,57 @@ def request_finished( transfer_id=params["transfer_id"], ) + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + """Free KV blocks from sends that never received a completion signal. + + Called every scheduler step. When a send is deferred (request_finished + returns True), blocks remain allocated until the worker reports + finished_sending. If that notification is lost (e.g. ibv_post_send + failure), blocks leak permanently. This method injects timed-out + entries into connector_output.finished_sending so the scheduler + frees them via the normal path. + """ + # Producer: unmap transfer_id<->request_id for sends that are now (async) + # reported as completed. This unmapping have to be deferred until now + # so get_finished can use it in any scheduler step. + # Consumer: unmapping already done in request_finished + if self.is_producer and connector_output.finished_sending: + for req_id in connector_output.finished_sending: + self.unmap_request_id(req_id) + + if not self._deferred_send_deadlines: + return + + # Remove entries the worker already reported as finished_sending, these will be + # freed anyways. + for req_id in connector_output.finished_sending or (): + self._deferred_send_deadlines.pop(req_id, None) + + now = time.monotonic() + expired_reqs = [ + req_id + for req_id, deadline in self._deferred_send_deadlines.items() + if now >= deadline + ] + if not expired_reqs: + return + + if connector_output.finished_sending is None: + connector_output.finished_sending = set() + # Register the expired requests as finished so the scheduler frees their blocks. + for req_id in expired_reqs: + connector_output.finished_sending.add(req_id) + del self._deferred_send_deadlines[req_id] + if self.is_producer: + self.unmap_request_id(req_id) + logger.warning( + "Reaped %d deferred sends with no finished_sending notification " + "after %.0fs. This indicates lost async KV completion " + "notifications from the KV connector.", + len(expired_reqs), + self._defer_timeout, + ) + class MoRIIOConnectorWorker: """Implementation of Worker side methods""" @@ -664,6 +737,9 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self._handle_request_thread = None self._ping_thread = None self._writer = MoRIIOWriter(self) + # Completions that arrived before transfer_id_to_request_id was populated. + # Retried each step until the mapping is established. + self._unmatched_write_completions: set[str] = set() role = "producer" if self.is_producer else "consumer" engine_suffix = ( @@ -696,7 +772,11 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): ) # Agent. - self.moriio_wrapper = MoRIIOWrapper(tp_rank=self.tp_rank, dp_rank=self.dp_rank) + self.moriio_wrapper = MoRIIOWrapper( + tp_rank=self.tp_rank, + dp_rank=self.dp_rank, + transfer_timeout=self.moriio_config.transfer_timeout, + ) self.moriio_wrapper.set_moriio_engine(self.moriio_engine) backend = ( BackendType.XGMI @@ -757,7 +837,8 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.dst_num_blocks: dict[EngineId, int] = {} # In progress transfers. self._recving_transfers: defaultdict[ReqId, list] = defaultdict(list) - self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str]] = {} + # Values are (remote_host, remote_notify_port, transfer_id). + self._recving_transfers_callback_addr: dict[ReqId, tuple[str, str, str]] = {} # Track the expiration time of requests that are waiting to be sent. self._reqs_to_send: dict[ReqId, float] = {} @@ -1257,13 +1338,30 @@ def get_finished(self) -> tuple[set[str], set[str]]: done_sending, done_recving = set(), set() if self.is_producer: - done_sending = self.moriio_wrapper.pop_finished_req_ids() - + # pop_finished_req_ids returns transfer_ids (the ZMQ payload sent + # by decode via send_notify); map back to req_ids for the scheduler. + finished_transfer_ids = self.moriio_wrapper.pop_finished_req_ids() + done_sending = { + self.transfer_id_to_request_id[xfer_id] + for xfer_id in finished_transfer_ids + if xfer_id in self.transfer_id_to_request_id + } else: if self.mode == MoRIIOMode.WRITE: - done_recving = self.moriio_wrapper.pop_finished_write_req_ids() + fresh = self.moriio_wrapper.pop_finished_write_req_ids() + # Accumulate with any completions that arrived before their + # transfer_id was registered in transfer_id_to_request_id. + self._unmatched_write_completions |= fresh + done_recving = self._unmatched_write_completions else: - done_recving = self._pop_done_transfers() + # READ mode: the scheduler treats KV loads as synchronous + # (load_kv_async=False), so requests go directly to RUNNING + # instead of WAITING_FOR_REMOTE_KVS. We still call + # _pop_done_transfers() to send the notify to the prefill + # side and clean up internal state, but we must NOT report + # these as done_recving because the scheduler doesn't + # expect a finished_recving signal for RUNNING requests. + self._pop_done_transfers() done_recving = { self.transfer_id_to_request_id[id] @@ -1271,6 +1369,14 @@ def get_finished(self) -> tuple[set[str], set[str]]: lambda id: id in self.transfer_id_to_request_id, done_recving ) } + if self.mode == MoRIIOMode.WRITE and not self.is_producer: + # Remove the ones we successfully matched; leave unmatched for retry. + matched_xfer_ids = { + id + for id in self._unmatched_write_completions + if id in self.transfer_id_to_request_id + } + self._unmatched_write_completions -= matched_xfer_ids return done_sending, done_recving @@ -1279,15 +1385,32 @@ def _pop_done_transfers(self) -> set[str]: with self.moriio_wrapper.lock: to_remove = [] for req_id, status_list in self._recving_transfers.items(): - if status_list[-1].Succeeded(): - done_req_ids.add(req_id) - - self.moriio_wrapper.send_notify( + last = status_list[-1] + if last.Succeeded(): + host, port, xfer_id = self._recving_transfers_callback_addr[req_id] + done_req_ids.add(xfer_id) + self.moriio_wrapper.send_notify(xfer_id, host, port) + to_remove.append(req_id) + elif last.Failed(): + logger.error( + "RDMA transfer failed for request %s: %s (code=%s). " + "Notifying prefill to free blocks; request will be " + "aborted by timeout.", req_id, - self._recving_transfers_callback_addr[req_id][0], - self._recving_transfers_callback_addr[req_id][1], + last.Message(), + last.Code(), ) + host, port, xfer_id = self._recving_transfers_callback_addr[req_id] + try: + self.moriio_wrapper.send_notify(xfer_id, host, port) + except Exception: + logger.exception( + "Failed to send error notification for request %s", + req_id, + ) to_remove.append(req_id) + # Do NOT add to done_req_ids: decode KV cache is incomplete. + # The request will expire via the normal request timeout. for req_id in to_remove: del self._recving_transfers[req_id] del self._recving_transfers_callback_addr[req_id] @@ -1328,11 +1451,22 @@ def save_kv_layer( continue self._write_blocks_for_req(req_id, meta, layer_name, kv_layer) + if remote_engine_id is None: + return + _deadline = time.monotonic() + self.moriio_config.transfer_timeout while True: if ( self._ready_requests.empty() and remote_engine_id not in self.write_ready_flags ): + if time.monotonic() > _deadline: + logger.warning( + "Timed out waiting for write_ready_flags[%s]; " + "adjust with kv_connector_extra_config.transfer_timeout", + remote_engine_id, + ) + break + time.sleep(0.001) continue elif not self._ready_requests.empty() and ( remote_engine_id in self.write_ready_flags @@ -1383,12 +1517,23 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): self._read_blocks_for_req(req_id, meta) # Start transfers for requests whose handshakes have now finished. + if remote_engine_id is None and not wait_handshake_readd_req: + return + _deadline = time.monotonic() + self.moriio_config.transfer_timeout while True: if ( self._ready_requests.empty() and remote_engine_id not in self.load_ready_flag and wait_handshake_readd_req ): + if time.monotonic() > _deadline: + logger.warning( + "Timed out waiting for load_ready_flag[%s]; " + "adjust with kv_connector_extra_config.transfer_timeout", + remote_engine_id, + ) + break + time.sleep(0.001) continue elif ( not self._ready_requests.empty() @@ -1409,6 +1554,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): ) self._read_blocks( request_id=req_id, + transfer_id=meta.transfer_id, dst_engine_id=meta.remote_engine_id, local_block_ids=meta.local_block_ids, remote_block_ids=meta.remote_block_ids, @@ -1566,6 +1712,7 @@ def _read_blocks( remote_block_ids: list[int], dst_engine_id: str, request_id: str, + transfer_id: str, remote_host: str, remote_notify_port: int, ) -> None: @@ -1593,4 +1740,5 @@ def _read_blocks( self._recving_transfers_callback_addr[request_id] = ( remote_host, str(remote_notify_port + self.tp_rank), + transfer_id, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 86f2533837ba..3ca5f37ca90f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import threading +import time from typing import TYPE_CHECKING, Any from weakref import ref as weakref_ref @@ -76,6 +77,7 @@ def __init__(self, worker: "MoRIIOConnectorWorker"): self._write_worker_started = False self._write_worker_lock = threading.Lock() self._deferred_tasks: list[WriteTask] = [] + self._defer_timeout = worker.moriio_config.defer_timeout @property def worker(self) -> "MoRIIOConnectorWorker": @@ -137,23 +139,55 @@ def _write_worker_loop(self) -> None: continue # Execute the task - - self._execute_write_task(task) + try: + self._execute_write_task(task) + except Exception: + logger.exception( + "Write task failed for request %s, marking done", + task.request_id, + ) + self._mark_request_done(task.transfer_id) def _process_deferred_tasks(self) -> None: """Process tasks that were previously deferred.""" if not self._deferred_tasks: return + defer_timeout = self._defer_timeout + now = time.perf_counter() still_deferred: list[WriteTask] = [] + for task in self._deferred_tasks: + if now - task.enqueue_time > defer_timeout: + logger.error( + "Deferred write task for request %s expired after %.1fs " + "(remote blocks never arrived), marking done", + task.request_id, + now - task.enqueue_time, + ) + self._mark_request_done(task.transfer_id) + continue if self._is_remote_ready(task): - self._execute_write_task(task) + try: + self._execute_write_task(task) + except Exception: + logger.exception( + "Deferred write task failed for request %s, marking done", + task.request_id, + ) + self._mark_request_done(task.transfer_id) else: still_deferred.append(task) self._deferred_tasks = still_deferred + def _mark_request_done(self, transfer_id: str) -> None: + """Mark a request done so its blocks are freed, even on transfer failure.""" + wrapper = self.worker.moriio_wrapper + with wrapper.lock: + wrapper.done_req_ids.append(transfer_id) + wrapper.done_remote_allocate_req_dict.pop(transfer_id, None) + def _is_remote_ready(self, task: WriteTask) -> bool: """Check if remote blocks are allocated for this task. @@ -323,7 +357,7 @@ def _finalize_if_complete( ) # mark request as done, then we can free the blocks with self.worker.moriio_wrapper.lock: - self.worker.moriio_wrapper.done_req_ids.append(task.request_id) + self.worker.moriio_wrapper.done_req_ids.append(task.transfer_id) del self.worker.moriio_wrapper.done_remote_allocate_req_dict[ task.transfer_id ] @@ -351,6 +385,7 @@ def __init__( moriio_engine: "IOEngine | None" = None, tp_rank: int = 0, dp_rank: int = 0, + transfer_timeout: float = MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT, ): self.tp_rank = tp_rank self.dp_rank = dp_rank @@ -365,6 +400,7 @@ def __init__( self.done_req_ids: list[str] = [] self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {} self.done_write_cache_req_ids: list[str] = [] + self._transfer_timeout = transfer_timeout self.notify_thread: threading.Thread | None = None self.sessions: list[IOEngine.Session] = [] self.paths: dict[str, zmq.Socket] = {} @@ -399,6 +435,14 @@ def set_backend_type( post_batch_size, num_workers, PollCqMode.POLLING, + # vLLM uses ZMQ for completion signaling + # and never calls PopInboundTransferStatus. + # With notifications enabled, ibv_post_send + # ENOMEM at high concurrency permanently + # poisons TransferStatus before the data WR + # completes, hanging requests in + # WAITING_FOR_REMOTE_KVS indefinitely. + enable_notification=False, ) self.moriio_engine.create_backend(backend_type, rdma_cfg) @@ -483,22 +527,43 @@ def waiting_for_transfer_complete(self): if not self.transfer_status: return - transfers_to_wait = [] with self.lock: transfers_to_wait = self.transfer_status[:] self.transfer_status.clear() - for status in transfers_to_wait: - try: - status.Wait() - if not status.Succeeded(): - logger.error( - "Transfer failed: %s, Code: %s", status.Message(), status.Code() + timeout = self._transfer_timeout + deadline = time.monotonic() + timeout + remaining = list(transfers_to_wait) + errors: list[str] = [] + + while remaining: + timed_out = time.monotonic() > deadline + still_waiting = [] + for status in remaining: + if status.Succeeded(): + continue + if status.Failed(): + errors.append( + f"RDMA transfer failed: {status.Message()} " + f"(code={status.Code()})" ) - raise TransferError("MoRIIO transfer failed!") - except Exception as e: - logger.error("Transfer %s failed: %s", status, e) - raise + continue + if timed_out: + errors.append( + f"RDMA transfer timed out after {timeout:.0f}s. " + f"Check NIC queue depth or reduce concurrency; " + f"adjust with kv_connector_extra_config.transfer_timeout." + ) + else: + still_waiting.append(status) + remaining = still_waiting + if remaining: + time.sleep(0.001) + if errors: + raise TransferError( + f"{len(errors)}/{len(transfers_to_wait)} transfers failed:\n" + + "\n".join(errors) + ) def async_wait_reqid(self): assert self.notify_port is not None, "Notify port cannot be None"