[Core][KV Connector] Release transfer-held KV before a cache reset

docs/training/async_rl.md

@@ -48,6 +48,9 @@ With `VLLM_SERVER_DEV_MODE=1`, the vLLM HTTP server exposes the same functionali
 !!! note "Data Parallelism"
     When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update.
 
+!!! note "Prefill/decode disaggregation"
+    Pause and resume every prefill and decode instance. Once in-flight work is done, a pause that clears the cache stops the KV transfers between prefill and decode instances: decode requests whose KV was still arriving recompute it after resume, and prefill instances drop the KV they still hold for decode instances. A pause that keeps the cache leaves those transfers to finish after resume. Releasing that KV needs the multiprocess engine. Set `kv_load_failure_policy="recompute"` so a decode request whose prefill KV was dropped is recomputed instead of failed.
+
 ## Typical Async RL Flow
 
 A typical async RL loop with weight syncing looks like this:

tests/v1/core/test_scheduler.py

@@ -5456,6 +5456,90 @@ def test_abort_request_finished_recving():
     assert not scheduler.finished_recving_kv_req_ids
 
 
+@pytest.mark.parametrize("then", ["reset", "resume"])
+@pytest.mark.parametrize("load_outcome", ["written", "failed-blocks", "failed-request"])
+def test_release_transfer_kv_frees_blocks_for_reset(load_outcome, then):
+    """Releasing transfer KV ends loads in flight and drops held sends; the
+    reset recomputes all waiting KV, and a load in flight once it ends, even
+    under the default fail policy and if a resume cancels the reset first."""
+    scheduler = create_scheduler(use_kv_connector=True)
+    running, loading, arrived, loaded, session, queued = create_requests(
+        num_requests=6,
+        req_ids=["running", "loading", "arrived", "loaded", "session", "queued"],
+    )
+    scheduler.add_request(running)
+    scheduler.schedule()
+    for request in (loading, arrived, loaded, session, queued):
+        scheduler.add_request(request)
+    # A remote load in flight, one arrived, one already promoted, and a
+    # streaming session between turns: all hold KV while waiting.
+    held = (loading, arrived, loaded, session)
+    scheduler.waiting.remove_requests(held)
+    for request, status in zip(
+        held,
+        (
+            RequestStatus.WAITING_FOR_REMOTE_KVS,
+            RequestStatus.WAITING_FOR_REMOTE_KVS,
+            RequestStatus.WAITING,
+            RequestStatus.WAITING_FOR_STREAMING_REQ,
+        ),
+    ):
+        scheduler.kv_cache_manager.allocate_slots(request, request.num_prompt_tokens)
+        request.num_computed_tokens = request.num_prompt_tokens
+        request.status = status
+        scheduler._enqueue_waiting_request(request)
+    scheduler.finished_recving_kv_req_ids.add(arrived.request_id)
+    scheduler.connector.abort_transfers = Mock()
+    blocks = scheduler.kv_cache_manager.get_block_ids
+
+    scheduler.set_pause_state(PauseState.PAUSED_ALL)
+    scheduler.release_transfer_kv()
+    scheduler.connector.abort_transfers.assert_called_once()
+    # The paused engine keeps stepping until the load in flight ends.
+    assert scheduler.has_requests()
+    assert running.status == RequestStatus.RUNNING
+    assert list(scheduler.waiting) == [queued]
+
+    # The load in flight keeps its blocks until its transfer ends; the session
+    # keeps its KV; loaded KV is recomputed.
+    assert not scheduler.reset_prefix_cache(reset_running_requests=True)
+    assert blocks(loading.request_id)[0] and blocks(session.request_id)[0]
+    assert session.num_computed_tokens == session.num_prompt_tokens
+    # Recomputed requests keep their place: running first, then waiting KV.
+    assert list(scheduler.waiting) == [running, arrived, loaded, queued]
+    for request in (arrived, loaded):
+        assert request.status == RequestStatus.WAITING
+        assert request.num_computed_tokens == 0
+        assert not blocks(request.request_id)[0]
+
+    if then == "resume":
+        scheduler.set_pause_state(PauseState.UNPAUSED)
+    output = KVConnectorOutput(finished_recving={loading.request_id})
+    if load_outcome == "failed-blocks":
+        output.invalid_block_ids = set(blocks(loading.request_id)[0])
+    elif load_outcome == "failed-request":
+        output.failed_recving = {loading.request_id}
+    scheduler.update_from_output(
+        SchedulerOutput.make_empty(),
+        ModelRunnerOutput(req_ids=[], req_id_to_index={}, kv_connector_output=output),
+    )
+    # Not failed, not even under the default fail policy.
+    assert scheduler.requests[loading.request_id] is loading
+    if then == "resume":
+        # Recomputed locally from scratch, without another remote load.
+        output = scheduler.schedule()
+        assert loading.status == RequestStatus.RUNNING
+        assert output.num_scheduled_tokens[loading.request_id] == len(
+            loading.prompt_token_ids
+        )
+    else:
+        scheduler.reset_prefix_cache(reset_running_requests=True)
+        assert loading.status == RequestStatus.WAITING
+        assert loading.num_computed_tokens == 0 and not blocks(loading.request_id)[0]
+        assert loading in scheduler.waiting
+    assert loading.request_id not in scheduler.failed_recving_kv_req_ids
+
+
 def test_delayed_kv_connector_free_keeps_scheduler_active():
     scheduler = create_scheduler(use_kv_connector=True)
     queued_request, request = create_requests(
@@ -5471,6 +5555,12 @@ def test_delayed_kv_connector_free_keeps_scheduler_active():
 
     assert scheduler.has_finished_requests()
     assert scheduler.has_requests()
+    # A paused engine waits for the held KV only while releasing it.
+    scheduler.set_pause_state(PauseState.PAUSED_ALL)
+    assert not scheduler.has_requests()
+    scheduler.release_transfer_kv()
+    assert scheduler.has_requests()
+    scheduler.set_pause_state(PauseState.UNPAUSED)
 
     scheduler_output = SchedulerOutput(
         scheduled_new_reqs=[],

tests/v1/engine/test_engine_core.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import copy
+import queue
 import time
 import uuid
 from collections import deque
@@ -698,6 +699,18 @@ def _pausable_engine_core_proc() -> EngineCoreProc:
     core.batch_queue = None
     core.engines_running = False
     core._idle_state_callbacks = []
+    core.is_running = lambda: True
+    core.process_input_queue_block = True
+    core.aborts_queue = MagicMock()
+    core.input_queue = MagicMock()
+    core.input_queue.empty.return_value = True
+
+    def get(block):
+        # Waiting for client input is only right when there is no work.
+        assert not core.has_work(), "blocked on input with work pending"
+        raise queue.Empty
+
+    core.input_queue.get.side_effect = get
     return core
 
 
@@ -724,26 +737,78 @@ def test_kv_cache_release_rejects_unsafe_state(pause_state, has_requests, has_ba
     core.model_executor.discard.assert_not_called()
 
 
-@pytest.mark.parametrize("deferred", [False, True])
-def test_pause_synchronizes_device_before_cache_reset(deferred: bool):
-    """A resolved pause promises an idle device: the barrier must run before
-    caches are cleared and before the caller is unblocked."""
+@pytest.mark.parametrize("held", [False, True], ids=["no-transfers", "transfers"])
+@pytest.mark.parametrize("deferred", [False, True], ids=["idle", "draining"])
+@pytest.mark.parametrize(
+    ("op", "clears"),
+    [
+        ("keep", True),
+        ("wait", True),
+        ("keep", False),
+        ("abort", False),
+        ("release", True),
+    ],
+)
+def test_cache_reset_waits_for_transfer_kv(op, clears, deferred, held):
+    """Once idle, an operation clearing the KV cache makes KV transfers give back
+    their blocks and waits for them; the device barrier precedes the reset."""
     core = _pausable_engine_core_proc()
+    core.scheduler.pause_state = PauseState.PAUSED_ALL
     core.engines_running = deferred
-    order: list[str] = []
-    core.model_executor.collective_rpc.side_effect = lambda method: order.append(method)
+    order: list = []
+    core.model_executor.collective_rpc.side_effect = order.append
+    core.model_executor.discard.side_effect = order.append
     core._reset_caches = lambda: order.append("reset_caches")
 
-    result = EngineCoreProc.pause_scheduler(core, mode="keep", clear_cache=True)
-    if deferred:
-        assert isinstance(result, Future)
-        assert not result.done() and order == []
-        core.engines_running = False
-        core._notify_idle_state_callbacks()
-        assert result.result(timeout=0) is None
+    def release_transfer_kv():
+        order.append("release")
+        core.scheduler.has_requests.return_value = held
+
+    core.scheduler.release_transfer_kv.side_effect = release_transfer_kv
+
+    if op == "release":
+        if deferred:
+            pytest.skip("release needs a completed pause")
+        core.model_executor.is_sleeping = False
+        result = core.release_kv_cache_memory()
+    else:
+        result = EngineCoreProc.pause_scheduler(core, mode=op, clear_cache=clears)
+    assert isinstance(result, Future) == (deferred or (clears and held))
+    for busy in (deferred, clears and held):
+        if busy:
+            assert not result.done() and "reset_caches" not in order
+            core.engines_running = False
+            core.scheduler.has_requests.return_value = False
+            core._process_input_queue()
+    if isinstance(result, Future):
+        result = result.result(timeout=0)
+    assert result is None
+
+    released = ["release"] * clears
+    if op == "release":
+        assert order == released + ["reset_caches", ("kv_cache",)]
     else:
-        assert result is None
-    assert order == ["synchronize_device", "reset_caches"]
+        assert order == released + ["synchronize_device"] + ["reset_caches"] * clears
+
+
+def test_resume_cancels_a_pending_cache_reset():
+    """Resume fails a reset still waiting for the engine to drain, at once, so
+    a later pause can never run it."""
+    core = _pausable_engine_core_proc()
+    core._reset_caches = MagicMock()
+    core.engines_running = True
+    core.scheduler.pause_state = PauseState.PAUSED_ALL
+    result = EngineCoreProc.pause_scheduler(core, mode="keep", clear_cache=True)
+
+    core.scheduler.set_pause_state.side_effect = lambda state: setattr(
+        core.scheduler, "pause_state", state
+    )
+    core.resume_scheduler()
+    with pytest.raises(RuntimeError, match="Resumed before"):
+        result.result(timeout=0)
+    assert not core._idle_state_callbacks
+    core.scheduler.release_transfer_kv.assert_not_called()
+    core._reset_caches.assert_not_called()
 
 
 def test_sleep_mode_refuses_unsupported_kv_connector_at_startup():

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -1112,10 +1112,11 @@ def test_scheduler_request_finished():
     assert "id-1" in scheduler_connector._reqs_not_processed
 
 
+@pytest.mark.parametrize("dropped_by", ["request-finished", "abort-transfers"])
 @pytest.mark.parametrize("load_done", [False, True])
-def test_consumer_releases_transfer_finished_mid_load(load_done: bool):
-    """D asks P to drop the transfer of a request that finishes while its KV
-    is still arriving, and only then."""
+def test_consumer_releases_transfer_finished_mid_load(load_done: bool, dropped_by: str):
+    """D asks P to drop the transfer of a request whose KV is still arriving,
+    once the request finishes or its transfers are aborted, and only then."""
     vllm_config = create_vllm_config(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
@@ -1135,8 +1136,11 @@ def test_consumer_releases_transfer_finished_mid_load(load_done: bool):
             KVConnectorOutput(finished_recving={request.request_id})
         )
 
-    request.status = RequestStatus.FINISHED_ABORTED
-    assert connector.request_finished(request, [5, 6]) == (False, None)
+    if dropped_by == "abort-transfers":
+        connector.abort_transfers()
+    else:
+        request.status = RequestStatus.FINISHED_ABORTED
+        assert connector.request_finished(request, [5, 6]) == (False, None)
     pulls = connector.build_connector_meta(SchedulerOutput.make_empty()).reqs_to_recv
     release = pulls.get("p", {}).get(request.request_id)
     assert (release is not None) == (not load_done)
@@ -1741,6 +1745,8 @@ async def test_kv_consumuer(monkeypatch):
         (["release", "claim", "ready", "pull"], False, True),
         (["claim", "release", "expire", "ready"], None, True),
         (["claim", "ready", "failing-pull"], None, True),
+        (["claim", "ready", "abort-all"], None, True),
+        (["claim", "pull", "abort-all"], None, False),
         (["claim", "ready", "busy", "pull", "abort", "free"], False, True),
         (["claim", "ready", "busy", "abort", "pull"], False, True),
         (["claim", "abort", "neighbor", "busy", "pull"], False, False),
@@ -1769,6 +1775,8 @@ async def test_kv_consumuer(monkeypatch):
         "release-before-claim",
         "release-outlives-expiry-while-claimed",
         "failed-send-frees-kv",
+        "abort-all-frees-ready-kv",
+        "abort-all-spares-unfinished",
         "pull-waiting-for-a-sender-stays-abortable",
         "aborted-pull-answered-while-senders-are-busy",
         "aborted-pull-answered-before-its-batch-writes",
@@ -1819,12 +1827,13 @@ def slow_write(*_) -> int:
         now = [time.perf_counter()]
         monkeypatch.setattr(mooncake_connector.time, "perf_counter", lambda: now[0])
 
-        def scheduler_meta(blocks=None, aborted=False):
+        def scheduler_meta(blocks=None, aborted=False, abort_all=False):
             meta = MooncakeConnectorMetadata()
             if blocks is not None:
                 meta.reqs_to_send["p-req"] = ("tx", blocks)
             if aborted:
                 meta.reqs_not_processed = {"tx"}
+            meta.abort_pending_sends = abort_all
             return meta
 
         with patch.object(worker, "_send_blocks", return_value=0) as send_blocks:
@@ -1859,6 +1868,8 @@ def scheduler_meta(blocks=None, aborted=False):
                     await asyncio.wait([pull], timeout=3)
                 elif event == "free":
                     worker._send_slots.release()
+                elif event == "abort-all":
+                    await worker.record_send_reqs(scheduler_meta(abort_all=True))
                 elif event in ("pull", "failing-pull"):
                     if event == "failing-pull":
                         send_blocks.side_effect = RuntimeError("transfer engine died")

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -662,6 +662,11 @@ def has_pending_push_work(self) -> bool:
         # scheduler alive (e.g. extend has_unfinished_requests).
         return False
 
+    def abort_transfers(self) -> None:
+        """Stop holding finished requests' KV for remote readers and end remote
+        loads in flight; their blocks come back through get_finished as usual."""
+        return
+
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
         """Get the required KV cache layout for this connector.

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -694,6 +694,7 @@ def __init__(self):
         self.reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict)
         self.reqs_to_send: dict[ReqId, tuple[TransferId, list[list[int]]]] = {}
         self.reqs_not_processed: set[TransferId] = set()
+        self.abort_pending_sends = False
 
     def add_new_req(
         self,
@@ -799,6 +800,10 @@ def update_connector_output(self, connector_output: KVConnectorOutput):
         assert self.connector_scheduler is not None
         self.connector_scheduler.update_connector_output(connector_output)
 
+    def abort_transfers(self) -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.abort_transfers()
+
     def request_finished(
         self,
         request: "Request",
@@ -940,6 +945,8 @@ def __init__(
         self._reqs_not_processed: set[TransferId] = set()
         # D: requests whose KV is still being pulled.
         self._reqs_loading: dict[ReqId, Request] = {}
+        # P: drop the KV held for D at the next step.
+        self._abort_pending_sends = False
 
         # Compute sliding window block counts per KV cache group.
         sw_sizes_tokens: list[tuple[int, int]] = [
@@ -1131,13 +1138,21 @@ def build_connector_meta(
             self._reqs_need_send.clear()
             meta.reqs_not_processed = self._reqs_not_processed
             self._reqs_not_processed = set()
+        meta.abort_pending_sends = self._abort_pending_sends
+        self._abort_pending_sends = False
 
         return meta
 
     def update_connector_output(self, connector_output: KVConnectorOutput):
         for req_id in connector_output.finished_recving or ():
             self._reqs_loading.pop(req_id, None)
 
+    def abort_transfers(self) -> None:
+        self._abort_pending_sends = True
+        # Ask P to drop each load in flight; the load ends with P's answer.
+        for req_id, request in self._reqs_loading.items():
+            self._reqs_need_recv[req_id] = (request, [])
+
     def request_finished(
         self,
         request: "Request",
@@ -2718,6 +2733,11 @@ async def record_send_reqs(self, metadata: MooncakeConnectorMetadata):
             send_meta = self._get_send_meta(transfer_id)
             send_meta.expire_time = expire_time
             self._end_send(send_meta)
+        if metadata.abort_pending_sends:
+            # Only finished requests' KV; P still owns its unfinished ones.
+            for send_meta in self.reqs_need_send.values():
+                if send_meta.state is SendState.READY:
+                    self._end_send(send_meta)
 
     def _end_send(self, send_meta: SendBlockMeta) -> None:
         """Send nothing more: pulls fail fast until the state expires, and P's
@@ -2735,7 +2755,9 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata):
             )
 
         if not self.is_kv_consumer and (
-            metadata.reqs_to_send or metadata.reqs_not_processed
+            metadata.reqs_to_send
+            or metadata.reqs_not_processed
+            or metadata.abort_pending_sends
         ):
             asyncio.run_coroutine_threadsafe(
                 self.record_send_reqs(metadata), self.sender_loop