[Bugfix][Mooncake] Offload producer partial tails on request finish (#52832)

Signed-off-by: Dao Le <Dao007forever@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
Co-authored-by: OpenAI Codex <noreply@openai.com>

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -12,6 +12,7 @@ import pytest
 import torch
 
 from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request
+from vllm.distributed.kv_transfer.kv_connector.v1.base import SupportsHMA
 from vllm.utils.hashing import sha256
 from vllm.v1.core.kv_cache_utils import (
     KVCacheBlockCopy,
@@ -735,11 +736,121 @@ def test_boundary_state_offloads_returns_cow_target():
     assert block_id != source_block_id
     # Draining clears it.
     assert manager.take_boundary_state_offloads() == {}
+    # CoW consumed the producer marker, so finishing cannot emit the same
+    # boundary a second time.
+    assert manager.finalize_partial_tail_offloads(req0) == []
 
     manager.block_pool.free_blocks(retained)
     manager.free(req0)
 
 
+def test_finished_partial_tail_uses_table_source_once():
+    """A finish without another allocation hands off the unchanged table block."""
+    hash_block_size = 2
+    block_size = 2 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=24,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=hash_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+
+    req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
+    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+    source_block_id = manager.get_blocks("0").get_block_ids()[1][1]
+
+    req0.num_computed_tokens = 6
+    offloads = manager.finalize_partial_tail_offloads(req0)
+
+    assert offloads == [(1, source_block_id, 6)]
+    assert manager.finalize_partial_tail_offloads(req0) == []
+    assert manager.take_boundary_state_offloads() == {}
+
+    # A table state that advanced beyond the boundary is unsafe and is
+    # consumed without producing a handoff.
+    advanced_manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+    req1 = make_request("1", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
+    computed_blocks, num_computed, _ = advanced_manager.get_computed_blocks(req1)
+    assert (
+        advanced_manager.allocate_slots(req1, 6, num_computed, computed_blocks)
+        is not None
+    )
+    req1.num_computed_tokens = 7
+    assert advanced_manager.finalize_partial_tail_offloads(req1) == []
+    assert advanced_manager.finalize_partial_tail_offloads(req1) == []
+
+
+def test_connector_finish_registers_partial_tail_before_cleanup():
+    class _Connector(SupportsHMA):
+        def register_finished_partial_tail(self, request, block_ids, offloads):
+            self.registered = (request, block_ids, offloads)
+            return True
+
+        def request_finished_all_groups(self, request, block_ids):
+            return False, None
+
+    scheduler = object.__new__(Scheduler)
+    scheduler.connector = _Connector()
+    scheduler.vllm_config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(is_kv_producer=True)
+    )
+    scheduler.kv_cache_manager = MagicMock()
+    scheduler.kv_cache_manager.finalize_partial_tail_offloads.return_value = [
+        (1, 9, 12)
+    ]
+    scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.return_value = (
+        [3],
+        [9],
+    )
+    request = SimpleNamespace(
+        request_id="0",
+        num_computed_tokens=12,
+        num_in_flight_tokens=0,
+        num_prompt_tokens=12,
+    )
+
+    delay_free, kv_xfer_params = scheduler._connector_finished(request)
+
+    assert delay_free is True
+    assert kv_xfer_params is None
+    assert scheduler.connector.registered == (
+        request,
+        ([3], [9]),
+        [(1, 9, 12)],
+    )
+    scheduler.kv_cache_manager.remove_skipped_blocks.assert_called_once()
+
+
 def test_block_pool_touch_pins_released_cow_target():
     """The connector can rescue an offered CoW target after its step-scoped
     retention is released by using the bound BlockPool's touch method."""

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

@@ -42,6 +42,7 @@ def _make_bare_scheduler(
     scheduler._unfinished_request_ids = {"req-0"}
     scheduler._unfinished_requests = {}
     scheduler._request_trackers = {}
+    scheduler._finished_partial_tail_metas = {}
     scheduler._gpu_block_pool = BlockPool(
         num_gpu_blocks=64, enable_caching=True, hash_block_size=hash_block_size
     )
@@ -971,6 +972,65 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():
     assert tracker.has_pending_offload is True
 
 
+def test_finished_partial_tail_is_pre_pinned_as_store_job():
+    scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    scheduler.client = SimpleNamespace(discard=lambda *_: None)
+    request = SimpleNamespace(
+        request_id="req-0",
+        block_hashes=[b"h0", b"h1", b"h2"],
+    )
+    scheduler._request_trackers["req-0"] = RequestTracker(
+        req_id="req-0",
+        token_len=12,
+        allocated_block_ids=([3], [9]),
+        token_ids=list(range(12)),
+        prefill_end_tokens=12,
+    )
+    block_ids = ([3], [9])
+
+    delay_free = scheduler.register_finished_partial_tail(
+        request,
+        block_ids,
+        [(1, 9, 12)],
+    )
+
+    # The exact source is pinned immediately, so the request can be freed
+    # before the next connector metadata build.
+    assert delay_free is False
+    assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1
+    assert scheduler._gpu_block_pool.blocks[3].ref_cnt == 0
+
+    out = SimpleNamespace(
+        finished_req_ids={"req-0"},
+        preempted_req_ids=set(),
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=SimpleNamespace(
+            req_ids=[],
+            new_block_ids=[],
+            num_computed_tokens=[],
+            resumed_req_ids=set(),
+        ),
+        num_scheduled_tokens={},
+        scheduled_spec_decode_tokens={},
+        kv_connector_block_state=_make_connector_block_state(),
+    )
+    meta = scheduler.build_connector_meta(out)
+
+    assert len(meta.requests) == 1
+    req_meta = meta.requests[0]
+    assert req_meta.req_id == "req-0"
+    assert req_meta.token_len_chunk == 0
+    assert req_meta.block_ids == block_ids
+    assert req_meta.block_hashes == request.block_hashes
+    assert req_meta.boundary_state_offloads == [(1, 9, 12)]
+    assert scheduler._pinned_saves[req_meta.store_job_id][0] == [9]
+    assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 1
+    assert scheduler._finished_partial_tail_metas == {}
+
+    scheduler.update_connector_output(_make_worker_output({req_meta.store_job_id: 1}))
+    assert scheduler._gpu_block_pool.blocks[9].ref_cnt == 0
+
+
 def test_decode_boundary_state_offload_dropped_unclaimed():
     # A hand-off past the prefill end can never complete a joint hybrid hit
     # (every other group stops saving there), so it is neither transferred nor

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -148,6 +148,26 @@ KVConnectorFactory.register_connector(
 )
 
 
+def test_register_finished_partial_tail_notifies_every_connector():
+    connector = object.__new__(MultiConnector)
+    first = MagicMock(spec_set=KVConnectorBase_V1)
+    second = MagicMock(spec_set=KVConnectorBase_V1)
+    first.register_finished_partial_tail.return_value = True
+    second.register_finished_partial_tail.return_value = False
+    connector._connectors = [first, second]
+    request = MagicMock()
+    block_ids = ([1], [2])
+    offloads = [(1, 2, 12)]
+
+    assert connector.register_finished_partial_tail(request, block_ids, offloads)
+    first.register_finished_partial_tail.assert_called_once_with(
+        request, block_ids, offloads
+    )
+    second.register_finished_partial_tail.assert_called_once_with(
+        request, block_ids, offloads
+    )
+
+
 @pytest.fixture
 def mc() -> MultiConnector:
     """MultiConnector using two mocked connectors"""

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

@@ -561,6 +561,19 @@ class KVConnectorBase_V1(ABC):
         """
         return False, None
 
+    def register_finished_partial_tail(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+        partial_tail_offloads: list[tuple[int, int, int]],
+    ) -> bool:
+        """Register finish-time partial-tail sources before block cleanup.
+
+        Returns True when the connector accepts responsibility for the sources
+        and the request's blocks must remain alive until ``get_finished()``.
+        """
+        return False
+
     def take_events(self) -> Iterable["KVCacheEvent"]:
         """
         Take the KV cache events from the connector.

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

@@ -265,6 +265,17 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
         # so a finishing request never has to defer freeing them.
         return False, None
 
+    def register_finished_partial_tail(
+        self,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+        partial_tail_offloads: list[tuple[int, int, int]],
+    ) -> bool:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.register_finished_partial_tail(
+            request, block_ids, partial_tail_offloads
+        )
+
     def reset_cache(self) -> bool | None:
         """Reset the external Mooncake store on prefix-cache reset.
 

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

@@ -102,6 +102,7 @@ class MooncakeStoreScheduler:
         self._request_trackers: dict[str, RequestTracker] = {}  # scheduled new requests
         self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}
         self._unfinished_request_ids: set[str] = set()
+        self._finished_partial_tail_metas: dict[str, ReqMeta] = {}
 
     def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:
         self._gpu_block_pool = gpu_block_pool
@@ -402,6 +403,14 @@ class MooncakeStoreScheduler:
             )
 
         self._apply_current_save_block_ids(meta, scheduler_output)
+
+        # Finish-time handoffs arrive after the producing step's metadata was
+        # built. Their exact blocks were pinned when they were registered, so
+        # they remain valid after request cleanup and need no current snapshot.
+        for req_meta in self._finished_partial_tail_metas.values():
+            meta.add_request(req_meta)
+        self._finished_partial_tail_metas.clear()
+
         self._reference_save_blocks(meta)
         return meta
 
@@ -440,6 +449,9 @@ class MooncakeStoreScheduler:
             assert pool is not None, (
                 "GPU block pool must be bound before any store job is emitted"
             )
+            if req_meta.store_job_id is not None:
+                assert req_meta.store_job_id in self._pinned_saves
+                continue
             req_meta.store_job_id = store_job_id = self._next_store_job_id
             self._next_store_job_id += 1
             block_ids: list[int] = []
@@ -470,6 +482,60 @@ class MooncakeStoreScheduler:
             self._pinned_saves[store_job_id] = (block_ids, self._num_workers)
             pool.touch([pool.blocks[block_id] for block_id in block_ids])
 
+    def register_finished_partial_tail(
+        self,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+        partial_tail_offloads: list[tuple[int, int, int]],
+    ) -> bool:
+        """Queue and pin a finish-time tail for the next connector step."""
+        if self.kv_role == "kv_consumer" or not partial_tail_offloads:
+            return False
+        tracker = self._request_trackers.get(request.request_id)
+        if tracker is None or not any(block_ids):
+            return False
+        boundaries = {boundary for _, _, boundary in partial_tail_offloads}
+        if len(boundaries) != 1:
+            raise ValueError(
+                "Partial-tail offloads for one request must share a boundary"
+            )
+        boundary_tokens = next(iter(boundaries))
+        if boundary_tokens > tracker.prefill_end_tokens:
+            return False
+
+        pinned_block_ids: list[int] = []
+        for group_id, block_id, _ in partial_tail_offloads:
+            if group_id not in self._boundary_state_group_ids:
+                return False
+            if block_id == NULL_BLOCK_ID:
+                return False
+            pinned_block_ids.append(block_id)
+        pinned_block_ids = list(dict.fromkeys(pinned_block_ids))
+
+        pool = self._gpu_block_pool
+        assert pool is not None, (
+            "GPU block pool must be bound before a finish-time handoff"
+        )
+        assert request.request_id not in self._finished_partial_tail_metas
+        store_job_id = self._next_store_job_id
+        self._next_store_job_id += 1
+        self._pinned_saves[store_job_id] = (pinned_block_ids, self._num_workers)
+        pool.touch([pool.blocks[block_id] for block_id in pinned_block_ids])
+
+        self._finished_partial_tail_metas[request.request_id] = ReqMeta(
+            req_id=request.request_id,
+            token_len_chunk=0,
+            block_ids=tuple(group.copy() for group in block_ids),
+            block_hashes=list(request.block_hashes),
+            can_save=True,
+            num_prompt_tokens=tracker.prefill_end_tokens,
+            store_job_id=store_job_id,
+            boundary_state_offloads=partial_tail_offloads,
+        )
+        tracker.has_pending_offload = True
+        # The store job owns exact block refs, so request cleanup need not wait.
+        return False
+
     def _handle_boundary_state_offloads(
         self,
         offloads: dict[str, list[tuple[int, int, int]]],

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

@@ -523,6 +523,18 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             lambda c: c.request_finished(request, blocks),
         )
 
+    def register_finished_partial_tail(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+        partial_tail_offloads: list[tuple[int, int, int]],
+    ) -> bool:
+        accepted = [
+            c.register_finished_partial_tail(request, block_ids, partial_tail_offloads)
+            for c in self._connectors
+        ]
+        return any(accepted)
+
     def request_finished_all_groups(
         self,
         request: "Request",

vllm/v1/core/kv_cache_manager.py

@@ -874,6 +874,29 @@ class KVCacheManager:
                 )
         return offloads
 
+    def finalize_partial_tail_offloads(
+        self, request: Request
+    ) -> list[tuple[int, int, int]]:
+        """Consume safe producer partial tails when a request finishes.
+
+        A mamba align table block is a valid boundary source only if no later
+        token was forwarded. The connector pins and queues the exact table
+        block before request cleanup, then releases the pin when every worker
+        reports the store job complete.
+        """
+        offloads: list[tuple[int, int, int]] = []
+        for mgr in self.coordinator.single_type_managers:
+            finalized = mgr.finalize_partial_tail_offload(
+                request.request_id,
+                request.num_computed_tokens,
+                request.num_in_flight_tokens,
+            )
+            if finalized is None:
+                continue
+            group_id, block, boundary_tokens = finalized
+            offloads.append((group_id, block.block_id, boundary_tokens))
+        return offloads
+
     def new_step_starts(self) -> None:
         """Notify the coordinator that a new step is starting."""
         self.coordinator.new_step_starts()

vllm/v1/core/sched/scheduler.py

@@ -2775,6 +2775,13 @@ class Scheduler(SchedulerInterface):
         if self.connector is None:
             return False, None
 
+        finished_partial_tails: list[tuple[int, int, int]] = []
+        kv_transfer_config = self.vllm_config.kv_transfer_config
+        if kv_transfer_config is not None and kv_transfer_config.is_kv_producer:
+            finished_partial_tails = (
+                self.kv_cache_manager.finalize_partial_tail_offloads(request)
+            )
+
         # Free any out-of-window prefix blocks before we hand the block table to
         # the connector, on the processed-token basis (see `allocate_slots`).
         self.kv_cache_manager.remove_skipped_blocks(
@@ -2789,6 +2796,13 @@ class Scheduler(SchedulerInterface):
             request_id=request.request_id,
             num_computed_tokens=request.num_computed_tokens,
         )
+        partial_tail_delay = False
+        if finished_partial_tails:
+            partial_tail_delay = self.connector.register_finished_partial_tail(
+                request,
+                block_ids,
+                finished_partial_tails,
+            )
 
         if not isinstance(self.connector, SupportsHMA):
             # NOTE(Kuntai): We should deprecate this code path after we enforce
@@ -2796,9 +2810,14 @@ class Scheduler(SchedulerInterface):
             # Hybrid memory allocator should be already turned off for this
             # code path, but let's double-check here.
             assert len(self.kv_cache_config.kv_cache_groups) == 1
-            return self.connector.request_finished(request, block_ids[0])
-
-        return self.connector.request_finished_all_groups(request, block_ids)
+            delay_free, kv_xfer_params = self.connector.request_finished(
+                request, block_ids[0]
+            )
+        else:
+            delay_free, kv_xfer_params = self.connector.request_finished_all_groups(
+                request, block_ids
+            )
+        return delay_free or partial_tail_delay, kv_xfer_params
 
     def _request_remaining_blocks(self, request: Request) -> int:
         """Blocks `request` still needs to allocate to hold its full sequence."""

vllm/v1/core/single_type_kv_cache_manager.py

@@ -407,6 +407,15 @@ class SingleTypeKVCacheManager(ABC):
         self._pending_boundary_state_offloads = []
         return pending
 
+    def finalize_partial_tail_offload(
+        self,
+        request_id: str,
+        num_computed_tokens: int,
+        num_in_flight_tokens: int,
+    ) -> tuple[int, KVCacheBlock, int] | None:
+        """Finalize a producer partial tail when its request finishes."""
+        return None
+
     def _apply_cow(
         self,
         request_id: str,
@@ -1383,10 +1392,10 @@ class MambaManager(SingleTypeKVCacheManager):
             # current allocation.
             self._num_checkpoint_blocks: dict[str, int] = {}
             # Requests that registered their own last-prompt-boundary partial
-            # tail (producers). On the next step's CoW the boundary state moves
-            # into a private cow_block; we record that block for connector
-            # offload (see _pending_boundary_state_offloads).
-            self._producer_partial_tail_reqs: dict[str, int] = {}
+            # tail (producers). A later CoW hands its private copy to the
+            # connector; a request that finishes first hands off this table
+            # source directly.
+            self._producer_partial_tail_reqs: dict[str, tuple[KVCacheBlock, int]] = {}
 
     @classmethod
     def find_longest_cache_hit(
@@ -1749,10 +1758,12 @@ class MambaManager(SingleTypeKVCacheManager):
                         self.block_pool.move_block_hashes(source_block, cow_block)
                         self._pending_cow_copies.append((source_block, cow_block))
                         source_block.ref_cnt += 1
-                        boundary_tokens = self._producer_partial_tail_reqs.pop(
+                        producer_tail = self._producer_partial_tail_reqs.pop(
                             request_id, None
                         )
-                        if boundary_tokens is not None:
+                        if producer_tail is not None:
+                            marker_block, boundary_tokens = producer_tail
+                            assert marker_block is source_block
                             # This CoW preserved a producer's own boundary
                             # state in cow_block; hand it to the connector for
                             # partial-tail offload once the copy has run.
@@ -1788,6 +1799,22 @@ class MambaManager(SingleTypeKVCacheManager):
         req_blocks.append(block)
         req_blocks[block_idx] = self._null_block
 
+    def finalize_partial_tail_offload(
+        self,
+        request_id: str,
+        num_computed_tokens: int,
+        num_in_flight_tokens: int,
+    ) -> tuple[int, KVCacheBlock, int] | None:
+        if self.mamba_cache_mode != "align":
+            return None
+        producer_tail = self._producer_partial_tail_reqs.pop(request_id, None)
+        if producer_tail is None:
+            return None
+        source_block, boundary_tokens = producer_tail
+        if num_in_flight_tokens != 0 or num_computed_tokens != boundary_tokens:
+            return None
+        return self.kv_cache_group_id, source_block, boundary_tokens
+
     def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
         if self.mamba_cache_mode == "align":
             self._allocated_block_reqs.discard(request_id)
@@ -1894,7 +1921,10 @@ class MambaManager(SingleTypeKVCacheManager):
             # in ``source_block`` but the next step's forward overwrites it. The
             # upcoming CoW copies it into a durable cow_block; record the req so
             # allocate_new_blocks hands that block to the connector for offload.
-            self._producer_partial_tail_reqs[request.request_id] = num_tokens
+            self._producer_partial_tail_reqs[request.request_id] = (
+                source_block,
+                num_tokens,
+            )
         return partial_hash