diff --git a/tests/v1/core/test_shared_external_prefix.py b/tests/v1/core/test_shared_external_prefix.py new file mode 100644 index 000000000000..0f041f11c675 --- /dev/null +++ b/tests/v1/core/test_shared_external_prefix.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.v1.core.kv_cache_utils import BlockHash +from vllm.v1.core.sched.shared_external_prefix import ( + SharedExternalPrefixKey, + SharedExternalPrefixLoadManager, +) +from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec +from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + + +def _key() -> SharedExternalPrefixKey: + return SharedExternalPrefixKey( + connector_namespace=("test", "source"), + hash_block_size=16, + num_local_tokens=0, + num_external_tokens=32, + local_end_hash=None, + external_end_hash=BlockHash(b"external-prefix"), + ) + + +def test_failed_load_stops_accepting_followers_without_blocking_retry(): + manager = SharedExternalPrefixLoadManager() + key = _key() + failed_load = manager.register_owner(key, "owner-1") + manager.add_follower(failed_load, "follower-1") + manager.add_follower(failed_load, "follower-2") + assert manager.detach_follower("follower-1") + assert not manager.is_follower("follower-1") + + manager.stop_accepting_followers("owner-1") + + assert manager.find(key) is None + assert manager.get_by_owner("owner-1") is failed_load + + retry_load = manager.register_owner(key, "owner-2") + assert manager.find(key) is retry_load + + # Draining the failed owner must not remove the replacement key entry. + assert manager.pop_owner("owner-1") is failed_load + assert manager.find(key) is retry_load + assert not manager.is_follower("follower-2") + + assert manager.pop_owner("owner-2") is retry_load + assert not manager.has_unresolved_loads() + + +def test_only_global_cpu_offloading_spec_opts_in(): + class CustomCPUOffloadingSpec(CPUOffloadingSpec): + pass + + cpu_spec = object.__new__(CPUOffloadingSpec) + tiering_spec = object.__new__(TieringOffloadingSpec) + custom_spec = object.__new__(CustomCPUOffloadingSpec) + + assert cpu_spec.shared_kv_load_namespace == ("cpu-offloading",) + assert tiering_spec.shared_kv_load_namespace is None + assert custom_spec.shared_kv_load_namespace is None diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 6c55b91d8da4..ef532631595f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -19,6 +19,7 @@ OffloadingConnectorScheduler, RequestOffloadState, ) +from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -425,6 +426,83 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: assert runner.connector_scheduler._block_id_to_pending_jobs == {} +@pytest.mark.parametrize("async_scheduling", [True, False]) +@pytest.mark.skip_global_cleanup +def test_core_coalesces_concurrent_partial_prefix_loads( + request_runner, async_scheduling: bool, monkeypatch +): + # RequestRunner explicitly exercises HMA. The production path supports it, + # but CPU/MPS test platforms disable it before connector validation. + monkeypatch.setattr(current_platform, "support_hybrid_kv_cache", lambda: True) + block_size = 4 + block_size_factor = 3 + offloaded_block_size = block_size * block_size_factor + runner = request_runner( + block_size=block_size, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + block_size_factor=block_size_factor, + extra_config_overrides={"enable_core_load_coalescing": True}, + ) + + common_prefix = [0] * offloaded_block_size + runner.new_request(token_ids=common_prefix) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0, 1, 2), + ) + + runner.scheduler.reset_prefix_cache() + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, context: 1 + + runner.new_request(token_ids=common_prefix + [1]) + runner.run(decoded_tokens=[], complete_transfers=False) + transfer_jobs = list(runner.offloading_spec.handler.transfer_specs) + + runner.new_request(token_ids=common_prefix + [2]) + runner.run(decoded_tokens=[], complete_transfers=False) + + owner_id, follower_id = "1", "2" + assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs) + assert runner.scheduler._shared_external_prefix_loads.is_follower(follower_id) + assert runner.scheduler.requests[owner_id].status == ( + RequestStatus.WAITING_FOR_REMOTE_KVS + ) + assert runner.scheduler.requests[follower_id].status == ( + RequestStatus.WAITING_FOR_REMOTE_KVS + ) + cache_manager = runner.scheduler.kv_cache_manager.coordinator.single_type_managers[ + 0 + ] + assert follower_id not in cache_manager.req_to_blocks + + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output([]) + ) + runner.run( + decoded_tokens=[], + complete_transfers=True, + expected_loaded=(0, 1, 2), + ) + # Promote the owner after transfer completion and the follower after the + # owner's blocks have been published into the ordinary GPU prefix cache. + runner.run(decoded_tokens=[], complete_transfers=True) + + owner_blocks = runner.scheduler.kv_cache_manager.get_block_ids(owner_id)[0] + follower_blocks = runner.scheduler.kv_cache_manager.get_block_ids(follower_id)[0] + assert owner_blocks[:block_size_factor] == follower_blocks[:block_size_factor] + assert all( + runner.scheduler.kv_cache_manager.block_pool.blocks[block_id].ref_cnt == 2 + for block_id in owner_blocks[:block_size_factor] + ) + + runner.run(decoded_tokens=[EOS_TOKEN_ID]) + assert not runner.scheduler._shared_external_prefix_loads.has_unresolved_loads() + + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_abort_loading_requests(request_runner, async_scheduling: bool): block_size = 4 diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 73ea5e2be1d2..9ba10bfb4da0 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -133,6 +133,12 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingWorker() + @property + def shared_kv_load_namespace(self) -> tuple[str] | None: + if self.extra_config.get("enable_core_load_coalescing", False): + return ("mock-offloading",) + return None + def get_manager(self) -> OffloadingManager: return self.manager diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 7d21d975c173..40e6c528fe01 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -740,3 +740,467 @@ def test_async_loads_both_admitted_when_pool_fits(): for req in reqs: assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def _make_external_prefix_dedupe_scheduler( + monkeypatch, kv_load_failure_policy="fail", matched_tokens=None +): + block_size = 16 + matched_tokens = matched_tokens or 2 * block_size + vllm_config = create_vllm_config( + block_size=block_size, + max_num_batched_tokens=128, + kv_connector="MockKVConnector", + kv_connector_extra_config={ + "matched_tokens": matched_tokens, + "is_async": True, + }, + kv_load_failure_policy=kv_load_failure_policy, + ) + scheduler = create_scheduler(vllm_config, num_blocks=64) + connector = scheduler.connector + assert connector is not None + + def get_num_new_matched_tokens(request, num_computed_tokens): + num_external_tokens = max(matched_tokens - num_computed_tokens, 0) + return num_external_tokens, num_external_tokens > 0 + + nonzero_allocs = [] + original_update_state_after_alloc = connector.update_state_after_alloc + + def record_update_state_after_alloc(request, blocks, num_external_tokens): + if num_external_tokens: + nonzero_allocs.append( + (request.request_id, blocks.get_block_ids(), num_external_tokens) + ) + return original_update_state_after_alloc(request, blocks, num_external_tokens) + + monkeypatch.setattr( + connector, + "get_num_new_matched_tokens", + get_num_new_matched_tokens, + ) + monkeypatch.setattr( + connector, + "get_shared_kv_load_namespace", + lambda request: ("mock",), + ) + monkeypatch.setattr( + connector, + "update_state_after_alloc", + record_update_state_after_alloc, + ) + return scheduler, block_size, matched_tokens, nonzero_allocs + + +def _make_external_prefix_requests(block_size, matched_tokens, same_prefix=True): + common_prefix_len = matched_tokens if same_prefix else 0 + return [ + create_request( + request_id=request_id, + block_size=block_size, + num_tokens=3 * block_size, + common_prefix_len=common_prefix_len, + ) + for request_id in (101, 102) + ] + + +def _finish_owner_load(scheduler, owner, initial_output): + scheduler.update_from_output(initial_output, EMPTY_MODEL_RUNNER_OUTPUT) + waiting_output = scheduler.schedule() + scheduler.update_from_output( + waiting_output, + create_model_runner_output( + reqs=[], + finished_recving={owner.request_id}, + ), + ) + return scheduler.schedule() + + +@pytest.mark.skip_global_cleanup +def test_shared_external_prefix_materializes_once(monkeypatch): + """Concurrent exact-prefix hits use one pending materialization.""" + ( + scheduler, + block_size, + matched_tokens, + nonzero_allocs, + ) = _make_external_prefix_dedupe_scheduler(monkeypatch) + owner, follower = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=True + ) + initial_num_free_blocks = ( + scheduler.kv_cache_manager.block_pool.get_num_free_blocks() + ) + + scheduler.add_request(owner) + scheduler.add_request(follower) + initial_output = scheduler.schedule() + + assert owner.status == RequestStatus.WAITING_FOR_REMOTE_KVS + assert follower.status == RequestStatus.WAITING_FOR_REMOTE_KVS + assert owner.num_computed_tokens == matched_tokens + assert follower.num_computed_tokens == 0 + assert owner in scheduler._inflight_prefills + assert follower not in scheduler._inflight_prefills + assert [alloc[0] for alloc in nonzero_allocs] == [owner.request_id] + + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] + assert owner.request_id in manager.req_to_blocks + assert follower.request_id not in manager.req_to_blocks + owner_prefix_ids = scheduler.kv_cache_manager.get_block_ids(owner.request_id)[0] + assert len(owner_prefix_ids) == matched_tokens // block_size + assert all( + scheduler.kv_cache_manager.block_pool.blocks[block_id].block_hash is None + for block_id in owner_prefix_ids + ) + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == ( + initial_num_free_blocks - len(owner_prefix_ids) + ) + + resumed_output = _finish_owner_load(scheduler, owner, initial_output) + + assert owner.status == RequestStatus.RUNNING + assert follower.status == RequestStatus.RUNNING + owner_block_ids = scheduler.kv_cache_manager.get_block_ids(owner.request_id)[0] + follower_block_ids = scheduler.kv_cache_manager.get_block_ids(follower.request_id)[ + 0 + ] + num_prefix_blocks = matched_tokens // block_size + assert owner_block_ids[:num_prefix_blocks] == owner_prefix_ids + assert follower_block_ids[:num_prefix_blocks] == owner_prefix_ids + assert set(owner_block_ids[num_prefix_blocks:]).isdisjoint( + follower_block_ids[num_prefix_blocks:] + ) + assert all( + scheduler.kv_cache_manager.block_pool.blocks[block_id].ref_cnt == 2 + for block_id in owner_prefix_ids + ) + assert set(resumed_output.num_scheduled_tokens) == { + owner.request_id, + follower.request_id, + } + assert len(nonzero_allocs) == 1 + + +@pytest.mark.skip_global_cleanup +def test_pending_shared_external_prefix_blocks_connector_reset(monkeypatch): + """A connector reset cannot invalidate a source with parked followers.""" + scheduler, block_size, matched_tokens, _ = _make_external_prefix_dedupe_scheduler( + monkeypatch + ) + owner, follower = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=True + ) + scheduler.add_request(owner) + scheduler.add_request(follower) + initial_output = scheduler.schedule() + + assert scheduler.connector is not None + + def unexpected_reset(): + raise AssertionError("connector reset must be rejected before delegation") + + monkeypatch.setattr(scheduler.connector, "reset_cache", unexpected_reset) + assert not scheduler.reset_connector_cache() + + resumed_output = _finish_owner_load(scheduler, owner, initial_output) + scheduler.update_from_output( + resumed_output, + create_model_runner_output([owner, follower]), + ) + decode_output = scheduler.schedule() + scheduler.update_from_output( + decode_output, + create_model_runner_output([owner, follower], use_eos=True), + ) + scheduler.schedule() + assert_scheduler_empty(scheduler) + + +@pytest.mark.skip_global_cleanup +def test_distinct_external_prefixes_do_not_dedupe(monkeypatch): + """Different exact-prefix hashes retain independent materializations.""" + ( + scheduler, + block_size, + matched_tokens, + nonzero_allocs, + ) = _make_external_prefix_dedupe_scheduler(monkeypatch) + requests = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=False + ) + initial_num_free_blocks = ( + scheduler.kv_cache_manager.block_pool.get_num_free_blocks() + ) + for request in requests: + scheduler.add_request(request) + + scheduler.schedule() + + assert all( + request.status == RequestStatus.WAITING_FOR_REMOTE_KVS for request in requests + ) + assert {alloc[0] for alloc in nonzero_allocs} == { + request.request_id for request in requests + } + block_id_sets = [ + set(scheduler.kv_cache_manager.get_block_ids(request.request_id)[0]) + for request in requests + ] + assert block_id_sets[0].isdisjoint(block_id_sets[1]) + assert all( + scheduler.kv_cache_manager.block_pool.blocks[block_id].ref_cnt == 1 + for block_ids in block_id_sets + for block_id in block_ids + ) + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == ( + initial_num_free_blocks - 2 * matched_tokens // block_size + ) + + +@pytest.mark.parametrize("case", ["unaligned", "namespace", "full_prompt"]) +@pytest.mark.skip_global_cleanup +def test_external_prefix_dedupe_fails_closed(monkeypatch, case): + """Unsupported spans and different connector namespaces never coalesce.""" + requested_match = {"unaligned": 17, "namespace": None, "full_prompt": 48}[case] + scheduler, block_size, matched_tokens, nonzero_allocs = ( + _make_external_prefix_dedupe_scheduler( + monkeypatch, matched_tokens=requested_match + ) + ) + if case == "namespace": + assert scheduler.connector is not None + monkeypatch.setattr( + scheduler.connector, + "get_shared_kv_load_namespace", + lambda request: ("mock", request.request_id), + ) + + requests = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=True + ) + for request in requests: + scheduler.add_request(request) + + scheduler.schedule() + + assert len(nonzero_allocs) == 2 + assert all( + request.status == RequestStatus.WAITING_FOR_REMOTE_KVS for request in requests + ) + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] + assert all(request.request_id in manager.req_to_blocks for request in requests) + + +@pytest.mark.skip_global_cleanup +def test_shared_external_prefix_follower_abort_does_not_leak(monkeypatch): + """A follower with no transfer or blocks detaches immediately on abort.""" + ( + scheduler, + block_size, + matched_tokens, + nonzero_allocs, + ) = _make_external_prefix_dedupe_scheduler(monkeypatch) + owner, follower = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=True + ) + initial_num_free_blocks = ( + scheduler.kv_cache_manager.block_pool.get_num_free_blocks() + ) + scheduler.add_request(owner) + scheduler.add_request(follower) + initial_output = scheduler.schedule() + scheduler.update_from_output(initial_output, EMPTY_MODEL_RUNNER_OUTPUT) + + scheduler.finish_requests( + follower.request_id, + RequestStatus.FINISHED_ABORTED, + ) + + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] + assert follower.request_id not in scheduler.requests + assert follower.request_id not in manager.req_to_blocks + assert owner.request_id in scheduler.requests + assert owner.status == RequestStatus.WAITING_FOR_REMOTE_KVS + owner_block_ids = scheduler.kv_cache_manager.get_block_ids(owner.request_id)[0] + assert all( + scheduler.kv_cache_manager.block_pool.blocks[block_id].ref_cnt == 1 + for block_id in owner_block_ids + ) + assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == ( + initial_num_free_blocks - len(owner_block_ids) + ) + assert len(nonzero_allocs) == 1 + + waiting_output = scheduler.schedule() + scheduler.update_from_output( + waiting_output, + create_model_runner_output( + reqs=[], + finished_recving={owner.request_id}, + ), + ) + owner_output = scheduler.schedule() + scheduler.update_from_output( + owner_output, + create_model_runner_output([owner]), + ) + decode_output = scheduler.schedule() + scheduler.update_from_output( + decode_output, + create_model_runner_output([owner], use_eos=True), + ) + scheduler.schedule() + assert_scheduler_empty(scheduler) + + +@pytest.mark.skip_global_cleanup +def test_shared_external_prefix_owner_abort_releases_follower(monkeypatch): + """Owner abort releases followers to recompute while its load drains.""" + ( + scheduler, + block_size, + matched_tokens, + nonzero_allocs, + ) = _make_external_prefix_dedupe_scheduler(monkeypatch) + owner, follower = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=True + ) + scheduler.add_request(owner) + scheduler.add_request(follower) + initial_output = scheduler.schedule() + scheduler.update_from_output(initial_output, EMPTY_MODEL_RUNNER_OUTPUT) + owner_block_ids = set(scheduler.kv_cache_manager.get_block_ids(owner.request_id)[0]) + + scheduler.finish_requests(owner.request_id, RequestStatus.FINISHED_ABORTED) + assert owner.request_id in scheduler.requests + + fallback_output = scheduler.schedule() + + assert follower.status == RequestStatus.RUNNING + assert fallback_output.num_scheduled_tokens[follower.request_id] == ( + follower.num_prompt_tokens + ) + follower_block_ids = set( + scheduler.kv_cache_manager.get_block_ids(follower.request_id)[0] + ) + assert owner_block_ids.isdisjoint(follower_block_ids) + assert len(nonzero_allocs) == 1 + + scheduler.update_from_output( + fallback_output, + create_model_runner_output( + [follower], + finished_recving={owner.request_id}, + ), + ) + assert owner.request_id not in scheduler.requests + decode_output = scheduler.schedule() + scheduler.update_from_output( + decode_output, + create_model_runner_output([follower], use_eos=True), + ) + scheduler.schedule() + assert_scheduler_empty(scheduler) + + +@pytest.mark.skip_global_cleanup +def test_shared_external_prefix_load_failure_fails_followers(monkeypatch): + """The fail policy applies one materialization outcome to every waiter.""" + scheduler, block_size, matched_tokens, _ = _make_external_prefix_dedupe_scheduler( + monkeypatch + ) + owner, follower = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=True + ) + scheduler.add_request(owner) + scheduler.add_request(follower) + initial_output = scheduler.schedule() + owner_block_ids = scheduler.kv_cache_manager.get_block_ids(owner.request_id)[0] + + outputs = scheduler.update_from_output( + initial_output, + create_model_runner_output( + reqs=[], + invalid_block_ids={owner_block_ids[-1]}, + ), + ) + + assert owner.status == RequestStatus.FINISHED_ERROR + assert follower.status == RequestStatus.FINISHED_ERROR + assert follower.request_id not in scheduler.requests + assert owner.request_id in scheduler.requests # Transfer target is draining. + assert { + output.request_id + for engine_output in outputs.values() + for output in engine_output.outputs + } == {owner.request_id, follower.request_id} + + drain_output = scheduler.schedule() + scheduler.update_from_output( + drain_output, + create_model_runner_output( + reqs=[], + finished_recving={owner.request_id}, + ), + ) + scheduler.schedule() + assert_scheduler_empty(scheduler) + + +@pytest.mark.skip_global_cleanup +def test_shared_external_prefix_load_failure_recomputes_safely(monkeypatch): + """Recompute publishes only the valid prefix, never shared failed blocks.""" + scheduler, block_size, matched_tokens, nonzero_allocs = ( + _make_external_prefix_dedupe_scheduler( + monkeypatch, kv_load_failure_policy="recompute" + ) + ) + owner, follower = _make_external_prefix_requests( + block_size, matched_tokens, same_prefix=True + ) + scheduler.add_request(owner) + scheduler.add_request(follower) + initial_output = scheduler.schedule() + owner_block_ids = scheduler.kv_cache_manager.get_block_ids(owner.request_id)[0] + + scheduler.update_from_output( + initial_output, + create_model_runner_output( + reqs=[], + invalid_block_ids={owner_block_ids[1]}, + finished_recving={owner.request_id}, + ), + ) + retry_output = scheduler.schedule() + + assert owner.status == RequestStatus.RUNNING + assert follower.status == RequestStatus.RUNNING + assert owner.num_computed_tokens == owner.num_prompt_tokens + assert follower.num_computed_tokens == follower.num_prompt_tokens + assert len(nonzero_allocs) == 1 + + owner_retry_ids = scheduler.kv_cache_manager.get_block_ids(owner.request_id)[0] + follower_retry_ids = scheduler.kv_cache_manager.get_block_ids(follower.request_id)[ + 0 + ] + # The valid block is published, but followers bypass APC for this retry so + # they cannot observe blocks optimistically hashed by the owner's recompute. + assert owner_retry_ids[0] == owner_block_ids[0] + assert scheduler.kv_cache_manager.block_pool.blocks[owner_block_ids[0]].ref_cnt == 1 + assert set(owner_retry_ids).isdisjoint(follower_retry_ids) + + scheduler.update_from_output( + retry_output, + create_model_runner_output([owner, follower]), + ) + decode_output = scheduler.schedule() + scheduler.update_from_output( + decode_output, + create_model_runner_output([owner, follower], use_eos=True), + ) + scheduler.schedule() + assert_scheduler_empty(scheduler) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 7df9e20e6a58..5cab0a3d8cc4 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -57,6 +57,9 @@ def assert_scheduler_empty(scheduler: Scheduler): assert len(scheduler.finished_req_ids) == 0 assert len(scheduler.finished_recving_kv_req_ids) == 0 assert len(scheduler._inflight_prefills) == 0 + assert not scheduler._shared_external_prefix_loads.has_unresolved_loads() + assert not scheduler._draining_shared_external_prefix_owners + assert not scheduler._force_local_recompute_once # EncoderCacheManager. assert len(scheduler.encoder_cache_manager.freed) == 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index c0eb24729793..a3fc03dc0b24 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -10,6 +10,8 @@ get_num_new_matched_tokens() - get number of new tokens that exist in the remote KV cache. Might be called multiple times for a given request and should be side-effect free. + get_shared_kv_load_namespace() - optionally identify async loads + that may share one exact local materialization. update_state_after_alloc() - update KVConnector state after temporary buffer alloc by the CacheManager. update_connector_output() - update KVConnector state after @@ -42,7 +44,7 @@ import enum from abc import ABC, abstractmethod -from collections.abc import Callable, Iterable +from collections.abc import Callable, Hashable, Iterable from typing import TYPE_CHECKING, Any, Literal import torch @@ -440,6 +442,22 @@ def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None: # Scheduler-side methods # ============================== + def get_shared_kv_load_namespace(self, request: "Request") -> Hashable | None: + """Return an identity namespace for coalescing exact async loads. + + Returning ``None`` disables load coalescing for this request. Connectors + that opt in must return equal, hashable values only when two requests' + external prefix loads use the same source and materialization semantics. + The scheduler additionally verifies the exact target block-hash range. + + This method is called after :meth:`get_num_new_matched_tokens` reports a + positive asynchronous hit and must not mutate connector state. For a + coalesced follower, the scheduler may defer + :meth:`update_state_after_alloc` until the request resumes through the + local prefix cache, or call :meth:`request_finished` first if it aborts. + """ + return None + def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None: """ Bind the GPU block pool to the connector for per-GPU block status tracking. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 75a9696b72cd..241ca28d470c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -48,7 +48,7 @@ make_offload_key, ) from vllm.v1.outputs import KVConnectorOutput -from vllm.v1.request import Request +from vllm.v1.request import Request, RequestStatus logger = init_logger(__name__) @@ -356,6 +356,7 @@ def __init__( spec: OffloadingSpec, ): self.config = SchedulerOffloadConfig.from_spec(spec) + self._shared_kv_load_namespace = spec.shared_kv_load_namespace self.manager: OffloadingManager = spec.get_manager() self._connector_stats = OffloadingConnectorStats() @@ -407,6 +408,37 @@ def _sliding_window_sort_key(i: int) -> int: self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) + def _can_coalesce_in_core( + self, req_status: RequestOffloadState, num_hit_tokens: int + ) -> bool: + """Whether the core's stricter exact-prefix gate will accept this hit.""" + namespace = self._shared_kv_load_namespace + if namespace is None: + return False + try: + hash(namespace) + except TypeError: + return False + if ( + req_status.req.status != RequestStatus.WAITING + or len(self.config.kv_group_configs) != 1 + ): + return False + + group_config = self.config.kv_group_configs[0] + num_local_tokens = req_status.num_locally_computed_tokens + num_total_tokens = num_local_tokens + num_hit_tokens + block_size = group_config.gpu_block_size + return ( + num_hit_tokens > 0 + and group_config.sliding_window_size_in_blocks is None + and not group_config.is_eagle_group + and num_total_tokens < req_status.req.num_tokens + and num_local_tokens % block_size == 0 + and num_hit_tokens % block_size == 0 + and num_total_tokens % block_size == 0 + ) + def _maybe_observe_lookup_async_delay( self, req_status: RequestOffloadState ) -> None: @@ -660,6 +692,10 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: if sliding_window_size_in_blocks is not None: offload_keys = offload_keys[-sliding_window_size_in_blocks:] if any(key in self._blocks_being_loaded for key in offload_keys): + if self._can_coalesce_in_core(req_status, num_hit_tokens): + # Let the core park this request behind the owner. The + # follower gets no GPU blocks and creates no load job. + continue # hit blocks are being loaded, delay request logger.debug( "Delaying request %s since some of its" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 197beca9aece..88b2f04be9e2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable +from collections.abc import Hashable, Iterable from typing import Any import torch @@ -57,6 +57,7 @@ def __init__( super().__init__(vllm_config, role, kv_cache_config) spec = OffloadingSpecFactory.create_spec(vllm_config, kv_cache_config) + self._shared_kv_load_namespace = spec.shared_kv_load_namespace self.connector_scheduler: OffloadingConnectorScheduler | None = None self.connector_worker: OffloadingConnectorWorker | None = None @@ -136,6 +137,9 @@ def get_num_new_matched_tokens( request, num_computed_tokens ) + def get_shared_kv_load_namespace(self, request: "Request") -> Hashable | None: + return self._shared_kv_load_namespace + def update_state_after_alloc( self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int ): diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 95071408876b..831f4d59cd2b 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -50,9 +50,14 @@ SchedulingPolicy, create_request_queue, ) +from vllm.v1.core.sched.shared_external_prefix import ( + PendingExternalPrefixLoad, + SharedExternalPrefixKey, + SharedExternalPrefixLoadManager, +) from vllm.v1.core.sched.utils import check_stop, remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig from vllm.v1.metrics.perf import ModelMetrics, PerfStats from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput @@ -259,6 +264,7 @@ def __init__( # Create the KV cache manager. if hash_block_size is None: hash_block_size = block_size + self.hash_block_size = hash_block_size self.kv_cache_manager = KVCacheManager( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, @@ -335,6 +341,118 @@ def __init__( # async KV loads). Their remaining-block reservation gates async loads. self._inflight_prefills: set[Request] = set() + # Exact external-prefix loads that can be coalesced. Pending followers + # own no blocks; they re-enter normal local APC lookup only after the + # owner's blocks have been published. + self._shared_external_prefix_loads = SharedExternalPrefixLoadManager() + # Aborted owners retain their transfer target until the connector + # reports completion, even though their followers have been released. + self._draining_shared_external_prefix_owners: set[str] = set() + # Followers of a failed or abandoned materialization bypass both local + # APC and the external source once. This prevents them from observing + # the owner's optimistic hashes or immediately retrying the same load. + self._force_local_recompute_once: set[str] = set() + + def _make_shared_external_prefix_key( + self, + request: Request, + num_local_tokens: int, + num_external_tokens: int, + ) -> SharedExternalPrefixKey | None: + """Build a fail-closed identity for a coalescible external load.""" + if ( + self.connector is None + or not self.kv_cache_manager.enable_caching + or request.skip_reading_prefix_cache + or request.status != RequestStatus.WAITING + or self.use_eagle + or len(self.kv_cache_config.kv_cache_groups) != 1 + or not isinstance( + self.kv_cache_config.kv_cache_groups[0].kv_cache_spec, + FullAttentionSpec, + ) + ): + return None + + num_total_tokens = num_local_tokens + num_external_tokens + if ( + num_external_tokens <= 0 + # A full-prompt hit must recompute its final token. APC may drop + # that entire block, so it cannot preserve exact block sharing. + or num_total_tokens >= request.num_tokens + or num_local_tokens % self.block_size != 0 + or num_external_tokens % self.block_size != 0 + or num_total_tokens % self.block_size != 0 + ): + return None + + namespace = self.connector.get_shared_kv_load_namespace(request) + if namespace is None: + return None + + local_hash_index = num_local_tokens // self.hash_block_size - 1 + external_hash_index = num_total_tokens // self.hash_block_size - 1 + if external_hash_index >= len(request.block_hashes): + return None + local_end_hash = ( + request.block_hashes[local_hash_index] if local_hash_index >= 0 else None + ) + key = SharedExternalPrefixKey( + connector_namespace=namespace, + hash_block_size=self.hash_block_size, + num_local_tokens=num_local_tokens, + num_external_tokens=num_external_tokens, + local_end_hash=local_end_hash, + external_end_hash=request.block_hashes[external_hash_index], + ) + try: + hash(key) + except TypeError: + logger.warning_once( + "KV connector returned an unhashable shared-load namespace; " + "external-prefix load coalescing is disabled for that request." + ) + return None + return key + + def _park_shared_external_prefix_follower( + self, + load: PendingExternalPrefixLoad, + request: Request, + ) -> None: + """Park a follower without allocating or exposing owner blocks.""" + assert request.num_computed_tokens == 0 + self._shared_external_prefix_loads.add_follower(load, request.request_id) + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + + def _release_shared_external_prefix_followers( + self, owner_request_id: str, force_local_recompute: bool = False + ) -> None: + """Publish barrier: release followers back through ordinary APC lookup.""" + load = self._shared_external_prefix_loads.pop_owner(owner_request_id) + if load is None: + return + + for follower_id in load.follower_request_ids: + follower = self.requests.get(follower_id) + if ( + follower is None + or follower.status != RequestStatus.WAITING_FOR_REMOTE_KVS + ): + continue + # Followers never owned blocks or a connector transfer. Keeping the + # count at zero forces a fresh local-prefix lookup after publication. + follower.num_computed_tokens = 0 + follower.status = ( + RequestStatus.PREEMPTED + if follower.num_preemptions + else RequestStatus.WAITING + ) + self.finished_recving_kv_req_ids.discard(follower_id) + self.failed_recving_kv_req_ids.discard(follower_id) + if force_local_recompute: + self._force_local_recompute_once.add(follower_id) + def _mamba_block_aligned_split( self, request: Request, @@ -682,11 +800,18 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: load_kv_async = False connector_prefix_cache_queries, connector_prefix_cache_hits = 0, 0 num_uncached_common_prefix_tokens = 0 + shared_external_prefix_key: SharedExternalPrefixKey | None = None + force_local_recompute = request_id in self._force_local_recompute_once # Get already-cached tokens. if request.num_computed_tokens == 0: # Get locally-cached tokens. - if ( + if force_local_recompute: + new_computed_blocks = ( + self.kv_cache_manager.empty_kv_cache_blocks + ) + num_new_local_computed_tokens = 0 + elif ( self.connector is not None and self.has_mamba_layers and isinstance( @@ -734,7 +859,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: ) # Get externally-cached tokens if using a KVConnector. - if self.connector is not None: + if self.connector is not None and not force_local_recompute: ext_tokens, load_kv_async = ( self.connector.get_num_new_matched_tokens( request, num_new_local_computed_tokens @@ -789,6 +914,25 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: num_new_local_computed_tokens = 0 num_computed_tokens = request.num_computed_tokens + if load_kv_async: + shared_external_prefix_key = self._make_shared_external_prefix_key( + request, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if shared_external_prefix_key is not None: + pending_load = self._shared_external_prefix_loads.find( + shared_external_prefix_key + ) + if pending_load is not None: + request = request_queue.pop_request() + self._force_local_recompute_once.discard(request_id) + self._park_shared_external_prefix_follower( + pending_load, request + ) + step_skipped_waiting.prepend_request(request) + continue + encoder_inputs_to_schedule = None external_load_encoder_input = [] new_encoder_compute_budget = encoder_compute_budget @@ -935,6 +1079,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.kv_cache_manager.get_blocks(request_id), num_external_computed_tokens, ) + if load_kv_async and shared_external_prefix_key is not None: + self._shared_external_prefix_loads.register_owner( + shared_external_prefix_key, request_id + ) if ( self.connector_prefix_cache_stats is not None and connector_prefix_cache_queries != 0 @@ -945,6 +1093,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: preempted=request.num_preemptions > 0, ) + self._force_local_recompute_once.discard(request_id) + request = request_queue.pop_request() if load_kv_async: # If loading async, allocate memory and put request @@ -2091,11 +2241,32 @@ def finish_requests( for request in valid_requests: delay_free_blocks = False if request.status == RequestStatus.WAITING_FOR_REMOTE_KVS: - delay_free_blocks = ( - request.request_id not in self.finished_recving_kv_req_ids + request_id = request.request_id + is_shared_follower = self._shared_external_prefix_loads.is_follower( + request_id ) - self.finished_recving_kv_req_ids.discard(request.request_id) - self.failed_recving_kv_req_ids.discard(request.request_id) + if is_shared_follower: + # A follower owns neither blocks nor a connector transfer, + # so waiting for its own completion would leak the request. + self._shared_external_prefix_loads.detach_follower(request_id) + else: + pending_load = self._shared_external_prefix_loads.get_by_owner( + request_id + ) + if pending_load is not None: + transfer_is_pending = ( + request_id not in self.finished_recving_kv_req_ids + ) + self._release_shared_external_prefix_followers( + request_id, force_local_recompute=True + ) + if transfer_is_pending: + self._draining_shared_external_prefix_owners.add(request_id) + delay_free_blocks = ( + request_id not in self.finished_recving_kv_req_ids + ) + self.finished_recving_kv_req_ids.discard(request_id) + self.failed_recving_kv_req_ids.discard(request_id) request.status = finished_status self._free_request(request, delay_free_blocks=delay_free_blocks) @@ -2108,6 +2279,7 @@ def _free_request( assert request.is_finished() self._inflight_prefills.discard(request) + self._force_local_recompute_once.discard(request.request_id) connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) self.encoder_cache_manager.free(request) request_id = request.request_id @@ -2262,6 +2434,16 @@ def reset_connector_cache(self) -> bool: ) return True + if ( + self._shared_external_prefix_loads.has_unresolved_loads() + or self._draining_shared_external_prefix_owners + ): + logger.warning( + "Cannot reset connector cache while shared external-prefix " + "loads are pending or draining." + ) + return False + if self.connector.reset_cache() is False: return False @@ -2422,7 +2604,8 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: """ assert self.connector is not None - if request.request_id in self.failed_recving_kv_req_ids: + load_failed = request.request_id in self.failed_recving_kv_req_ids + if load_failed: # Request had KV load failures; num_computed_tokens was already # updated in _update_requests_with_invalid_blocks if request.num_computed_tokens: @@ -2444,6 +2627,11 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: if request.num_computed_tokens == request.num_tokens: request.num_computed_tokens = request.num_tokens - 1 + # Publication must happen before followers are released. Successful + # followers use ordinary APC; failed loads bypass APC for one retry. + self._release_shared_external_prefix_followers( + request.request_id, force_local_recompute=load_failed + ) self.finished_recving_kv_req_ids.remove(request.request_id) def _try_promote_blocked_waiting_request(self, request: Request) -> bool: @@ -2502,10 +2690,12 @@ def _update_from_kv_xfer_finished(self, kv_connector_output: KVConnectorOutput): self.finished_recving_kv_req_ids.add(req_id) else: assert RequestStatus.is_finished(req.status) + self._draining_shared_external_prefix_owners.discard(req_id) 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) assert req_id in self.requests + self._draining_shared_external_prefix_owners.discard(req_id) self._free_blocks(self.requests[req_id]) def _update_requests_with_invalid_blocks( @@ -2637,6 +2827,24 @@ def _handle_invalid_blocks( ) ) + for owner_id in async_failed_req_ids: + self._shared_external_prefix_loads.stop_accepting_followers(owner_id) + + if should_fail: + # Followers have no blocks to intersect with invalid_block_ids, but + # under the fail policy they depend on the same materialization and + # must observe the same terminal outcome. + coalesced_followers: set[str] = set() + for owner_id in tuple(async_failed_req_ids): + load = self._shared_external_prefix_loads.get_by_owner(owner_id) + if load is None: + continue + coalesced_followers.update(load.follower_request_ids) + num_failed_tokens += len(load.follower_request_ids) * ( + load.key.num_local_tokens + load.key.num_external_tokens + ) + async_failed_req_ids |= coalesced_followers + total_failed_requests = len(async_failed_req_ids) total_failed_tokens = num_failed_tokens diff --git a/vllm/v1/core/sched/shared_external_prefix.py b/vllm/v1/core/sched/shared_external_prefix.py new file mode 100644 index 000000000000..a55ee376a814 --- /dev/null +++ b/vllm/v1/core/sched/shared_external_prefix.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bookkeeping for coalesced asynchronous external-prefix loads.""" + +from collections.abc import Hashable +from dataclasses import dataclass, field + +from vllm.v1.core.kv_cache_utils import BlockHash + + +@dataclass(frozen=True) +class SharedExternalPrefixKey: + """Identity of one exact, block-aligned external-prefix load.""" + + connector_namespace: Hashable + hash_block_size: int + num_local_tokens: int + num_external_tokens: int + local_end_hash: BlockHash | None + external_end_hash: BlockHash + + +@dataclass +class PendingExternalPrefixLoad: + """An owner load and requests waiting for its APC publication.""" + + key: SharedExternalPrefixKey + owner_request_id: str + follower_request_ids: set[str] = field(default_factory=set) + accepting_followers: bool = True + + +class SharedExternalPrefixLoadManager: + """Indexes pending exact-prefix loads without owning any KV blocks.""" + + def __init__(self) -> None: + self._loads_by_key: dict[ + SharedExternalPrefixKey, PendingExternalPrefixLoad + ] = {} + self._loads_by_owner: dict[str, PendingExternalPrefixLoad] = {} + self._follower_to_owner: dict[str, str] = {} + + def find(self, key: SharedExternalPrefixKey) -> PendingExternalPrefixLoad | None: + load = self._loads_by_key.get(key) + return load if load is not None and load.accepting_followers else None + + def register_owner( + self, key: SharedExternalPrefixKey, request_id: str + ) -> PendingExternalPrefixLoad: + assert key not in self._loads_by_key + assert request_id not in self._loads_by_owner + assert request_id not in self._follower_to_owner + load = PendingExternalPrefixLoad(key=key, owner_request_id=request_id) + self._loads_by_key[key] = load + self._loads_by_owner[request_id] = load + return load + + def add_follower(self, load: PendingExternalPrefixLoad, request_id: str) -> None: + assert self._loads_by_key.get(load.key) is load + assert request_id not in self._loads_by_owner + assert request_id not in self._follower_to_owner + load.follower_request_ids.add(request_id) + self._follower_to_owner[request_id] = load.owner_request_id + + def get_by_owner(self, request_id: str) -> PendingExternalPrefixLoad | None: + return self._loads_by_owner.get(request_id) + + def is_follower(self, request_id: str) -> bool: + return request_id in self._follower_to_owner + + def stop_accepting_followers(self, request_id: str) -> None: + load = self.get_by_owner(request_id) + if load is not None: + load.accepting_followers = False + if self._loads_by_key.get(load.key) is load: + del self._loads_by_key[load.key] + + def detach_follower(self, request_id: str) -> bool: + owner_id = self._follower_to_owner.pop(request_id, None) + if owner_id is None: + return False + load = self.get_by_owner(owner_id) + if load is not None: + load.follower_request_ids.discard(request_id) + return True + + def pop_owner(self, request_id: str) -> PendingExternalPrefixLoad | None: + load = self._loads_by_owner.pop(request_id, None) + if load is None: + return None + if self._loads_by_key.get(load.key) is load: + del self._loads_by_key[load.key] + for follower_id in load.follower_request_ids: + self._follower_to_owner.pop(follower_id, None) + return load + + def has_unresolved_loads(self) -> bool: + return bool(self._loads_by_owner) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 5a2e3c184d39..aa7deb73aef0 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -5,7 +5,7 @@ """ from abc import ABC, abstractmethod -from collections.abc import Collection, Iterable, Sequence +from collections.abc import Collection, Hashable, Iterable, Sequence from dataclasses import dataclass from enum import Enum, auto from typing import TYPE_CHECKING, Any, NamedTuple, NewType @@ -475,6 +475,16 @@ def shutdown(self) -> None: class OffloadingSpec(ABC): """Spec for an offloading connector""" + @property + def shared_kv_load_namespace(self) -> Hashable | None: + """Identity of a global source safe for exact load coalescing. + + Specs return ``None`` by default. An opt-in namespace must describe a + source whose lookup and materialization semantics are identical for all + requests using the spec. + """ + return None + @classmethod def build_metric_definitions( cls, extra_config: dict[str, Any] diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 26ea3728191f..43553babcb29 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -26,6 +26,15 @@ class CPUOffloadingSpec(OffloadingSpec): BLOCK_SIZE_ALIGNMENT = 1 + @property + @override + def shared_kv_load_namespace(self) -> tuple[str] | None: + # Subclasses may replace the global CPU manager with request-scoped or + # remote sources and must opt in independently. + if type(self) is CPUOffloadingSpec: + return ("cpu-offloading",) + return None + @classmethod def build_metric_definitions( cls, extra_config: dict[str, Any] diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 3dc31a3622ee..1952f3f5c79b 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -71,6 +71,12 @@ class TieringOffloadingSpec(CPUOffloadingSpec): BLOCK_SIZE_ALIGNMENT = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + @property + @override + def shared_kv_load_namespace(self) -> None: + # A tiering lookup may select a request-scoped P2P source. + return None + @classmethod @override def build_metric_definitions(