vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector ¶

class MoRIIOConnector(KVConnectorBase_V1):
    def __init__(
        self,
        vllm_config: VllmConfig,
        role: KVConnectorRole,
        kv_cache_config: "KVCacheConfig",
    ):
        super().__init__(vllm_config, role, kv_cache_config)
        assert vllm_config.kv_transfer_config is not None, (
            "kv_transfer_config must be set for MoRIIOConnector"
        )

        self.kv_transfer_config = vllm_config.kv_transfer_config
        self._set_port_defaults(vllm_config)

        self.engine_id = (
            str(resolve_host_ip(self.kv_transfer_config.kv_connector_extra_config))
            + ":"
            + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
        )
        self.mode = get_moriio_mode(self.kv_transfer_config)
        if (
            self.mode == MoRIIOMode.READ
            and self.kv_transfer_config.is_kv_consumer
            and vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
        ):
            # warn only; kv-read barrier requires PIECEWISE cudagraph mode
            logger.warning_once(
                "MoRIIO READ mode is running with %s CUDA graphs: per-layer "
                "KV-read barrier can't fire inside full graph; accuracy may "
                "degrade at high concurrency. Set cudagraph_mode=PIECEWISE "
                "in --compilation-config.",
                vllm_config.compilation_config.cudagraph_mode.name,
            )
        if role == KVConnectorRole.SCHEDULER:
            self.connector_scheduler: MoRIIOConnectorScheduler | None = (
                MoRIIOConnectorScheduler(vllm_config, self.engine_id)
            )
            self.connector_worker: MoRIIOConnectorWorker | None = None
        elif role == KVConnectorRole.WORKER:
            self.connector_scheduler = None
            self.connector_worker = MoRIIOConnectorWorker(
                vllm_config, self.engine_id, kv_cache_config
            )
        logger.info(
            "Initialized MoRIIO Connector,engine_id:%s,role: %s",
            self.engine_id,
            role.value,
        )

    ############################################################
    # Scheduler Side Methods
    ############################################################

    def _set_port_defaults(self, vllm_config: VllmConfig):
        assert vllm_config.kv_transfer_config is not None, (
            "kv_transfer_config must be set for MoRIIOConnector"
        )
        kv_transfer_config = vllm_config.kv_transfer_config
        extra_config = kv_transfer_config.kv_connector_extra_config

        if "handshake_port" not in extra_config or not extra_config["handshake_port"]:
            extra_config["handshake_port"] = MoRIIOConstants.DEFAULT_HANDSHAKE_PORT

        if "notify_port" not in extra_config or not extra_config["notify_port"]:
            extra_config["notify_port"] = MoRIIOConstants.DEFAULT_NOTIFY_PORT

    def get_num_new_matched_tokens(
        self, request: "Request", num_computed_tokens: int
    ) -> tuple[int, bool]:
        assert self.connector_scheduler is not None
        return self.connector_scheduler.get_num_new_matched_tokens(
            request, num_computed_tokens
        )

    def update_state_after_alloc(
        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
    ):
        assert self.connector_scheduler is not None
        return self.connector_scheduler.update_state_after_alloc(
            request, blocks, num_external_tokens, self.connector_worker
        )

    def build_connector_meta(
        self,
        scheduler_output: SchedulerOutput,
    ) -> KVConnectorMetadata:
        assert self.connector_scheduler is not None
        return self.connector_scheduler.build_connector_meta(scheduler_output)

    def request_finished(
        self,
        request: "Request",
        block_ids: list[int],
    ) -> tuple[bool, dict[str, Any] | None]:
        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
    ############################################################
    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
        assert self.connector_worker is not None
        self.connector_worker.register_kv_caches(kv_caches)

    def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
        """Get the finished recving and sending requests."""
        assert self.connector_worker is not None
        return self.connector_worker.get_finished()

    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
        assert self.connector_worker is not None
        if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.CONSUMER:
            self.connector_worker.moriio_wrapper.async_wait_reqid()

        assert isinstance(self._connector_metadata, MoRIIOConnectorMetadata)
        self.connector_worker.start_load_kv(self._connector_metadata)

    def wait_for_layer_load(self, layer_name: str) -> None:
        assert self.connector_worker is not None
        self.connector_worker.wait_for_layer_load(layer_name)

    def save_kv_layer(
        self,
        layer_name: str,
        kv_layer: torch.Tensor,
        attn_metadata: "AttentionMetadata",
        **kwargs,
    ) -> None:
        # Only producer/prefill saves KV Cache
        if get_role() == ROLE.CONSUMER:
            return
        assert self.connector_worker is not None, (
            "save_kv_layer called on scheduler role"
        )

        assert isinstance(self._connector_metadata, MoRIIOConnectorMetadata), (
            "Connector metadata not initialized yet"
        )
        self.connector_worker.save_kv_layer(
            self._connector_metadata, layer_name, kv_layer, attn_metadata, **kwargs
        )

        return None

    def wait_for_save(self):
        if self.mode != MoRIIOMode.WRITE or get_role() != ROLE.PRODUCER:
            return
        assert self.connector_worker is not None
        assert isinstance(self._connector_metadata, MoRIIOConnectorMetadata), (
            "Connector metadata not initialized yet"
        )
        self.connector_worker.wait_for_save(self._connector_metadata)

    def shutdown(self):
        if self.connector_worker is not None:
            self.connector_worker.shutdown()
        if self.connector_scheduler is not None:
            self.connector_scheduler.shutdown()

    def has_connector_metadata(self) -> bool:
        """Check whether the connector metadata is currently set.

        Returns:
            bool: True if connector metadata exists, False otherwise.
        """
        try:
            return self._connector_metadata is not None
        except AttributeError:
            return False

class MoRIIOConnectorScheduler:
    """Implementation of Scheduler side methods"""

    def __init__(self, vllm_config: VllmConfig, engine_id: str):
        self.vllm_config = vllm_config

        assert vllm_config.kv_transfer_config is not None, (
            "kv_transfer_config must be set for MoRIIOConnector"
        )
        self.kv_transfer_config = vllm_config.kv_transfer_config
        self.block_size = vllm_config.cache_config.block_size
        self.engine_id: EngineId = engine_id
        self.mode = get_moriio_mode(self.kv_transfer_config)
        self.host_ip = resolve_host_ip(
            self.kv_transfer_config.kv_connector_extra_config
        )
        self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[
            "handshake_port"
        ]
        logger.info("Initializing MoRIIO Scheduler engine_id = %s", engine_id)

        self.side_notify_port = self.kv_transfer_config.kv_connector_extra_config[
            "notify_port"
        ]
        self.tp_size = self.vllm_config.parallel_config.tensor_parallel_size
        # Local DP rank (0..dp_local-1) for port arithmetic; fold_local_rank
        # handles the external-DP sentinel (dp_local==0) so the arithmetic
        # can't ZeroDivisionError at init.
        _pc = self.vllm_config.parallel_config
        self.dp_rank = fold_local_rank(
            _pc.data_parallel_rank, _pc.data_parallel_size_local
        )
        # Only first-pod ranks originate notify to avoid duplicates. Guarding
        # via pod_index keeps _is_kv_master True for the external-DP sentinel
        # (single pod) instead of being silently stuck False.
        self._is_kv_master = (
            pod_index(_pc.data_parallel_rank, _pc.data_parallel_size_local) == 0
        )
        # Global DP rank for pinned request ownership check.
        self._global_dp_rank = self.vllm_config.parallel_config.data_parallel_rank
        self.is_producer = self.kv_transfer_config.kv_role == "kv_producer"
        # Requests that need to start recv/send.
        # New requests are added by update_state_after_alloc in
        # the scheduler. Used to make metadata passed to Worker.
        self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
        self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
        # Snapshot of kv_transfer_params for chunked prefill recovery.
        self._req_kv_params: dict[ReqId, dict] = {}

        # For chunked prefill, we perform layer-wise access within the final chunk.
        # TODO: Perform transfer at end chunk.
        self._reqs_need_pending_save: dict[ReqId, tuple[Request, list[int]]] = {}

        if self.is_producer:
            set_role(ROLE.PRODUCER)
        else:
            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.
        # Value: (deadline, transfer_id) for unmapping after mutation.
        self._deferred_send_deadlines: dict[ReqId, tuple[float, TransferId | None]] = {}
        self._defer_timeout = float(
            self.kv_transfer_config.kv_connector_extra_config.get(
                "defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT
            )
        )
        # Buffer for early ACKs that arrive before request_finished.
        self._pending_sent_acks: dict[ReqId, float] = {}
        self.paths: dict[str, zmq.Socket] = {}
        self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
        self.request_id_to_transfer_id: dict[ReqId, TransferId] = {}

    def map_request_id(self, request_id: ReqId, transfer_id: TransferId):
        self.transfer_id_to_request_id[transfer_id] = request_id
        self.request_id_to_transfer_id[request_id] = transfer_id

    def unmap_request_id(
        self, request_id: ReqId, transfer_id: TransferId | None = None
    ):
        """Unmap request_id/transfer_id. Uses transfer_id for lookup if
        exact request_id match fails (handles input_processor mutation)."""
        if request_id in self.request_id_to_transfer_id:
            tid = self.request_id_to_transfer_id[request_id]
            del self.request_id_to_transfer_id[request_id]
            if tid in self.transfer_id_to_request_id:
                del self.transfer_id_to_request_id[tid]
            return

        # Fallback: use transfer_id to find original request_id.
        if transfer_id is not None and transfer_id in self.transfer_id_to_request_id:
            original_rid = self.transfer_id_to_request_id[transfer_id]
            if original_rid != request_id:
                logger.debug(
                    "MoRI-IO unmap via transfer_id: %r -> %r", request_id, original_rid
                )
            if original_rid in self.request_id_to_transfer_id:
                del self.request_id_to_transfer_id[original_rid]
            del self.transfer_id_to_request_id[transfer_id]
            return

        logger.warning(
            "MoRI-IO unmap MISS: rid=%r transfer_id=%r table_size=%d",
            request_id,
            transfer_id,
            len(self.request_id_to_transfer_id),
        )

    def get_num_new_matched_tokens(
        self,
        request: "Request",
        num_computed_tokens: int,
    ) -> tuple[int, bool]:
        """
        For remote prefill, pull all prompt blocks from remote
        asynchronously relative to engine execution.

        Args:
            request (Request): the request object.
            num_computed_tokens (int): the number of locally
                computed tokens for this request
        Returns:
            * the number of tokens that can be loaded from the
              external KV cache beyond what is already computed.
            * true if the external KV cache tokens will be loaded
              asynchronously (between scheduler steps).
        """
        if self.is_producer:
            return 0, False

        token_ids = request.prompt_token_ids or []
        if self.mode == MoRIIOMode.WRITE:
            # MoriiO in write mode, no remote prefill

            return len(token_ids) - num_computed_tokens, True

        return len(token_ids) - 1 - num_computed_tokens, False

    def send_notify_block(
        self,
        req_id: ReqId,
        transfer_id: TransferId,
        block_notify_list: list[int],
        host=None,
        port=None,
    ):
        path = make_zmq_path("tcp", host, port)
        if path not in self.paths:
            ctx = zmq.Context.instance()
            sock = make_zmq_socket(
                ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False
            )
            self.paths[path] = sock

        data = {
            "req_id": req_id,
            "transfer_id": transfer_id,
            "block_notify_list": block_notify_list or [],
            # GLOBAL decode dp rank: producer derives the per-pod notify offset
            # (% dp_local), owning pod index (// dp_local), and write-target
            # from it. Sending the LOCAL rank made child-pod consumers look
            # like master rank 0 and hang in WAITING_FOR_REMOTE_KVS. Single-pod:
            # local == global.
            "decode_rank": self._global_dp_rank,
            "type": "remote_blocks",
        }
        serialized_data = msgpack.dumps(data)
        self.paths[path].send(serialized_data)

    def _send_transfer_release(self, transfer_id: TransferId, host: str, port: int):
        path = make_zmq_path("tcp", host, port)
        if path not in self.paths:
            ctx = zmq.Context.instance()
            sock = make_zmq_socket(
                ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False
            )
            self.paths[path] = sock

        # Advertise the consumer (decode) TP size so prefill counts the right
        # number of ACKs (get_moriio_expected_ack_count; see the upstream
        # READ-mode release in _pop_done_transfers). Resolves to 1 for our
        # homogeneous TP=1 configs; only matters under heterogeneous-TP fan-in.
        self.paths[path].send(
            msgpack.dumps(
                {
                    "type": "release",
                    "transfer_id": transfer_id,
                    "consumer_tp_size": self.tp_size,
                }
            )
        )

    def _release_write_prefill_blocks(self, request_id: ReqId, params: dict[str, Any]):
        transfer_id = params.get("transfer_id")
        if transfer_id is None:
            logger.warning(
                "Cannot release WRITE prefill blocks for request %s: "
                "missing transfer_id",
                request_id,
            )
            return

        remote_dp_rank = params.get("remote_dp_rank", 0)
        remote_host = params.get("remote_host")
        remote_notify_port = params.get("remote_notify_port")
        if remote_host is None or remote_notify_port is None:
            try:
                peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=False)
                if peer_zmq is None:
                    raise ValueError("no peer zmq address for request")
                remote_host, _, remote_notify_port = parse_moriio_zmq_address(peer_zmq)
            except ValueError:
                logger.warning(
                    "Cannot release WRITE prefill blocks for request %s: "
                    "missing remote notify address",
                    request_id,
                )
                return

        remote_notify_port = int(remote_notify_port)
        for tp_index in range(self.tp_size):
            target_port = remote_notify_port + get_port_offset(remote_dp_rank, tp_index)
            self._send_transfer_release(transfer_id, remote_host, target_port)

    def update_state_after_alloc(
        self,
        request: "Request",
        blocks: "KVCacheBlocks",
        num_external_tokens: int,
        connector_worker: "MoRIIOConnectorWorker | None" = None,
    ):
        """Scheduler-side post-allocation hook (decode leg in WRITE mode).

        In WRITE mode this fires the decode->prefill "blocks ready" notify
        that lets the producer RDMA-Write its KV into the freshly allocated
        decode blocks.

        DP-rank routing contract (router-authoritative)
        ------------------------------------------------
        Both legs of a disagg pair must agree on a single prefill DP rank,
        otherwise the notify lands on a rank that never handshook and the
        request hangs until ``VLLM_MORIIO_DEFERRED_TIMEOUT_S``.

        The routing authority is the external router/sidecar, NOT the
        connector. The llm-d routing sidecar computes ``H = pickDPRank(uuid,
        dp_size)`` (see ``dp_rank.go``) once and pins BOTH legs to ``H`` two
        ways: the ``X-data-parallel-rank`` dispatch header (which vLLM engine
        the leg lands on) and ``kv_transfer_params.remote_dp_rank=H`` (which
        prefill rank this notify targets), with ``remote_dp_rank_override=True``.

        The connector consumes ``remote_dp_rank`` verbatim and does NOT
        self-derive a rank: an independent hash here could disagree with the
        router's dispatch pin and misroute the notify. For the returnable paths
        (READ / serial WRITE) the prefill leg also echoes back the rank it ran
        on via ``request_finished`` so routing is pure propagation.

        The owning rank is the only one that originates the notify
        (``remote_dp_rank_override`` -> global-rank match; otherwise the
        ``_is_kv_master`` anti-duplicate gate). See the inline comments
        below for the exactly-once-across-pods reasoning.
        """
        params = request.kv_transfer_params
        if not params:
            return
        # LLM-D sidecar compat: the nixlv2 routing sidecar emits NIXL-shaped
        # kv_transfer_params without MoRI-IO's transfer_id. Synthesize one
        # deterministically from request_id so producer and consumer agree
        # without a sidecar wire-protocol change.
        transfer_id = params.get("transfer_id") or f"sidecar-{request.request_id}"
        params.setdefault("transfer_id", transfer_id)
        request_id = request.request_id
        self.map_request_id(request_id, transfer_id)
        if params.get("do_remote_decode"):
            local_block_ids = blocks.get_block_ids()[0]
            self._reqs_need_save[request.request_id] = (request, local_block_ids)
            # Snapshot params now so chunked-prefill build_connector_meta
            # can recover them on the final chunk even if the live
            # request.kv_transfer_params has been mutated/cleared.
            self._req_kv_params[request.request_id] = dict(params)

        if params is not None and params.get("do_remote_prefill"):
            if self.mode == MoRIIOMode.READ:
                if remote_block_ids := params.get("remote_block_ids"):
                    # remote_engine_id is returned by the prefill's request_finished.
                    # host/ports come from the request_id (parsed in add_new_req).
                    if "remote_engine_id" in params:
                        if num_external_tokens > 0:
                            # Get unhashed blocks to pull from remote.
                            local_block_ids = blocks.get_block_ids()[0]
                            assert len(local_block_ids) <= len(remote_block_ids)
                            if len(local_block_ids) != len(remote_block_ids):
                                local_block_ids = remote_block_ids[
                                    -len(local_block_ids) :
                                ]
                        else:
                            # If remote_blocks and num_external_tokens = 0, we have
                            # a full prefix cache hit on the D worker. We need to call
                            # send_notify in _read_blocks to free the memory on the P.
                            local_block_ids = []

                        self._reqs_need_recv[request.request_id] = (
                            request,
                            local_block_ids,
                        )
                        # Snapshot params for chunked prefill consumption
                        # in build_connector_meta (see comment above).
                        self._req_kv_params[request.request_id] = dict(params)
                    else:
                        logger.warning(
                            "Got invalid KVTransferParams: %s. This "
                            "request will not utilize KVTransfer",
                            params,
                        )

            else:
                # 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"
                )

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

                # Effective DP fan-out for the per-pod port/host math below
                # (see _remote_dp_rank_for_port / _pod_idx). Capped to the
                # per-pod local size when the router advertises it (Wide-EP).
                _dp_size = int(request.kv_transfer_params.get("remote_dp_size", 1) or 1)
                try:
                    _dp_local = int(
                        request.kv_transfer_params.get("remote_dp_size_local", 0) or 0
                    )
                    if _dp_local > 0:
                        _dp_size = min(_dp_size, _dp_local)
                except (TypeError, ValueError):
                    _dp_local = 0

                # Rank routing is ROUTER-AUTHORITATIVE: honor the router-pinned
                # remote_dp_rank (matched to the X-data-parallel-rank dispatch
                # pin). Self-deriving a rank could disagree and notify a rank
                # that never served the request. If unpinned in a multi-rank
                # deployment, warn instead of guessing.
                if (
                    _dp_size > 1
                    and "remote_dp_rank" not in request.kv_transfer_params
                    and "is_request_leader" not in request.kv_transfer_params
                ):
                    logger.warning(
                        "MoRI-IO decode notify: remote_dp_size=%d but the router "
                        "did not pin remote_dp_rank for request %s; defaulting to "
                        "rank 0. The router must set remote_dp_rank (and the "
                        "X-data-parallel-rank dispatch header) so the prefill and "
                        "decode legs agree on the same rank.",
                        _dp_size,
                        request.request_id,
                    )

                # Only the owning rank originates notify (exactly-once).
                # Priority: is_request_leader > remote_dp_rank_override > _is_kv_master
                _leader_flag = request.kv_transfer_params.get("is_request_leader")
                if _leader_flag is not None:
                    _should_notify = bool(_leader_flag)
                elif request.kv_transfer_params.get("remote_dp_rank_override"):
                    _should_notify = self._global_dp_rank == remote_dp_rank
                else:
                    _should_notify = self._is_kv_master
                if _should_notify:
                    peer_zmq = get_peer_zmq_from_request_id(
                        request.request_id, is_producer=False
                    )
                    if peer_zmq is not None:
                        remote_host, _, remote_notify_port = parse_moriio_zmq_address(
                            peer_zmq
                        )
                    else:
                        # Sidecar fallback: use explicit params fields.
                        params = request.kv_transfer_params or {}
                        remote_host = params.get("remote_host") or ""
                        try:
                            remote_notify_port = int(
                                params.get("remote_notify_port") or 0
                            )
                        except (TypeError, ValueError):
                            remote_notify_port = 0
                        if not remote_host or not remote_notify_port:
                            raise ValueError(
                                f"request {request.request_id!r}: "
                                f"request_id has no embedded peer "
                                f"zmq_address and kv_transfer_params is "
                                f"missing remote_host / remote_notify_port "
                                f"(got remote_host={remote_host!r}, "
                                f"remote_notify_port={remote_notify_port!r})"
                            )

                    # num_external_tokens == 0: nothing to push, so don't tell
                    # the producer to write into these blocks.
                    block_notify_list = (
                        blocks.get_block_ids()[0] if num_external_tokens > 0 else []
                    )

                    # Wide-EP multi-pod: a pod binds notify sockets only for
                    # its LOCAL ranks, so the port offset must use the per-pod
                    # local rank (% dp_local), not the global rank. Single-pod
                    # is bit-identical (modulus is a no-op).
                    _remote_dp_rank_for_port = fold_local_rank(
                        remote_dp_rank, _dp_local
                    )
                    # The target rank may live on a child pod at a different IP,
                    # so resolve the per-pod host (pod_idx = global // dp_local).
                    # Otherwise a notify for child ranks lands on the master
                    # pod and the request hangs in WAITING_FOR_REMOTE_KVS.
                    _notify_host = remote_host
                    _kvp = request.kv_transfer_params or {}
                    _remote_hosts = _kvp.get("remote_hosts") or []
                    if _dp_local > 0 and _remote_hosts:
                        _pod_idx = pod_index(remote_dp_rank, _dp_local)
                        if 0 <= _pod_idx < len(_remote_hosts):
                            _notify_host = _remote_hosts[_pod_idx]
                    for tp_index in range(self.tp_size):
                        target_port = remote_notify_port + get_port_offset(
                            _remote_dp_rank_for_port, tp_index
                        )

                        self.send_notify_block(
                            req_id=request.request_id,
                            transfer_id=request.kv_transfer_params["transfer_id"],
                            block_notify_list=block_notify_list,
                            host=_notify_host,
                            port=target_port,
                        )

            # Only trigger 1 KV transfer per request.

            params["do_remote_prefill"] = False

    def build_connector_meta(
        self,
        scheduler_output: SchedulerOutput,
    ) -> KVConnectorMetadata:
        meta = MoRIIOConnectorMetadata()
        meta.transfer_id_to_request_id = self.transfer_id_to_request_id

        if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.PRODUCER:
            # This is the logic for checking against chunked prefill.
            # When the last chunk is identified,
            # It places the request metadata into the saving queue.

            for i, req_id in enumerate(scheduler_output.scheduled_cached_reqs.req_ids):
                new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i]

                if new_block_ids is not None:
                    block_ids = new_block_ids[0]
                    # TODO : hybrid attn, etc
                    # A non-disagg request (no kv_transfer_params, e.g. smoke
                    # test) is never registered in _reqs_need_pending_save;
                    # indexing it unconditionally would KeyError and crash the
                    # EngineCore. Skip it silently.
                    if req_id not in self._reqs_need_pending_save:
                        continue
                    req, existing_blocks = self._reqs_need_pending_save[req_id]
                    updated_blocks = list(existing_blocks) + (block_ids)
                    self._reqs_need_pending_save[req_id] = (req, updated_blocks)
                    if (
                        len(self._reqs_need_pending_save[req_id][1]) * self.block_size
                        >= req.num_prompt_tokens
                    ):
                        # Final chunk: live kv_transfer_params may be cleared,
                        # so prefer the snapshot from update_state_after_alloc.
                        kv_params = self._req_kv_params.pop(
                            req_id, req.kv_transfer_params or {}
                        )
                        meta.add_new_req(
                            request_id=req_id,
                            local_block_ids=self._reqs_need_pending_save[req_id][1],
                            kv_transfer_params=kv_params,
                            write_mode=True,
                        )
                        del self._reqs_need_pending_save[req_id]

        # Loop through scheduled reqs and convert to ReqMeta.
        for req_id, (req, block_ids) in self._reqs_need_recv.items():
            kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})
            meta.add_new_req(
                request_id=req_id,
                local_block_ids=block_ids,
                kv_transfer_params=kv_params,
            )

        for req_id, (req, block_ids) in self._reqs_need_save.items():
            kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})
            if req.num_prompt_tokens > len(block_ids) * self.block_size:
                # not last chunk prefill
                self._reqs_need_pending_save[req_id] = (req, block_ids)
                continue
            meta.add_new_req(
                request_id=req_id,
                local_block_ids=block_ids,
                kv_transfer_params=kv_params,
                write_mode=True,
            )
        # Clear the list once workers start the transfers

        meta.reqs_to_send = self._reqs_need_send

        # Reclaim snapshot cache entries that completed this step. Keep
        # entries that are still pending (chunked prefill not yet at the
        # final chunk) — those will be popped above on the final chunk.
        for req_id in self._reqs_need_recv:
            self._req_kv_params.pop(req_id, None)
        for req_id in self._reqs_need_save:
            if req_id not in self._reqs_need_pending_save:
                self._req_kv_params.pop(req_id, None)

        self._reqs_need_recv.clear()
        self._reqs_need_save.clear()
        self._reqs_need_send = {}

        return meta

    def shutdown(self):
        for path, sock in self.paths.items():
            try:
                sock.close(linger=0)
                logger.debug("Closed ZMQ socket for path: %s", path)
            except Exception as e:
                logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
        self.paths.clear()

    def request_finished(
        self,
        request: "Request",
        block_ids: list[int],
    ) -> tuple[bool, dict[str, Any] | None]:
        """
        Once a request is finished, determine whether request blocks
        should be freed now or will be sent asynchronously and freed later.
        """

        request_id = request.request_id
        params = request.kv_transfer_params
        # 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:
            transfer_id = params.get("transfer_id") if params else None
            self.unmap_request_id(request_id, transfer_id=transfer_id)
        logger.debug(
            "MoriioConnector request_finished, request_status=%s, "
            "kv_transfer_params=%s",
            request.status,
            params,
        )
        if not params:
            return False, None

        if params.get("do_remote_prefill"):
            # If do_remote_prefill is still True when the request is finished,
            # update_state_after_alloc must not have been called (the request
            # must have been aborted before it was scheduled).
            # To avoid stranding the prefill blocks in the prefill instance,
            # READ mode adds empty block_ids to _reqs_need_recv so the worker
            # side notifies the prefill instance. WRITE mode should notify the
            # producer directly: there is no decode allocation for the producer
            # to write into, and a plain request_id may not contain router-
            # embedded MoRIIO ZMQ addresses.
            if self.mode == MoRIIOMode.WRITE:
                self._release_write_prefill_blocks(request.request_id, params)
            else:
                self._reqs_need_recv[request.request_id] = (request, [])
            params["do_remote_prefill"] = False
            return False, None

        if (
            not params.get("do_remote_decode")
            or request.status != RequestStatus.FINISHED_LENGTH_CAPPED
        ):
            return False, None

        # computed_block_ids = block_ids if all_full else block_ids[:-1]
        computed_block_ids = block_ids
        # If prompt < block_size, no xfer so free blocks immediately.
        delay_free_blocks = len(computed_block_ids) > 0

        if delay_free_blocks:
            # Prefill request on remote. It will be read from D upon completion
            self._reqs_need_send[request.request_id] = (
                time.perf_counter()
                + MoRIIOConstants.VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT
            )
            self._deferred_send_deadlines[request.request_id] = (
                time.monotonic() + self._defer_timeout,
                params.get("transfer_id") if params else None,
            )

        # Return KV transfer params forwarded verbatim to the decode instance by
        # the router. remote_dp_rank is the rank THIS prefill leg actually ran
        # on: on the returnable paths (READ / serial WRITE) the router forwards
        # it to the decode leg, so the decode->prefill notify is routed by pure
        # propagation of the producer's real rank -- no hashing, no reliance on
        # the router independently pinning the same rank. remote_dp_rank_override
        # makes the decode side honor it via the global-rank match gate.
        return delay_free_blocks, dict(
            do_remote_prefill=True,
            do_remote_decode=False,
            remote_block_ids=computed_block_ids,
            remote_engine_id=self.engine_id,
            remote_host=self.host_ip,
            remote_handshake_port=self.handshake_port,
            remote_notify_port=self.side_notify_port,
            remote_dp_rank=self._global_dp_rank,
            remote_dp_rank_override=True,
            remote_dp_size=self.vllm_config.parallel_config.data_parallel_size,
            remote_dp_size_local=(
                self.vllm_config.parallel_config.data_parallel_size_local
            ),
            tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
            transfer_id=params["transfer_id"],
        )

    def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
        """Reconcile worker finished_sending ACKs with the producer lifecycle.

        Called every scheduler step (when there is KV connector output),
        BEFORE the scheduler's finished_sending free loop. We rewrite
        ``connector_output.finished_sending`` in place so it contains
        exactly the producer requests that are safe for the scheduler to
        free this step, letting the scheduler keep the plain upstream
        ``assert req_id in self.requests`` + ``_free_blocks`` path:

        * An ACK whose request is in ``_deferred_send_deadlines``
          (request_finished ran with delay_free_blocks=True, so the
          scheduler is holding its blocks) is surfaced -> freed now.
        * An ACK that arrives BEFORE its request finished is parked in
          ``_pending_sent_acks`` and released on a later step once the
          request enters ``_deferred_send_deadlines``.
        * A deferred send whose ACK never arrives is reaped after
          ``_defer_timeout`` and surfaced, so leaked blocks are freed.
        * A parked ACK that never matches a deferral before its own
          deadline is a stale duplicate (e.g. a real ACK landing after the
          send was already reaped) and is dropped.

        Consumers never populate finished_sending (they report
        finished_recving), and they unmap in request_finished, so this is a
        no-op for them.
        """
        if not self.is_producer:
            return

        incoming = set(connector_output.finished_sending or ())
        now = time.monotonic()

        # Surface ACKs whose request is already finished (blocks held for
        # delayed free); park the rest until their request finishes.
        safe: set[ReqId] = set()
        for req_id in incoming:
            if req_id in self._deferred_send_deadlines:
                safe.add(req_id)
            else:
                self._pending_sent_acks.setdefault(req_id, now + self._defer_timeout)

        # Release previously parked ACKs whose request has since finished.
        for req_id in self._pending_sent_acks:
            if req_id in self._deferred_send_deadlines:
                safe.add(req_id)

        # Reap deferred sends whose ACK never arrived (avoid leaking blocks).
        expired = [
            req_id
            for req_id, (deadline, _) in self._deferred_send_deadlines.items()
            if now >= deadline
        ]
        if expired:
            safe.update(expired)
            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),
                self._defer_timeout,
            )

        # Finalize the requests we are surfacing: drop their deferral/park
        # bookkeeping and unmap their transfer ids.
        for req_id in safe:
            deferred_info = self._deferred_send_deadlines.pop(req_id, None)
            transfer_id = deferred_info[1] if deferred_info else None
            self._pending_sent_acks.pop(req_id, None)
            self.unmap_request_id(req_id, transfer_id=transfer_id)

        # Drop stale parked ACKs that never matched a deferral in time.
        stale = [
            req_id
            for req_id, deadline in self._pending_sent_acks.items()
            if now >= deadline
        ]
        for req_id in stale:
            self._pending_sent_acks.pop(req_id, None)

        connector_output.finished_sending = safe or None