From 16a7115d085444054c6a7776f72926e598013804 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 12:00:23 +0200
Subject: [PATCH 01/26] Initial commit

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

From 4beb0f22d31904b9568b73fee1bdbc39e1a2d8dd Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 12:01:23 +0200
Subject: [PATCH 02/26] fix: remove mori notifications and fix hang issues and
 indefinite waits

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../v1/moriio/moriio_connector.py             | 71 ++++++++--------
 .../kv_connector/v1/moriio/moriio_engine.py   | 81 +++++++++++++++----
 vllm/envs.py                                  | 10 +++
 vllm/v1/core/sched/scheduler.py               | 28 ++++++-
 4 files changed, 140 insertions(+), 50 deletions(-)

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 dcde7665f344..29134c9f1cfc 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
@@ -1246,15 +1246,38 @@ 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():
+                last = status_list[-1]
+                if last.Succeeded():
                     done_req_ids.add(req_id)
-
                     self.moriio_wrapper.send_notify(
                         req_id,
                         self._recving_transfers_callback_addr[req_id][0],
                         self._recving_transfers_callback_addr[req_id][1],
                     )
                     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,
+                        last.Message(),
+                        last.Code(),
+                    )
+                    try:
+                        self.moriio_wrapper.send_notify(
+                            req_id,
+                            self._recving_transfers_callback_addr[req_id][0],
+                            self._recving_transfers_callback_addr[req_id][1],
+                        )
+                    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]
@@ -1295,21 +1318,13 @@ def save_kv_layer(
                         continue
             self._write_blocks_for_req(req_id, meta, layer_name, kv_layer)
 
-        while True:
-            if (
-                self._ready_requests.empty()
-                and remote_engine_id not in self.write_ready_flags
-            ):
-                continue
-            elif not self._ready_requests.empty() and (
-                remote_engine_id in self.write_ready_flags
-            ):
-                self._write_blocks_for_req(
-                    *self._ready_requests.get_nowait(), layer_name, kv_layer
-                )
-                break
-            else:
-                break
+        # Drain any requests whose background handshakes just completed.
+        # Requests whose handshakes are still in progress will be retried
+        # on the next scheduler step when save_kv_layer is called again.
+        while not self._ready_requests.empty():
+            req_id, meta = self._ready_requests.get_nowait()
+            if meta.remote_engine_id in self.write_ready_flags:
+                self._write_blocks_for_req(req_id, meta, layer_name, kv_layer)
 
     def get_engine_name_with_dp(self, engine_name, dp_rank):
         return f"{engine_name}_dp{dp_rank}"
@@ -1350,21 +1365,13 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
             self._read_blocks_for_req(req_id, meta)
         # Start transfers for requests whose handshakes have now finished.
 
-        while True:
-            if (
-                self._ready_requests.empty()
-                and remote_engine_id not in self.load_ready_flag
-                and wait_handshake_readd_req
-            ):
-                continue
-            elif (
-                not self._ready_requests.empty()
-                and remote_engine_id in self.load_ready_flag
-            ):
-                self._read_blocks_for_req(*self._ready_requests.get_nowait())
-                break
-            else:
-                break
+        # Drain any requests whose background handshakes just completed.
+        # Requests whose handshakes are still in progress will be retried
+        # on the next scheduler step when start_load_kv is called again.
+        while not self._ready_requests.empty():
+            req_id, meta = self._ready_requests.get_nowait()
+            if meta.remote_engine_id in self.load_ready_flag:
+                self._read_blocks_for_req(req_id, meta)
 
         self._reqs_to_send.update(metadata.reqs_to_send)
 
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 973c0bb801c8..653c396bdb8e 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,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import time
 import threading
 from typing import TYPE_CHECKING, Any
 from weakref import ref as weakref_ref
@@ -136,23 +137,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.request_id, task.transfer_id)
 
     def _process_deferred_tasks(self) -> None:
         """Process tasks that were previously deferred."""
         if not self._deferred_tasks:
             return
 
+        defer_timeout = envs.VLLM_MORIIO_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.request_id, 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.request_id, task.transfer_id)
             else:
                 still_deferred.append(task)
 
         self._deferred_tasks = still_deferred
 
+    def _mark_request_done(self, request_id: str, 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(request_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.
 
@@ -385,6 +418,13 @@ def set_backend_type(self, backend_type):
             post_batch_size,
             num_worker_threads,
             poll_mode,
+            enable_notification=False,  # 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.
         )
         self.moriio_engine.create_backend(backend_type, rdma_cfg)
 
@@ -469,22 +509,35 @@ 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 = envs.VLLM_MORIIO_TRANSFER_TIMEOUT
+        deadline = time.monotonic() + timeout
+        remaining = list(transfers_to_wait)
+
+        while remaining:
+            if time.monotonic() > deadline:
+                raise TransferError(
+                    f"RDMA transfer timed out after {timeout:.0f}s "
+                    f"({len(remaining)}/{len(transfers_to_wait)} transfers still "
+                    f"in progress). Check NIC queue depth or reduce concurrency; "
+                    f"adjust with VLLM_MORIIO_TRANSFER_TIMEOUT."
+                )
+            still_waiting = []
+            for status in remaining:
+                if status.Succeeded():
+                    continue
+                if status.Failed():
+                    raise TransferError(
+                        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
+                still_waiting.append(status)
+            remaining = still_waiting
+            if remaining:
+                time.sleep(0.001)
 
     def async_wait_reqid(self):
         assert self.notify_port is not None, "Notify port cannot be None"
diff --git a/vllm/envs.py b/vllm/envs.py
index 8ed1d33434cb..7602d249689b 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -203,6 +203,8 @@
     VLLM_MORIIO_QP_PER_TRANSFER: int = 1
     VLLM_MORIIO_POST_BATCH_SIZE: int = -1
     VLLM_MORIIO_NUM_WORKERS: int = 1
+    VLLM_MORIIO_TRANSFER_TIMEOUT: float = 30.0
+    VLLM_MORIIO_DEFER_TIMEOUT: float = 60.0
     VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480
     VLLM_ENABLE_CUDAGRAPH_GC: bool = False
     VLLM_LOOPBACK_IP: str = ""
@@ -1443,6 +1445,14 @@ def _get_or_set_default() -> str:
     ),
     # Controls the number of workers for Mori operations for the Mori-IO connector
     "VLLM_MORIIO_NUM_WORKERS": lambda: int(os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")),
+    # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError
+    "VLLM_MORIIO_TRANSFER_TIMEOUT": lambda: float(
+        os.getenv("VLLM_MORIIO_TRANSFER_TIMEOUT", "30.0")
+    ),
+    # Timeout (seconds) before a deferred WriteTask with no remote blocks is dropped
+    "VLLM_MORIIO_DEFER_TIMEOUT": lambda: float(
+        os.getenv("VLLM_MORIIO_DEFER_TIMEOUT", "60.0")
+    ),
     # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
     "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
         os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480")
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 40b5899f0457..0a4999bf18b6 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -2117,16 +2117,36 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
         # KV Connector:: update recv and send status from last step.
         for req_id in kv_connector_output.finished_recving or ():
             logger.debug("Finished recving KV transfer for request %s", req_id)
-            assert req_id in self.requests
+            if req_id not in self.requests:
+                # Request was removed (finished/aborted) before the async
+                # RDMA completion was reported. Nothing left to do.
+                logger.debug(
+                    "Request %s no longer tracked, skipping recv", req_id
+                )
+                continue
             req = self.requests[req_id]
             if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
                 self.finished_recving_kv_req_ids.add(req_id)
+            elif RequestStatus.is_finished(req.status):
+                self._free_blocks(req)
             else:
-                assert RequestStatus.is_finished(req.status)
-                self._free_blocks(self.requests[req_id])
+                # Connectors that poll for transfer completion asynchronously
+                # (rather than completing within the same scheduler step) can
+                # report finished_recving after the request has already
+                # advanced past WAITING_FOR_REMOTE_KVS. Blocks will be freed
+                # when the request finishes normally.
+                logger.debug(
+                    "Request %s recv finished in status=%s, skipping block free",
+                    req_id,
+                    req.status.name,
+                )
         for req_id in kv_connector_output.finished_sending or ():
             logger.debug("Finished sending KV transfer for request %s", req_id)
-            assert req_id in self.requests
+            if req_id not in self.requests:
+                logger.debug(
+                    "Request %s no longer tracked, skipping send", req_id
+                )
+                continue
             self._free_blocks(self.requests[req_id])
 
     def _update_requests_with_invalid_blocks(

From 300893e9b4f69736cdf0f334dee42025a1837a00 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 12:20:03 +0200
Subject: [PATCH 03/26] fix: collect all errors in
 waiting_for_transfer_complete before raising

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_engine.py   | 26 ++++++++++++-------
 1 file changed, 17 insertions(+), 9 deletions(-)

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 653c396bdb8e..5fdf8a75939c 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
@@ -516,26 +516,34 @@ def waiting_for_transfer_complete(self):
         timeout = envs.VLLM_MORIIO_TRANSFER_TIMEOUT
         deadline = time.monotonic() + timeout
         remaining = list(transfers_to_wait)
+        errors: list[str] = []
 
         while remaining:
-            if time.monotonic() > deadline:
-                raise TransferError(
-                    f"RDMA transfer timed out after {timeout:.0f}s "
-                    f"({len(remaining)}/{len(transfers_to_wait)} transfers still "
-                    f"in progress). Check NIC queue depth or reduce concurrency; "
-                    f"adjust with VLLM_MORIIO_TRANSFER_TIMEOUT."
-                )
+            timed_out = time.monotonic() > deadline
             still_waiting = []
             for status in remaining:
                 if status.Succeeded():
                     continue
                 if status.Failed():
-                    raise TransferError(
+                    errors.append(
                         f"RDMA transfer failed: {status.Message()} "
                         f"(code={status.Code()})"
                     )
-                still_waiting.append(status)
+                    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 VLLM_MORIIO_TRANSFER_TIMEOUT."
+                    )
+                else:
+                    still_waiting.append(status)
             remaining = still_waiting
+            if errors:
+                raise TransferError(
+                    f"{len(errors)}/{len(transfers_to_wait)} transfers failed:\n"
+                    + "\n".join(errors)
+                )
             if remaining:
                 time.sleep(0.001)
 

From ebde8ca287ccb10681ebe6d66f2f54035f9e0923 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 12:24:49 +0200
Subject: [PATCH 04/26] fix: send_notify using transfer id instead of request
 id

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_connector.py         | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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 29134c9f1cfc..e9ba08bcd4c4 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
@@ -1249,8 +1249,11 @@ def _pop_done_transfers(self) -> set[str]:
                 last = status_list[-1]
                 if last.Succeeded():
                     done_req_ids.add(req_id)
+                    transfer_id = self.request_id_to_transfer_id.get(
+                        req_id, req_id
+                    )
                     self.moriio_wrapper.send_notify(
-                        req_id,
+                        transfer_id,
                         self._recving_transfers_callback_addr[req_id][0],
                         self._recving_transfers_callback_addr[req_id][1],
                     )
@@ -1264,9 +1267,12 @@ def _pop_done_transfers(self) -> set[str]:
                         last.Message(),
                         last.Code(),
                     )
+                    transfer_id = self.request_id_to_transfer_id.get(
+                        req_id, req_id
+                    )
                     try:
                         self.moriio_wrapper.send_notify(
-                            req_id,
+                            transfer_id,
                             self._recving_transfers_callback_addr[req_id][0],
                             self._recving_transfers_callback_addr[req_id][1],
                         )

From 38eb6179799ca3d872dd9696d2dd54070021bb0f Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 12:30:38 +0200
Subject: [PATCH 05/26] chore: pre-commit

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_connector.py  | 10 ++--------
 .../kv_connector/v1/moriio/moriio_engine.py     | 17 +++++++++--------
 vllm/v1/core/sched/scheduler.py                 |  8 ++------
 3 files changed, 13 insertions(+), 22 deletions(-)

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 e9ba08bcd4c4..ecbb549f9e56 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
@@ -1249,9 +1249,7 @@ def _pop_done_transfers(self) -> set[str]:
                 last = status_list[-1]
                 if last.Succeeded():
                     done_req_ids.add(req_id)
-                    transfer_id = self.request_id_to_transfer_id.get(
-                        req_id, req_id
-                    )
+                    transfer_id = self.request_id_to_transfer_id.get(req_id, req_id)
                     self.moriio_wrapper.send_notify(
                         transfer_id,
                         self._recving_transfers_callback_addr[req_id][0],
@@ -1267,9 +1265,7 @@ def _pop_done_transfers(self) -> set[str]:
                         last.Message(),
                         last.Code(),
                     )
-                    transfer_id = self.request_id_to_transfer_id.get(
-                        req_id, req_id
-                    )
+                    transfer_id = self.request_id_to_transfer_id.get(req_id, req_id)
                     try:
                         self.moriio_wrapper.send_notify(
                             transfer_id,
@@ -1363,8 +1359,6 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
                         self._background_moriio_handshake(
                             req_id, remote_engine_id, meta
                         )
-                        wait_handshake_readd_req = True
-
                         continue
 
             # Handshake already completed, start async read xfer.
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 5fdf8a75939c..ff051f25c40b 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,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import time
 import threading
+import time
 from typing import TYPE_CHECKING, Any
 from weakref import ref as weakref_ref
 
@@ -418,13 +418,14 @@ def set_backend_type(self, backend_type):
             post_batch_size,
             num_worker_threads,
             poll_mode,
-            enable_notification=False,  # 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.
+            # 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)
 
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 0a4999bf18b6..23368b33b28e 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -2120,9 +2120,7 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
             if req_id not in self.requests:
                 # Request was removed (finished/aborted) before the async
                 # RDMA completion was reported. Nothing left to do.
-                logger.debug(
-                    "Request %s no longer tracked, skipping recv", req_id
-                )
+                logger.debug("Request %s no longer tracked, skipping recv", req_id)
                 continue
             req = self.requests[req_id]
             if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
@@ -2143,9 +2141,7 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
         for req_id in kv_connector_output.finished_sending or ():
             logger.debug("Finished sending KV transfer for request %s", req_id)
             if req_id not in self.requests:
-                logger.debug(
-                    "Request %s no longer tracked, skipping send", req_id
-                )
+                logger.debug("Request %s no longer tracked, skipping send", req_id)
                 continue
             self._free_blocks(self.requests[req_id])
 

From 69d4bebd649f9487937228288f2475d9f7dcefbe Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 12:37:06 +0200
Subject: [PATCH 06/26] fix: add transfer_id to
 _recving_transfers_callback_addr

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../v1/moriio/moriio_connector.py             | 22 ++++++++-----------
 1 file changed, 9 insertions(+), 13 deletions(-)

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 ecbb549f9e56..b2e174d83ef1 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
@@ -735,7 +735,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] = {}
@@ -1249,12 +1250,8 @@ def _pop_done_transfers(self) -> set[str]:
                 last = status_list[-1]
                 if last.Succeeded():
                     done_req_ids.add(req_id)
-                    transfer_id = self.request_id_to_transfer_id.get(req_id, req_id)
-                    self.moriio_wrapper.send_notify(
-                        transfer_id,
-                        self._recving_transfers_callback_addr[req_id][0],
-                        self._recving_transfers_callback_addr[req_id][1],
-                    )
+                    host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
+                    self.moriio_wrapper.send_notify(xfer_id, host, port)
                     to_remove.append(req_id)
                 elif last.Failed():
                     logger.error(
@@ -1265,13 +1262,9 @@ def _pop_done_transfers(self) -> set[str]:
                         last.Message(),
                         last.Code(),
                     )
-                    transfer_id = self.request_id_to_transfer_id.get(req_id, req_id)
+                    host, port, xfer_id = self._recving_transfers_callback_addr[req_id]
                     try:
-                        self.moriio_wrapper.send_notify(
-                            transfer_id,
-                            self._recving_transfers_callback_addr[req_id][0],
-                            self._recving_transfers_callback_addr[req_id][1],
-                        )
+                        self.moriio_wrapper.send_notify(xfer_id, host, port)
                     except Exception:
                         logger.exception(
                             "Failed to send error notification for request %s",
@@ -1383,6 +1376,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,
@@ -1540,6 +1534,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:
@@ -1567,4 +1562,5 @@ def _read_blocks(
                 self._recving_transfers_callback_addr[request_id] = (
                     remote_host,
                     str(remote_notify_port + self.tp_rank),
+                    transfer_id,
                 )

From 6d9687d17759c2fa1b0ae44ca9e1daf74c8d5469 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 12:37:22 +0200
Subject: [PATCH 07/26] chore: remove unused variable

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/moriio/moriio_connector.py       | 1 -
 1 file changed, 1 deletion(-)

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 b2e174d83ef1..9ab30ea84a47 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
@@ -1336,7 +1336,6 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
         if self.mode == MoRIIOMode.WRITE:
             return
 
-        wait_handshake_readd_req = False
         remote_engine_id = None
 
         for req_id, meta in metadata.reqs_to_recv.items():

From 1be872383dffb96b351365ae8114d791c0c16471 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 14:12:05 +0200
Subject: [PATCH 08/26] fix: revert while-true loop to retain correctness

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../v1/moriio/moriio_connector.py             | 47 +++++++++++++------
 1 file changed, 33 insertions(+), 14 deletions(-)

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 9ab30ea84a47..1e6f095a4045 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
@@ -1313,13 +1313,21 @@ def save_kv_layer(
                         continue
             self._write_blocks_for_req(req_id, meta, layer_name, kv_layer)
 
-        # Drain any requests whose background handshakes just completed.
-        # Requests whose handshakes are still in progress will be retried
-        # on the next scheduler step when save_kv_layer is called again.
-        while not self._ready_requests.empty():
-            req_id, meta = self._ready_requests.get_nowait()
-            if meta.remote_engine_id in self.write_ready_flags:
-                self._write_blocks_for_req(req_id, meta, layer_name, kv_layer)
+        while True:
+            if (
+                self._ready_requests.empty()
+                and remote_engine_id not in self.write_ready_flags
+            ):
+                continue
+            elif not self._ready_requests.empty() and (
+                remote_engine_id in self.write_ready_flags
+            ):
+                self._write_blocks_for_req(
+                    *self._ready_requests.get_nowait(), layer_name, kv_layer
+                )
+                break
+            else:
+                break
 
     def get_engine_name_with_dp(self, engine_name, dp_rank):
         return f"{engine_name}_dp{dp_rank}"
@@ -1336,6 +1344,7 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
         if self.mode == MoRIIOMode.WRITE:
             return
 
+        wait_handshake_readd_req = False
         remote_engine_id = None
 
         for req_id, meta in metadata.reqs_to_recv.items():
@@ -1351,19 +1360,29 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
                         self._background_moriio_handshake(
                             req_id, remote_engine_id, meta
                         )
+                        wait_handshake_readd_req = True
+
                         continue
 
             # Handshake already completed, start async read xfer.
             self._read_blocks_for_req(req_id, meta)
         # Start transfers for requests whose handshakes have now finished.
 
-        # Drain any requests whose background handshakes just completed.
-        # Requests whose handshakes are still in progress will be retried
-        # on the next scheduler step when start_load_kv is called again.
-        while not self._ready_requests.empty():
-            req_id, meta = self._ready_requests.get_nowait()
-            if meta.remote_engine_id in self.load_ready_flag:
-                self._read_blocks_for_req(req_id, meta)
+        while True:
+            if (
+                self._ready_requests.empty()
+                and remote_engine_id not in self.load_ready_flag
+                and wait_handshake_readd_req
+            ):
+                continue
+            elif (
+                not self._ready_requests.empty()
+                and remote_engine_id in self.load_ready_flag
+            ):
+                self._read_blocks_for_req(*self._ready_requests.get_nowait())
+                break
+            else:
+                break
 
         self._reqs_to_send.update(metadata.reqs_to_send)
 

From 562d8e99e0e94fc433ed79bf6ea13630779529eb Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 14:24:35 +0200
Subject: [PATCH 09/26] fix: raise on failed transfers after all requests have
 been waited on

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_engine.py            | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

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 ff051f25c40b..5de03901c619 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
@@ -540,13 +540,13 @@ def waiting_for_transfer_complete(self):
                 else:
                     still_waiting.append(status)
             remaining = still_waiting
-            if errors:
-                raise TransferError(
-                    f"{len(errors)}/{len(transfers_to_wait)} transfers failed:\n"
-                    + "\n".join(errors)
-                )
             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"

From adadc15de23b7dacccf7976569225cd0c179859a Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 14:27:25 +0200
Subject: [PATCH 10/26] fix: properly map from transfer_id to req_id

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_connector.py         | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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 1e6f095a4045..760bc910eaa2 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
@@ -1225,8 +1225,14 @@ 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.
+            raw_done_sending = self.moriio_wrapper.pop_finished_req_ids()
+            done_sending = {
+                self.transfer_id_to_request_id[xfer_id]
+                for xfer_id in raw_done_sending
+                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()

From bd6ef0aeae6afc44503de377eae8518d81b5da13 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 19:25:22 +0200
Subject: [PATCH 11/26] fix: add 1ms sleeps to yield main thread control to
 engine core

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../v1/moriio/moriio_connector.py             | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

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 760bc910eaa2..90b6c436a0b9 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
@@ -15,6 +15,7 @@
 import torch
 import zmq
 
+from vllm import envs
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorBase_V1,
@@ -1319,11 +1320,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() + envs.VLLM_MORIIO_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 VLLM_MORIIO_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
@@ -1374,12 +1386,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() + envs.VLLM_MORIIO_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 VLLM_MORIIO_TRANSFER_TIMEOUT",
+                        remote_engine_id,
+                    )
+                    break
+                time.sleep(0.001)
                 continue
             elif (
                 not self._ready_requests.empty()

From 75a7a893366a37f733f0c19eb722ed050dd04787 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 20 Apr 2026 20:45:23 +0200
Subject: [PATCH 12/26] fix: add reaping of leaked kv blocks

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 vllm/v1/core/sched/scheduler.py | 37 +++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 23368b33b28e..ffa6f17a1c23 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -181,6 +181,9 @@ def __init__(
         # KV Connector: requests in process of async KV loading or recving
         self.finished_recving_kv_req_ids: set[str] = set()
         self.failed_recving_kv_req_ids: set[str] = set()
+        # Monotonic deadline after which _reap_leaked_kv_send_blocks runs.
+        # Reset to now+idle_threshold whenever active requests are present.
+        self._kv_send_block_reap_deadline: float = 0.0
 
         # Encoder-related.
         # Calculate encoder cache size if applicable
@@ -2145,6 +2148,40 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
                 continue
             self._free_blocks(self.requests[req_id])
 
+        if self.connector is not None:
+            self._reap_leaked_kv_send_blocks()
+
+    def _reap_leaked_kv_send_blocks(self) -> None:
+        """Free KV cache blocks for finished requests on the prefill side.
+
+        Under high concurrency, ibv_post_send failures can cause
+        finished_sending notifications to be lost, leaving KV blocks
+        permanently allocated. This method is called after a configurable
+        idle period (VLLM_MORIIO_DEFER_TIMEOUT seconds with no running
+        requests) to sweep up any such leaks.
+        """
+        now = time.monotonic()
+        if self.running:
+            self._kv_send_block_reap_deadline = now + envs.VLLM_MORIIO_DEFER_TIMEOUT
+            return
+        if now < self._kv_send_block_reap_deadline:
+            return
+
+        leaked = [req for req in self.requests.values() if req.is_finished()]
+        if not leaked:
+            return
+
+        for req in leaked:
+            self._free_blocks(req)
+        logger.warning(
+            "Reaped %d finished requests with leaked KV send blocks "
+            "after %.0fs idle. This indicates lost finished_sending "
+            "notifications (typically from RDMA errors at high concurrency).",
+            len(leaked),
+            envs.VLLM_MORIIO_DEFER_TIMEOUT,
+        )
+        self._kv_send_block_reap_deadline = now + envs.VLLM_MORIIO_DEFER_TIMEOUT
+
     def _update_requests_with_invalid_blocks(
         self,
         requests: Iterable[Request],

From b5a6fee69763f9866c45489025a989ce0c1cb3f8 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Tue, 21 Apr 2026 13:49:51 +0200
Subject: [PATCH 13/26] refactor: improve reaper impl

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/base.py       | 18 ++++++++++
 .../v1/moriio/moriio_connector.py             |  6 ++++
 vllm/v1/core/sched/scheduler.py               | 34 ++++++++++++-------
 3 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
index ef143cba7fb5..f3f7947a1666 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
@@ -610,6 +610,24 @@ def get_finished_count(self) -> int | None:
 
         return None
 
+    def get_block_leak_reap_timeout(self) -> float | None:
+        """
+        Return the idle period (seconds) after which the scheduler should
+        reap KV cache blocks held by finished requests that never signalled
+        completion via get_finished().
+
+        Connectors that free blocks synchronously (request_finished returns
+        False) should return None to disable reaping entirely.  Connectors
+        that defer freeing asynchronously (request_finished returns True)
+        should return a positive timeout long enough to let normal completion
+        signals arrive before the reaper runs.
+
+        Returns:
+            float: idle seconds before reaping leaked blocks.
+            None: reaping is disabled (default).
+        """
+        return None
+
     @classmethod
     def build_kv_connector_stats(
         cls, data: dict[str, Any] | None = None
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 90b6c436a0b9..3611172b0c65 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
@@ -170,6 +170,12 @@ def request_finished(
         assert self.connector_scheduler is not None
         return self.connector_scheduler.request_finished(request, block_ids)
 
+    def get_block_leak_reap_timeout(self) -> float | None:
+        # ibv_post_send failures under high concurrency can silently drop the
+        # finished_sending notification, leaving KV blocks permanently allocated
+        # on the prefill side. The reaper frees them after this idle period.
+        return envs.VLLM_MORIIO_DEFER_TIMEOUT
+
     ############################################################
     # Worker Side Methods
     ############################################################
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index ffa6f17a1c23..f8f91f0df017 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -183,6 +183,12 @@ def __init__(
         self.failed_recving_kv_req_ids: set[str] = set()
         # Monotonic deadline after which _reap_leaked_kv_send_blocks runs.
         # Reset to now+idle_threshold whenever active requests are present.
+        # None means the connector has disabled reaping (synchronous freeing).
+        self._kv_block_reap_timeout: float | None = (
+            self.connector.get_block_leak_reap_timeout()
+            if self.connector is not None
+            else None
+        )
         self._kv_send_block_reap_deadline: float = 0.0
 
         # Encoder-related.
@@ -2148,21 +2154,23 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
                 continue
             self._free_blocks(self.requests[req_id])
 
-        if self.connector is not None:
+        if self._kv_block_reap_timeout is not None:
             self._reap_leaked_kv_send_blocks()
 
     def _reap_leaked_kv_send_blocks(self) -> None:
-        """Free KV cache blocks for finished requests on the prefill side.
-
-        Under high concurrency, ibv_post_send failures can cause
-        finished_sending notifications to be lost, leaving KV blocks
-        permanently allocated. This method is called after a configurable
-        idle period (VLLM_MORIIO_DEFER_TIMEOUT seconds with no running
-        requests) to sweep up any such leaks.
+        """Free KV cache blocks for finished requests that never signalled
+        completion via get_finished().
+
+        Connectors that defer block freeing asynchronously (request_finished
+        returns True) may fail to deliver completion signals under certain
+        error conditions, leaving blocks permanently allocated.  This method
+        runs after the connector-specified idle period with no running requests
+        to sweep up any such leaks.
         """
+        assert self._kv_block_reap_timeout is not None
         now = time.monotonic()
         if self.running:
-            self._kv_send_block_reap_deadline = now + envs.VLLM_MORIIO_DEFER_TIMEOUT
+            self._kv_send_block_reap_deadline = now + self._kv_block_reap_timeout
             return
         if now < self._kv_send_block_reap_deadline:
             return
@@ -2175,12 +2183,12 @@ def _reap_leaked_kv_send_blocks(self) -> None:
             self._free_blocks(req)
         logger.warning(
             "Reaped %d finished requests with leaked KV send blocks "
-            "after %.0fs idle. This indicates lost finished_sending "
-            "notifications (typically from RDMA errors at high concurrency).",
+            "after %.0fs idle. This indicates lost async KV completion "
+            "notifications from the KV connector.",
             len(leaked),
-            envs.VLLM_MORIIO_DEFER_TIMEOUT,
+            self._kv_block_reap_timeout,
         )
-        self._kv_send_block_reap_deadline = now + envs.VLLM_MORIIO_DEFER_TIMEOUT
+        self._kv_send_block_reap_deadline = now + self._kv_block_reap_timeout
 
     def _update_requests_with_invalid_blocks(
         self,

From f50078cafdf4667956c3bfead2ab18c5d190ff25 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Tue, 21 Apr 2026 15:49:42 +0200
Subject: [PATCH 14/26] fix: add transfer_id to done_req_ids (misnomer)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_engine.py            | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

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 5de03901c619..eb4a9a70505c 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
@@ -144,7 +144,7 @@ def _write_worker_loop(self) -> None:
                     "Write task failed for request %s, marking done",
                     task.request_id,
                 )
-                self._mark_request_done(task.request_id, task.transfer_id)
+                self._mark_request_done(task.transfer_id)
 
     def _process_deferred_tasks(self) -> None:
         """Process tasks that were previously deferred."""
@@ -163,7 +163,7 @@ def _process_deferred_tasks(self) -> None:
                     task.request_id,
                     now - task.enqueue_time,
                 )
-                self._mark_request_done(task.request_id, task.transfer_id)
+                self._mark_request_done(task.transfer_id)
                 continue
             if self._is_remote_ready(task):
                 try:
@@ -173,17 +173,17 @@ def _process_deferred_tasks(self) -> None:
                         "Deferred write task failed for request %s, marking done",
                         task.request_id,
                     )
-                    self._mark_request_done(task.request_id, task.transfer_id)
+                    self._mark_request_done(task.transfer_id)
             else:
                 still_deferred.append(task)
 
         self._deferred_tasks = still_deferred
 
-    def _mark_request_done(self, request_id: str, transfer_id: str) -> None:
+    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(request_id)
+            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:

From 21ac6f9e534b570aa8bcf4b191772810b69e7f3f Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Tue, 21 Apr 2026 15:59:21 +0200
Subject: [PATCH 15/26] fix: add xfer id to done_req_ids (misnomer)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/moriio/moriio_connector.py      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 3611172b0c65..1b95eec4fd21 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
@@ -1262,8 +1262,8 @@ def _pop_done_transfers(self) -> set[str]:
             for req_id, status_list in self._recving_transfers.items():
                 last = status_list[-1]
                 if last.Succeeded():
-                    done_req_ids.add(req_id)
                     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():

From f56157b13a581019e05d29b6258aac772176672f Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Mon, 27 Apr 2026 11:39:45 +0200
Subject: [PATCH 16/26] fix: set is_producer=False ccorrectly for decode path

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/moriio/moriio_connector.py  | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

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 4f44c26ebd16..1a37fc4ce2a9 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
@@ -415,9 +415,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"
                 )
@@ -425,7 +423,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)
 

From 17a852d8fa4e12cb5477ec870e4c9ed81afecb08 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Wed, 29 Apr 2026 14:19:43 +0000
Subject: [PATCH 17/26] fix: fix WRITE mode race condition

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_connector.py   | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

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 1a37fc4ce2a9..092c1d10adb3 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
@@ -666,6 +666,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 = (
@@ -1260,7 +1263,11 @@ def get_finished(self) -> tuple[set[str], set[str]]:
             }
         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()
 
@@ -1270,6 +1277,13 @@ 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
 

From f5a7c4c30ab13c920c5acda1fdfc7234aa1df3cf Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Wed, 29 Apr 2026 16:30:21 +0200
Subject: [PATCH 18/26] chore: pre-commit

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/moriio/moriio_connector.py     | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

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 092c1d10adb3..c326b4bdb6ae 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
@@ -1280,7 +1280,8 @@ def get_finished(self) -> tuple[set[str], set[str]]:
         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
+                id
+                for id in self._unmatched_write_completions
                 if id in self.transfer_id_to_request_id
             }
             self._unmatched_write_completions -= matched_xfer_ids

From a436497134b5edb2d655c96763141994fab99b70 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Thu, 7 May 2026 11:23:11 +0200
Subject: [PATCH 19/26] fix: cache envvars in moriio engine and connector

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_connector.py    |  8 +++++---
 .../kv_connector/v1/moriio/moriio_engine.py       | 15 ++++++++++-----
 2 files changed, 15 insertions(+), 8 deletions(-)

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 c326b4bdb6ae..aa337c0baf9f 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
@@ -110,6 +110,7 @@ def __init__(
             + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
         )
         self.mode = get_moriio_mode()
+        self._defer_timeout = envs.VLLM_MORIIO_DEFER_TIMEOUT
         if role == KVConnectorRole.SCHEDULER:
             self.connector_scheduler: MoRIIOConnectorScheduler | None = (
                 MoRIIOConnectorScheduler(vllm_config, self.engine_id)
@@ -176,7 +177,7 @@ def get_block_leak_reap_timeout(self) -> float | None:
         # ibv_post_send failures under high concurrency can silently drop the
         # finished_sending notification, leaving KV blocks permanently allocated
         # on the prefill side. The reaper frees them after this idle period.
-        return envs.VLLM_MORIIO_DEFER_TIMEOUT
+        return self._defer_timeout
 
     ############################################################
     # Worker Side Methods
@@ -665,6 +666,7 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         self.moriio_engine = None
         self._handle_request_thread = None
         self._ping_thread = None
+        self._transfer_timeout = envs.VLLM_MORIIO_TRANSFER_TIMEOUT
         self._writer = MoRIIOWriter(self)
         # Completions that arrived before transfer_id_to_request_id was populated.
         # Retried each step until the mapping is established.
@@ -1361,7 +1363,7 @@ def save_kv_layer(
 
         if remote_engine_id is None:
             return
-        _deadline = time.monotonic() + envs.VLLM_MORIIO_TRANSFER_TIMEOUT
+        _deadline = time.monotonic() + self._transfer_timeout
         while True:
             if (
                 self._ready_requests.empty()
@@ -1427,7 +1429,7 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
 
         if remote_engine_id is None and not wait_handshake_readd_req:
             return
-        _deadline = time.monotonic() + envs.VLLM_MORIIO_TRANSFER_TIMEOUT
+        _deadline = time.monotonic() + self._transfer_timeout
         while True:
             if (
                 self._ready_requests.empty()
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 eb4a9a70505c..282a276c8eac 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
@@ -76,6 +76,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 = envs.VLLM_MORIIO_DEFER_TIMEOUT
 
     @property
     def worker(self) -> "MoRIIOConnectorWorker":
@@ -151,7 +152,7 @@ def _process_deferred_tasks(self) -> None:
         if not self._deferred_tasks:
             return
 
-        defer_timeout = envs.VLLM_MORIIO_DEFER_TIMEOUT
+        defer_timeout = self._defer_timeout
         now = time.perf_counter()
         still_deferred: list[WriteTask] = []
 
@@ -397,6 +398,10 @@ 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 = envs.VLLM_MORIIO_TRANSFER_TIMEOUT
+        self._qp_per_transfer = envs.VLLM_MORIIO_QP_PER_TRANSFER
+        self._post_batch_size = envs.VLLM_MORIIO_POST_BATCH_SIZE
+        self._num_worker_threads = envs.VLLM_MORIIO_NUM_WORKERS
         self.notify_thread: threading.Thread | None = None
         self.sessions: list[IOEngine.Session] = []
         self.paths: dict[str, zmq.Socket] = {}
@@ -409,9 +414,9 @@ def set_moriio_engine(self, moriio_engine):
 
     def set_backend_type(self, backend_type):
         assert self.moriio_engine is not None, "MoRIIO engine must be set first"
-        qp_per_transfer = envs.VLLM_MORIIO_QP_PER_TRANSFER
-        post_batch_size = envs.VLLM_MORIIO_POST_BATCH_SIZE
-        num_worker_threads = envs.VLLM_MORIIO_NUM_WORKERS
+        qp_per_transfer = self._qp_per_transfer
+        post_batch_size = self._post_batch_size
+        num_worker_threads = self._num_worker_threads
         poll_mode = PollCqMode.POLLING
         rdma_cfg = RdmaBackendConfig(
             qp_per_transfer,
@@ -514,7 +519,7 @@ def waiting_for_transfer_complete(self):
             transfers_to_wait = self.transfer_status[:]
             self.transfer_status.clear()
 
-        timeout = envs.VLLM_MORIIO_TRANSFER_TIMEOUT
+        timeout = self._transfer_timeout
         deadline = time.monotonic() + timeout
         remaining = list(transfers_to_wait)
         errors: list[str] = []

From 28eb79fd09ff6911c886b421e1f824a0c12a1418 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Thu, 7 May 2026 17:43:18 +0200
Subject: [PATCH 20/26] fix: append to done_req_ids transfer_id rather than
 request_id

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/moriio/moriio_engine.py         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 282a276c8eac..0e89d9738e48 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
@@ -356,7 +356,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
             ]

From 455e5cc9fab571ab2d9daafaaa06b48c66cca2f2 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Thu, 7 May 2026 17:44:13 +0200
Subject: [PATCH 21/26] chore: rename raw_done_sending to finished_transfer_ids

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/moriio/moriio_connector.py    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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 aa337c0baf9f..ce0ad526574e 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
@@ -1257,10 +1257,10 @@ def get_finished(self) -> tuple[set[str], set[str]]:
         if self.is_producer:
             # pop_finished_req_ids returns transfer_ids (the ZMQ payload sent
             # by decode via send_notify); map back to req_ids for the scheduler.
-            raw_done_sending = self.moriio_wrapper.pop_finished_req_ids()
+            finished_transfer_ids = self.moriio_wrapper.pop_finished_req_ids()
             done_sending = {
                 self.transfer_id_to_request_id[xfer_id]
-                for xfer_id in raw_done_sending
+                for xfer_id in finished_transfer_ids
                 if xfer_id in self.transfer_id_to_request_id
             }
         else:

From 77e8fffa6474d7f84d4354a8ac7b72617b24c11d Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Fri, 8 May 2026 10:18:42 +0200
Subject: [PATCH 22/26] fix: move kv reaper logic to connector through
 update_connector_output

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_transfer/kv_connector/v1/base.py       | 18 ------
 .../v1/moriio/moriio_connector.py             | 61 +++++++++++++++++--
 vllm/v1/core/sched/scheduler.py               | 45 --------------
 3 files changed, 55 insertions(+), 69 deletions(-)

diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
index f3f7947a1666..ef143cba7fb5 100644
--- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py
+++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py
@@ -610,24 +610,6 @@ def get_finished_count(self) -> int | None:
 
         return None
 
-    def get_block_leak_reap_timeout(self) -> float | None:
-        """
-        Return the idle period (seconds) after which the scheduler should
-        reap KV cache blocks held by finished requests that never signalled
-        completion via get_finished().
-
-        Connectors that free blocks synchronously (request_finished returns
-        False) should return None to disable reaping entirely.  Connectors
-        that defer freeing asynchronously (request_finished returns True)
-        should return a positive timeout long enough to let normal completion
-        signals arrive before the reaper runs.
-
-        Returns:
-            float: idle seconds before reaping leaked blocks.
-            None: reaping is disabled (default).
-        """
-        return None
-
     @classmethod
     def build_kv_connector_stats(
         cls, data: dict[str, Any] | None = None
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 ce0ad526574e..6661f6773501 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
@@ -61,6 +61,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:
@@ -110,7 +111,6 @@ def __init__(
             + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
         )
         self.mode = get_moriio_mode()
-        self._defer_timeout = envs.VLLM_MORIIO_DEFER_TIMEOUT
         if role == KVConnectorRole.SCHEDULER:
             self.connector_scheduler: MoRIIOConnectorScheduler | None = (
                 MoRIIOConnectorScheduler(vllm_config, self.engine_id)
@@ -173,11 +173,9 @@ def request_finished(
         assert self.connector_scheduler is not None
         return self.connector_scheduler.request_finished(request, block_ids)
 
-    def get_block_leak_reap_timeout(self) -> float | None:
-        # ibv_post_send failures under high concurrency can silently drop the
-        # finished_sending notification, leaving KV blocks permanently allocated
-        # on the prefill side. The reaper frees them after this idle period.
-        return self._defer_timeout
+    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
@@ -287,6 +285,13 @@ 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 = envs.VLLM_MORIIO_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] = {}
@@ -598,6 +603,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.
@@ -610,6 +618,47 @@ 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.
+        """
+        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]
+        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"""
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 9a8ff48e909e..4ad9f988d247 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -177,15 +177,6 @@ def __init__(
         # KV Connector: requests in process of async KV loading or recving
         self.finished_recving_kv_req_ids: set[str] = set()
         self.failed_recving_kv_req_ids: set[str] = set()
-        # Monotonic deadline after which _reap_leaked_kv_send_blocks runs.
-        # Reset to now+idle_threshold whenever active requests are present.
-        # None means the connector has disabled reaping (synchronous freeing).
-        self._kv_block_reap_timeout: float | None = (
-            self.connector.get_block_leak_reap_timeout()
-            if self.connector is not None
-            else None
-        )
-        self._kv_send_block_reap_deadline: float = 0.0
 
         # Encoder-related.
         # Calculate encoder cache size if applicable
@@ -2082,42 +2073,6 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
                 continue
             self._free_blocks(self.requests[req_id])
 
-        if self._kv_block_reap_timeout is not None:
-            self._reap_leaked_kv_send_blocks()
-
-    def _reap_leaked_kv_send_blocks(self) -> None:
-        """Free KV cache blocks for finished requests that never signalled
-        completion via get_finished().
-
-        Connectors that defer block freeing asynchronously (request_finished
-        returns True) may fail to deliver completion signals under certain
-        error conditions, leaving blocks permanently allocated.  This method
-        runs after the connector-specified idle period with no running requests
-        to sweep up any such leaks.
-        """
-        assert self._kv_block_reap_timeout is not None
-        now = time.monotonic()
-        if self.running:
-            self._kv_send_block_reap_deadline = now + self._kv_block_reap_timeout
-            return
-        if now < self._kv_send_block_reap_deadline:
-            return
-
-        leaked = [req for req in self.requests.values() if req.is_finished()]
-        if not leaked:
-            return
-
-        for req in leaked:
-            self._free_blocks(req)
-        logger.warning(
-            "Reaped %d finished requests with leaked KV send blocks "
-            "after %.0fs idle. This indicates lost async KV completion "
-            "notifications from the KV connector.",
-            len(leaked),
-            self._kv_block_reap_timeout,
-        )
-        self._kv_send_block_reap_deadline = now + self._kv_block_reap_timeout
-
     def _update_requests_with_invalid_blocks(
         self,
         requests: Iterable[Request],

From 973f2cb8c1b6c8c4e82369e9c82eeff3e77b0b9a Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Wed, 20 May 2026 18:41:54 +0200
Subject: [PATCH 23/26] fix: keep producer transfer_id mapping alive until
 finished_sending fire

Before the transfer_id<->request_id mapping was freed on the prefiller
as soon as prefill finished. When a "done transferring the KV"
notification subsequently came in from the decoder, the prefiller
couldn't tell the scheduler to now free the blocks because it couldn't
find the corresponding request id for that completed transfer in its
map. The effect was that every request's blocks were freed by the reaper
rather than the connector telling the scheduler when it was time for
freeing. Instead, we now defer the unmapping until the transfer has
actually completed. The result is that we no longer get lots of "reaper
logs".

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_connector.py  | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

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 279ca166e1d2..09fde5afa8ab 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
@@ -563,7 +563,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(
@@ -628,6 +633,14 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
         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
 
@@ -651,6 +664,8 @@ def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
         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 "

From 66228fd0cf500566639871da3f6ed0981fb3bab7 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Wed, 20 May 2026 19:09:37 +0200
Subject: [PATCH 24/26] fix: convert new envvars to config arguments

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_common.py   | 18 +++++++++++++++
 .../v1/moriio/moriio_connector.py             | 22 ++++++++++++-------
 .../kv_connector/v1/moriio/moriio_engine.py   |  7 +++---
 vllm/envs.py                                  | 10 ---------
 4 files changed, 36 insertions(+), 21 deletions(-)

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 b843c5b59308..e8f8973b98f8 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
@@ -189,6 +189,8 @@ class MoRIIOConfig:
     dp_rank: int
     dp_size: int
     tp_size: int
+    transfer_timeout: float
+    defer_timeout: float
 
     @classmethod
     def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
@@ -227,6 +229,14 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
             dp_rank=dp_rank,
             dp_size=dp_size,
             tp_size=tp_size,
+            transfer_timeout=float(
+                extra_config.get(
+                    "transfer_timeout", MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT
+                )
+            ),
+            defer_timeout=float(
+                extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT)
+            ),
         )
 
 
@@ -247,6 +257,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 09fde5afa8ab..2c72ffb919dc 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
@@ -15,7 +15,6 @@
 import torch
 import zmq
 
-from vllm import envs
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorBase_V1,
@@ -291,7 +290,11 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         # 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 = envs.VLLM_MORIIO_DEFER_TIMEOUT
+        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] = {}
@@ -730,7 +733,6 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         self.moriio_engine = None
         self._handle_request_thread = None
         self._ping_thread = None
-        self._transfer_timeout = envs.VLLM_MORIIO_TRANSFER_TIMEOUT
         self._writer = MoRIIOWriter(self)
         # Completions that arrived before transfer_id_to_request_id was populated.
         # Retried each step until the mapping is established.
@@ -767,7 +769,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)
         self.moriio_wrapper.set_backend_type(BackendType.RDMA)
         self.moriio_wrapper.notify_port = self.moriio_config.notify_port
@@ -1427,7 +1433,7 @@ def save_kv_layer(
 
         if remote_engine_id is None:
             return
-        _deadline = time.monotonic() + self._transfer_timeout
+        _deadline = time.monotonic() + self.moriio_config.transfer_timeout
         while True:
             if (
                 self._ready_requests.empty()
@@ -1436,7 +1442,7 @@ def save_kv_layer(
                 if time.monotonic() > _deadline:
                     logger.warning(
                         "Timed out waiting for write_ready_flags[%s]; "
-                        "adjust with VLLM_MORIIO_TRANSFER_TIMEOUT",
+                        "adjust with kv_connector_extra_config.transfer_timeout",
                         remote_engine_id,
                     )
                     break
@@ -1493,7 +1499,7 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
 
         if remote_engine_id is None and not wait_handshake_readd_req:
             return
-        _deadline = time.monotonic() + self._transfer_timeout
+        _deadline = time.monotonic() + self.moriio_config.transfer_timeout
         while True:
             if (
                 self._ready_requests.empty()
@@ -1503,7 +1509,7 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata):
                 if time.monotonic() > _deadline:
                     logger.warning(
                         "Timed out waiting for load_ready_flag[%s]; "
-                        "adjust with VLLM_MORIIO_TRANSFER_TIMEOUT",
+                        "adjust with kv_connector_extra_config.transfer_timeout",
                         remote_engine_id,
                     )
                     break
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 0e89d9738e48..f5dd1e5a2bfb 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
@@ -76,7 +76,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 = envs.VLLM_MORIIO_DEFER_TIMEOUT
+        self._defer_timeout = worker.moriio_config.defer_timeout
 
     @property
     def worker(self) -> "MoRIIOConnectorWorker":
@@ -384,6 +384,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
@@ -398,7 +399,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 = envs.VLLM_MORIIO_TRANSFER_TIMEOUT
+        self._transfer_timeout = transfer_timeout
         self._qp_per_transfer = envs.VLLM_MORIIO_QP_PER_TRANSFER
         self._post_batch_size = envs.VLLM_MORIIO_POST_BATCH_SIZE
         self._num_worker_threads = envs.VLLM_MORIIO_NUM_WORKERS
@@ -540,7 +541,7 @@ def waiting_for_transfer_complete(self):
                     errors.append(
                         f"RDMA transfer timed out after {timeout:.0f}s. "
                         f"Check NIC queue depth or reduce concurrency; "
-                        f"adjust with VLLM_MORIIO_TRANSFER_TIMEOUT."
+                        f"adjust with kv_connector_extra_config.transfer_timeout."
                     )
                 else:
                     still_waiting.append(status)
diff --git a/vllm/envs.py b/vllm/envs.py
index ce7450952205..8d0240898074 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -213,8 +213,6 @@
     VLLM_MORIIO_QP_PER_TRANSFER: int = 1
     VLLM_MORIIO_POST_BATCH_SIZE: int = -1
     VLLM_MORIIO_NUM_WORKERS: int = 1
-    VLLM_MORIIO_TRANSFER_TIMEOUT: float = 30.0
-    VLLM_MORIIO_DEFER_TIMEOUT: float = 60.0
     VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480
     VLLM_ENABLE_CUDAGRAPH_GC: bool = False
     VLLM_LOOPBACK_IP: str = ""
@@ -1519,14 +1517,6 @@ def _get_or_set_default() -> str:
     ),
     # Controls the number of workers for Mori operations for the Mori-IO connector
     "VLLM_MORIIO_NUM_WORKERS": lambda: int(os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")),
-    # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError
-    "VLLM_MORIIO_TRANSFER_TIMEOUT": lambda: float(
-        os.getenv("VLLM_MORIIO_TRANSFER_TIMEOUT", "30.0")
-    ),
-    # Timeout (seconds) before a deferred WriteTask with no remote blocks is dropped
-    "VLLM_MORIIO_DEFER_TIMEOUT": lambda: float(
-        os.getenv("VLLM_MORIIO_DEFER_TIMEOUT", "60.0")
-    ),
     # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
     "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
         os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480")

From e0c29714e583d93ed14f5b9621ec4665c22483b8 Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Wed, 20 May 2026 19:16:44 +0000
Subject: [PATCH 25/26] revert: last scheduler changes

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 vllm/v1/core/sched/scheduler.py | 24 ++++--------------------
 1 file changed, 4 insertions(+), 20 deletions(-)

diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 958353357df2..45a5169a3efe 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -2141,32 +2141,16 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput):
         # KV Connector:: update recv and send status from last step.
         for req_id in kv_connector_output.finished_recving or ():
             logger.debug("Finished recving KV transfer for request %s", req_id)
-            if req_id not in self.requests:
-                # Request was removed (finished/aborted) before the async
-                # RDMA completion was reported. Nothing left to do.
-                logger.debug("Request %s no longer tracked, skipping recv", req_id)
-                continue
+            assert req_id in self.requests
             req = self.requests[req_id]
             if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
                 self.finished_recving_kv_req_ids.add(req_id)
-            elif RequestStatus.is_finished(req.status):
-                self._free_blocks(req)
             else:
-                # Connectors that poll for transfer completion asynchronously
-                # (rather than completing within the same scheduler step) can
-                # report finished_recving after the request has already
-                # advanced past WAITING_FOR_REMOTE_KVS. Blocks will be freed
-                # when the request finishes normally.
-                logger.debug(
-                    "Request %s recv finished in status=%s, skipping block free",
-                    req_id,
-                    req.status.name,
-                )
+                assert RequestStatus.is_finished(req.status)
+                self._free_blocks(self.requests[req_id])
         for req_id in kv_connector_output.finished_sending or ():
             logger.debug("Finished sending KV transfer for request %s", req_id)
-            if req_id not in self.requests:
-                logger.debug("Request %s no longer tracked, skipping send", req_id)
-                continue
+            assert req_id in self.requests
             self._free_blocks(self.requests[req_id])
 
     def _update_requests_with_invalid_blocks(

From 09059d178060a69e744078aba58d183d0a1a21de Mon Sep 17 00:00:00 2001
From: simondanielsson <simon.danielsson99@hotmail.com>
Date: Wed, 20 May 2026 19:31:01 +0000
Subject: [PATCH 26/26] fix: don't report done_recving for READ mode transfers
 in MoRI

In READ mode, the scheduler puts requests directly in a RUNNING state
rather than into a WAITING_FOR_REMOTE_KVS. However, the connector's
get_finished still reported these requests as done_recving after the
D requests had transferred (read) the KVs synchronously from P instance.
This triggered the scheduler to see these requests in
_update_from_kv_xfer_finished, but since they were already in RUNNING
state they triggered an assertion error as that code path expects only
asynchronously freed request blocks. This change avoids reporting these
sync requests as needing an async scheduling.

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
---
 .../kv_connector/v1/moriio/moriio_connector.py           | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

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 2c72ffb919dc..ba5b17fb4b56 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
@@ -1341,7 +1341,14 @@ def get_finished(self) -> tuple[set[str], set[str]]:
                 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]