diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 2dcc0e1184..9558409ae9 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/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.""" diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 9cdaa9666c..110a83db48 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/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 diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 3f2d5da5b0..0c13b1f578 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/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""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 36e53828c0..fa764571ac 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/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. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index a73c6dea43..2c8d118cdd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/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. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index bc5fdcd5d8..b36f31e727 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/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]]], diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 2b567da792..348f26acfe 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/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", diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index ef9f8da7be..50f185c6ac 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/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() diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b0fd3d25bf..cbef879b29 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/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.""" diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2091be31c6..6da0b8c070 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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