[Core] Coalesce concurrent external-prefix KV loads

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

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

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