Add vLLM v0.18.1 source tree with KV transfer abort fix

third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:

  vllm/v1/core/sched/scheduler.py:
    Replace fatal assert with graceful skip when KV transfer callback
    arrives for an already-aborted request during PD disaggregated serving.

Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

third_party/vllm/tests/v1/kv_connector/unit/test_backwards_compatibility.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for backwards compatibility with external KV connector implementations.
+
+This test ensures that external connectors (loaded via kv_connector_module_path)
+implemented with the old signature continue to work:
+- Old signature: __init__(self, vllm_config, role)
+- New signature: __init__(self, vllm_config, role, kv_cache_config)
+"""
+
+from typing import TYPE_CHECKING
+from unittest.mock import patch
+
+import pytest
+
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.distributed.kv_transfer.kv_connector.v1 import (
+    KVConnectorBase_V1,
+    KVConnectorRole,
+)
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+
+from .utils import create_scheduler, create_vllm_config
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+
+class OldStyleTestConnector(KVConnectorBase_V1):
+    """
+    Test connector using the old signature with 2 required arguments.
+    This simulates external connectors that haven't been updated yet.
+    """
+
+    def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
+        # Old-style call to super().__init__ with only 2 arguments
+        super().__init__(vllm_config=vllm_config, role=role)
+
+    def get_num_new_matched_tokens(
+        self, request: "Request", num_computed_tokens: int
+    ) -> tuple[int | None, bool]:
+        return 0, False
+
+    def update_state_after_alloc(
+        self,
+        request: "Request",
+        blocks: "KVCacheBlocks",
+        num_external_tokens: int,
+    ):
+        pass
+
+    def build_connector_meta(self, scheduler_output: SchedulerOutput):
+        return None
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer,
+        attn_metadata: AttentionMetadata,
+        **kwargs,
+    ) -> None:
+        pass
+
+    def wait_for_save(self):
+        pass
+
+
+class NewStyleTestConnector(KVConnectorBase_V1):
+    """
+    Test connector using the new signature with 3 required arguments.
+    """
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        # New-style call to super().__init__ with all 3 arguments
+        super().__init__(
+            vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
+        )
+
+    def get_num_new_matched_tokens(
+        self, request: "Request", num_computed_tokens: int
+    ) -> tuple[int | None, bool]:
+        return 0, False
+
+    def update_state_after_alloc(
+        self,
+        request: "Request",
+        blocks: "KVCacheBlocks",
+        num_external_tokens: int,
+    ):
+        pass
+
+    def build_connector_meta(self, scheduler_output: SchedulerOutput):
+        return None
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer,
+        attn_metadata: AttentionMetadata,
+        **kwargs,
+    ) -> None:
+        pass
+
+    def wait_for_save(self):
+        pass
+
+
+@pytest.mark.parametrize("role", [KVConnectorRole.SCHEDULER, KVConnectorRole.WORKER])
+def test_external_old_signature_factory_instantiation(role):
+    """
+    Test that external connectors with old signature (2 required args) loaded
+    via kv_connector_module_path are correctly instantiated with backwards
+    compatibility support.
+    """
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_connector = "OldStyleTestConnector"
+    vllm_config.kv_transfer_config.kv_connector_module_path = (
+        "tests.v1.kv_connector.unit.test_backwards_compatibility"
+    )
+
+    scheduler = create_scheduler(vllm_config)
+    kv_cache_config = scheduler.kv_cache_config
+
+    connector = KVConnectorFactory.create_connector(vllm_config, role, kv_cache_config)
+
+    assert connector is not None
+    assert isinstance(connector, OldStyleTestConnector)
+    assert connector.role == role
+    assert connector._kv_cache_config is None
+
+
+@pytest.mark.parametrize("role", [KVConnectorRole.SCHEDULER, KVConnectorRole.WORKER])
+def test_external_new_signature_factory_instantiation(role):
+    """
+    Test that external connectors with new signature (3 required args) loaded
+    via kv_connector_module_path are correctly instantiated.
+    """
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_connector = "NewStyleTestConnector"
+    vllm_config.kv_transfer_config.kv_connector_module_path = (
+        "tests.v1.kv_connector.unit.test_backwards_compatibility"
+    )
+
+    scheduler = create_scheduler(vllm_config)
+    kv_cache_config = scheduler.kv_cache_config
+
+    connector = KVConnectorFactory.create_connector(vllm_config, role, kv_cache_config)
+
+    assert connector is not None
+    assert isinstance(connector, NewStyleTestConnector)
+    assert connector.role == role
+    assert connector._kv_cache_config is not None
+    assert connector._kv_cache_config == kv_cache_config
+
+
+@pytest.mark.parametrize("role", [KVConnectorRole.SCHEDULER, KVConnectorRole.WORKER])
+def test_old_signature_super_init(role):
+    """
+    Test that old-style connectors can call super().__init__() without
+    kv_cache_config parameter.
+    """
+    vllm_config = create_vllm_config()
+
+    connector = OldStyleTestConnector(vllm_config, role)
+
+    assert connector is not None
+    assert connector.role == role
+    assert connector._kv_cache_config is None
+
+
+def test_old_signature_super_init_with_kwargs():
+    """
+    Test that old-style connectors can call super().__init__() with keyword
+    arguments in different orders.
+    """
+    vllm_config = create_vllm_config()
+
+    # Test with vllm_config= and role= kwargs
+    connector1 = OldStyleTestConnector(
+        vllm_config=vllm_config, role=KVConnectorRole.SCHEDULER
+    )
+    assert connector1 is not None
+    assert connector1._kv_cache_config is None
+
+    # Test with role= and vllm_config= in reversed order
+    connector2 = OldStyleTestConnector(
+        role=KVConnectorRole.WORKER, vllm_config=vllm_config
+    )
+    assert connector2 is not None
+    assert connector2._kv_cache_config is None
+
+
+def test_internal_connector_uses_new_signature():
+    """
+    Test that internal connectors (registered in factory) always use the new
+    signature and get kv_cache_config.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (
+        ExampleConnector,
+    )
+
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_connector = "ExampleConnector"
+
+    scheduler = create_scheduler(vllm_config)
+    kv_cache_config = scheduler.kv_cache_config
+
+    connector = KVConnectorFactory.create_connector(
+        vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+    )
+
+    assert connector is not None
+    assert isinstance(connector, ExampleConnector)
+    assert connector._kv_cache_config is not None
+    assert connector._kv_cache_config == kv_cache_config
+
+
+def test_signature_detection_with_mocking():
+    """
+    Test that the factory correctly applies compat_sig flag returned from
+    _get_connector_class_with_compat.
+    """
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+    kv_cache_config = scheduler.kv_cache_config
+
+    # Mock _get_connector_class_with_compat to return old-style connector
+    with patch.object(
+        KVConnectorFactory,
+        "_get_connector_class_with_compat",
+        return_value=(OldStyleTestConnector, True),
+    ):
+        old_connector = KVConnectorFactory.create_connector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )
+        assert old_connector is not None
+        assert isinstance(old_connector, OldStyleTestConnector)
+        assert old_connector._kv_cache_config is None
+
+    # Mock _get_connector_class_with_compat to return new-style connector
+    with patch.object(
+        KVConnectorFactory,
+        "_get_connector_class_with_compat",
+        return_value=(NewStyleTestConnector, False),
+    ):
+        new_connector = KVConnectorFactory.create_connector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )
+        assert new_connector is not None
+        assert isinstance(new_connector, NewStyleTestConnector)
+        assert new_connector._kv_cache_config is not None
+        assert new_connector._kv_cache_config == kv_cache_config

third_party/vllm/tests/v1/kv_connector/unit/test_cache_pollution_prevention.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+test that invalid blocks are evicted from prefix cache to prevent pollution.
+
+verifies that when sync-loading fails, invalid blocks are removed from the
+prefix cache hash table so future requests cannot match and reuse corrupted data.
+"""
+
+from collections.abc import Callable
+from unittest.mock import Mock
+
+import pytest
+
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.request import Request, RequestStatus
+
+from .utils import (
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def _make_get_num_new_matched_tokens(
+    req_num_new_matched_tokens: dict[str, int],
+    async_load: bool,
+) -> Callable[[Request, int], tuple[int, bool]]:
+    def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
+        value = req_num_new_matched_tokens.get(request.request_id, 0)
+        return value, async_load
+
+    return get_num_new_matched_tokens
+
+
+@pytest.fixture
+def fail_scheduler():
+    """scheduler with kv_load_failure_policy='fail'"""
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
+    return create_scheduler(vllm_config)
+
+
+def test_invalid_blocks_evicted_prevents_cache_pollution(
+    fail_scheduler: Scheduler,
+):
+    """
+    verify invalid blocks are evicted to prevent future cache hits.
+
+    scenario:
+    1. request 1 loads externally-computed blocks (sync mode)
+    2. some blocks fail to load and are marked invalid
+    3. with fail policy, invalid blocks should be evicted from prefix cache
+    4. request is marked as FINISHED_ERROR
+    """
+    num_prompt_blocks = 100
+    num_external_computed_blocks = 99
+    invalid_block_idx = 50
+
+    num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
+    num_external_computed_tokens = (
+        num_external_computed_blocks * fail_scheduler.block_size
+    )
+
+    # request 1: will have invalid blocks
+    request1 = create_request(num_tokens=num_prompt_tokens, request_id=1)
+    fail_scheduler.add_request(request=request1)
+
+    req_num_new_matched_tokens = {
+        request1.request_id: num_external_computed_tokens,
+    }
+
+    # mock connector indicating sync load
+    fail_scheduler.connector = Mock()
+    fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
+    )
+    fail_scheduler.connector.request_finished.return_value = (False, None)
+    fail_scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = fail_scheduler.schedule()
+
+    # request should be running with sync KV load
+    assert len(fail_scheduler.running) == 1
+    assert request1.status == RequestStatus.RUNNING
+
+    # get allocated block IDs
+    req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
+    invalid_block_id = req_block_ids[invalid_block_idx]
+    invalid_block_ids = {invalid_block_id}
+
+    # get the block object to verify eviction later
+    block = fail_scheduler.kv_cache_manager.block_pool.blocks[invalid_block_id]
+
+    # cache the blocks to simulate they've been computed and cached
+    # (in real scenario blocks would be cached after compute)
+    fail_scheduler.kv_cache_manager.cache_blocks(request1, num_external_computed_tokens)
+
+    # verify block has a hash (is cached) before reporting invalid blocks
+    assert block.block_hash is not None, (
+        f"block {invalid_block_id} should be cached (have a hash) before "
+        f"eviction test, but hash is None"
+    )
+
+    # report invalid blocks
+    model_runner_output = create_model_runner_output(
+        [request1],
+        invalid_block_ids=invalid_block_ids,
+        use_eos=False,
+    )
+
+    fail_scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # verify request finished with error (fail policy)
+    assert request1.status == RequestStatus.FINISHED_ERROR
+
+    # critical assertion: invalid block and all subsequent blocks should be evicted
+    # all blocks from invalid_block_idx onwards become invalid since they were
+    # computed based on the failed block
+    for idx in range(invalid_block_idx, len(req_block_ids)):
+        block_id = req_block_ids[idx]
+        block_obj = fail_scheduler.kv_cache_manager.block_pool.blocks[block_id]
+        assert block_obj.block_hash is None, (
+            f"block {block_id} at index {idx} should have been evicted "
+            f"(hash reset to None), but hash is {block_obj.block_hash}. "
+            f"All blocks from index {invalid_block_idx} onwards should be evicted "
+            f"since they depend on the invalid block at index {invalid_block_idx}."
+        )
+
+    # verify cache contains exactly the valid blocks (before first affected block)
+    # and none of the invalid blocks (from first affected block onwards)
+
+    # valid blocks: all blocks before invalid_block_idx should be cached
+    for idx in range(invalid_block_idx):
+        block_id = req_block_ids[idx]
+        block_obj = fail_scheduler.kv_cache_manager.block_pool.blocks[block_id]
+        assert block_obj.block_hash is not None, (
+            f"valid block {block_id} at index {idx} should still be cached "
+            f"(have a hash), but hash is None. Only blocks from index "
+            f"{invalid_block_idx} onwards should be evicted."
+        )
+
+    # invalid blocks: verify they're not in the cached_block_hash_to_block map
+    cached_blocks = (
+        fail_scheduler.kv_cache_manager.block_pool.cached_block_hash_to_block
+    )
+    cached_block_ids = {
+        b.block_id
+        for blocks_val in cached_blocks._cache.values()
+        for b in (
+            [blocks_val] if not isinstance(blocks_val, dict) else blocks_val.values()
+        )
+    }
+
+    for idx in range(invalid_block_idx, len(req_block_ids)):
+        block_id = req_block_ids[idx]
+        assert block_id not in cached_block_ids, (
+            f"invalid block {block_id} at index {idx} should not be in cache hash table"
+        )

third_party/vllm/tests/v1/kv_connector/unit/test_config.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for KV cache offloading configuration."""
+
+import pytest
+
+from vllm.config import CacheConfig, KVTransferConfig, ParallelConfig, VllmConfig
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.mark.parametrize(
+    "kv_offloading_backend,kv_offloading_size,tp,pp,expected_backend,expected_bytes",
+    [
+        ("native", 4.0, 1, 1, "OffloadingConnector", 4.0 * (1 << 30)),
+        # bytes per rank: 8.0 GiB / (2 * 2) = 2.0 GiB
+        ("native", 8.0, 2, 2, "OffloadingConnector", 8.0 * (1 << 30)),
+        ("lmcache", 4.0, 1, 1, "LMCacheConnectorV1", 4.0),
+        # size per rank: 8.0 GiB / (2 * 2) = 2.0 GiB
+        ("lmcache", 8.0, 2, 2, "LMCacheConnectorV1", 2.0),
+        # When kv_offloading_size is None, offloading is disabled (backend is ignored)
+        ("native", None, 1, 1, None, None),
+    ],
+)
+def test_kv_connector(
+    kv_offloading_backend, kv_offloading_size, tp, pp, expected_backend, expected_bytes
+):
+    kv_transfer_config = (
+        KVTransferConfig(kv_connector_extra_config={"existing_key": "existing_value"})
+        if expected_backend is not None
+        else None
+    )
+
+    vllm_config = VllmConfig(
+        cache_config=CacheConfig(
+            kv_offloading_backend=kv_offloading_backend,
+            kv_offloading_size=kv_offloading_size,
+        ),
+        kv_transfer_config=kv_transfer_config,
+        parallel_config=ParallelConfig(
+            tensor_parallel_size=tp, pipeline_parallel_size=pp
+        ),
+    )
+
+    # No KV transfer config expected
+    if expected_backend is None:
+        assert vllm_config.kv_transfer_config is expected_backend
+        return
+
+    kv_transfer_config = vllm_config.kv_transfer_config
+    kv_connector_extra_config = kv_transfer_config.kv_connector_extra_config
+
+    assert kv_transfer_config.kv_connector == expected_backend
+    assert kv_transfer_config.kv_role == "kv_both"
+
+    if kv_offloading_backend == "native":
+        assert kv_connector_extra_config["cpu_bytes_to_use"] == expected_bytes
+        # Existing config should be preserved
+        assert kv_connector_extra_config["existing_key"] == "existing_value"
+    elif kv_offloading_backend == "lmcache":
+        assert kv_connector_extra_config["lmcache.local_cpu"] is True
+        assert kv_connector_extra_config["lmcache.max_local_cpu_size"] == expected_bytes
+        # Existing config should be replaced
+        assert "existing_key" not in kv_connector_extra_config
+
+
+def test_kv_offloading_size_only_uses_native_default():
+    """Test that setting only kv_offloading_size enables native offloading."""
+    vllm_config = VllmConfig(
+        cache_config=CacheConfig(
+            kv_offloading_size=4.0,
+            # kv_offloading_backend not set, should default to "native"
+        ),
+    )
+
+    kv_transfer_config = vllm_config.kv_transfer_config
+    kv_connector_extra_config = kv_transfer_config.kv_connector_extra_config
+    assert kv_transfer_config.kv_connector == "OffloadingConnector"
+    assert kv_transfer_config.kv_role == "kv_both"
+    assert kv_connector_extra_config["cpu_bytes_to_use"] == 4.0 * (1 << 30)

third_party/vllm/tests/v1/kv_connector/unit/test_decode_bench_connector.py

@@ -0,0 +1,417 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for DecodeBenchConnector.
+
+Tests the functionality of the DecodeBenchConnector which fills KV cache
+with dummy values for decode performance benchmarking.
+"""
+
+import pytest
+import torch
+
+from vllm import SamplingParams
+from vllm.config import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
+
+# ruff: noqa: E501
+from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import (
+    DecodeBenchConnector,
+    DecodeBenchConnectorMetadata,
+)
+from vllm.forward_context import ForwardContext
+from vllm.utils.hashing import sha256
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.request import Request
+
+from .utils import (
+    EOS_TOKEN_ID,
+    create_model_runner_output,
+    create_scheduler,
+    create_vllm_config,
+)
+
+
+class DecodeBenchTestRunner:
+    """Test runner for DecodeBenchConnector."""
+
+    def __init__(self, block_size: int, num_gpu_blocks: int):
+        self.block_size = block_size
+        self.num_gpu_blocks = num_gpu_blocks
+
+        self.req_id = -1
+
+        # Create vllm config with DecodeBenchConnector
+        vllm_config = create_vllm_config(
+            block_size=block_size, max_num_batched_tokens=1000
+        )
+        vllm_config.kv_transfer_config = KVTransferConfig(
+            kv_connector="DecodeBenchConnector",
+            kv_role="kv_both",
+        )
+
+        self.vllm_config = vllm_config
+        self.scheduler: Scheduler = create_scheduler(
+            vllm_config, num_blocks=num_gpu_blocks
+        )
+
+        # Create worker-side connector
+        self.worker_connector = DecodeBenchConnector(
+            vllm_config, KVConnectorRole.WORKER
+        )
+
+        # Create dummy KV caches for testing
+        # Shape: [num_blocks, 2, num_heads, block_size, head_dim]
+        # Using simplified shape for testing
+        num_heads = 4
+        head_dim = 64
+        self.kv_caches = {
+            f"layer_{i}": torch.zeros(
+                num_gpu_blocks, 2, num_heads, block_size, head_dim
+            )
+            for i in range(2)  # 2 layers for testing
+        }
+
+        # Register KV caches with worker connector
+        self.worker_connector.register_kv_caches(self.kv_caches)
+
+        # Extract scheduler-side connector
+        scheduler_connector = self.scheduler.connector
+        assert scheduler_connector is not None
+        assert isinstance(scheduler_connector, DecodeBenchConnector)
+        self.scheduler_connector: DecodeBenchConnector = scheduler_connector
+
+        init_none_hash(sha256)
+        self._block_hasher = get_request_block_hasher(block_size, sha256)
+
+        self._dummy_ctx: ForwardContext = ForwardContext(
+            no_compile_layers={}, attn_metadata={}, virtual_engine=0, slot_mapping={}
+        )
+
+    def new_request(self, token_ids: list[int]) -> Request:
+        """Create a new request with given token IDs."""
+        self.req_id += 1
+
+        sampling_params = SamplingParams(max_tokens=100)
+        sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
+
+        req = Request(
+            request_id=str(self.req_id),
+            prompt_token_ids=token_ids,
+            sampling_params=sampling_params,
+            pooling_params=None,
+            block_hasher=self._block_hasher,
+        )
+
+        self.scheduler.add_request(req)
+        return req
+
+    def run_single_step(self, token_id: int = 0):
+        """Run a single scheduler + worker step."""
+        scheduler_output = self.scheduler.schedule()
+
+        # Get connector metadata
+        kv_connector_metadata = scheduler_output.kv_connector_metadata
+        assert kv_connector_metadata is not None
+        assert isinstance(kv_connector_metadata, DecodeBenchConnectorMetadata)
+
+        # Bind metadata and load KV
+        self.worker_connector.bind_connector_metadata(kv_connector_metadata)
+        self.worker_connector.start_load_kv(self._dummy_ctx)
+
+        if scheduler_output.total_num_scheduled_tokens > 0:
+            self.worker_connector.wait_for_save()
+
+        self.worker_connector.clear_connector_metadata()
+
+        # Create model runner output
+        model_runner_output = create_model_runner_output(
+            reqs=self.scheduler.running,
+            token_id=token_id,
+        )
+
+        self.scheduler.update_from_output(scheduler_output, model_runner_output)
+
+        return scheduler_output, kv_connector_metadata
+
+
+def test_decode_bench_connector_basic():
+    """Test basic functionality of DecodeBenchConnector."""
+    block_size = 16
+    num_gpu_blocks = 100
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # Create a request with multiple blocks worth of tokens
+    num_tokens = block_size * 3  # 3 blocks
+    token_ids = [1] * num_tokens
+
+    req = runner.new_request(token_ids)
+
+    # Run first step - should fill KV cache with dummy values
+    scheduler_output, metadata = runner.run_single_step()
+
+    # Check that get_num_new_matched_tokens returned correct value
+    # Should be num_tokens - 1 (all except the last token for decode)
+    expected_fill_tokens = num_tokens - 1
+
+    # Check metadata has the request to fill
+    assert len(metadata.reqs_to_fill) == 1
+    assert req.request_id in metadata.reqs_to_fill
+
+    block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
+    assert num_tokens_to_fill == expected_fill_tokens
+
+    # For standard attention, there's only one group
+    assert len(block_ids_per_group) == 1
+    block_ids = block_ids_per_group[0]
+
+    # Calculate expected number of blocks
+    expected_num_blocks = (expected_fill_tokens + block_size - 1) // block_size
+    assert len(block_ids) == expected_num_blocks
+
+    # Verify KV caches were filled with constant value
+    for layer_name, kv_cache in runner.kv_caches.items():
+        for block_id in block_ids:
+            # Check that the block was filled
+            block_data = kv_cache[block_id]
+            # Should be filled with constant value 0.015
+            assert torch.allclose(block_data, torch.tensor(0.015))
+
+
+def test_decode_bench_connector_no_refill():
+    """Test that DecodeBenchConnector only fills once per request."""
+    block_size = 16
+    num_gpu_blocks = 100
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # Create a request
+    num_tokens = block_size * 2
+    token_ids = [1] * num_tokens
+
+    runner.new_request(token_ids)
+
+    # Run first step - should fill KV cache
+    _, metadata1 = runner.run_single_step()
+    assert len(metadata1.reqs_to_fill) == 1
+
+    # Run second step - should NOT fill again (already filled)
+    _, metadata2 = runner.run_single_step()
+    assert len(metadata2.reqs_to_fill) == 0
+
+
+def test_decode_bench_connector_single_token():
+    """Test DecodeBenchConnector with single token request."""
+    block_size = 16
+    num_gpu_blocks = 100
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # Create a request with just 1 token
+    # Should not fill anything (need at least 2 tokens: 1 to fill, 1 to decode)
+    token_ids = [1]
+
+    runner.new_request(token_ids)
+
+    # Run step - should NOT fill KV cache
+    _, metadata = runner.run_single_step()
+    assert len(metadata.reqs_to_fill) == 0
+
+
+def test_decode_bench_connector_two_tokens():
+    """Test DecodeBenchConnector with two token request."""
+    block_size = 16
+    num_gpu_blocks = 100
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # Create a request with 2 tokens
+    # Should fill 1 token (first token), decode the second
+    token_ids = [1, 2]
+
+    req = runner.new_request(token_ids)
+
+    # Run step
+    _, metadata = runner.run_single_step()
+
+    assert len(metadata.reqs_to_fill) == 1
+    assert req.request_id in metadata.reqs_to_fill
+
+    block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
+    assert num_tokens_to_fill == 1
+    # For standard attention, there's only one group
+    assert len(block_ids_per_group) == 1
+    assert len(block_ids_per_group[0]) == 1  # 1 token needs 1 block
+
+
+def test_decode_bench_connector_large_context():
+    """Test DecodeBenchConnector with large context size."""
+    block_size = 16
+    num_gpu_blocks = 1000
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # Create a request with many blocks
+    num_blocks = 20
+    num_tokens = block_size * num_blocks
+    token_ids = list(range(num_tokens))
+
+    req = runner.new_request(token_ids)
+
+    # Run step
+    _, metadata = runner.run_single_step()
+
+    assert len(metadata.reqs_to_fill) == 1
+    assert req.request_id in metadata.reqs_to_fill
+
+    block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
+
+    # Should fill all tokens except the last one
+    expected_fill_tokens = num_tokens - 1
+    assert num_tokens_to_fill == expected_fill_tokens
+
+    # For standard attention, there's only one group
+    assert len(block_ids_per_group) == 1
+    block_ids = block_ids_per_group[0]
+
+    # Calculate expected number of blocks
+    expected_num_blocks = (expected_fill_tokens + block_size - 1) // block_size
+    assert len(block_ids) == expected_num_blocks
+
+    # Verify blocks were filled
+    for layer_name, kv_cache in runner.kv_caches.items():
+        for block_id in block_ids:
+            block_data = kv_cache[block_id]
+            assert torch.allclose(block_data, torch.tensor(0.015))
+
+
+def test_decode_bench_connector_multiple_requests():
+    """Test DecodeBenchConnector with multiple sequential requests."""
+    block_size = 16
+    num_gpu_blocks = 100
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # First request
+    req1 = runner.new_request([1] * (block_size * 2))
+    _, metadata1 = runner.run_single_step()
+
+    assert len(metadata1.reqs_to_fill) == 1
+    assert req1.request_id in metadata1.reqs_to_fill
+
+    # Complete first request
+    while runner.scheduler.running:
+        runner.run_single_step()
+
+    # Add EOS to finish
+    scheduler_output = runner.scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=runner.scheduler.running,
+        token_id=EOS_TOKEN_ID,
+        use_eos=True,
+    )
+    runner.scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # Second request - should also get filled
+    req2 = runner.new_request([2] * (block_size * 3))
+    _, metadata2 = runner.run_single_step()
+
+    assert len(metadata2.reqs_to_fill) == 1
+    assert req2.request_id in metadata2.reqs_to_fill
+
+    # Different request should have different metadata
+    _, num_tokens1 = metadata1.reqs_to_fill[req1.request_id]
+    _, num_tokens2 = metadata2.reqs_to_fill[req2.request_id]
+
+    assert num_tokens1 == block_size * 2 - 1
+    assert num_tokens2 == block_size * 3 - 1
+
+
+def test_decode_bench_connector_partial_block():
+    """Test DecodeBenchConnector with partial block filling."""
+    block_size = 16
+    num_gpu_blocks = 100
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # Create a request that doesn't align to block boundaries
+    # e.g., 2.5 blocks worth of tokens
+    num_tokens = block_size * 2 + block_size // 2
+    token_ids = [1] * num_tokens
+
+    req = runner.new_request(token_ids)
+
+    # Run step
+    _, metadata = runner.run_single_step()
+
+    assert len(metadata.reqs_to_fill) == 1
+    assert req.request_id in metadata.reqs_to_fill
+
+    block_ids_per_group, num_tokens_to_fill = metadata.reqs_to_fill[req.request_id]
+
+    # Should fill all tokens except the last one
+    expected_fill_tokens = num_tokens - 1
+    assert num_tokens_to_fill == expected_fill_tokens
+
+    # For standard attention, there's only one group
+    assert len(block_ids_per_group) == 1
+    block_ids = block_ids_per_group[0]
+
+    # Should allocate 3 blocks to hold the partial data
+    expected_num_blocks = 3
+    assert len(block_ids) == expected_num_blocks
+
+
+def test_decode_bench_connector_concurrent_requests():
+    """Test DecodeBenchConnector with multiple concurrent requests in the same batch."""
+    block_size = 16
+    num_gpu_blocks = 1000
+
+    runner = DecodeBenchTestRunner(block_size=block_size, num_gpu_blocks=num_gpu_blocks)
+
+    # Create multiple requests that will be batched together
+    req1 = runner.new_request([1] * (block_size * 2))
+    req2 = runner.new_request([2] * (block_size * 3))
+    req3 = runner.new_request([3] * (block_size * 1))
+
+    # Run first step - all requests should be filled concurrently
+    _, metadata = runner.run_single_step()
+
+    # All three requests should be in the metadata
+    assert len(metadata.reqs_to_fill) == 3
+    assert req1.request_id in metadata.reqs_to_fill
+    assert req2.request_id in metadata.reqs_to_fill
+    assert req3.request_id in metadata.reqs_to_fill
+
+    # Verify each request has correct fill info
+    block_ids_per_group1, num_tokens1 = metadata.reqs_to_fill[req1.request_id]
+    block_ids_per_group2, num_tokens2 = metadata.reqs_to_fill[req2.request_id]
+    block_ids_per_group3, num_tokens3 = metadata.reqs_to_fill[req3.request_id]
+
+    # Verify token counts (all tokens except last one)
+    assert num_tokens1 == block_size * 2 - 1
+    assert num_tokens2 == block_size * 3 - 1
+    assert num_tokens3 == block_size * 1 - 1
+
+    # Verify block counts for each request
+    assert len(block_ids_per_group1[0]) == 2  # 2 blocks
+    assert len(block_ids_per_group2[0]) == 3  # 3 blocks
+    assert len(block_ids_per_group3[0]) == 1  # 1 block
+
+    # Verify all blocks are filled in KV cache
+    for req_id, (block_ids_per_group, _) in metadata.reqs_to_fill.items():
+        block_ids = block_ids_per_group[0]
+        for layer_name, kv_cache in runner.kv_caches.items():
+            for block_id in block_ids:
+                block_data = kv_cache[block_id]
+                assert torch.allclose(block_data, torch.tensor(0.015))
+
+    # Run second step - should NOT fill again (already filled)
+    _, metadata2 = runner.run_single_step()
+    assert len(metadata2.reqs_to_fill) == 0
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

third_party/vllm/tests/v1/kv_connector/unit/test_error_propagation.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from unittest.mock import Mock
+
+import pytest
+
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.request import FinishReason, Request, RequestStatus
+
+from .utils import (
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def _make_get_num_new_matched_tokens(
+    req_num_new_matched_tokens: dict[str, int],
+    async_load: bool,
+) -> Callable[[Request, int], tuple[int, bool]]:
+    def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
+        value = req_num_new_matched_tokens.get(request.request_id, 0)
+        return value, async_load
+
+    return get_num_new_matched_tokens
+
+
+@pytest.fixture
+def fail_scheduler():
+    """scheduler with kv_load_failure_policy='fail'"""
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
+    return create_scheduler(vllm_config)
+
+
+def test_error_propagation_sync_load(fail_scheduler: Scheduler):
+    """test invalid_block_ids with fail policy -> FINISHED_ERROR (sync load)"""
+    num_prompt_blocks = 100
+    num_external_computed_blocks = 99
+    invalid_block_idx = 50
+
+    num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
+    num_external_computed_tokens = (
+        num_external_computed_blocks * fail_scheduler.block_size
+    )
+
+    request = create_request(num_tokens=num_prompt_tokens)
+    fail_scheduler.add_request(request=request)
+
+    req_num_new_matched_tokens = {
+        request.request_id: num_external_computed_tokens,
+    }
+
+    fail_scheduler.connector = Mock()
+    fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
+    )
+    fail_scheduler.connector.request_finished.return_value = (False, None)
+    fail_scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = fail_scheduler.schedule()
+
+    assert len(fail_scheduler.running) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 1
+    assert fail_scheduler.connector.get_num_new_matched_tokens.call_count == 1
+
+    req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
+    invalid_block_ids = {req_block_ids[invalid_block_idx]}
+    model_runner_output = create_model_runner_output(
+        [request],
+        invalid_block_ids=invalid_block_ids,
+        use_eos=True,
+    )
+
+    outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    assert request.status == RequestStatus.FINISHED_ERROR
+    assert request.get_finished_reason() == FinishReason.ERROR
+
+    assert len(outputs) == 1
+    engine_outputs = next(iter(outputs.values()))
+    assert len(engine_outputs.outputs) == 1
+    output = engine_outputs.outputs[0]
+    assert output.request_id == request.request_id
+    assert output.finish_reason == FinishReason.ERROR
+
+    assert len(fail_scheduler.running) == 0
+
+
+def test_error_propagation_async_load(fail_scheduler: Scheduler):
+    """test invalid_block_ids with fail policy -> FINISHED_ERROR (async load)"""
+    num_prompt_blocks = 100
+    num_external_computed_blocks = 99
+    invalid_block_idx = 50
+
+    num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
+    num_external_computed_tokens = (
+        num_external_computed_blocks * fail_scheduler.block_size
+    )
+
+    request = create_request(num_tokens=num_prompt_tokens)
+    fail_scheduler.add_request(request=request)
+
+    req_num_new_matched_tokens = {
+        request.request_id: num_external_computed_tokens,
+    }
+
+    fail_scheduler.connector = Mock()
+    fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True)
+    )
+    fail_scheduler.connector.request_finished.return_value = (False, None)
+    fail_scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = fail_scheduler.schedule()
+
+    assert len(fail_scheduler.skipped_waiting) == 1
+    assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert request.num_computed_tokens == num_external_computed_tokens
+
+    (req_block_ids,) = fail_scheduler.kv_cache_manager.get_block_ids(request.request_id)
+    invalid_block_ids = {req_block_ids[invalid_block_idx]}
+    model_runner_output = create_model_runner_output(
+        reqs=[],
+        finished_recving=set(),
+        invalid_block_ids=invalid_block_ids,
+        use_eos=True,
+    )
+
+    outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    assert request.status == RequestStatus.FINISHED_ERROR
+    assert request.get_finished_reason() == FinishReason.ERROR
+
+    assert len(outputs) == 1
+    engine_outputs = next(iter(outputs.values()))
+    assert len(engine_outputs.outputs) == 1
+    output = engine_outputs.outputs[0]
+    assert output.request_id == request.request_id
+    assert output.finish_reason == FinishReason.ERROR
+
+    assert len(fail_scheduler.waiting) == 0
+    assert len(fail_scheduler.skipped_waiting) == 0

third_party/vllm/tests/v1/kv_connector/unit/test_example_connector.py

@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import asdict
+from typing import NamedTuple
+
+import pytest
+from PIL import Image
+
+from vllm import LLM, EngineArgs, SamplingParams
+from vllm.assets.image import ImageAsset
+from vllm.config import AttentionConfig, KVTransferConfig
+from vllm.multimodal.utils import encode_image_url
+from vllm.platforms import current_platform
+
+MODEL_NAME = "RedHatAI/Qwen2.5-VL-3B-Instruct-quantized.w8a8"
+
+SAMPLING_PARAMS = SamplingParams(temperature=0.0, top_k=1, max_tokens=128)
+
+TEXT_PROMPTS = [
+    "What's in the image(s)? Around 30 words. What's special in 2nd image?",
+    "The future of AI is",
+]
+
+
+class InputCase(NamedTuple):
+    text: str
+    img: list[Image]
+    expected_len: int
+    info: str
+
+
+def _check_path_len(path):
+    """Return the latest length in path"""
+    return len(list(path.iterdir()))
+
+
+def _list_path(path):
+    """Return the list of foldername (hashes generated) under the path"""
+    return list(path.iterdir())
+
+
+def run_test(
+    tmp_path,
+    processor,
+    llm: LLM,
+    question: str,
+    image_urls: list[Image],
+    expected_len: int,
+    info: str,
+):
+    """
+    One individual test to process the prompt and output base on 1 set of input
+    Then check if the length in the storage path matches the expected length
+    `info` introduces details or purpose of the individual test
+    """
+    print(f"***info: {info}***")
+    print(f"**Expected storage path length after llm generate: {expected_len}**")
+    process_prompt(processor, llm, question, image_urls)
+
+    print(f"Path matched expected length: {_check_path_len(tmp_path)}")
+    print(f"Hashes under the storage path: {_list_path(tmp_path)}")
+
+    assert _check_path_len(tmp_path) == expected_len, (
+        f"Expect storage path length {expected_len} ;",
+        f"but end up {_check_path_len(tmp_path)} instead. ",
+        f"Info: {info}",
+    )
+
+
+def process_prompt(processor, llm: LLM, question: str, image_urls: list[Image]):
+    """
+    Form the prompt based on the text and image input, then llm generate output
+    """
+    placeholders = [
+        {
+            "type": "image_url",
+            "image_url": {"url": encode_image_url(image_pil)},
+        }
+        for image_pil in image_urls
+    ]
+
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {
+            "role": "user",
+            "content": [
+                *placeholders,
+                {"type": "text", "text": question},
+            ],
+        },
+    ]
+
+    prompt = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+
+    outputs = llm.generate(
+        {
+            "prompt": prompt,
+            **({"multi_modal_data": {"image": [*image_urls]}} if image_urls else {}),
+        },
+        sampling_params=SAMPLING_PARAMS,
+    )
+
+    print("-" * 50)
+    print("Output:")
+    for o in outputs:
+        generated_text = o.outputs[0].text
+        print(generated_text)
+        print("-" * 50)
+
+
+@pytest.mark.parametrize(
+    "attn_backend",
+    (
+        ["FLASH_ATTN", "TRITON_ATTN"]
+        if current_platform.is_cuda()
+        else ["TRITON_ATTN"]
+        if current_platform.is_rocm()
+        else []
+    ),
+)
+def test_shared_storage_connector_hashes(tmp_path, attn_backend):
+    """
+    Tests that ExampleConnector saves KV to the storage locations
+    with proper hashes; that are unique for inputs with identical text but
+    different images (same size), or same multiple images but different orders.
+    """
+    # Using tmp_path as the storage path to store KV
+    print(f"KV storage path at: {str(tmp_path)}")
+
+    # Configure the ExampleConnector
+    kv_transfer_config = KVTransferConfig(
+        kv_connector="ExampleConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={"shared_storage_path": str(tmp_path)},
+    )
+
+    engine_args = EngineArgs(
+        model=MODEL_NAME,
+        max_model_len=8192,
+        max_num_seqs=1,
+        gpu_memory_utilization=0.4,
+        attention_config=AttentionConfig(backend=attn_backend),
+        enforce_eager=True,
+        kv_transfer_config=kv_transfer_config,
+        limit_mm_per_prompt={"image": 2},
+    )
+
+    # don't put this import at the top level
+    # it will call torch.accelerator.device_count()
+    from transformers import AutoProcessor
+
+    # Create processor to handle the chat prompt
+    processor = AutoProcessor.from_pretrained(MODEL_NAME)
+
+    # Prepare images for the tests
+    # Resize to the same size to check hashes correctness
+    image_1 = ImageAsset("stop_sign").pil_image.resize((1280, 720))
+    image_2 = ImageAsset("cherry_blossom").pil_image.resize((1280, 720))
+
+    # Make sure that they are not the same picture
+    assert image_1 != image_2, "The images should not be identical"
+
+    # Create the LLM instance
+    engine_args = asdict(engine_args)
+    llm = LLM(**engine_args)
+
+    # Prepare the input cases
+    input_cases = [
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_1],
+            expected_len=1,
+            info="image_1 single input the first time.",
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_2],
+            expected_len=2,
+            info=(
+                "image_2 single input the first time. "
+                "It is in same pixel size with image_1, yet it "
+                "should be able to form a new unique hash."
+            ),
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_1],
+            expected_len=2,
+            info=(
+                "image_1 single input the 2nd time. "
+                "It should not form another new hash."
+            ),
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_2],
+            expected_len=2,
+            info=(
+                "image_2 single input the 2nd time. "
+                "It should not form another new hash."
+            ),
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_1, image_2],
+            expected_len=3,
+            info="image_1 with image_2 input the first time.",
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_2, image_1],
+            expected_len=4,
+            info="The image order is swapped. Should form new hash.",
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_1, image_2],
+            expected_len=4,
+            info=(
+                "[image_1, image_2] input the 2nd time. "
+                "It should not form another new hash."
+            ),
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[image_2, image_1],
+            expected_len=4,
+            info=(
+                "[image_2, image_1] input the 2nd time. "
+                "It should not form another new hash."
+            ),
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[],
+            expected_len=5,
+            info="Pure text input test as a case-control",
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[0],
+            img=[],
+            expected_len=5,
+            info="Identical pure text input as a case-control",
+        ),
+        InputCase(
+            text=TEXT_PROMPTS[1],
+            img=[],
+            expected_len=6,
+            info="Another pure text input as a case-control",
+        ),
+    ]
+
+    # Run tests
+    for case_id, (text, img, expected_len, info) in enumerate(input_cases):
+        print("\n", "=" * 25, f"Below running input case: {case_id}", "=" * 25)
+        run_test(tmp_path, processor, llm, text, img, expected_len, info)
+
+    print("All tests passed successfully!")

third_party/vllm/tests/v1/kv_connector/unit/test_flexkv_connector.py

@@ -0,0 +1,232 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for FlexKVConnectorV1.
+
+These tests mock the ``flexkv`` package so they can run without a real FlexKV
+installation.  They verify:
+
+1. That ``FlexKVConnectorV1`` raises a helpful ``ImportError`` when FlexKV is
+   not installed.
+2. That all public methods are correctly delegated to the underlying
+   ``FlexKVConnectorV1Impl``.
+"""
+
+import sys
+import types
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm.config import KVTransferConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+from .utils import create_vllm_config
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_vllm_config(
+    kv_connector: str = "FlexKVConnectorV1",
+    kv_role: str = "kv_both",
+) -> VllmConfig:
+    """Return a minimal VllmConfig with a KVTransferConfig attached."""
+    vllm_config = create_vllm_config(block_size=16, max_num_batched_tokens=512)
+    vllm_config.kv_transfer_config = KVTransferConfig(
+        kv_connector=kv_connector,
+        kv_role=kv_role,
+    )
+    return vllm_config
+
+
+def _make_kv_cache_config() -> KVCacheConfig:
+    return MagicMock(spec=KVCacheConfig)
+
+
+def _make_flexkv_module(
+    impl_mock: MagicMock,
+) -> tuple[types.ModuleType, types.ModuleType]:
+    """Build a fake ``flexkv`` package hierarchy that returns *impl_mock*
+    when ``FlexKVConnectorV1Impl`` is instantiated."""
+    flexkv_mod = types.ModuleType("flexkv")
+    integration_mod = types.ModuleType("flexkv.integration")
+    vllm_mod = types.ModuleType("flexkv.integration.vllm")
+    adapter_mod = types.ModuleType("flexkv.integration.vllm.vllm_v1_adapter")
+
+    # Make FlexKVConnectorV1Impl() return our mock instance.
+    # The "# type: ignore" markers below are needed because ModuleType does
+    # not declare these attributes statically; they are set dynamically.
+    FlexKVConnectorV1ImplCls = MagicMock(return_value=impl_mock)
+    adapter_mod.FlexKVConnectorV1Impl = FlexKVConnectorV1ImplCls  # type: ignore
+
+    flexkv_mod.integration = integration_mod  # type: ignore
+    integration_mod.vllm = vllm_mod  # type: ignore
+    vllm_mod.vllm_v1_adapter = adapter_mod  # type: ignore
+
+    return flexkv_mod, adapter_mod
+
+
+def _install_flexkv_mock(impl_mock: MagicMock):
+    """Insert fake flexkv modules into sys.modules and return a context that
+    cleans them up afterwards."""
+    flexkv_mod, adapter_mod = _make_flexkv_module(impl_mock)
+    mods = {
+        "flexkv": flexkv_mod,
+        "flexkv.integration": flexkv_mod.integration,
+        "flexkv.integration.vllm": flexkv_mod.integration.vllm,
+        "flexkv.integration.vllm.vllm_v1_adapter": adapter_mod,
+    }
+    return patch.dict(sys.modules, mods)
+
+
+def _build_connector(vllm_config: VllmConfig, impl_mock: MagicMock):
+    """Instantiate FlexKVConnectorV1 with faked flexkv modules."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import (
+        FlexKVConnectorV1,
+    )
+
+    with _install_flexkv_mock(impl_mock):
+        connector = FlexKVConnectorV1(
+            vllm_config=vllm_config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=_make_kv_cache_config(),
+        )
+    return connector
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestFlexKVConnectorImportError:
+    """FlexKVConnectorV1 should fail with a helpful message when flexkv is
+    absent."""
+
+    def test_import_error_message(self):
+        from vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector import (
+            FlexKVConnectorV1,
+        )
+
+        # Ensure flexkv is NOT in sys.modules
+        for key in list(sys.modules):
+            if key.startswith("flexkv"):
+                del sys.modules[key]
+
+        with pytest.raises(ImportError, match="(?i)flexkv") as exc_info:
+            FlexKVConnectorV1(
+                vllm_config=_make_vllm_config(),
+                role=KVConnectorRole.WORKER,
+                kv_cache_config=_make_kv_cache_config(),
+            )
+
+        assert "https://github.com/taco-project/FlexKV" in str(exc_info.value)
+
+
+class TestFlexKVConnectorDelegation:
+    """All public API methods should be forwarded to the impl."""
+
+    @pytest.fixture()
+    def connector_and_impl(self):
+        impl = MagicMock()
+        cfg = _make_vllm_config()
+        connector = _build_connector(cfg, impl)
+        return connector, impl
+
+    def test_shutdown(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        connector.shutdown()
+        impl.shutdown.assert_called_once()
+
+    def test_start_load_kv(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        ctx = MagicMock()
+        connector.start_load_kv(ctx, extra_arg="x")
+        impl.start_load_kv.assert_called_once_with(ctx, extra_arg="x")
+
+    def test_save_kv_layer(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        kv_layer = torch.zeros(4, 4)
+        attn_meta = MagicMock()
+        connector.save_kv_layer("layer_0", kv_layer, attn_meta)
+        impl.save_kv_layer.assert_called_once_with("layer_0", kv_layer, attn_meta)
+
+    def test_wait_for_save(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        connector.wait_for_save()
+        impl.wait_for_save.assert_called_once()
+
+    def test_get_finished(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        impl.get_finished.return_value = ({"req1"}, None)
+        result = connector.get_finished({"req1"})
+        impl.get_finished.assert_called_once_with({"req1"})
+        assert result == ({"req1"}, None)
+
+    def test_register_kv_caches(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        kv_caches = {"layer_0": torch.zeros(1)}
+        connector.register_kv_caches(kv_caches)
+        impl.register_kv_caches.assert_called_once_with(kv_caches)
+
+    def test_get_num_new_matched_tokens(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        req = MagicMock()
+        impl.get_num_new_matched_tokens.return_value = (10, False)
+        result = connector.get_num_new_matched_tokens(req, 5)
+        impl.get_num_new_matched_tokens.assert_called_once_with(req, 5)
+        assert result == (10, False)
+
+    def test_update_state_after_alloc(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        req = MagicMock()
+        blocks = MagicMock()
+        connector.update_state_after_alloc(req, blocks, 4)
+        impl.update_state_after_alloc.assert_called_once_with(req, blocks, 4)
+
+    def test_build_connector_meta(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        sched_out = MagicMock()
+        connector.build_connector_meta(sched_out)
+        impl.build_connector_meta.assert_called_once_with(sched_out)
+
+    def test_update_connector_output(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        out = MagicMock()
+        connector.update_connector_output(out)
+        impl.update_connector_output.assert_called_once_with(out)
+
+    def test_request_finished(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        req = MagicMock()
+        impl.request_finished.return_value = (True, {"key": "val"})
+        result = connector.request_finished(req, [1, 2, 3])
+        impl.request_finished.assert_called_once_with(req, [1, 2, 3])
+        assert result == (True, {"key": "val"})
+
+    def test_take_events(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        impl.take_events.return_value = iter([])
+        list(connector.take_events())
+        impl.take_events.assert_called_once()
+
+    def test_get_kv_connector_stats(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        impl.get_kv_connector_stats.return_value = None
+        result = connector.get_kv_connector_stats()
+        impl.get_kv_connector_stats.assert_called_once()
+        assert result is None
+
+    def test_get_block_ids_with_load_errors(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        impl.get_block_ids_with_load_errors.return_value = {7, 8}
+        result = connector.get_block_ids_with_load_errors()
+        assert result == {7, 8}
+
+    def test_wait_for_layer_load(self, connector_and_impl):
+        connector, impl = connector_and_impl
+        connector.wait_for_layer_load("layer_0")
+        impl.wait_for_layer_load.assert_called_once_with("layer_0")

third_party/vllm/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py

@@ -0,0 +1,480 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Tests for correctness in invalid block handling.
+
+These tests verify correct behavior in three scenarios:
+1. Sync recompute case: Blocks should not be freed for running requests
+   that need to recompute invalid blocks
+2. Sync fail case: Invalid blocks must be evicted from cache when request fails
+3. Async recompute case: Invalid blocks should not be cached after transfer
+"""
+
+from collections.abc import Callable
+from unittest.mock import Mock
+
+import pytest
+
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.request import FinishReason, Request, RequestStatus
+
+from .utils import (
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def _make_get_num_new_matched_tokens(
+    req_num_new_matched_tokens: dict[str, int],
+    async_load: bool,
+) -> Callable[[Request, int], tuple[int, bool]]:
+    def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
+        value = req_num_new_matched_tokens.get(request.request_id, 0)
+        return value, async_load
+
+    return get_num_new_matched_tokens
+
+
+@pytest.fixture
+def fail_scheduler():
+    """scheduler with kv_load_failure_policy='fail'"""
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_load_failure_policy = "fail"
+    return create_scheduler(vllm_config)
+
+
+@pytest.fixture
+def recompute_scheduler():
+    """scheduler with kv_load_failure_policy='recompute'"""
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_load_failure_policy = "recompute"
+    return create_scheduler(vllm_config)
+
+
+def test_sync_recompute_blocks_not_freed_for_running_requests(
+    recompute_scheduler: Scheduler,
+):
+    """
+    Test sync recompute case - blocks must not be freed for running requests.
+
+    When a running request has invalid blocks and retry_policy is 'recompute':
+    1. Request should remain in RUNNING state
+    2. num_computed_tokens should be truncated to invalid block boundary
+    3. Blocks should NOT be freed (request still needs them for recomputation)
+    4. Request should remain in scheduler.requests and scheduler.running
+    """
+    num_prompt_blocks = 100
+    num_external_computed_blocks = 99
+    invalid_block_idx = 50
+
+    num_prompt_tokens = num_prompt_blocks * recompute_scheduler.block_size
+    num_external_computed_tokens = (
+        num_external_computed_blocks * recompute_scheduler.block_size
+    )
+
+    request = create_request(num_tokens=num_prompt_tokens)
+    recompute_scheduler.add_request(request=request)
+
+    req_num_new_matched_tokens = {
+        request.request_id: num_external_computed_tokens,
+    }
+
+    # mock connector indicating sync load
+    recompute_scheduler.connector = Mock()
+    recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
+    )
+    recompute_scheduler.connector.request_finished.return_value = (False, None)
+    recompute_scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = recompute_scheduler.schedule()
+
+    # request should be running with sync KV load
+    assert len(recompute_scheduler.running) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 1
+    assert request.status == RequestStatus.RUNNING
+
+    # get the allocated block IDs before invalid blocks are reported
+    req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
+    invalid_block_ids = {req_block_ids[invalid_block_idx]}
+
+    # store original num_computed_tokens for comparison
+    original_num_computed_tokens = request.num_computed_tokens
+
+    model_runner_output = create_model_runner_output(
+        [request],
+        invalid_block_ids=invalid_block_ids,
+        use_eos=False,  # not finished - should continue running
+    )
+
+    outputs = recompute_scheduler.update_from_output(
+        scheduler_output, model_runner_output
+    )
+
+    # critical assertions for recompute case:
+
+    # 1. request should still be RUNNING (not finished, not aborted)
+    assert request.status == RequestStatus.RUNNING, (
+        f"Request should remain RUNNING for recompute, got {request.status}"
+    )
+
+    # 2. num_computed_tokens should be truncated to first invalid block
+    expected_truncated_tokens = invalid_block_idx * recompute_scheduler.block_size
+    assert request.num_computed_tokens == expected_truncated_tokens, (
+        f"num_computed_tokens should be truncated to {expected_truncated_tokens}, "
+        f"got {request.num_computed_tokens}"
+    )
+    assert request.num_computed_tokens < original_num_computed_tokens, (
+        "num_computed_tokens should be reduced after invalid block detection"
+    )
+
+    # 3. no output should be generated (request is still running)
+    # the request should be skipped in the output loop
+    assert len(outputs) == 0 or request.request_id not in [
+        out.request_id for outs in outputs.values() for out in outs.outputs
+    ], "No output should be generated for recompute requests"
+
+    # 4. request should still be in running queue
+    assert request in recompute_scheduler.running, (
+        "Request should remain in running queue for recomputation"
+    )
+
+    # 5. request should still be in scheduler.requests (not deleted)
+    assert request.request_id in recompute_scheduler.requests, (
+        "Request should not be deleted from scheduler.requests"
+    )
+
+    # 6. blocks should NOT be freed - verify blocks are still allocated
+    try:
+        allocated_blocks = recompute_scheduler.kv_cache_manager.get_block_ids(
+            request.request_id
+        )
+        assert allocated_blocks is not None
+        assert len(allocated_blocks[0]) > 0, (
+            "Blocks should still be allocated for recomputation"
+        )
+    except KeyError:
+        pytest.fail(
+            "Blocks were freed incorrectly! Running requests need their blocks "
+            "to recompute invalid portions."
+        )
+
+    # 7. verify request can be rescheduled in next step
+    scheduler_output_2 = recompute_scheduler.schedule()
+
+    # request should appear in the new schedule to recompute invalid blocks
+    scheduled_req_ids = [
+        req.request_id for req in scheduler_output_2.scheduled_new_reqs
+    ]
+    if scheduler_output_2.num_scheduled_tokens:
+        scheduled_req_ids.extend(scheduler_output_2.num_scheduled_tokens.keys())
+
+    assert (
+        request.request_id in scheduled_req_ids or len(recompute_scheduler.running) > 0
+    ), "Request should be reschedulable for recomputation"
+
+
+def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler):
+    """
+    Test sync fail case - invalid blocks must be evicted from cache.
+
+    When a request fails with policy='fail' and has invalid blocks from sync loading:
+    1. Request should be finished with FINISHED_ERROR
+    2. Invalid blocks should be evicted from the KV cache
+    3. Valid blocks (if shared) should remain in cache
+    4. Future requests should not reuse the invalid blocks
+
+    This test verifies that invalid blocks are properly evicted to prevent
+    cache corruption and reuse of invalid data.
+    """
+    num_prompt_blocks = 100
+    num_external_computed_blocks = 99
+    invalid_block_idx = 50
+
+    num_prompt_tokens = num_prompt_blocks * fail_scheduler.block_size
+    num_external_computed_tokens = (
+        num_external_computed_blocks * fail_scheduler.block_size
+    )
+
+    request = create_request(num_tokens=num_prompt_tokens)
+    fail_scheduler.add_request(request=request)
+
+    req_num_new_matched_tokens = {
+        request.request_id: num_external_computed_tokens,
+    }
+
+    # mock connector indicating sync load
+    fail_scheduler.connector = Mock()
+    fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, False)
+    )
+    fail_scheduler.connector.request_finished.return_value = (False, None)
+    fail_scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = fail_scheduler.schedule()
+
+    # request should be running with sync KV load
+    assert len(fail_scheduler.running) == 1
+    assert request.status == RequestStatus.RUNNING
+
+    # get allocated block IDs
+    req_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
+    invalid_block_id = req_block_ids[invalid_block_idx]
+    invalid_block_ids = {invalid_block_id}
+
+    # verify the block is in the block pool before we report it as invalid
+    block = fail_scheduler.kv_cache_manager.block_pool.blocks[invalid_block_id]
+    assert block is not None
+
+    # report invalid blocks - request should fail
+    model_runner_output = create_model_runner_output(
+        [request],
+        invalid_block_ids=invalid_block_ids,
+        use_eos=True,
+    )
+
+    outputs = fail_scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # verify request is finished with error
+    assert request.status == RequestStatus.FINISHED_ERROR
+    assert request.get_finished_reason() == FinishReason.ERROR
+
+    # verify output is generated
+    assert len(outputs) == 1
+    engine_outputs = next(iter(outputs.values()))
+    assert len(engine_outputs.outputs) == 1
+    output = engine_outputs.outputs[0]
+    assert output.request_id == request.request_id
+    assert output.finish_reason == FinishReason.ERROR
+
+    # verify the request was removed from scheduler
+    assert request.request_id not in fail_scheduler.requests
+    assert len(fail_scheduler.running) == 0
+
+    # critical: verify invalid block was actually freed from cache
+    # this is the key assertion - the invalid block should no longer be
+    # tracked by the KV cache manager for this request
+    # if it's still there, a future request could reuse the invalid data
+    try:
+        block_ids = fail_scheduler.kv_cache_manager.get_block_ids(request.request_id)
+        # if we get here, check if blocks were actually freed
+        if block_ids is not None and len(block_ids[0]) > 0:
+            pytest.fail(
+                f"Invalid blocks still tracked for finished request! "
+                f"Request {request.request_id} should have been freed but "
+                f"still has {len(block_ids[0])} blocks allocated."
+            )
+        # blocks list exists but is empty - this is fine, they were freed
+    except KeyError:
+        # expected - request completely removed from tracking
+        pass
+
+    # critical: verify invalid block was evicted from prefix cache
+    # the block should no longer have a hash (hash is reset on eviction)
+    assert block.block_hash is None, (
+        f"Invalid block {invalid_block_id} should have been evicted from cache "
+        f"(hash should be None), but hash is still {block.block_hash}"
+    )
+
+    # Verify connector prefix cache stats:
+    # - queries = num_prompt_tokens (total tokens not in local cache)
+    # - hits = num_external_computed_tokens (tokens loaded externally)
+    assert engine_outputs.scheduler_stats is not None
+    stats = engine_outputs.scheduler_stats
+    assert stats.connector_prefix_cache_stats is not None
+    conn_stats = stats.connector_prefix_cache_stats
+    assert conn_stats.requests == 1
+    assert conn_stats.queries == num_prompt_tokens
+    assert conn_stats.hits == num_external_computed_tokens
+
+
+def test_async_recompute_blocks_not_cached_when_invalid(
+    recompute_scheduler: Scheduler,
+):
+    """
+    Test async recompute case - invalid blocks not cached after transfer.
+
+    When async KV loading has invalid blocks and retry_policy is 'recompute':
+    1. Blocks are allocated but not cached yet
+    2. When async transfer completes, only valid blocks should be cached
+    3. Invalid blocks should never enter the prefix cache
+
+    This test verifies correctness, the failed_recving_kv_req_ids protection
+    ensures only valid blocks are cached when the transfer completes, and we
+    only evict blocks from cache that are already hashed in the block table.
+    """
+    from unittest.mock import patch
+
+    num_prompt_blocks = 100
+    num_external_computed_blocks = 99
+    invalid_block_idx = 50
+
+    num_prompt_tokens = num_prompt_blocks * recompute_scheduler.block_size
+    num_external_computed_tokens = (
+        num_external_computed_blocks * recompute_scheduler.block_size
+    )
+
+    request = create_request(num_tokens=num_prompt_tokens)
+    recompute_scheduler.add_request(request=request)
+
+    req_num_new_matched_tokens = {
+        request.request_id: num_external_computed_tokens,
+    }
+
+    # mock connector indicating async load
+    recompute_scheduler.connector = Mock()
+    recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True)
+    )
+    recompute_scheduler.connector.request_finished.return_value = (False, None)
+    recompute_scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = recompute_scheduler.schedule()
+
+    # request should be waiting for remote KVs
+    assert len(recompute_scheduler.skipped_waiting) == 1
+    assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert request.num_computed_tokens == num_external_computed_tokens
+
+    # get the allocated block IDs
+    (req_block_ids,) = recompute_scheduler.kv_cache_manager.get_block_ids(
+        request.request_id
+    )
+    invalid_block_id = req_block_ids[invalid_block_idx]
+    invalid_block_ids = {invalid_block_id}
+
+    # get the block object to verify it's not cached yet and stays uncached
+    block = recompute_scheduler.kv_cache_manager.block_pool.blocks[invalid_block_id]
+
+    # verify block has no hash before invalid blocks are reported
+    assert block.block_hash is None, (
+        "Async loading blocks should not be cached yet (no hash)"
+    )
+
+    # report invalid blocks (transfer not finished yet)
+    model_runner_output = create_model_runner_output(
+        reqs=[],
+        finished_recving=None,  # transfer NOT finished
+        invalid_block_ids=invalid_block_ids,
+        use_eos=False,
+    )
+
+    # critical: spy on evict_blocks to verify it's NOT called for async blocks
+    original_evict_blocks = recompute_scheduler.kv_cache_manager.evict_blocks
+    evict_blocks_calls = []
+
+    def evict_blocks_spy(block_ids):
+        evict_blocks_calls.append(set(block_ids))
+        return original_evict_blocks(block_ids)
+
+    with patch.object(
+        recompute_scheduler.kv_cache_manager, "evict_blocks", evict_blocks_spy
+    ):
+        outputs = recompute_scheduler.update_from_output(
+            scheduler_output, model_runner_output
+        )
+
+    # verify evict_blocks was NOT called (async blocks excluded from eviction)
+    assert len(evict_blocks_calls) == 0, (
+        f"evict_blocks should not be called for async-only invalid blocks, "
+        f"but was called {len(evict_blocks_calls)} time(s) with {evict_blocks_calls}"
+    )
+
+    # request should still be waiting (not finished with error due to recompute policy)
+    assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert request.request_id in recompute_scheduler.failed_recving_kv_req_ids
+
+    # verify num_computed_tokens was truncated to before invalid block
+    expected_valid_tokens = invalid_block_idx * recompute_scheduler.block_size
+    assert request.num_computed_tokens == expected_valid_tokens
+
+    # verify invalid block still has no hash (was not evicted)
+    assert block.block_hash is None, (
+        f"Async loading blocks shouldn't be cached or evicted. "
+        f"Block {invalid_block_id} hash should be None but is {block.block_hash}"
+    )
+
+    # Verify connector prefix cache stats:
+    # - queries = num_prompt_tokens (total tokens not in local cache)
+    # - hits = num_external_computed_tokens (tokens loaded externally)
+    assert len(outputs) == 1
+    engine_outputs = next(iter(outputs.values()))
+    assert engine_outputs.scheduler_stats is not None
+    stats = engine_outputs.scheduler_stats
+    assert stats.connector_prefix_cache_stats is not None
+    conn_stats = stats.connector_prefix_cache_stats
+    assert conn_stats.requests == 1
+    assert conn_stats.queries == num_prompt_tokens
+    assert conn_stats.hits == num_external_computed_tokens
+
+    # now simulate async transfer completing
+    model_runner_output_2 = create_model_runner_output(
+        reqs=[],
+        finished_recving={request.request_id},
+        invalid_block_ids=None,
+        use_eos=False,
+    )
+
+    recompute_scheduler.update_from_output(scheduler_output, model_runner_output_2)
+
+    # verify request is now marked as finished receiving and ready to be processed
+    assert request.request_id in recompute_scheduler.finished_recving_kv_req_ids
+    assert request.request_id in recompute_scheduler.failed_recving_kv_req_ids
+
+    # critical: verify invalid block still has no hash before recompute
+    # the async transfer invalid data was never cached
+    assert block.block_hash is None, (
+        f"Invalid block {invalid_block_id} should not be cached before recompute "
+        f"(hash should be None), but hash is {block.block_hash}"
+    )
+
+    # critical end-to-end test: spy on cache_blocks to verify it's called with
+    # the truncated num_computed_tokens value
+    original_cache_blocks = recompute_scheduler.kv_cache_manager.cache_blocks
+    cache_blocks_calls = []
+
+    def cache_blocks_spy(req, num_tokens):
+        cache_blocks_calls.append((req.request_id, num_tokens))
+        return original_cache_blocks(req, num_tokens)
+
+    with patch.object(
+        recompute_scheduler.kv_cache_manager, "cache_blocks", cache_blocks_spy
+    ):
+        # call schedule() again - this triggers _update_waiting_for_remote_kv()
+        # which should call cache_blocks with the truncated value
+        recompute_scheduler.schedule()
+
+    # verify cache_blocks was called with the truncated value
+    assert len(cache_blocks_calls) == 1, (
+        f"cache_blocks should be called exactly once, "
+        f"got {len(cache_blocks_calls)} calls"
+    )
+    cached_req_id, cached_num_tokens = cache_blocks_calls[0]
+    assert cached_req_id == request.request_id
+    assert cached_num_tokens == expected_valid_tokens, (
+        f"cache_blocks should be called with truncated value {expected_valid_tokens}, "
+        f"but was called with {cached_num_tokens}"
+    )
+
+    # request should now be RUNNING (scheduled immediately after transfer completes)
+    # the flow is: WAITING_FOR_REMOTE_KVS -> WAITING -> RUNNING in same schedule() call
+    assert request.status == RequestStatus.RUNNING
+
+    # num_computed_tokens should be >= expected_valid_tokens because the scheduler
+    # will schedule additional new tokens (up to max_num_batched_tokens) for the request
+    assert request.num_computed_tokens >= expected_valid_tokens, (
+        f"num_computed_tokens should be at least {expected_valid_tokens}, "
+        f"got {request.num_computed_tokens}"
+    )
+
+    # request should no longer be in the failed/finished receiving sets
+    assert request.request_id not in recompute_scheduler.failed_recving_kv_req_ids
+    assert request.request_id not in recompute_scheduler.finished_recving_kv_req_ids
+
+    # request should be in the running queue
+    assert request in recompute_scheduler.running

third_party/vllm/tests/v1/kv_connector/unit/test_kv_cache_layout.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def test_mla_backend_rejects_cross_layer_kv_cache():
+    """MLA backends return identity permutation (layers dim first)
+    to signal cross-layer KV cache is unsupported."""
+    from vllm.model_executor.layers.attention.mla_attention import (
+        MLACommonBackend,
+    )
+
+    stride_order = MLACommonBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=True
+    )
+    assert stride_order == (0, 1, 2, 3)
+    assert stride_order[0] == 0  # layers dim first => no cross-layer
+    assert MLACommonBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=False
+    ) == (0, 1, 2)
+
+
+def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
+    """DeepseekV32Indexer returns identity permutation (layers dim first)
+    to signal cross-layer KV cache is unsupported."""
+    from vllm.v1.attention.backends.mla.indexer import (
+        DeepseekV32IndexerBackend,
+    )
+
+    stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=True
+    )
+    assert stride_order == (0, 1, 2, 3)
+    assert stride_order[0] == 0  # layers dim first => no cross-layer
+    assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=False
+    ) == (0, 1, 2)

third_party/vllm/tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (  # noqa: E501
+    ExampleConnectorMetadata,
+)
+from vllm.distributed.kv_transfer.kv_transfer_state import (
+    ensure_kv_transfer_initialized,
+    get_kv_transfer_group,
+)
+from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
+
+# Importing utils registers TestExampleConnector with the factory
+from .utils import create_vllm_config
+
+
+def _make_empty_scheduler_output():
+    return SchedulerOutput(
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        num_scheduled_tokens={},
+        total_num_scheduled_tokens=0,
+        scheduled_spec_decode_tokens={},
+        scheduled_encoder_inputs={},
+        num_common_prefix_blocks=[],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+        kv_connector_metadata=ExampleConnectorMetadata(),
+    )
+
+
+def test_kv_connector_mixin_clears_metadata():
+    vllm_config = create_vllm_config()
+    vllm_config.kv_transfer_config.kv_connector = "TestExampleConnector"
+    vllm_config.kv_transfer_config.kv_role = "kv_both"
+    vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = "unit"
+
+    # Initialize the global connector instance
+    ensure_kv_transfer_initialized(vllm_config)
+
+    try:
+        # Minimal scheduler output with empty metadata; mixin should still
+        # bind/clear metadata even if no loads happen
+        scheduler_output = _make_empty_scheduler_output()
+
+        # Invoke the no-forward path which uses the mixin context manager
+        KVConnectorModelRunnerMixin.kv_connector_no_forward(
+            scheduler_output, vllm_config
+        )
+
+        # Verify clear_connector_metadata was called on the connector
+        connector = get_kv_transfer_group()
+        assert connector._connector_metadata is None
+        # Test connector wrapper records method calls
+        assert connector.call_record.get("bind_connector_metadata", 0) == 1
+        assert connector.call_record.get("clear_connector_metadata", 0) == 1
+    finally:
+        # Ensure we clean up the global connector between tests
+        KVConnectorModelRunnerMixin.ensure_kv_transfer_shutdown()

third_party/vllm/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py

@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from unittest.mock import Mock
+
+import pytest
+
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.request import Request, RequestStatus
+
+from .utils import (
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+
+
+def _make_get_num_new_matched_tokens(
+    req_num_new_matched_tokens: dict[str, int],
+    async_load,
+) -> Callable[[Request, int], tuple[int, bool]]:
+    def get_num_new_matched_tokens(request: Request, _: int) -> tuple[int, bool]:
+        value = req_num_new_matched_tokens.get(request.request_id, 0)
+        return value, async_load
+
+    return get_num_new_matched_tokens
+
+
+@pytest.fixture
+def scheduler():
+    vllm_config = create_vllm_config(kv_load_failure_policy="recompute")
+    return create_scheduler(vllm_config)
+
+
+@pytest.mark.parametrize(
+    "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs",
+    [
+        (100, 99, {0, 98}),
+        (100, 99, {50, 98}),
+        (100, 99, {98}),
+    ],
+)
+def test_async_load_failure(
+    scheduler: Scheduler,
+    num_prompt_blocks: int,
+    num_external_computed_blocks: int,
+    invalid_block_idxs: set[int],
+):
+    assert num_prompt_blocks >= num_external_computed_blocks
+
+    num_prompt_tokens = num_prompt_blocks * scheduler.block_size
+    num_external_computed_tokens = num_external_computed_blocks * scheduler.block_size
+
+    request1 = create_request(num_tokens=num_prompt_tokens)
+    scheduler.add_request(request=request1)
+    request2 = create_request(num_tokens=num_prompt_tokens)
+    scheduler.add_request(request=request2)
+    request3 = create_request(num_tokens=num_prompt_tokens)
+    scheduler.add_request(request=request3)
+
+    # Mock KV connector method.
+    # req_id -> num_external_computed_tokens
+    req_num_new_matched_tokens = {
+        request1.request_id: num_external_computed_tokens,
+        request2.request_id: num_external_computed_tokens,
+        request3.request_id: num_external_computed_tokens,
+    }
+
+    scheduler.connector = Mock()
+    scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True)
+    )
+    scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = scheduler.schedule()
+
+    assert len(scheduler.waiting) == 0
+    assert len(scheduler.skipped_waiting) == 3
+    for request in scheduler.skipped_waiting:
+        assert request.num_computed_tokens == num_external_computed_tokens
+        assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
+
+    # Simulate a failure in loading some of request2 blocks.
+    (req2_block_ids,) = scheduler.kv_cache_manager.get_block_ids(request2.request_id)
+    invalid_block_ids = {req2_block_ids[i] for i in invalid_block_idxs}
+    model_runner_output = create_model_runner_output(
+        reqs=[],
+        finished_recving={request1.request_id, request3.request_id},
+        invalid_block_ids=invalid_block_ids,
+        use_eos=True,
+    )
+
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    min_invalid_block_idx = min(invalid_block_idxs)
+
+    assert len(scheduler.waiting) == 0
+    assert len(scheduler.skipped_waiting) == 3
+    for request in scheduler.skipped_waiting:
+        if request.request_id == request2.request_id:
+            assert request.num_computed_tokens == (
+                min_invalid_block_idx * scheduler.block_size
+            )
+        else:
+            assert request.num_computed_tokens == num_external_computed_tokens
+        assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert scheduler.failed_recving_kv_req_ids == {request2.request_id}
+    assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
+
+
+@pytest.mark.parametrize(
+    "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs",
+    [
+        (100, 99, {0, 98}),
+        (100, 99, {50, 98}),
+        (100, 99, {98}),
+    ],
+)
+def test_sync_load_failure(
+    scheduler: Scheduler,
+    num_prompt_blocks: int,
+    num_external_computed_blocks: int,
+    invalid_block_idxs: set[int],
+):
+    assert num_prompt_blocks >= num_external_computed_blocks
+
+    num_prompt_tokens = num_prompt_blocks * scheduler.block_size
+    num_external_computed_tokens = num_external_computed_blocks * scheduler.block_size
+
+    request1 = create_request(num_tokens=num_prompt_tokens)
+    scheduler.add_request(request=request1)
+    request2 = create_request(num_tokens=num_prompt_tokens)
+    scheduler.add_request(request=request2)
+    request3 = create_request(num_tokens=num_prompt_tokens)
+    scheduler.add_request(request=request3)
+
+    # Mock KV connector method.
+    # req_id -> num_external_computed_tokens
+    req_num_new_matched_tokens = {
+        request1.request_id: num_external_computed_tokens,
+        request2.request_id: num_external_computed_tokens,
+        request3.request_id: num_external_computed_tokens,
+    }
+
+    scheduler.connector = Mock()
+    scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=False)
+    )
+    scheduler.connector.request_finished.return_value = (False, None)
+    scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = scheduler.schedule()
+
+    # req_id -> num_computed_tokens
+    expected_computed_tokens = {
+        request1.request_id: num_external_computed_tokens,
+        request2.request_id: num_external_computed_tokens,
+        request3.request_id: num_external_computed_tokens,
+    }
+
+    assert len(scheduler.running) == 3
+    assert len(scheduler_output.scheduled_new_reqs) == 3
+    for request in scheduler_output.scheduled_new_reqs:
+        assert request.num_computed_tokens == expected_computed_tokens[request.req_id]
+    assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
+
+    # Simulate a failure in loading some of request2 blocks.
+    req2_block_ids = scheduler_output.scheduled_new_reqs[1].block_ids[0]
+    invalid_block_ids = {req2_block_ids[i] for i in invalid_block_idxs}
+    model_runner_output = create_model_runner_output(
+        [request1, request2, request3],
+        invalid_block_ids=invalid_block_ids,
+        use_eos=True,
+    )
+
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    assert len(scheduler.running) == 1
+    assert scheduler.running[0].request_id == request2.request_id
+    assert scheduler.running[0].num_computed_tokens == (
+        min(invalid_block_idxs) * scheduler.block_size
+    )
+    assert scheduler.connector.get_num_new_matched_tokens.call_count == 3
+    assert scheduler.connector.request_finished.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "num_prompt_blocks,"
+    "num_external_computed_blocks,"
+    "num_common_prefix_blocks,"
+    "invalid_block_idxs",
+    [
+        (100, 99, 50, {0, 49}),
+        (100, 99, 50, {25, 49}),
+        (100, 99, 50, {49}),
+    ],
+)
+def test_sync_load_failure_with_shared_blocks(
+    scheduler: Scheduler,
+    num_prompt_blocks: int,
+    num_external_computed_blocks: int,
+    num_common_prefix_blocks: int,
+    invalid_block_idxs: set[int],
+):
+    assert num_prompt_blocks >= num_external_computed_blocks >= num_common_prefix_blocks
+
+    num_prompt_tokens = num_prompt_blocks * scheduler.block_size
+    num_external_computed_tokens = num_external_computed_blocks * scheduler.block_size
+    common_prefix_len = num_common_prefix_blocks * scheduler.block_size
+
+    request1 = create_request(
+        num_tokens=num_prompt_tokens, common_prefix_len=common_prefix_len
+    )
+    scheduler.add_request(request=request1)
+    request2 = create_request(
+        num_tokens=num_prompt_tokens, common_prefix_len=common_prefix_len
+    )
+    scheduler.add_request(request=request2)
+
+    # Mock KV connector method.
+    # req_id -> num_external_computed_tokens
+    req_num_new_matched_tokens = {
+        request1.request_id: num_external_computed_tokens,
+    }
+
+    scheduler.connector = Mock()
+    scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=False)
+    )
+    scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = scheduler.schedule()
+
+    # req_id -> num_computed_tokens
+    expected_computed_tokens = {
+        request1.request_id: num_external_computed_tokens,
+        request2.request_id: common_prefix_len,
+    }
+
+    assert len(scheduler.running) == 2
+    assert len(scheduler_output.scheduled_new_reqs) == 2
+    for request in scheduler_output.scheduled_new_reqs:
+        assert request.num_computed_tokens == expected_computed_tokens[request.req_id]
+    assert scheduler.connector.get_num_new_matched_tokens.call_count == 2
+
+    # Simulate a failure in loading some of the shared blocks.
+    req1_block_ids = scheduler_output.scheduled_new_reqs[0].block_ids[0]
+    invalid_block_ids = {req1_block_ids[i] for i in invalid_block_idxs}
+    model_runner_output = create_model_runner_output(
+        [request1, request2], invalid_block_ids=invalid_block_ids, use_eos=True
+    )
+
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # req_id -> num_computed_tokens
+    # all the common prefix blocks will be computed by request1
+    expected_computed_tokens = {
+        request1.request_id: min(invalid_block_idxs) * scheduler.block_size,
+        request2.request_id: common_prefix_len,
+    }
+
+    assert len(scheduler.running) == 2
+    for request in scheduler.running:
+        assert (
+            request.num_computed_tokens == expected_computed_tokens[request.request_id]
+        )
+    assert scheduler.connector.get_num_new_matched_tokens.call_count == 2
+
+
+@pytest.mark.parametrize(
+    "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs",
+    [
+        (100, 99, {0, 50, 98}),
+        (100, 99, {98, 50, 0}),
+    ],
+)
+def test_async_progressive_load_failure(
+    scheduler: Scheduler,
+    num_prompt_blocks: int,
+    num_external_computed_blocks: int,
+    invalid_block_idxs: set[int],
+):
+    assert num_prompt_blocks >= num_external_computed_blocks
+
+    num_prompt_tokens = num_prompt_blocks * scheduler.block_size
+    num_external_computed_tokens = num_external_computed_blocks * scheduler.block_size
+
+    request = create_request(num_tokens=num_prompt_tokens)
+    scheduler.add_request(request=request)
+
+    # Mock KV connector method.
+    # req_id -> num_external_computed_tokens
+    req_num_new_matched_tokens = {
+        request.request_id: num_external_computed_tokens,
+    }
+
+    scheduler.connector = Mock()
+    scheduler.connector.get_num_new_matched_tokens.side_effect = (
+        _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True)
+    )
+    scheduler.connector.take_events.return_value = ()
+
+    scheduler_output = scheduler.schedule()
+
+    assert len(scheduler.waiting) == 0
+    assert len(scheduler.skipped_waiting) == 1
+    assert scheduler.skipped_waiting.peek_request().request_id == request.request_id
+    assert request.num_computed_tokens == num_external_computed_tokens
+    assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert scheduler.connector.get_num_new_matched_tokens.call_count == 1
+
+    min_invalid_block_idx = max(invalid_block_idxs) + 1
+    # Simulate failures when progressively loading request blocks.
+    for invalid_block_idx in invalid_block_idxs:
+        (req_block_ids,) = scheduler.kv_cache_manager.get_block_ids(request.request_id)
+        invalid_block_ids = {req_block_ids[invalid_block_idx]}
+        model_runner_output = create_model_runner_output(
+            reqs=[],
+            finished_recving=set(),
+            invalid_block_ids=invalid_block_ids,
+            use_eos=True,
+        )
+
+        scheduler.update_from_output(scheduler_output, model_runner_output)
+
+        min_invalid_block_idx = min(min_invalid_block_idx, invalid_block_idx)
+
+        assert len(scheduler.waiting) == 0
+        assert len(scheduler.skipped_waiting) == 1
+        assert scheduler.skipped_waiting.peek_request().request_id == request.request_id
+        assert request.num_computed_tokens == (
+            min_invalid_block_idx * scheduler.block_size
+        )
+        assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+        assert scheduler.failed_recving_kv_req_ids == {request.request_id}
+        assert scheduler.connector.get_num_new_matched_tokens.call_count == 1

third_party/vllm/tests/v1/kv_connector/unit/test_lmcache_connector.py

@@ -0,0 +1,786 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm.distributed.kv_events import BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector import (
+    LMCacheConnectorV1,
+    LMCacheKVEvents,
+)
+from vllm.v1.outputs import KVConnectorOutput
+
+
+@pytest.fixture
+def mock_lmcache_engine_event():
+    """Create a mock event object that mimics what the lmcache engine returns."""
+
+    class MockEvent:
+        def __init__(
+            self,
+            block_hashes,
+            parent_block_hash,
+            token_ids,
+            lora_id,
+            block_size,
+            medium,
+            lora_name,
+        ):
+            self.block_hashes = block_hashes
+            self.parent_block_hash = parent_block_hash
+            self.token_ids = token_ids
+            self.lora_id = lora_id
+            self.block_size = block_size
+            self.medium = medium
+            self.lora_name = lora_name
+
+    return MockEvent(
+        block_hashes=["hash1", "hash2"],
+        parent_block_hash="parent_hash",
+        token_ids=[1, 2, 3, 4],
+        lora_id=None,
+        block_size=16,
+        medium="GPU",
+        lora_name=None,
+    )
+
+
+@pytest.fixture
+def mock_connector():
+    """Create a mock LMCacheConnectorV1 instance with mocked dependencies."""
+    connector = MagicMock(spec=LMCacheConnectorV1)
+    connector._kv_cache_events = None
+    connector._lmcache_engine = MagicMock()
+
+    # Make the methods use the real implementation
+    connector.get_kv_connector_kv_cache_events = (
+        LMCacheConnectorV1.get_kv_connector_kv_cache_events.__get__(
+            connector, LMCacheConnectorV1
+        )
+    )
+    connector.update_connector_output = (
+        LMCacheConnectorV1.update_connector_output.__get__(
+            connector, LMCacheConnectorV1
+        )
+    )
+    connector.take_events = LMCacheConnectorV1.take_events.__get__(
+        connector, LMCacheConnectorV1
+    )
+
+    return connector
+
+
+class TestGetKVConnectorKVCacheEvents:
+    """Test get_kv_connector_kv_cache_events method."""
+
+    def test_returns_none_when_no_events(self, mock_connector):
+        """Test that None is returned when lmcache engine has no events."""
+        mock_connector._lmcache_engine.get_kv_events.return_value = None
+
+        result = mock_connector.get_kv_connector_kv_cache_events()
+
+        assert result is None
+        mock_connector._lmcache_engine.get_kv_events.assert_called_once()
+
+    def test_returns_none_when_empty_list(self, mock_connector):
+        """Test that None is returned when lmcache engine returns empty list."""
+        mock_connector._lmcache_engine.get_kv_events.return_value = []
+
+        result = mock_connector.get_kv_connector_kv_cache_events()
+
+        assert result is None
+
+    def test_converts_single_event(self, mock_connector, mock_lmcache_engine_event):
+        """Test conversion of a single event from lmcache engine format."""
+        mock_connector._lmcache_engine.get_kv_events.return_value = [
+            mock_lmcache_engine_event
+        ]
+
+        result = mock_connector.get_kv_connector_kv_cache_events()
+
+        assert result is not None
+        assert isinstance(result, LMCacheKVEvents)
+        assert result.get_number_of_workers() == 1
+
+        events = result.get_all_events()
+        assert len(events) == 1
+        assert isinstance(events[0], BlockStored)
+        assert events[0].block_hashes == ["hash1", "hash2"]
+        assert events[0].parent_block_hash == "parent_hash"
+        assert events[0].token_ids == [1, 2, 3, 4]
+        assert events[0].lora_id is None
+        assert events[0].block_size == 16
+        assert events[0].medium == "GPU"
+        assert events[0].lora_name is None
+
+    def test_converts_multiple_events(self, mock_connector):
+        """Test conversion of multiple events from lmcache engine format."""
+
+        class MockEvent:
+            def __init__(self, i):
+                self.block_hashes = [f"hash{i}"]
+                self.parent_block_hash = f"parent{i}"
+                self.token_ids = [i]
+                self.lora_id = None
+                self.block_size = 16
+                self.medium = "GPU"
+                self.lora_name = None
+
+        events = [MockEvent(i) for i in range(5)]
+        mock_connector._lmcache_engine.get_kv_events.return_value = events
+
+        result = mock_connector.get_kv_connector_kv_cache_events()
+
+        assert result is not None
+        assert isinstance(result, LMCacheKVEvents)
+
+        converted_events = result.get_all_events()
+        assert len(converted_events) == 5
+
+        for i, event in enumerate(converted_events):
+            assert isinstance(event, BlockStored)
+            assert event.block_hashes == [f"hash{i}"]
+            assert event.parent_block_hash == f"parent{i}"
+            assert event.token_ids == [i]
+
+    def test_preserves_event_attributes(self, mock_connector):
+        """Test that all event attributes are correctly preserved."""
+
+        class MockEventWithLora:
+            def __init__(self):
+                self.block_hashes = ["hash_a", "hash_b", "hash_c"]
+                self.parent_block_hash = "parent_xyz"
+                self.token_ids = [100, 200, 300]
+                self.lora_id = 42
+                self.block_size = 32
+                self.medium = "DISK"
+                self.lora_name = "lora_example"
+
+        mock_connector._lmcache_engine.get_kv_events.return_value = [
+            MockEventWithLora()
+        ]
+
+        result = mock_connector.get_kv_connector_kv_cache_events()
+
+        events = result.get_all_events()
+        event = events[0]
+
+        assert event.block_hashes == ["hash_a", "hash_b", "hash_c"]
+        assert event.parent_block_hash == "parent_xyz"
+        assert event.token_ids == [100, 200, 300]
+        assert event.lora_id == 42
+        assert event.block_size == 32
+        assert event.medium == "DISK"
+        assert event.lora_name == "lora_example"
+
+    def test_handles_none_parent_block_hash(self, mock_connector):
+        """Test handling of events with None parent_block_hash."""
+
+        class MockEventNoParent:
+            def __init__(self):
+                self.block_hashes = ["hash1"]
+                self.parent_block_hash = None
+                self.token_ids = [1, 2]
+                self.lora_id = None
+                self.block_size = 16
+                self.medium = "GPU"
+                self.lora_name = None
+
+        mock_connector._lmcache_engine.get_kv_events.return_value = [
+            MockEventNoParent()
+        ]
+
+        result = mock_connector.get_kv_connector_kv_cache_events()
+
+        events = result.get_all_events()
+        assert events[0].parent_block_hash is None
+
+
+class TestUpdateConnectorOutput:
+    """Test update_connector_output method."""
+
+    def test_does_nothing_when_kv_cache_events_is_none(self, mock_connector):
+        """Test that method returns early when kv_cache_events is None."""
+        connector_output = KVConnectorOutput(kv_cache_events=None)
+
+        mock_connector.update_connector_output(connector_output)
+
+        assert mock_connector._kv_cache_events is None
+
+    def test_does_nothing_when_kv_cache_events_is_not_lmcache_kv_events(
+        self, mock_connector
+    ):
+        """Test that method returns early when kv_cache_events is not
+        LMCacheKVEvents."""
+        # Create a mock object that is not LMCacheKVEvents
+        fake_events = MagicMock()
+        connector_output = KVConnectorOutput(kv_cache_events=fake_events)
+
+        mock_connector.update_connector_output(connector_output)
+
+        assert mock_connector._kv_cache_events is None
+
+    def test_sets_kv_cache_events_when_none(self, mock_connector):
+        """Test that _kv_cache_events is set when it was None."""
+        kv_events = LMCacheKVEvents(num_workers=1)
+        event = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1, 2],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        kv_events.add_events([event])
+
+        connector_output = KVConnectorOutput(kv_cache_events=kv_events)
+
+        mock_connector.update_connector_output(connector_output)
+
+        assert mock_connector._kv_cache_events is kv_events
+
+    def test_adds_events_when_kv_cache_events_already_exists(self, mock_connector):
+        """Test that events are added when _kv_cache_events already exists."""
+        # Set up existing events
+        existing_events = LMCacheKVEvents(num_workers=2)
+        event1 = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        existing_events.add_events([event1])
+        existing_events.add_events([event1])  # Simulate 2 workers reporting
+
+        mock_connector._kv_cache_events = existing_events
+
+        # Create new events to add
+        new_events = LMCacheKVEvents(num_workers=1)
+        event2 = BlockStored(
+            block_hashes=["hash2"],
+            parent_block_hash=None,
+            token_ids=[2],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        new_events.add_events([event2])
+
+        connector_output = KVConnectorOutput(kv_cache_events=new_events)
+
+        mock_connector.update_connector_output(connector_output)
+
+        # Check that events were added
+        all_events = mock_connector._kv_cache_events.get_all_events()
+        assert len(all_events) == 3  # 2 from existing + 1 from new
+        assert event1 in all_events
+        assert event2 in all_events
+
+    def test_increments_workers_when_kv_cache_events_already_exists(
+        self, mock_connector
+    ):
+        """Test that worker count is incremented correctly."""
+        # Set up existing events with 2 workers
+        existing_events = LMCacheKVEvents(num_workers=2)
+        mock_connector._kv_cache_events = existing_events
+
+        # Create new events from 3 workers
+        new_events = LMCacheKVEvents(num_workers=3)
+        event = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        new_events.add_events([event])
+
+        connector_output = KVConnectorOutput(kv_cache_events=new_events)
+
+        mock_connector.update_connector_output(connector_output)
+
+        # Worker count should be 2 + 3 = 5
+        assert mock_connector._kv_cache_events.get_number_of_workers() == 5
+
+    def test_multiple_updates(self, mock_connector):
+        """Test multiple consecutive updates."""
+        # First update
+        events1 = LMCacheKVEvents(num_workers=1)
+        event1 = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        events1.add_events([event1])
+        output1 = KVConnectorOutput(kv_cache_events=events1)
+        mock_connector.update_connector_output(output1)
+
+        # Second update
+        events2 = LMCacheKVEvents(num_workers=2)
+        event2 = BlockStored(
+            block_hashes=["hash2"],
+            parent_block_hash=None,
+            token_ids=[2],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        events2.add_events([event2])
+        output2 = KVConnectorOutput(kv_cache_events=events2)
+        mock_connector.update_connector_output(output2)
+
+        # Third update
+        events3 = LMCacheKVEvents(num_workers=1)
+        event3 = BlockStored(
+            block_hashes=["hash3"],
+            parent_block_hash=None,
+            token_ids=[3],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        events3.add_events([event3])
+        output3 = KVConnectorOutput(kv_cache_events=events3)
+        mock_connector.update_connector_output(output3)
+
+        # Check final state
+        all_events = mock_connector._kv_cache_events.get_all_events()
+        assert len(all_events) == 3
+        assert mock_connector._kv_cache_events.get_number_of_workers() == 4  # 1+2+1
+
+    def test_updates_with_empty_events(self, mock_connector):
+        """Test updating with empty event lists."""
+        # First update with actual events
+        events1 = LMCacheKVEvents(num_workers=1)
+        event1 = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        events1.add_events([event1])
+        output1 = KVConnectorOutput(kv_cache_events=events1)
+        mock_connector.update_connector_output(output1)
+
+        # Second update with empty events
+        events2 = LMCacheKVEvents(num_workers=2)
+        # No events added
+        output2 = KVConnectorOutput(kv_cache_events=events2)
+        mock_connector.update_connector_output(output2)
+
+        # Should still have the original event
+        all_events = mock_connector._kv_cache_events.get_all_events()
+        assert len(all_events) == 1
+        assert mock_connector._kv_cache_events.get_number_of_workers() == 3
+
+
+class TestTakeEvents:
+    """Test take_events method."""
+
+    def test_yields_nothing_when_kv_cache_events_is_none(self, mock_connector):
+        """Test that nothing is yielded when _kv_cache_events is None."""
+        mock_connector._kv_cache_events = None
+
+        events = list(mock_connector.take_events())
+
+        assert events == []
+
+    def test_yields_events_and_clears(self, mock_connector):
+        """Test that events are yielded and then cleared."""
+        # Set up events
+        kv_events = LMCacheKVEvents(num_workers=1)
+        event1 = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        event2 = BlockStored(
+            block_hashes=["hash2"],
+            parent_block_hash=None,
+            token_ids=[2],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        kv_events.add_events([event1, event2])
+        mock_connector._kv_cache_events = kv_events
+
+        # Take events
+        events = list(mock_connector.take_events())
+
+        # Check that events were yielded
+        assert len(events) == 2
+        assert event1 in events
+        assert event2 in events
+
+        # Check that _kv_cache_events was cleared
+        assert mock_connector._kv_cache_events is None
+
+    def test_aggregates_before_yielding(self, mock_connector):
+        """Test that events are aggregated before yielding."""
+        # Set up events from multiple workers
+        kv_events = LMCacheKVEvents(num_workers=3)
+        common_event = BlockStored(
+            block_hashes=["hash_common"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        uncommon_event = BlockStored(
+            block_hashes=["hash_uncommon"],
+            parent_block_hash=None,
+            token_ids=[2],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+
+        # All 3 workers report common_event
+        kv_events.add_events([common_event])
+        kv_events.add_events([common_event])
+        kv_events.add_events([common_event])
+
+        # Only 1 worker reports uncommon_event
+        kv_events.add_events([uncommon_event])
+
+        mock_connector._kv_cache_events = kv_events
+
+        # Take events
+        events = list(mock_connector.take_events())
+
+        # Only the common event should be yielded
+        assert len(events) == 1
+        assert events[0] == common_event
+
+    def test_multiple_take_events_calls(self, mock_connector):
+        """Test calling take_events multiple times."""
+        # First call with events
+        kv_events1 = LMCacheKVEvents(num_workers=1)
+        event1 = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        kv_events1.add_events([event1])
+        mock_connector._kv_cache_events = kv_events1
+
+        events1 = list(mock_connector.take_events())
+        assert len(events1) == 1
+        assert events1[0] == event1
+        assert mock_connector._kv_cache_events is None
+
+        # Second call with no events
+        events2 = list(mock_connector.take_events())
+        assert events2 == []
+
+        # Third call after adding new events
+        kv_events2 = LMCacheKVEvents(num_workers=1)
+        event2 = BlockStored(
+            block_hashes=["hash2"],
+            parent_block_hash=None,
+            token_ids=[2],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        kv_events2.add_events([event2])
+        mock_connector._kv_cache_events = kv_events2
+
+        events3 = list(mock_connector.take_events())
+        assert len(events3) == 1
+        assert events3[0] == event2
+
+    def test_yields_empty_after_aggregation_removes_all(self, mock_connector):
+        """Test that nothing is yielded if aggregation removes all events."""
+        # Set up events from 2 workers with no common events
+        kv_events = LMCacheKVEvents(num_workers=2)
+        event1 = BlockStored(
+            block_hashes=["hash1"],
+            parent_block_hash=None,
+            token_ids=[1],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+        event2 = BlockStored(
+            block_hashes=["hash2"],
+            parent_block_hash=None,
+            token_ids=[2],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+
+        # Worker 1 reports event1
+        kv_events.add_events([event1])
+        # Worker 2 reports event2
+        kv_events.add_events([event2])
+
+        mock_connector._kv_cache_events = kv_events
+
+        # Take events
+        events = list(mock_connector.take_events())
+
+        # No common events, so nothing should be yielded
+        assert events == []
+        assert mock_connector._kv_cache_events is None
+
+
+class TestIntegrationScenarios:
+    """Test integration scenarios."""
+
+    def test_full_workflow(self, mock_connector, mock_lmcache_engine_event):
+        """Test a complete workflow from getting events to taking them."""
+        # Step 1: Get events from lmcache engine
+        mock_connector._lmcache_engine.get_kv_events.return_value = [
+            mock_lmcache_engine_event
+        ]
+        kv_events = mock_connector.get_kv_connector_kv_cache_events()
+
+        assert kv_events is not None
+        assert len(kv_events.get_all_events()) == 1
+
+        # Step 2: Update connector output (simulate receiving from worker)
+        output1 = KVConnectorOutput(kv_cache_events=kv_events)
+        mock_connector.update_connector_output(output1)
+
+        assert mock_connector._kv_cache_events is not None
+
+        # Step 3: Take events
+        taken_events = list(mock_connector.take_events())
+
+        assert len(taken_events) == 1
+        assert mock_connector._kv_cache_events is None
+
+    def test_multiple_workers_workflow(self, mock_connector):
+        """Test workflow with multiple workers."""
+
+        class MockEvent:
+            def __init__(self, hash_val):
+                self.block_hashes = [hash_val]
+                self.parent_block_hash = None
+                self.token_ids = [1]
+                self.lora_id = None
+                self.block_size = 16
+                self.medium = "GPU"
+                self.lora_name = None
+
+        # Worker 1
+        mock_connector._lmcache_engine.get_kv_events.return_value = [
+            MockEvent("hash_common"),
+            MockEvent("hash_worker1"),
+        ]
+        kv_events1 = mock_connector.get_kv_connector_kv_cache_events()
+        output1 = KVConnectorOutput(kv_cache_events=kv_events1)
+        mock_connector.update_connector_output(output1)
+
+        # Worker 2
+        mock_connector._lmcache_engine.get_kv_events.return_value = [
+            MockEvent("hash_common"),
+            MockEvent("hash_worker2"),
+        ]
+        kv_events2 = mock_connector.get_kv_connector_kv_cache_events()
+        output2 = KVConnectorOutput(kv_cache_events=kv_events2)
+        mock_connector.update_connector_output(output2)
+
+        # Take events (should only get common events)
+        taken_events = list(mock_connector.take_events())
+
+        # With aggregation, only events reported by both workers should be present
+        # In this case, hash_common was reported by both
+        event_hashes = [e.block_hashes[0] for e in taken_events]
+        assert "hash_common" in event_hashes
+
+    def test_empty_workflow(self, mock_connector):
+        """Test workflow when there are no events at any stage."""
+        # Get events returns None
+        mock_connector._lmcache_engine.get_kv_events.return_value = None
+        kv_events = mock_connector.get_kv_connector_kv_cache_events()
+
+        assert kv_events is None
+
+        # Update with None
+        output = KVConnectorOutput(kv_cache_events=None)
+        mock_connector.update_connector_output(output)
+
+        # Take events
+        taken_events = list(mock_connector.take_events())
+
+        assert taken_events == []
+        assert mock_connector._kv_cache_events is None
+
+    def test_repeated_cycles(self, mock_connector):
+        """Test multiple cycles of the complete workflow."""
+
+        class MockEvent:
+            def __init__(self, cycle_num):
+                self.block_hashes = [f"hash_cycle_{cycle_num}"]
+                self.parent_block_hash = None
+                self.token_ids = [cycle_num]
+                self.lora_id = None
+                self.block_size = 16
+                self.medium = "GPU"
+                self.lora_name = None
+
+        for cycle in range(3):
+            # Get events
+            mock_connector._lmcache_engine.get_kv_events.return_value = [
+                MockEvent(cycle)
+            ]
+            kv_events = mock_connector.get_kv_connector_kv_cache_events()
+
+            # Update
+            output = KVConnectorOutput(kv_cache_events=kv_events)
+            mock_connector.update_connector_output(output)
+
+            # Take
+            taken_events = list(mock_connector.take_events())
+
+            # Verify
+            assert len(taken_events) == 1
+            assert taken_events[0].block_hashes[0] == f"hash_cycle_{cycle}"
+            assert mock_connector._kv_cache_events is None
+
+    def test_lmcache_kv_events_aggregation(self):
+        """
+        Test LMCacheKVEvents aggregation across TP ranks using
+        KVOutputAggregator (used by MultiprocExecutor).
+        """
+        from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
+        from vllm.v1.outputs import ModelRunnerOutput
+
+        # Create KVOutputAggregator for 3 workers (simulating TP=3)
+        aggregator = KVOutputAggregator(expected_finished_count=3)
+
+        # Define common and unique events
+        common_event = BlockStored(
+            block_hashes=["hash_common"],
+            parent_block_hash="parent_common",
+            token_ids=[1, 2, 3],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+
+        worker1_unique_event = BlockStored(
+            block_hashes=["hash_worker1"],
+            parent_block_hash="parent_w1",
+            token_ids=[4, 5],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+
+        worker2_unique_event = BlockStored(
+            block_hashes=["hash_worker2"],
+            parent_block_hash="parent_w2",
+            token_ids=[6, 7],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+
+        worker3_unique_event = BlockStored(
+            block_hashes=["hash_worker3"],
+            parent_block_hash="parent_w3",
+            token_ids=[8, 9],
+            block_size=16,
+            lora_id=None,
+            medium="GPU",
+            lora_name=None,
+        )
+
+        # Create events for each worker
+        # Worker 0: reports common event and its unique event
+        worker0_events = LMCacheKVEvents(num_workers=1)
+        worker0_events.add_events([common_event, worker1_unique_event])
+
+        # Worker 1: reports common event and its unique event
+        worker1_events = LMCacheKVEvents(num_workers=1)
+        worker1_events.add_events([common_event, worker2_unique_event])
+
+        # Worker 2: reports common event and its unique event
+        worker2_events = LMCacheKVEvents(num_workers=1)
+        worker2_events.add_events([common_event, worker3_unique_event])
+
+        # Create ModelRunnerOutput instances for each worker
+        worker_outputs = []
+        for i, worker_events in enumerate(
+            [worker0_events, worker1_events, worker2_events]
+        ):
+            output = ModelRunnerOutput(
+                req_ids=[f"req_{i}"],
+                req_id_to_index={f"req_{i}": 0},
+                sampled_token_ids=[[123]],  # dummy token
+                logprobs=None,
+                prompt_logprobs_dict={},
+                pooler_output=[None],
+                kv_connector_output=KVConnectorOutput(
+                    finished_sending=set([f"req_{i}_send"])
+                    if i < 2
+                    else None,  # Workers 0,1 finished sending
+                    finished_recving=set([f"req_{i}_recv"])
+                    if i > 0
+                    else None,  # Workers 1,2 finished receiving
+                    kv_cache_events=worker_events,
+                ),
+            )
+            worker_outputs.append(output)
+
+        # Use the real aggregation mechanism (like MultiprocExecutor.execute_model)
+        aggregated_output = aggregator.aggregate(worker_outputs, output_rank=0)
+        kv_cache_events = aggregated_output.kv_connector_output.kv_cache_events
+
+        assert isinstance(kv_cache_events, LMCacheKVEvents)
+
+        # After aggregation, events should be combined from all workers
+        # The aggregator doesn't automatically aggregate events, so we need to call
+        # aggregate() to get only common events
+        kv_cache_events.aggregate()
+        aggregated_events = kv_cache_events.get_all_events()
+
+        # Only the common event should remain after aggregation
+        # because it's the only event reported by all 3 workers
+        assert len(aggregated_events) == 1
+        assert aggregated_events[0] == common_event
+
+        # Verify the common event properties
+        assert aggregated_events[0].block_hashes == ["hash_common"]
+        assert aggregated_events[0].parent_block_hash == "parent_common"
+        assert aggregated_events[0].token_ids == [1, 2, 3]

third_party/vllm/tests/v1/kv_connector/unit/test_lmcache_integration.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# NOTE: if your PR has broken one of the tests here (sorry),
+# kindly patch the corresponding integration in
+# /vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+# or reach out to @aposataC for assistance
+
+# Assumption vs. Correctness Tests:
+# these unit tests do *not* test correctness of LMCache-side or vLLM-side logic
+# it is to ensure that assumptions LMCache makes about vLLM's interface are stable
+
+import pytest
+
+from vllm.platforms import current_platform
+
+
+def assumes(obj, attr, is_callable=False, is_instance_of=None):
+    import inspect
+    from dataclasses import is_dataclass
+
+    assumption_msg = (
+        f"LMCache connector currently assumes that {obj} has a(n) {attr} attribute"
+    )
+    if hasattr(obj, attr):
+        attr_value = getattr(obj, attr)
+    elif is_dataclass(obj) and attr in getattr(obj, "__dataclass_fields__", {}):
+        field = obj.__dataclass_fields__[attr]
+        field_type = field.type
+        origin = getattr(field_type, "__origin__", None)
+        if origin is not None:
+            field_type = origin
+        attr_value = field_type
+    else:
+        raise AssertionError(assumption_msg)
+    if is_callable:
+        assumption_msg += f" and that {obj}.{attr} is a callable"
+        assert callable(attr_value), assumption_msg
+    if is_instance_of:
+        assumption_msg += f" and that {obj}.{attr} is an instance of {is_instance_of}"
+        if isinstance(attr_value, property):
+            fget = attr_value.fget
+            assert fget is not None, f"Property {obj}.{attr} has no fget"
+            sig = inspect.signature(fget)
+            ret_anno = sig.return_annotation
+            assert ret_anno is not inspect._empty, (
+                f"Property {obj}.{attr} has no return annotation"
+            )
+            assert ret_anno == is_instance_of, assumption_msg
+        else:
+            if isinstance(attr_value, type):
+                assert attr_value is is_instance_of, assumption_msg
+            else:
+                assert isinstance(attr_value, is_instance_of), assumption_msg
+
+
+@pytest.mark.skipif(
+    current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
+)
+def test_multimodal_interface():
+    # protect against interface changes
+    from vllm.multimodal.inputs import PlaceholderRange
+
+    assumes(PlaceholderRange, "offset")
+    assumes(PlaceholderRange, "length")
+
+
+@pytest.mark.skipif(
+    current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
+)
+def test_config_interface():
+    # protect against interface changes
+    from vllm.config import VllmConfig
+    from vllm.config.cache import CacheConfig
+    from vllm.config.kv_transfer import KVTransferConfig
+    from vllm.config.model import ModelConfig
+    from vllm.config.parallel import ParallelConfig
+
+    assumes(VllmConfig, "model_config")
+    assumes(VllmConfig, "cache_config")
+    assumes(VllmConfig, "parallel_config")
+    assumes(VllmConfig, "kv_transfer_config")
+
+    assumes(KVTransferConfig, "kv_role")
+    assumes(KVTransferConfig, "kv_connector_extra_config")
+
+    assumes(ModelConfig, "use_mla", is_instance_of=bool)
+    assumes(ModelConfig, "dtype")
+    assumes(ModelConfig, "max_model_len")
+    assumes(ModelConfig, "get_vocab_size", is_callable=True)
+    assumes(ModelConfig, "get_num_attention_heads", is_callable=True)
+    assumes(ModelConfig, "get_num_kv_heads", is_callable=True)
+    assumes(ModelConfig, "get_head_size", is_callable=True)
+    assumes(ModelConfig, "get_num_layers", is_callable=True)
+    assumes(ModelConfig, "get_num_kv_heads", is_callable=True)
+    assumes(ModelConfig, "model")
+
+    assumes(ParallelConfig, "world_size")
+    assumes(ParallelConfig, "rank")
+    assumes(ParallelConfig, "tensor_parallel_size")
+    assumes(ParallelConfig, "pipeline_parallel_size")
+    assumes(ParallelConfig, "data_parallel_size_local")
+    assumes(ParallelConfig, "data_parallel_rank_local")
+
+    assumes(CacheConfig, "cache_dtype")
+    assumes(CacheConfig, "block_size")
+    assumes(CacheConfig, "gpu_memory_utilization")
+
+    # kv metadata minimal case
+    from vllm.utils.torch_utils import get_kv_cache_torch_dtype
+
+    model_config = ModelConfig(dtype="bfloat16")
+    parallel_config = ParallelConfig()
+    cache_config = CacheConfig(cache_dtype="bfloat16")
+    kv_dtype = get_kv_cache_torch_dtype(cache_config.cache_dtype, model_config.dtype)
+    use_mla = False
+    chunk_size = 256
+    num_layer = model_config.get_num_layers(parallel_config)
+    num_kv_head = model_config.get_num_kv_heads(parallel_config)
+    head_size = model_config.get_head_size()
+    kv_shape = (num_layer, 1 if use_mla else 2, chunk_size, num_kv_head, head_size)
+
+    # dummy lmcache metadata creation example
+    _ = (
+        model_config.model,
+        parallel_config.world_size,
+        parallel_config.rank,
+        "vllm",
+        kv_dtype,
+        kv_shape,
+        use_mla,
+    )
+
+
+@pytest.mark.skipif(
+    current_platform.is_rocm(), reason="Requires libcudart.so, not available on ROCm"
+)
+def test_request_interface():
+    # protect against interface changes
+    from types import NoneType
+
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.request import Request
+
+    sampling_params = SamplingParams(max_tokens=10)
+    sampling_params.update_from_generation_config({}, eos_token_id=100)
+
+    req = Request(
+        request_id="test_request",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=sampling_params,
+        pooling_params=None,
+        lora_request=None,
+    )
+    assumes(req, "mm_features", is_instance_of=(list, NoneType))
+    assumes(req, "request_id")
+    assumes(req, "priority")
+    assumes(req, "prompt_token_ids")
+    assumes(req, "sampling_params")
+    assumes(req, "num_tokens")
+    assumes(req, "kv_transfer_params", is_instance_of=(dict, NoneType))
+
+    from vllm.multimodal.inputs import MultiModalFeatureSpec
+
+    assumes(MultiModalFeatureSpec, "identifier")
+    assumes(MultiModalFeatureSpec, "mm_position")
+
+
+def test_new_request_interface():
+    # protect against interface changes
+    from vllm.v1.core.sched.output import NewRequestData
+
+    assumes(NewRequestData, "req_id")
+    assumes(NewRequestData, "block_ids")
+    assumes(NewRequestData, "prompt_token_ids")
+    assumes(NewRequestData, "sampling_params")
+
+
+def test_sampling_params_interface():
+    # protect against interface changes
+    from vllm.sampling_params import SamplingParams
+
+    assumes(SamplingParams, "extra_args")
+
+    # dumb example use case in LMCache
+    kv_transfer_params = {
+        "lmcache.tag.user": "example_user_1",
+        "lmcache.ttl": 60,
+    }
+    sampling_params = SamplingParams(
+        extra_args={"kv_transfer_params": kv_transfer_params}
+    )
+    assert sampling_params.extra_args["kv_transfer_params"] == kv_transfer_params
+
+
+def test_tp_interface():
+    # protect against interface changes
+    import inspect
+
+    from vllm.distributed.parallel_state import get_tp_group
+
+    sig = inspect.signature(get_tp_group)
+    GroupCoordinator = sig.return_annotation
+
+    assumes(GroupCoordinator, "broadcast", is_callable=True)
+    assumes(GroupCoordinator, "broadcast_object", is_callable=True)
+
+
+def test_forward_context_interface():
+    # protect against interface changes
+    from vllm.forward_context import ForwardContext
+
+    assumes(ForwardContext, "no_compile_layers", is_instance_of=dict)
+    assumes(ForwardContext, "virtual_engine")
+    assumes(ForwardContext, "attn_metadata")
+
+
+def test_scheduler_output_interface():
+    # protect against interface changes
+    from vllm.v1.core.sched.output import SchedulerOutput
+
+    assumes(SchedulerOutput, "finished_req_ids")
+    assumes(SchedulerOutput, "scheduled_new_reqs", is_instance_of=list)
+    assumes(SchedulerOutput, "num_scheduled_tokens", is_instance_of=dict)
+    assumes(SchedulerOutput, "scheduled_cached_reqs")
+
+    from vllm.v1.core.sched.output import CachedRequestData
+
+    assumes(CachedRequestData, "req_ids", is_instance_of=list)
+    assumes(CachedRequestData, "new_block_ids", is_instance_of=list)

third_party/vllm/tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -0,0 +1,566 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import importlib.util
+import os
+import subprocess
+from unittest.mock import MagicMock, patch
+
+import msgspec
+import pytest
+import torch
+import zmq
+
+from tests.conftest import _find_free_port
+from vllm.config import (
+    CacheConfig,
+    DeviceConfig,
+    KVTransferConfig,
+    ModelConfig,
+    SchedulerConfig,
+    VllmConfig,
+    set_current_vllm_config,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
+    MoRIIOAgentMetadata,
+    MoRIIOConnectorMetadata,
+    MoRIIOConstants,
+    zmq_ctx,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
+    KVConnectorRole,
+    MoRIIOConnector,
+    MoRIIOConnectorWorker,
+)
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import (
+    get_ip,
+    make_zmq_path,
+)
+
+from .utils import create_request, create_scheduler
+
+aiter_available = importlib.util.find_spec("aiter") is not None
+mori_available = importlib.util.find_spec("mori") is not None
+
+
+def _rdma_available() -> bool:
+    """Check if RDMA devices are available."""
+    try:
+        result = subprocess.run(["ibv_devinfo"], capture_output=True, text=True)
+        return "No IB devices found" not in result.stderr
+    except FileNotFoundError:
+        return False
+
+
+rdma_available = _rdma_available()
+
+pytestmark = pytest.mark.skipif(
+    not (current_platform.is_rocm() and mori_available),
+    reason="MoRIIOs are only available on ROCm with aiter package installed",
+)
+
+
+@pytest.fixture
+def mock_parallel_groups():
+    """Mock tensor/data parallel group functions for single-rank tests."""
+    mock_group = MagicMock()
+    mock_group.rank = 0
+    mock_group.local_rank = 0
+    mock_group.world_size = 1
+
+    with (
+        patch.multiple(
+            "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common",
+            get_tensor_model_parallel_rank=MagicMock(return_value=0),
+            get_tensor_model_parallel_world_size=MagicMock(return_value=0),
+        ),
+        patch.multiple(
+            "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector",
+            get_tensor_model_parallel_world_size=MagicMock(return_value=0),
+            get_world_group=MagicMock(return_value=mock_group),
+            get_tp_group=MagicMock(return_value=mock_group),
+        ),
+    ):
+        yield mock_group
+
+
+def _setup_kv_transfer_request(
+    request, remote_host="127.0.0.1", fake_port=4789, fake_transfer_id="0"
+):
+    """Setup KV transfer parameters for a request."""
+    request.kv_transfer_params.update(
+        {
+            "transfer_id": fake_transfer_id,
+            "remote_notify_port": fake_port,
+            "remote_block_ids": None,
+            "remote_host": remote_host,
+            "remote_port": fake_port,
+            "remote_handshake_port": fake_port,
+            "remote_engine_id": "test_engine",
+        }
+    )
+    return request
+
+
+class FakeMoRIIOWrapper:
+    # A fake MoRIIOWrapper for testing purposes
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def set_moriio_engine(self, moriio_engine):
+        pass
+
+    def set_backend_type(self, backend_type):
+        pass
+
+    def get_agent_metadata(self):
+        pass
+
+    def register_remote_engine(self, remote_packed_engine_metadata):
+        pass
+
+    def register_local_tensor(self, tensor: torch.Tensor):
+        pass
+
+    def get_unpack_memory_metadata(self, packed_memory_metadata):
+        pass
+
+    def build_session(self, local_memory_metadata, remote_memory_metadata):
+        pass
+
+    def read_remote_data(
+        self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
+    ):
+        pass
+
+    def write_remote_data(
+        self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
+    ):
+        pass
+
+    def write_remote_data_single(
+        self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0
+    ):
+        pass
+
+    def waiting_for_transfer_complete(self):
+        pass
+
+    def async_wait_reqid(self):
+        pass
+
+    def _handle_message(self, msg: bytes):
+        pass
+
+    def _handle_structured_message(self, data: dict):
+        pass
+
+    def _handle_completion_message(self, msg: str):
+        pass
+
+    def send_notify(self, req_ids, remote_ip, remote_port):
+        pass
+
+    def pop_finished_req_ids(self):
+        pass
+
+    def pop_finished_write_req_ids(self):
+        pass
+
+    def shutdown(self):
+        pass
+
+
+class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker):
+    # Define a fake remote engine id for testing
+    REMOTE_ENGINE_ID = "remote_engine"
+
+    def __init__(
+        self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+
+
+def create_vllm_config(
+    model: str = "facebook/opt-125m",
+    max_num_seqs: int = 16,
+    max_num_batched_tokens: int = 64,
+    block_size: int = 16,
+    max_model_len: int = 10000,
+    enable_chunked_prefill: bool = True,
+    enable_permute_local_kv: bool = False,
+    role="kv_consumer",
+) -> VllmConfig:
+    """Initialize VllmConfig for testing."""
+    scheduler_config = SchedulerConfig(
+        max_num_seqs=max_num_seqs,
+        max_num_batched_tokens=max_num_batched_tokens,
+        max_model_len=max_model_len,
+        enable_chunked_prefill=enable_chunked_prefill,
+        is_encoder_decoder=False,
+    )
+    model_config = ModelConfig(
+        model=model,
+        trust_remote_code=True,
+        dtype="bfloat16",
+        seed=42,
+    )
+    # Cache config, optionally force APC
+    cache_config = CacheConfig(
+        block_size=block_size,
+        gpu_memory_utilization=0.9,
+        cache_dtype="auto",
+        enable_prefix_caching=True,
+    )
+    kv_transfer_config = KVTransferConfig(
+        kv_connector="MoRIIOConnector",
+        kv_role=role,
+        enable_permute_local_kv=enable_permute_local_kv,
+    )
+    return VllmConfig(
+        scheduler_config=scheduler_config,
+        model_config=model_config,
+        cache_config=cache_config,
+        kv_transfer_config=kv_transfer_config,
+        device_config=DeviceConfig("cpu"),
+    )
+
+
+@pytest.fixture
+def moriio_read_mode():
+    """Force the connector into read mode via env for tests."""
+    os.environ["VLLM_MORIIO_CONNECTOR_READ_MODE"] = "True"
+    yield
+    # Cleanup after test
+    os.environ.pop("VLLM_MORIIO_CONNECTOR_READ_MODE", None)
+
+
+def test_write_mode_saves_local_block_ids():
+    """Write mode records local block ids in MoRIIOConnectorMetadata.reqs_to_save."""
+
+    # Setup Scheduler and Request
+    vllm_config = create_vllm_config(role="kv_producer")
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=True,
+        do_remote_prefill=False,
+    )
+    request_id = request.request_id
+
+    scheduler.add_request(request)
+
+    # Fake Config
+    request = _setup_kv_transfer_request(request)
+
+    # Remote Prefill, triggers MoRIIOConnectorMetadata.
+    scheduler_output = scheduler.schedule()
+    kv_connector_metadata = scheduler_output.kv_connector_metadata
+    assert kv_connector_metadata is not None, "kv_connector_metadata is None"
+    assert isinstance(kv_connector_metadata, MoRIIOConnectorMetadata)
+
+    assert len(kv_connector_metadata.reqs_to_save) == 1, (
+        "Unexpected number of reqs_to_save"
+    )
+    assert len(kv_connector_metadata.reqs_to_recv) == 0, (
+        "Unexpected number of reqs_to_recv"
+    )
+    assert len(kv_connector_metadata.reqs_to_send) == 0, (
+        "Unexpected number of reqs_to_send"
+    )
+    assert request_id in kv_connector_metadata.reqs_to_save, (
+        "Request ID not in reqs_to_save"
+    )
+    req_meta = kv_connector_metadata.reqs_to_save[request_id]
+
+    for block_id, block in zip(
+        req_meta.local_block_ids,
+        scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
+            request_id
+        ],
+    ):
+        assert block_id == block.block_id, f"{block_id} != {block.block_id}"
+
+
+def test_write_mode_with_chunked_prefill_saves_local_block_ids():
+    """Write mode with chunked prefill still records correct local block ids."""
+    # Setup Scheduler and Request
+    MAX_NUM_BATCHED_TOKENS = 64
+    NUM_TOKENS = MAX_NUM_BATCHED_TOKENS * 2 + MAX_NUM_BATCHED_TOKENS // 2
+
+    vllm_config = create_vllm_config(
+        max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS, role="kv_producer"
+    )
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=True,
+        do_remote_prefill=False,
+    )
+    request_id = request.request_id
+
+    scheduler.add_request(request)
+
+    # Fake Config
+    request = _setup_kv_transfer_request(request)
+
+    # Remote Prefill with chunked prefill, triggers multiple schedules.
+    expected_counts = [(0, 0, 0), (0, 0, 0), (1, 0, 0)]
+    kv_connector_metadata = None
+    for _, (expected_save, expected_recv, expected_send) in enumerate(expected_counts):
+        scheduler_output = scheduler.schedule()
+        kv_connector_metadata = scheduler_output.kv_connector_metadata
+
+        assert len(kv_connector_metadata.reqs_to_save) == expected_save
+        assert len(kv_connector_metadata.reqs_to_recv) == expected_recv
+        assert len(kv_connector_metadata.reqs_to_send) == expected_send
+    assert kv_connector_metadata is not None, "kv_connector_metadata is None"
+    assert request_id in kv_connector_metadata.reqs_to_save, (
+        "Request ID not in reqs_to_save"
+    )
+    req_meta = kv_connector_metadata.reqs_to_save[request_id]
+
+    for block_id, block in zip(
+        req_meta.local_block_ids,
+        scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
+            request_id
+        ],
+    ):
+        assert block_id == block.block_id, f"{block_id} != {block.block_id}"
+
+
+def test_read_mode_loads_remote_block_ids(moriio_read_mode):
+    """Read mode loads remote block ids into local cache mapping."""
+
+    # Setup Scheduler and Request
+    vllm_config = create_vllm_config(role="kv_consumer")
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=False,
+        do_remote_prefill=True,
+    )
+    request_id = request.request_id
+
+    scheduler.add_request(request)
+    block_list = scheduler.kv_cache_manager.coordinator.single_type_managers[
+        0
+    ].req_to_blocks[request_id]
+
+    request = _setup_kv_transfer_request(request)
+
+    # Set remote block ids to be fetched.
+    request.kv_transfer_params["remote_block_ids"] = block_list
+
+    # Remote Prefill, triggers MoRIIOConnectorMetadata.
+
+    scheduler_output = scheduler.schedule()
+    kv_connector_metadata = scheduler_output.kv_connector_metadata
+    assert kv_connector_metadata is not None, "kv_connector_metadata is None"
+    assert isinstance(kv_connector_metadata, MoRIIOConnectorMetadata), (
+        "kv_connector_metadata is not MoRIIOConnectorMetadata"
+    )
+    assert len(kv_connector_metadata.reqs_to_save) == 0, (
+        "Unexpected number of reqs_to_save"
+    )
+    assert len(kv_connector_metadata.reqs_to_recv) == 1, (
+        "Unexpected number of reqs_to_recv"
+    )
+    assert len(kv_connector_metadata.reqs_to_send) == 0, (
+        "Unexpected number of reqs_to_send"
+    )
+    assert request_id in kv_connector_metadata.reqs_to_recv, (
+        "Request ID not in reqs_to_recv"
+    )
+    req_meta = kv_connector_metadata.reqs_to_recv[request_id]
+
+    for block_id, block in zip(
+        req_meta.local_block_ids,
+        scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
+            request_id
+        ],
+    ):
+        assert block_id == block.block_id, f"{block_id} != {block.block_id}"
+
+
+@pytest.mark.skipif(
+    not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend"
+)
+@pytest.mark.skipif(not rdma_available, reason="No RDMA devices available")
+def test_register_kv_caches(mock_parallel_groups):
+    """Test that MoRIIOConnector.register_kv_caches correctly registers kv caches."""
+    ROLE = "kv_consumer"
+    IP = get_ip()
+    vllm_config = create_vllm_config(role=ROLE)
+    DEFAULT_PORT = 6301
+    TP_RANK = 0
+    DP_RANK = 0
+    from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
+
+    backend_cls = AiterFlashAttentionBackend
+
+    # Create test kv cache tensors using proper backend shape
+    kv_cache_shape = backend_cls.get_kv_cache_shape(
+        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+    )
+    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    kv_caches = {
+        "layer0": shared_tensor,
+        "layer1": unique_tensor,
+        "layer2": shared_tensor,
+    }
+
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.threading.Event"
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector.threading.Thread"
+        ),
+    ):
+        # Create connector
+        vllm_config.kv_transfer_config.kv_connector_extra_config.update(
+            {
+                "proxy_ip": "127.0.0.1",
+                "proxy_ping_port": 12345,
+                "http_port": 12346,
+            }
+        )
+
+        with set_current_vllm_config(vllm_config):
+            connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)
+            connector.connector_worker = FakeMoRIIOConnectorWorker(
+                vllm_config, connector.engine_id, hand_shake_latency=0
+            )
+
+        from mori.io import (
+            MemoryDesc,
+        )
+
+        # Execute register_kv_caches
+        connector.register_kv_caches(kv_caches)
+
+        # Verify that the MemoryDesc stored in layer_name_to_local_kv_cache_metadata
+        assert (
+            shared_tensor.data_ptr()
+            == MemoryDesc.unpack(
+                connector.connector_worker.layer_name_to_local_kv_cache_metadata[
+                    "layer0"
+                ][0]
+            ).data
+        )
+        assert (
+            unique_tensor.data_ptr()
+            == MemoryDesc.unpack(
+                connector.connector_worker.layer_name_to_local_kv_cache_metadata[
+                    "layer1"
+                ][0]
+            ).data
+        )
+        assert (
+            shared_tensor.data_ptr()
+            == MemoryDesc.unpack(
+                connector.connector_worker.layer_name_to_local_kv_cache_metadata[
+                    "layer2"
+                ][0]
+            ).data
+        )
+
+        # Verify engine keys
+        expected_engine_key = f"{ROLE[3:]}:{IP}:{DEFAULT_PORT}:tp{TP_RANK}:dp{DP_RANK}"
+        assert (
+            MemoryDesc.unpack(
+                connector.connector_worker.layer_name_to_local_kv_cache_metadata[
+                    "layer0"
+                ][0]
+            ).engine_key
+            == expected_engine_key
+        )
+
+
+@pytest.mark.skipif(
+    not aiter_available, reason="Requires aiter package for ROCm FlashAttention backend"
+)
+@pytest.mark.skipif(not rdma_available, reason="No RDMA devices available")
+def test_moriio_handshake_returns_metadata(mock_parallel_groups):
+    """MoRIIO handshake socket returns valid agent metadata over ZMQ."""
+
+    ROLE = "kv_consumer"
+    vllm_config = create_vllm_config(role=ROLE)
+    from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
+
+    backend_cls = AiterFlashAttentionBackend
+
+    # Create test kv cache tensors using proper backend shape
+    kv_cache_shape = backend_cls.get_kv_cache_shape(
+        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+    )
+    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    kv_caches = {
+        "layer0": shared_tensor,
+        "layer1": unique_tensor,
+        "layer2": shared_tensor,
+    }
+
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine.MoRIIOWrapper",
+            FakeMoRIIOWrapper,
+        ),
+    ):
+        handshake_port = _find_free_port()
+        # Create connector
+        vllm_config.kv_transfer_config.kv_connector_extra_config.update(
+            {
+                "proxy_ip": "127.0.0.1",
+                "proxy_ping_port": 12345,
+                "http_port": 12346,
+                "handshake_port": handshake_port,
+            }
+        )
+        with set_current_vllm_config(vllm_config):
+            connector = MoRIIOConnector(vllm_config, KVConnectorRole.WORKER)
+
+        # Execute register_kv_caches
+        connector.register_kv_caches(kv_caches)
+
+        # Connect to handshake socket and request metadata
+        path = make_zmq_path("tcp", "127.0.0.1", handshake_port)
+        with zmq_ctx(zmq.DEALER, path) as sock:
+            sock.send(MoRIIOConstants.GET_META_MSG)
+            received_frame = sock.recv_multipart()
+
+            if len(received_frame) != 2 or received_frame[0] != b"":
+                raise ValueError(f"Unexpected frame! {received_frame = }")
+
+            metadata_bytes = received_frame[1]
+            decoder = msgspec.msgpack.Decoder(MoRIIOAgentMetadata)
+            metadata = decoder.decode(metadata_bytes)
+            assert isinstance(metadata, MoRIIOAgentMetadata), (
+                "Decoded metadata is not MoRIIOAgentMetadata"
+            )

third_party/vllm/tests/v1/kv_connector/unit/test_multi_connector.py

@@ -0,0 +1,920 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import filecmp
+import shutil
+import tempfile
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from tests.v1.kv_connector.unit.utils import create_vllm_config
+from vllm import LLM, SamplingParams
+from vllm.config import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
+from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
+    MultiConnector,
+    MultiKVConnectorStats,
+    MultiKVConnectorWorkerMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+    NixlKVConnectorStats,
+)
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import KVConnectorOutput, KVConnectorWorkerMetadata
+
+MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
+
+PROMPT_CONTEXT = "Hi " * 100
+PROMPTS = [
+    PROMPT_CONTEXT + "Hello, my name is",
+    PROMPT_CONTEXT + "The capital of France is",
+]
+
+SAMPLING_PARAMS = SamplingParams(temperature=0, max_tokens=20)
+
+
+# Test connector with custom stats for testing MultiConnector
+class MockConnectorStats(KVConnectorStats):
+    """Mock stats class for testing."""
+
+    pass
+
+
+class MockConnector(KVConnectorBase_V1):
+    """Mock connector for testing."""
+
+    def __new__(cls, *args, **kwargs):
+        # mock all KVConnectorBase_V1 functions
+        mock = MagicMock(spec_set=KVConnectorBase_V1)
+        # Override just build_kv_connector_stats
+        mock.build_kv_connector_stats = cls.build_kv_connector_stats
+        return mock
+
+    @classmethod
+    def build_kv_connector_stats(
+        cls, data: dict[str, Any] | None = None
+    ) -> KVConnectorStats | None:
+        return MockConnectorStats(data=data) if data is not None else None
+
+    def start_load_kv(self, forward_context, **kwargs):
+        pass
+
+    def wait_for_layer_load(self, layer_name):
+        pass
+
+    def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
+        pass
+
+    def wait_for_save(self):
+        pass
+
+    def build_connector_meta(self, scheduler_output):
+        return None
+
+    def get_num_new_matched_tokens(self, request, num_computed_tokens):
+        return (0, False)
+
+    def update_state_after_alloc(self, request, blocks, num_tokens) -> None:
+        pass
+
+
+# Register the mock connector
+KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
+
+
+@pytest.fixture
+def mc() -> MultiConnector:
+    """MultiConnector using two mocked connectors"""
+    vllm_config = create_vllm_config()
+
+    mock_connector_config = {
+        "kv_connector": "MockConnector",
+        "kv_role": "kv_both",
+        "kv_connector_module_path": "tests.v1.kv_connector.unit.test_multi_connector",
+    }
+
+    vllm_config.kv_transfer_config = KVTransferConfig(
+        kv_connector="MultiConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={
+            "connectors": [mock_connector_config, mock_connector_config],
+        },
+    )
+
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
+    )
+
+    mc = MultiConnector(
+        vllm_config=vllm_config,
+        role=KVConnectorRole.WORKER,
+        kv_cache_config=kv_cache_config,
+    )
+
+    return mc
+
+
+# Helper function to compare directories recursively
+def _compare_directories(dir1: Path, dir2: Path) -> bool:
+    """Compares two directories recursively for identical content."""
+    dcmp = filecmp.dircmp(dir1, dir2)
+    if dcmp.left_only or dcmp.right_only or dcmp.diff_files:
+        print(f"Differences found between {dir1} and {dir2}:")
+        print(f"  Left only: {dcmp.left_only}")
+        print(f"  Right only: {dcmp.right_only}")
+        print(f"  Different files: {dcmp.diff_files}")
+        return False
+    for sub_dir in dcmp.common_dirs:
+        if not _compare_directories(dir1 / sub_dir, dir2 / sub_dir):
+            return False
+    return True
+
+
+def test_multi_example_connector_consistency():
+    """
+    Tests that MultiConnector with two ExampleConnectors saves
+    identical KV cache data to separate storage locations.
+    """
+    storage_1_path = Path("storage_1/")
+    storage_2_path = Path("storage_2/")
+    shutil.rmtree(storage_1_path, ignore_errors=True)
+    shutil.rmtree(storage_2_path, ignore_errors=True)
+    storage_1_path.mkdir()
+    storage_2_path.mkdir()
+
+    # Configure MultiConnector with two ExampleConnectors
+    kv_transfer_config = KVTransferConfig(
+        kv_connector="MultiConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={
+            "connectors": [
+                {
+                    "kv_connector": "TestExampleConnector",
+                    "kv_role": "kv_both",
+                    "kv_connector_extra_config": {
+                        "shared_storage_path": str(storage_1_path),
+                        "name": "storage1",
+                    },
+                    "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",
+                },
+                {
+                    "kv_connector": "TestExampleConnector",
+                    "kv_role": "kv_both",
+                    "kv_connector_extra_config": {
+                        "shared_storage_path": str(storage_2_path),
+                        "name": "storage2",
+                    },
+                    "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",
+                },
+            ]
+        },
+    )
+
+    llm = LLM(
+        model=MODEL_NAME,
+        enforce_eager=True,
+        gpu_memory_utilization=0.5,
+        kv_transfer_config=kv_transfer_config,
+    )
+    # Run generation - this should trigger saving KV cache
+    # Use a single prompt to avoid race conditions depending on the order of scheduling
+    _ = llm.generate(PROMPTS[0], SAMPLING_PARAMS)
+
+    # --- Verification ---
+
+    # Check that both storage directories were populated
+    local_subdirs = list(storage_1_path.iterdir())
+    external_subdirs = list(storage_2_path.iterdir())
+
+    assert len(local_subdirs) > 0, (
+        f"Local storage path {storage_1_path} is empty after generation."
+    )
+    assert len(external_subdirs) > 0, (
+        f"External storage path {storage_2_path} is empty after generation."
+    )
+    assert len(local_subdirs) == len(external_subdirs), (
+        f"Mismatch in number of cache entries: "
+        f"Local={len(local_subdirs)}, External={len(external_subdirs)}"
+    )
+
+    # The subdirectories should correspond to the prompt hashes
+    # Since prompts are the same, the hash directories should be the same name
+    local_subdir_names = sorted([d.name for d in local_subdirs])
+    external_subdir_names = sorted([d.name for d in external_subdirs])
+    assert local_subdir_names == external_subdir_names, (
+        "Cache directory names do not match between local and external storage"
+    )
+
+    # Compare the contents of each corresponding cache directory
+    for subdir_name in local_subdir_names:
+        print(f"Comparing contents of cache directory: {subdir_name}")
+        assert _compare_directories(
+            storage_1_path / subdir_name, storage_2_path / subdir_name
+        ), (
+            f"Contents differ for cache directory '{subdir_name}' between "
+            f"{storage_1_path} and {storage_2_path}"
+        )
+
+    events = get_connector_events()
+    # First event is set_xfer_handshake_metadata from initialization, then
+    # get_num_new_matched_tokens and update_state_after_alloc from generate().
+    assert events["storage1-SCHEDULER"][:4] == [
+        "set_xfer_handshake_metadata",
+        "get_num_new_matched_tokens 0",
+        "update_state_after_alloc num_blocks=[0] 0",
+        "build_connector_meta",
+    ]
+    # First three events are from initialization (register_kv_caches,
+    # set_host_xfer_buffer_ops, get_handshake_metadata), then generate() events.
+    assert events["storage1-WORKER"][:7] == [
+        "register_kv_caches",
+        "set_host_xfer_buffer_ops",
+        "get_handshake_metadata",
+        "bind_connector_metadata",
+        "start_load_kv",
+        "wait_for_layer_load",
+        "save_kv_layer",
+    ]
+    assert events["storage2-SCHEDULER"][:4] == [
+        "set_xfer_handshake_metadata",
+        "get_num_new_matched_tokens 0",
+        "update_state_after_alloc num_blocks=[0] 0",
+        "build_connector_meta",
+    ]
+    assert events["storage2-WORKER"][:7] == [
+        "register_kv_caches",
+        "set_host_xfer_buffer_ops",
+        "get_handshake_metadata",
+        "bind_connector_metadata",
+        "start_load_kv",
+        "wait_for_layer_load",
+        "save_kv_layer",
+    ]
+
+    # Reset prefix cache or else we'll just get the tokens back from there.
+    llm.reset_prefix_cache()
+
+    # Run generation again - this should trigger loading from the first
+    # connector.
+    _ = llm.generate(PROMPTS[1], SAMPLING_PARAMS)
+
+    events = get_connector_events()
+    # get_num_new_matched_tokens will return new tokens from the first
+    # connector so update_state_after_alloc will be with allocated blocks
+    # on that one but with zero blocks for others (first nonzero match is
+    # chosen).
+    assert events["storage1-SCHEDULER"][:3] == [
+        "get_num_new_matched_tokens 0",
+        "update_state_after_alloc num_blocks=[7] 96",
+        "build_connector_meta",
+    ]
+    assert events["storage2-SCHEDULER"][:3] == [
+        "get_num_new_matched_tokens 0",
+        "update_state_after_alloc num_blocks=[0] 0",
+        "build_connector_meta",
+    ]
+
+    # Delete storage1 connector state
+    shutil.rmtree(storage_1_path)
+
+    # Reset prefix cache or else we'll just get the tokens back from there.
+    llm.reset_prefix_cache()
+
+    # Run generation again - this should trigger loading from the first
+    # connector.
+    _ = llm.generate(PROMPTS[0], SAMPLING_PARAMS)
+
+    events = get_connector_events()
+    # get_num_new_matched_tokens will be called for both connectors but will
+    # return 0 from the first connector, but the second connector should have
+    # a hit, so update_state_after_alloc will only be called with allocated
+    # blocks for the second connector.
+    assert events["storage1-SCHEDULER"][:3] == [
+        "get_num_new_matched_tokens 0",
+        "update_state_after_alloc num_blocks=[0] 0",
+        "build_connector_meta",
+    ]
+    assert events["storage2-SCHEDULER"][:3] == [
+        "get_num_new_matched_tokens 0",
+        "update_state_after_alloc num_blocks=[7] 96",
+        "build_connector_meta",
+    ]
+
+    # Clean up
+    shutil.rmtree(storage_1_path)
+    shutil.rmtree(storage_2_path)
+
+
+def get_connector_events() -> dict[str, list[str]]:
+    # Read in connector events and reset the files.
+    import glob
+
+    event_files = glob.glob(tempfile.gettempdir() + "/connector_*_events.log")
+    connector_events = {}
+    for fname in event_files:
+        name = fname.split("connector_")[1].split("_events.log")[0]
+        try:
+            with open(fname, "r+") as f:
+                connector_events[name] = [line.strip() for line in f if line.strip()]
+                f.truncate(0)
+        except Exception as e:
+            print(f"[ERROR] Could not read connector events for {name}: {e}")
+
+    return connector_events
+
+
+def test_engine_id_conflict():
+    configs = [KVTransferConfig() for _ in range(2)]
+    ids = [config.engine_id for config in configs]
+    assert ids[0] != ids[1], (
+        f"Engine IDs should be different for different configs. Got {ids}"
+    )
+
+
+def test_multi_connector_handle_preemptions_integration():
+    """
+    Integration test: verify MultiConnector delegates handle_preemptions
+    to all sub-connectors.
+
+    Uses TestExampleConnector which logs all method calls to temp files.
+    This test directly calls handle_preemptions on a MultiConnector with
+    TestExampleConnector sub-connectors and verifies the calls are logged.
+    """
+    from tests.v1.kv_connector.unit.utils import (
+        create_scheduler,
+        create_vllm_config,
+    )
+
+    storage_path = Path(tempfile.mkdtemp())
+
+    try:
+        # Configure MultiConnector with two TestExampleConnectors
+        kv_transfer_config = KVTransferConfig(
+            kv_connector="MultiConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "connectors": [
+                    {
+                        "kv_connector": "TestExampleConnector",
+                        "kv_role": "kv_both",
+                        "kv_connector_extra_config": {
+                            "shared_storage_path": str(storage_path / "s1"),
+                            "name": "preempt1",
+                        },
+                        "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",
+                    },
+                    {
+                        "kv_connector": "TestExampleConnector",
+                        "kv_role": "kv_both",
+                        "kv_connector_extra_config": {
+                            "shared_storage_path": str(storage_path / "s2"),
+                            "name": "preempt2",
+                        },
+                        "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",
+                    },
+                ]
+            },
+        )
+
+        vllm_config = create_vllm_config(
+            block_size=16,
+            max_num_batched_tokens=100,
+            kv_connector_extra_config=kv_transfer_config.kv_connector_extra_config,
+        )
+        vllm_config.kv_transfer_config = kv_transfer_config
+
+        # Create scheduler - this initializes the MultiConnector with SCHEDULER role
+        scheduler = create_scheduler(vllm_config, num_blocks=10)
+
+        # Clear any events from initialization
+        get_connector_events()
+
+        # Directly call handle_preemptions on the scheduler's connector
+        # Note: handle_preemptions is normally a worker-side method, but we're
+        # testing the delegation behavior of MultiConnector here.
+        # The connector attribute contains the KV connector.
+        assert scheduler.connector is not None, "Scheduler should have a connector"
+        preempted_req_ids = {"req-1", "req-2", "req-3"}
+        scheduler.connector.handle_preemptions(preempted_req_ids)
+
+        # Verify both connectors received the handle_preemptions call
+        events = get_connector_events()
+
+        # Both SCHEDULER-role connectors should have logged handle_preemptions
+        assert "handle_preemptions" in events.get("preempt1-SCHEDULER", []), (
+            f"preempt1-SCHEDULER should have handle_preemptions call. "
+            f"Got events: {events}"
+        )
+        assert "handle_preemptions" in events.get("preempt2-SCHEDULER", []), (
+            f"preempt2-SCHEDULER should have handle_preemptions call. "
+            f"Got events: {events}"
+        )
+
+    finally:
+        # Cleanup
+        shutil.rmtree(storage_path, ignore_errors=True)
+
+
+class TestMultiConnectorStats:
+    """Tests for MultiConnector stats reconstruction and operations."""
+
+    def test_build_kv_connector_stats_with_none(self):
+        """Test that build_kv_connector_stats returns empty stats when given None."""
+        stats = MultiConnector.build_kv_connector_stats(data=None)
+
+        assert stats is not None
+        assert isinstance(stats, MultiKVConnectorStats)
+        assert len(stats.data) == 0
+        assert stats.is_empty()
+
+    def test_build_kv_connector_stats_with_empty_dict(self):
+        """Test that build_kv_connector_stats returns empty stats with empty dict."""
+        stats = MultiConnector.build_kv_connector_stats(data={})
+
+        assert stats is not None
+        assert isinstance(stats, MultiKVConnectorStats)
+        assert len(stats.data) == 0
+        assert stats.is_empty()
+
+    def test_build_kv_connector_stats_reconstructs_nixl_stats(self):
+        """Test that NixlConnector stats are properly reconstructed with
+        correct data."""
+        serialized_data = {
+            "NixlConnector": {
+                "data": {
+                    "transfer_duration": [1.5, 2.3],
+                    "post_duration": [0.1, 0.2],
+                    "bytes_transferred": [1024, 2048],
+                    "num_descriptors": [10, 20],
+                    "num_failed_transfers": [],
+                    "num_failed_notifications": [],
+                }
+            }
+        }
+
+        stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
+
+        assert "NixlConnector" in stats.data
+        nixl_stats = stats.data["NixlConnector"]
+        assert isinstance(nixl_stats, NixlKVConnectorStats)
+        assert nixl_stats.data["transfer_duration"] == [1.5, 2.3]
+        assert nixl_stats.data["post_duration"] == [0.1, 0.2]
+        assert nixl_stats.data["bytes_transferred"] == [1024, 2048]
+        assert nixl_stats.data["num_descriptors"] == [10, 20]
+
+    def test_build_kv_connector_stats_with_multiple_connectors(self):
+        """Test reconstruction with multiple connector types that have custom stats."""
+        serialized_data = {
+            "NixlConnector": {
+                "data": {
+                    "transfer_duration": [1.5],
+                    "post_duration": [0.1],
+                    "bytes_transferred": [1024],
+                    "num_descriptors": [10],
+                    "num_failed_transfers": [],
+                    "num_failed_notifications": [],
+                }
+            },
+            "MockConnector": {"data": {"mock_field": [1, 2, 3]}},
+        }
+
+        stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
+
+        assert stats is not None
+        assert isinstance(stats, MultiKVConnectorStats)
+        # Both connectors should be reconstructed
+        assert len(stats.data) == 2
+        assert "NixlConnector" in stats.data
+        assert "MockConnector" in stats.data
+        assert isinstance(stats.data["NixlConnector"], NixlKVConnectorStats)
+        assert isinstance(stats.data["MockConnector"], MockConnectorStats)
+        # Verify data is preserved
+        assert stats.data["MockConnector"].data == {"mock_field": [1, 2, 3]}
+
+    def test_build_kv_connector_stats_raises_error_for_unknown_connector(self):
+        """Test that unknown connectors raise an error."""
+        serialized_data = {
+            "UnknownConnector": {"data": {"some_field": [1, 2, 3]}},
+            "NixlConnector": {
+                "data": {
+                    "transfer_duration": [1.5],
+                    "post_duration": [0.1],
+                    "bytes_transferred": [1024],
+                    "num_descriptors": [10],
+                    "num_failed_transfers": [],
+                    "num_failed_notifications": [],
+                }
+            },
+        }
+
+        with pytest.raises(
+            ValueError, match="Connector 'UnknownConnector' is not registered."
+        ):
+            MultiConnector.build_kv_connector_stats(data=serialized_data)
+
+    def test_build_kv_connector_stats_with_already_instantiated_objects(self):
+        """Test that already-instantiated stats objects are preserved (same process)."""
+        # This simulates the in-process case where stats are not serialized
+        nixl_stats = NixlKVConnectorStats(
+            data={
+                "transfer_duration": [1.5],
+                "post_duration": [0.1],
+                "bytes_transferred": [1024],
+                "num_descriptors": [10],
+                "num_failed_transfers": [],
+                "num_failed_notifications": [],
+            }
+        )
+        mock_stats = MockConnectorStats(data={"mock_field": [1, 2, 3]})
+
+        data_with_objects = {
+            "NixlConnector": nixl_stats,
+            "MockConnector": mock_stats,
+        }
+
+        stats = MultiConnector.build_kv_connector_stats(data=data_with_objects)
+
+        assert stats is not None
+        assert isinstance(stats, MultiKVConnectorStats)
+        assert len(stats.data) == 2
+        # Verify objects are preserved as-is
+        assert stats.data["NixlConnector"] is nixl_stats
+        assert stats.data["MockConnector"] is mock_stats
+
+    def test_build_kv_connector_stats_with_mixed_objects_and_dicts(self):
+        """Test handling mixed already-instantiated and serialized stats."""
+        # This can happen during transition or partial serialization
+        nixl_stats = NixlKVConnectorStats(
+            data={
+                "transfer_duration": [1.5],
+                "post_duration": [0.1],
+                "bytes_transferred": [1024],
+                "num_descriptors": [10],
+                "num_failed_transfers": [],
+                "num_failed_notifications": [],
+            }
+        )
+
+        mixed_data = {
+            "NixlConnector": nixl_stats,  # Already instantiated
+            "MockConnector": {"data": {"mock_field": [1, 2, 3]}},  # Serialized
+        }
+
+        stats = MultiConnector.build_kv_connector_stats(data=mixed_data)
+
+        assert stats is not None
+        assert isinstance(stats, MultiKVConnectorStats)
+        assert len(stats.data) == 2
+        # Instantiated object preserved
+        assert stats.data["NixlConnector"] is nixl_stats
+        # Serialized object reconstructed
+        assert isinstance(stats.data["MockConnector"], MockConnectorStats)
+        assert stats.data["MockConnector"].data == {"mock_field": [1, 2, 3]}
+
+    def test_build_kv_connector_stats_skips_connectors_without_custom_stats(self):
+        """Test that connectors without custom stats (return None) are skipped."""
+        # ExampleConnector doesn't override build_kv_connector_stats,
+        # so it returns None and should be skipped
+        serialized_data = {
+            "NixlConnector": {
+                "data": {
+                    "transfer_duration": [1.5],
+                    "post_duration": [0.1],
+                    "bytes_transferred": [1024],
+                    "num_descriptors": [10],
+                    "num_failed_transfers": [],
+                    "num_failed_notifications": [],
+                }
+            },
+            "ExampleConnector": {"data": {"some_field": [1, 2, 3]}},
+        }
+
+        stats = MultiConnector.build_kv_connector_stats(data=serialized_data)
+
+        assert stats is not None
+        assert isinstance(stats, MultiKVConnectorStats)
+        # Only NixlConnector should be reconstructed
+        assert len(stats.data) == 1
+        assert "NixlConnector" in stats.data
+        assert isinstance(stats.data["NixlConnector"], NixlKVConnectorStats)
+        # ExampleConnector should be skipped (returns None)
+        assert "ExampleConnector" not in stats.data
+
+    def test_build_kv_connector_stats_handles_malformed_data(self):
+        """Test that malformed data raises appropriate errors."""
+        serialized_data = {
+            "NixlConnector": {"wrong_field": {"transfer_duration": [1.5]}}
+        }
+
+        with pytest.raises(AssertionError, match="Expected a dict with a 'data' field"):
+            MultiConnector.build_kv_connector_stats(data=serialized_data)
+
+    def test_aggregate_same_connector(self):
+        """Test aggregating stats from the same connector type."""
+        stats1 = MultiKVConnectorStats(
+            data={
+                "NixlConnector": NixlKVConnectorStats(
+                    data={
+                        "transfer_duration": [1.0],
+                        "post_duration": [0.1],
+                        "bytes_transferred": [1024],
+                        "num_descriptors": [10],
+                        "num_failed_transfers": [],
+                        "num_failed_notifications": [],
+                    }
+                )
+            }
+        )
+
+        stats2 = MultiKVConnectorStats(
+            data={
+                "NixlConnector": NixlKVConnectorStats(
+                    data={
+                        "transfer_duration": [2.0],
+                        "post_duration": [0.2],
+                        "bytes_transferred": [2048],
+                        "num_descriptors": [20],
+                        "num_failed_transfers": [],
+                        "num_failed_notifications": [],
+                    }
+                )
+            }
+        )
+
+        result = stats1.aggregate(stats2)
+
+        assert result is stats1  # Should return self
+        assert "NixlConnector" in result.data
+        nixl_stats = result.data["NixlConnector"]
+        assert nixl_stats.data["transfer_duration"] == [1.0, 2.0]
+        assert nixl_stats.data["post_duration"] == [0.1, 0.2]
+        assert nixl_stats.data["bytes_transferred"] == [1024, 2048]
+        assert nixl_stats.data["num_descriptors"] == [10, 20]
+
+    def test_aggregate_new_connector(self):
+        """Test aggregating stats when a new connector type appears."""
+        from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+            KVConnectorStats,
+        )
+
+        stats1 = MultiKVConnectorStats(
+            data={
+                "NixlConnector": NixlKVConnectorStats(
+                    data={
+                        "transfer_duration": [1.0],
+                        "post_duration": [0.1],
+                        "bytes_transferred": [1024],
+                        "num_descriptors": [10],
+                        "num_failed_transfers": [],
+                        "num_failed_notifications": [],
+                    }
+                )
+            }
+        )
+
+        stats2 = MultiKVConnectorStats(
+            data={"ExampleConnector": KVConnectorStats(data={"field": [1, 2]})}
+        )
+
+        result = stats1.aggregate(stats2)
+
+        assert "NixlConnector" in result.data
+        assert "ExampleConnector" in result.data
+
+    def test_reduce(self):
+        """Test that reduce() correctly reduces all nested connector stats."""
+        stats = MultiKVConnectorStats(
+            data={
+                "NixlConnector": NixlKVConnectorStats(
+                    data={
+                        "transfer_duration": [1.0, 2.0],
+                        "post_duration": [0.1, 0.2],
+                        "bytes_transferred": [1024, 2048],
+                        "num_descriptors": [10, 20],
+                        "num_failed_transfers": [],
+                        "num_failed_notifications": [],
+                    }
+                )
+            }
+        )
+
+        reduced = stats.reduce()
+
+        assert "NixlConnector" in reduced
+        assert isinstance(reduced["NixlConnector"], dict)
+        # Check that the stats were reduced (should have aggregated values)
+        assert "Num successful transfers" in reduced["NixlConnector"]
+        assert reduced["NixlConnector"]["Num successful transfers"] == 2
+
+    def test_reset(self):
+        """Test that reset() resets all nested connector stats."""
+        stats = MultiKVConnectorStats(
+            data={
+                "NixlConnector": NixlKVConnectorStats(
+                    data={
+                        "transfer_duration": [1.0, 2.0],
+                        "post_duration": [0.1, 0.2],
+                        "bytes_transferred": [1024, 2048],
+                        "num_descriptors": [10, 20],
+                        "num_failed_transfers": [],
+                        "num_failed_notifications": [],
+                    }
+                )
+            }
+        )
+
+        assert not stats.is_empty()
+
+        stats.reset()
+
+        # After reset, stats should be empty
+        assert stats.is_empty()
+        nixl_stats = stats.data["NixlConnector"]
+        assert len(nixl_stats.data["transfer_duration"]) == 0
+
+    def test_is_empty_with_multiple_connectors(self):
+        """Test is_empty() returns correct value with multiple connectors."""
+        # All empty
+        stats = MultiKVConnectorStats(
+            data={
+                "NixlConnector": NixlKVConnectorStats(data={}),
+            }
+        )
+        # Initialize empty stats
+        stats.data["NixlConnector"].reset()
+        assert stats.is_empty()
+
+        # One non-empty
+        stats.data["NixlConnector"].data["transfer_duration"].append(1.0)
+        assert not stats.is_empty()
+
+
+def test_multi_connector_overrides_all_base_methods():
+    """
+    Ensure MultiConnector overrides all public methods from KVConnectorBase_V1.
+    """
+    # These are fine to inherit from KVConnectorBase_V1
+    # TODO(https://github.com/vllm-project/vllm/pull/31811): Remove
+    # get_kv_connector_kv_cache_events from INHERITED_OK once implemented.
+    INHERITED_OK = {
+        "role",
+        "has_connector_metadata",
+        "get_kv_connector_kv_cache_events",
+    }
+
+    base_members = {
+        name for name in dir(KVConnectorBase_V1) if not name.startswith("_")
+    } - KVConnectorBase_V1.__abstractmethods__
+
+    missing = [
+        name
+        for name in sorted(base_members)
+        if name not in INHERITED_OK and name not in MultiConnector.__dict__
+    ]
+
+    if missing:
+        pytest.fail(f"""
+MultiConnector does not override these KVConnectorBase_V1 methods: {missing}
+
+MultiConnector wraps other connectors and must delegate all methods.
+Please add overrides that delegate to self._connectors.
+
+Options:
+  1. Add delegation in MultiConnector (preferred)
+  2. Add to INHERITED_OK if the base implementation works correctly
+""")
+
+
+def test_multi_connector_prefer_cross_layer_blocks(mc):
+    mc._connectors[0].prefer_cross_layer_blocks = False
+    mc._connectors[1].prefer_cross_layer_blocks = True
+    assert mc.prefer_cross_layer_blocks is False
+
+    mc._connectors[0].prefer_cross_layer_blocks = True
+    mc._connectors[1].prefer_cross_layer_blocks = True
+    assert mc.prefer_cross_layer_blocks is True
+
+
+def test_multi_connector_worker_metadata(mc):
+    class MockConnectorWorkerMetadata(KVConnectorWorkerMetadata):
+        def __init__(self, data: set[str]):
+            self.data = data
+
+    class MockConnectorWorkerMetadata0(MockConnectorWorkerMetadata):
+        def aggregate(
+            self, other: KVConnectorWorkerMetadata
+        ) -> KVConnectorWorkerMetadata:
+            assert isinstance(other, MockConnectorWorkerMetadata)
+            return MockConnectorWorkerMetadata0(data=self.data | other.data)
+
+    class MockConnectorWorkerMetadata1(MockConnectorWorkerMetadata):
+        def aggregate(
+            self, other: KVConnectorWorkerMetadata
+        ) -> KVConnectorWorkerMetadata:
+            assert isinstance(other, MockConnectorWorkerMetadata)
+            return MockConnectorWorkerMetadata1(data=self.data | other.data)
+
+    # -------------------- test build_worker_connector_meta -------------------
+
+    # both connectors return None
+    mc._connectors[0].build_connector_worker_meta.return_value = None
+    mc._connectors[1].build_connector_worker_meta.return_value = None
+    assert mc.build_connector_worker_meta() is None
+
+    # only first connector returns None
+    worker_meta1a = MockConnectorWorkerMetadata1({"1a"})
+    mc._connectors[0].build_connector_worker_meta.return_value = None
+    mc._connectors[1].build_connector_worker_meta.return_value = worker_meta1a
+    mc_worker_meta_none_1a = mc.build_connector_worker_meta()
+    assert isinstance(mc_worker_meta_none_1a, MultiKVConnectorWorkerMetadata)
+    assert mc_worker_meta_none_1a.metadata == (None, worker_meta1a)
+
+    # only second connector returns None
+    worker_meta0a = MockConnectorWorkerMetadata0({"0a"})
+    mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0a
+    mc._connectors[1].build_connector_worker_meta.return_value = None
+    mc_worker_meta_0a_none = mc.build_connector_worker_meta()
+    assert isinstance(mc_worker_meta_0a_none, MultiKVConnectorWorkerMetadata)
+    assert mc_worker_meta_0a_none.metadata == (worker_meta0a, None)
+
+    # both connectors do not return None
+    worker_meta0b = MockConnectorWorkerMetadata0({"0b"})
+    worker_meta1b = MockConnectorWorkerMetadata1({"1b"})
+    mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0b
+    mc._connectors[1].build_connector_worker_meta.return_value = worker_meta1b
+    mc_worker_meta_0b_1b = mc.build_connector_worker_meta()
+    assert isinstance(mc_worker_meta_0b_1b, MultiKVConnectorWorkerMetadata)
+    assert mc_worker_meta_0b_1b.metadata == (worker_meta0b, worker_meta1b)
+
+    # ----------------------------- test aggregate ----------------------------
+
+    # aggregate ({"0a"}, None) and (None, {"1a"}) -> ({"0a"}, {"1a"})
+    mc_worker_meta_0a_1a = mc_worker_meta_0a_none.aggregate(mc_worker_meta_none_1a)
+    assert isinstance(mc_worker_meta_0a_1a, MultiKVConnectorWorkerMetadata)
+    assert mc_worker_meta_0a_1a.metadata == (worker_meta0a, worker_meta1a)
+
+    # aggregate ({"0a"}, None) and ({"0b"}, None) -> ({"0a", "0b"}, None)
+    mc._connectors[0].build_connector_worker_meta.return_value = worker_meta0b
+    mc._connectors[1].build_connector_worker_meta.return_value = None
+    mc_worker_meta_0b_none = mc.build_connector_worker_meta()
+    mc_worker_meta_0a_0b = mc_worker_meta_0a_none.aggregate(mc_worker_meta_0b_none)
+    assert isinstance(mc_worker_meta_0a_0b, MultiKVConnectorWorkerMetadata)
+    assert mc_worker_meta_0a_0b.metadata[1] is None
+    connector0_md = mc_worker_meta_0a_0b.metadata[0]
+    assert isinstance(connector0_md, MockConnectorWorkerMetadata0)
+    assert connector0_md.data == {"0a", "0b"}
+
+    # aggregate ({"0a"}, {"1a"}) and ({"0b"}, {"1b"}) -> ({"0a", "0b"}, {"1a", "1b"})
+    mc_worker_meta_01a_01b = mc_worker_meta_0a_1a.aggregate(mc_worker_meta_0b_1b)
+    assert isinstance(mc_worker_meta_01a_01b, MultiKVConnectorWorkerMetadata)
+    metadata = mc_worker_meta_01a_01b.metadata
+    assert len(metadata) == 2
+    connector0_md, connector1_md = metadata
+    assert isinstance(connector0_md, MockConnectorWorkerMetadata0)
+    assert isinstance(connector1_md, MockConnectorWorkerMetadata1)
+    assert connector0_md.data == {"0a", "0b"}
+    assert connector1_md.data == {"1a", "1b"}
+
+    # ---------------------- test update_connector_output ---------------------
+
+    def verify_worker_metadata(expected_metadata: MockConnectorWorkerMetadata | None):
+        def _verify_worker_metadata(connector_output: KVConnectorOutput):
+            worker_meta = connector_output.kv_connector_worker_meta
+            if expected_metadata is None:
+                assert worker_meta is None
+                return
+
+            assert isinstance(worker_meta, MockConnectorWorkerMetadata)
+            assert type(worker_meta) is type(expected_metadata)
+            assert expected_metadata.data == worker_meta.data
+
+        return _verify_worker_metadata
+
+    def assert_update_connector_output_called(mc: MultiConnector):
+        for c in mc._connectors:
+            c.update_connector_output.assert_called_once()
+            c.update_connector_output.reset_mock()
+
+    # no worker meta
+    kv_connector_output = KVConnectorOutput()
+    mc._connectors[0].update_connector_output.side_effect = verify_worker_metadata(None)
+    mc._connectors[1].update_connector_output.side_effect = verify_worker_metadata(None)
+    mc.update_connector_output(kv_connector_output)
+    assert_update_connector_output_called(mc)
+
+    # multi worker meta
+    kv_connector_output.kv_connector_worker_meta = mc_worker_meta_01a_01b
+    mc._connectors[0].update_connector_output.side_effect = verify_worker_metadata(
+        connector0_md
+    )
+    mc._connectors[1].update_connector_output.side_effect = verify_worker_metadata(
+        connector1_md
+    )
+    mc.update_connector_output(kv_connector_output)
+    assert_update_connector_output_called(mc)
+    assert kv_connector_output.kv_connector_worker_meta == mc_worker_meta_01a_01b

third_party/vllm/tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -0,0 +1,315 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for NixlConnectorScheduler sw_sizes calculation with HMA."""
+
+from unittest.mock import patch
+
+import pytest
+
+from vllm import LLM, SamplingParams
+from vllm.config import KVTransferConfig
+from vllm.v1.core.single_type_kv_cache_manager import (
+    FullAttentionManager,
+    SlidingWindowManager,
+)
+
+from .utils import (
+    create_vllm_config,
+    make_kv_cache_config,
+)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "hma_enabled,expected_sw_sizes",
+    [
+        # HMA enabled: FullAttentionSpec (0) + SlidingWindowSpec (2048/16=128)
+        (True, [0, 128 + 1]),
+        # HMA disabled: only FullAttentionSpec (0)
+        (False, [0]),
+    ],
+)
+@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")
+def test_sw_sizes(mock_platform, hma_enabled, expected_sw_sizes):
+    """Test sw_sizes is correctly computed based on HMA enabled/disabled."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorScheduler,
+    )
+
+    mock_platform.device_type = "cpu"
+
+    block_size = 16
+    vllm_config = create_vllm_config(block_size=block_size)
+    # SW 2048 tokens=>128 blocks
+    kv_cache_config = make_kv_cache_config(
+        block_size=block_size, hma_enabled=hma_enabled, sw_size=2048
+    )
+
+    scheduler = NixlConnectorScheduler(
+        vllm_config=vllm_config,
+        engine_id="test-engine",
+        kv_cache_config=kv_cache_config,
+    )
+    # in number of blocks
+    assert scheduler.blocks_per_sw == expected_sw_sizes, (
+        f"Expected sw_sizes={expected_sw_sizes}, got {scheduler.blocks_per_sw}"
+    )
+
+
+@pytest.mark.cpu_test
+def test_logical_to_kernel_block_ids_with_hma():
+    """Test _logical_to_kernel_block_ids expands blocks when HMA is enabled.
+
+    When HMA is enabled, the logical block size may differ from the kernel
+    block size. Each logical block maps to multiple kernel blocks.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorWorker,
+    )
+
+    # Create a mock worker with just the required attributes
+    # (use __new__ to skip __init__)
+    worker = object.__new__(NixlConnectorWorker)
+
+    # Simulate HMA scenario: logical block size = 32, kernel block size = 16
+    # So each logical block maps to 2 kernel blocks eg [0]->[0,1]
+    worker._physical_blocks_per_logical_kv_block = 2
+    # FA + SW groups (neither is MambaSpec, so both get expanded)
+    worker.kv_cache_config = make_kv_cache_config(block_size=16, hma_enabled=True)
+
+    # Test conversion: FA + SW group
+    logical_block_ids = [[0, 1, 2], [3, 4]]
+    kernel_block_ids = worker._logical_to_kernel_block_ids(logical_block_ids)
+
+    expected_kernel_block_ids = [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]]
+    assert kernel_block_ids == expected_kernel_block_ids, (
+        f"Expected {expected_kernel_block_ids}, got {kernel_block_ids}"
+    )
+
+
+@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
+def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
+    """Test that a prefill instance returns fewer "remote blocks" for the SWA groups
+    when sequence exceeds the sliding window.
+    """
+    kv_transfer_config = KVTransferConfig(
+        kv_connector="NixlConnector",
+        kv_role="kv_both",
+    )
+    block_size = 16
+    llm_kwargs = {
+        "model": model_name,
+        "enforce_eager": True,
+        "gpu_memory_utilization": 0.5,
+        "kv_transfer_config": kv_transfer_config,
+        "max_model_len": 2048,
+        # NOTE: Make sure HMA is enabled
+        "disable_hybrid_kv_cache_manager": False,
+        "max_num_batched_tokens": 1024,
+        "enable_prefix_caching": False,
+        "block_size": block_size,
+    }
+
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+    def run_hma_test(llm: LLM):
+        remote_prefill_opts = {
+            "do_remote_decode": True,
+            "do_remote_prefill": False,
+            "remote_engine_id": None,
+            "remote_block_ids": None,
+            "remote_host": None,
+            "remote_port": None,
+        }
+        # Simulate sidecar request
+        sampling_params = SamplingParams(
+            temperature=0.0,
+            max_tokens=1,
+            extra_args={"kv_transfer_params": remote_prefill_opts},
+        )
+        scheduler = llm.llm_engine.engine_core.engine_core.scheduler
+        kv_managers = scheduler.kv_cache_manager.coordinator.single_type_managers
+        # HMA enabled with FA + SWA groups
+        assert len(kv_managers) > 2
+        for kv_manager in kv_managers:
+            assert isinstance(kv_manager, (SlidingWindowManager, FullAttentionManager))
+        req_to_blocks = kv_managers[0].req_to_blocks
+        assert len(req_to_blocks) == 0
+
+        # Process some request with length exceeding the sliding window
+        outputs = llm.generate(["hi" * 1401], sampling_params)
+        kv_params = outputs[0].kv_transfer_params
+
+        # +1 to account for overlapping window across blocks.
+        expected_num_remote_blocks = sw_size // block_size + 1
+        remote_block_ids = kv_params["remote_block_ids"]
+        assert (
+            len(remote_block_ids[0])
+            == expected_num_remote_blocks
+            < len(remote_block_ids[-1])
+        )
+        for group_block_ids in remote_block_ids[:-1]:
+            assert len(group_block_ids) == expected_num_remote_blocks
+
+    def run_test_and_cleanup():
+        llm = LLM(**llm_kwargs)
+        try:
+            run_hma_test(llm)
+        finally:
+            llm.llm_engine.engine_core.shutdown()
+
+    run_test_and_cleanup()
+
+
+@pytest.mark.cpu_test
+def test_nixl_metadata_hma_block_ids_structure():
+    """
+    Test that NixlConnectorMetadata correctly stores block IDs for multiple
+    KV cache groups when HMA is enabled.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorMetadata,
+    )
+
+    metadata = NixlConnectorMetadata()
+
+    # Add request with block IDs for 2 groups (FA + SW)
+    fa_blocks = [0, 1, 2, 3, 4, 5, 6, 7]  # 8 blocks for FA
+    sw_blocks = [8, 9, 10, 11]  # 4 blocks for SW (clipped)
+
+    metadata.add_new_req_to_recv(
+        request_id="test-req-hma",
+        local_block_ids=(fa_blocks, sw_blocks),
+        kv_transfer_params={
+            "remote_block_ids": ([10, 11, 12, 13, 14, 15, 16, 17], [18, 19, 20, 21]),
+            "remote_engine_id": "remote-engine",
+            "remote_request_id": "prefill-test-req-hma",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "tp_size": 1,
+        },
+    )
+
+    assert "test-req-hma" in metadata.reqs_to_recv
+    req_meta = metadata.reqs_to_recv["test-req-hma"]
+
+    # Verify local block IDs structure
+    assert len(req_meta.local_block_ids) == 2
+    assert list(req_meta.local_block_ids[0]) == fa_blocks
+    assert list(req_meta.local_block_ids[1]) == sw_blocks
+
+    # Verify remote block IDs structure
+    assert req_meta.remote is not None
+    assert len(req_meta.remote.block_ids) == 2
+    assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]
+    assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21]
+
+
+@pytest.mark.cpu_test
+def test_get_block_descs_ids_hybrid_ssm():
+    """Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM
+    when ratio=1 (no kernel block size mismatch)."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorWorker,
+    )
+
+    worker = object.__new__(NixlConnectorWorker)
+
+    num_blocks = 100
+    engine_id = "test-engine"
+    worker.num_regions = 2
+    worker.dst_num_blocks = {engine_id: num_blocks}
+    worker._has_mamba = True
+    worker._is_mamba_group = [False, True]
+    worker._physical_blocks_per_logical_kv_block = 1
+    # num_descs = num_regions * num_blocks (no blocks_first doubling)
+    worker.num_descs = 2 * num_blocks
+
+    fa_blocks = [3, 5]
+    ssm_blocks = [1, 2]
+    result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
+
+    # FA group: stride=num_blocks=100, offset=0
+    #   region0: [3, 5],  region1: [103, 105]
+    # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1),
+    #   offset=num_descs=200
+    #   region0: [201, 202],  region1: [301, 302]
+    expected = [3, 5, 103, 105, 201, 202, 301, 302]
+    assert list(result) == expected, f"Expected {expected}, got {list(result)}"
+
+
+@pytest.mark.cpu_test
+def test_get_block_descs_ids_kernel_block_mismatch():
+    """Test _get_block_descs_ids uses different strides for FA (kernel blocks)
+    vs SSM (logical blocks) when ratio > 1."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorWorker,
+    )
+
+    worker = object.__new__(NixlConnectorWorker)
+
+    ratio = 4
+    logical_blocks = 100
+    num_blocks = logical_blocks * ratio  # 400 kernel blocks
+    engine_id = "test-engine"
+    worker.num_regions = 2
+    worker.dst_num_blocks = {engine_id: num_blocks}
+    worker._has_mamba = True
+    worker._is_mamba_group = [False, True]
+    worker._physical_blocks_per_logical_kv_block = ratio
+    worker.num_descs = 2 * num_blocks  # 800
+
+    fa_blocks = [3, 7]  # kernel-level block IDs
+    ssm_blocks = [1, 2]  # logical block IDs
+    result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
+
+    # FA group: stride=num_blocks=400, offset=0
+    #   region0: [3, 7],  region1: [403, 407]
+    # SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800
+    #   region0: [801, 802],  region1: [901, 902]
+    expected = [3, 7, 403, 407, 801, 802, 901, 902]
+    assert list(result) == expected, f"Expected {expected}, got {list(result)}"
+
+
+@pytest.mark.cpu_test
+def test_nixl_metadata_hybrid_ssm_block_ids():
+    """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM
+    groups with different block counts (kernel mismatch active)."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorMetadata,
+    )
+
+    metadata = NixlConnectorMetadata()
+
+    # FA: 8 kernel blocks (2 logical * ratio=4), SSM: 2 logical blocks
+    fa_blocks = [0, 1, 2, 3, 4, 5, 6, 7]
+    ssm_blocks = [0, 1]
+
+    metadata.add_new_req_to_recv(
+        request_id="test-req-hybrid",
+        local_block_ids=(fa_blocks, ssm_blocks),
+        kv_transfer_params={
+            "remote_block_ids": ([10, 11, 12, 13, 14, 15, 16, 17], [20, 21]),
+            "remote_engine_id": "remote-engine",
+            "remote_request_id": "prefill-test-req-hybrid",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "tp_size": 1,
+        },
+    )
+
+    assert "test-req-hybrid" in metadata.reqs_to_recv
+    req_meta = metadata.reqs_to_recv["test-req-hybrid"]
+
+    # Verify local block IDs: different lengths per group
+    assert len(req_meta.local_block_ids) == 2
+    assert list(req_meta.local_block_ids[0]) == fa_blocks
+    assert list(req_meta.local_block_ids[1]) == ssm_blocks
+    assert len(req_meta.local_block_ids[0]) != len(req_meta.local_block_ids[1])
+
+    # Verify remote block IDs: same asymmetry preserved
+    assert req_meta.remote is not None
+    assert len(req_meta.remote.block_ids) == 2
+    assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]
+    assert list(req_meta.remote.block_ids[1]) == [20, 21]
+    assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1])

third_party/vllm/tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -0,0 +1,991 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm import SamplingParams
+from vllm.config import KVTransferConfig, VllmConfig
+from vllm.distributed.kv_events import BlockRemoved, BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+    OffloadingConnector,
+    OffloadingConnectorMetadata,
+    OffloadingConnectorStats,
+)
+from vllm.forward_context import ForwardContext
+from vllm.utils.hashing import sha256
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+from vllm.v1.core.kv_cache_utils import (
+    BlockHash,
+    get_request_block_hasher,
+    init_none_hash,
+)
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+)
+from vllm.v1.kv_offload.abstract import (
+    LoadStoreSpec,
+    OffloadingEvent,
+    OffloadingManager,
+    PrepareStoreOutput,
+)
+from vllm.v1.kv_offload.mediums import GPULoadStoreSpec
+from vllm.v1.kv_offload.spec import OffloadingSpec
+from vllm.v1.kv_offload.worker.worker import (
+    OffloadingHandler,
+    TransferResult,
+    TransferSpec,
+)
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
+from vllm.v1.request import Request, RequestStatus
+from vllm.v1.structured_output import StructuredOutputManager
+
+from .utils import (
+    EOS_TOKEN_ID,
+    create_model_runner_output,
+    create_vllm_config,
+)
+
+
+class MockLoadStoreSpec(LoadStoreSpec):
+    def __init__(self, block_hashes: Iterable[BlockHash]):
+        self.block_hashes: list[BlockHash] = list(block_hashes)
+
+    @staticmethod
+    def medium() -> str:
+        return "Mock"
+
+    def __repr__(self) -> str:
+        return repr(self.block_hashes)
+
+
+class MockOffloadingHandler(OffloadingHandler):
+    def __init__(self):
+        self.transfer_specs: dict[int, TransferSpec] = {}
+        self.completed_transfers: list[TransferResult] = []
+        self.waiting_jobs: set[int] = set()
+        self.completed_jobs: list[int] = []
+        self.flushed_jobs: set[int] = set()
+
+    def get_finished(self) -> list[TransferResult]:
+        finished = self.completed_transfers
+        self.completed_transfers = []
+        return finished
+
+    def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:
+        self.transfer_specs[job_id] = spec
+        self.waiting_jobs.add(job_id)
+        return True
+
+    def complete_jobs(self, job_ids: set[int]) -> None:
+        for job_id in job_ids:
+            if job_id in self.waiting_jobs:
+                self.waiting_jobs.remove(job_id)
+                self.completed_jobs.append(job_id)
+                result = TransferResult(
+                    job_id=job_id,
+                    success=True,
+                    transfer_size=None,
+                    transfer_time=None,
+                    transfer_type=None,
+                )
+                self.completed_transfers.append(result)
+
+    def wait(self, job_ids: set[int]) -> None:
+        self.flushed_jobs |= job_ids
+        self.complete_jobs(job_ids)
+
+
+class MockOffloadingSpec(OffloadingSpec):
+    def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
+        super().__init__(vllm_config, kv_cache_config)
+
+        self.manager = MagicMock(spec=OffloadingManager)
+        self.manager.lookup.return_value = 0
+        self.manager.prepare_load = lambda block_hashes: (
+            MockLoadStoreSpec(block_hashes)
+        )
+        self.handler = MockOffloadingHandler()
+
+    def get_manager(self) -> OffloadingManager:
+        return self.manager
+
+    def get_handlers(
+        self, _, __
+    ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
+        yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler
+        yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler
+
+    def complete_transfers(self):
+        self.handler.complete_jobs(self.handler.waiting_jobs.copy())
+
+    def get_completed_transfers(self) -> list[TransferSpec]:
+        specs = [
+            self.handler.transfer_specs[job_id]
+            for job_id in self.handler.completed_jobs
+        ]
+        self.handler.completed_jobs.clear()
+        return specs
+
+    def get_flushed_transfers(self):
+        specs = [
+            self.handler.transfer_specs[job_id] for job_id in self.handler.flushed_jobs
+        ]
+        self.handler.flushed_jobs.clear()
+        return specs
+
+
+@dataclass
+class TransferSummary:
+    gpu_block_indices: list[int]
+    offload_addresses: list[Any]
+
+
+class RequestRunner:
+    def __init__(
+        self,
+        offloaded_block_size: int,
+        gpu_block_size: int,
+        num_gpu_blocks: int,
+        async_scheduling: bool = True,
+    ):
+        self.offloaded_block_size: int = offloaded_block_size
+        self.gpu_block_size: int = gpu_block_size
+        self.num_gpu_blocks: int = num_gpu_blocks
+        self.async_scheduling: bool = async_scheduling
+
+        self.req_id: int = -1
+
+        vllm_config = create_vllm_config(
+            block_size=gpu_block_size, max_num_batched_tokens=1000
+        )
+        vllm_config.scheduler_config.async_scheduling = async_scheduling
+        vllm_config.kv_transfer_config = KVTransferConfig(
+            kv_connector="OffloadingConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "spec_name": "MockOffloadingSpec",
+                "spec_module_path": "tests.v1.kv_connector.unit.test_offloading_connector",  # noqa: E501
+                "block_size": offloaded_block_size,
+            },
+        )
+
+        block_size = vllm_config.cache_config.block_size
+        kv_cache_config = KVCacheConfig(
+            num_blocks=num_gpu_blocks,
+            kv_cache_tensors=[],
+            kv_cache_groups=[
+                KVCacheGroupSpec(
+                    ["layer"],
+                    FullAttentionSpec(
+                        block_size=block_size,
+                        num_kv_heads=1,
+                        head_size=1,
+                        dtype=torch.float32,
+                    ),
+                )
+            ],
+        )
+        vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks
+        self.num_kv_groups = len(kv_cache_config.kv_cache_groups)
+
+        scheduler_cls = AsyncScheduler if async_scheduling else Scheduler
+        self.scheduler = scheduler_cls(
+            vllm_config=vllm_config,
+            kv_cache_config=kv_cache_config,
+            log_stats=True,
+            structured_output_manager=StructuredOutputManager(vllm_config),
+            block_size=block_size,
+        )
+
+        self.worker_connector = OffloadingConnector(
+            vllm_config, KVConnectorRole.WORKER, kv_cache_config
+        )
+
+        # register worker kv_caches to enable OffloadingWorker creations
+        self.worker_connector.register_cross_layers_kv_cache(
+            kv_cache=torch.empty(0),
+            attn_backend=FlashAttentionBackend,
+        )
+
+        # extract connector of scheduler
+        scheduler_connector = self.scheduler.connector
+        assert scheduler_connector is not None
+        assert isinstance(scheduler_connector, OffloadingConnector)
+        self.scheduler_connector: OffloadingConnector = scheduler_connector
+
+        # extract mocked OffloadingManager of scheduler connector
+        connector_scheduler = scheduler_connector.connector_scheduler
+        assert connector_scheduler is not None
+        manager = connector_scheduler.manager
+        assert isinstance(manager, MagicMock)
+        self.manager: MagicMock = manager
+
+        assert connector_scheduler.gpu_block_size == gpu_block_size
+        assert connector_scheduler.offloaded_block_size == offloaded_block_size
+
+        # extract OffloadingSpec of worker_connector
+        connector_worker = self.worker_connector.connector_worker
+        assert connector_worker is not None
+        offloading_spec = connector_worker.spec
+        assert isinstance(offloading_spec, MockOffloadingSpec)
+        self.offloading_spec: MockOffloadingSpec = offloading_spec
+
+        # mapping (offloading address) -> gpu_block_index
+        self.offloaded: dict[Any, int] = {}
+
+        self.completed_loads: list[TransferSummary] = []
+        self.completed_stores: list[TransferSummary] = []
+        self.flushed_gpu_block_indexes: set[int] = set()
+
+        # maps {block_id: block_offset}
+        self.gpu_block_index: dict[int, int] = {}
+
+        init_none_hash(sha256)
+        self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
+
+        self._dummy_ctx: ForwardContext = ForwardContext(
+            no_compile_layers={},
+            attn_metadata={},
+            virtual_engine=0,
+            slot_mapping={},
+        )
+
+    def new_request(self, token_ids: list[int]):
+        self.req_id += 1
+
+        sampling_params = SamplingParams(max_tokens=1000)
+        sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
+
+        req = Request(
+            request_id=str(self.req_id),
+            prompt_token_ids=token_ids,
+            sampling_params=sampling_params,
+            pooling_params=None,
+            block_hasher=self._block_hasher,
+        )
+
+        self.scheduler.add_request(req)
+
+    def _parse_transfers(self):
+        for transfer_spec in self.offloading_spec.get_flushed_transfers():
+            src_spec, dst_spec = transfer_spec
+            assert isinstance(src_spec, GPULoadStoreSpec)
+
+            for block_id in src_spec.block_ids:
+                self.flushed_gpu_block_indexes.add(
+                    self.gpu_block_index[block_id.item()]
+                )
+
+        block_size_factor = self.offloaded_block_size // self.gpu_block_size
+
+        for transfer_spec in self.offloading_spec.get_completed_transfers():
+            src_spec, dst_spec = transfer_spec
+
+            if isinstance(src_spec, GPULoadStoreSpec):
+                store = True
+                gpu_spec = src_spec
+                offload_spec = dst_spec
+            else:
+                store = False
+                gpu_spec = dst_spec
+                offload_spec = src_spec
+
+            assert isinstance(offload_spec, MockLoadStoreSpec)
+            assert isinstance(gpu_spec, GPULoadStoreSpec)
+
+            gpu_block_indices: list[int] = []
+            for block_id in gpu_spec.block_ids:
+                gpu_block_indices.append(self.gpu_block_index[block_id.item()])
+
+            # list of (block_hash, sub_block_offset)
+            offload_addresses: list[Any] = []
+            for block_hash in offload_spec.block_hashes:
+                for sub_block_idx in range(block_size_factor):
+                    offload_addresses.append((block_hash, sub_block_idx))
+
+            if store:
+                assert len(gpu_block_indices) == len(offload_addresses)
+
+                self.completed_stores.append(
+                    TransferSummary(gpu_block_indices, offload_addresses)
+                )
+            else:
+                remainder_sub_block_count = len(offload_addresses) - len(
+                    gpu_block_indices
+                )
+                assert remainder_sub_block_count >= 0
+                assert remainder_sub_block_count < block_size_factor
+                offload_addresses = offload_addresses[remainder_sub_block_count:]
+
+                self.completed_loads.append(
+                    TransferSummary(gpu_block_indices, offload_addresses)
+                )
+
+    def _update_gpu_block_idx(self):
+        for blocks in self.scheduler.kv_cache_manager.coordinator.single_type_managers[
+            0
+        ].req_to_blocks.values():
+            for block_idx, block in enumerate(blocks):
+                self.gpu_block_index[block.block_id] = block_idx
+
+    def _run(self, decoded_tokens: list[int], complete_transfers: bool):
+        """
+        Runs multiple engine (scheduler + worker) steps.
+        Assumes a single request is running.
+
+        Args:
+            decoded_tokens: the tokens to yield at each step.
+            complete_transfers: complete transfers immediately
+        """
+
+        tokens_iter = iter(decoded_tokens)
+        token_id = next(tokens_iter, None)
+        prev_scheduler_output = None
+        prev_model_runner_output = None
+        while True:
+            assert self.scheduler.requests
+
+            scheduler_output = self.scheduler.schedule()
+            self._update_gpu_block_idx()
+
+            kv_connector_metadata = scheduler_output.kv_connector_metadata
+            assert kv_connector_metadata is not None
+            assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)
+
+            if scheduler_output.preempted_req_ids:
+                self.worker_connector.handle_preemptions(
+                    scheduler_output.preempted_req_ids
+                )
+
+            self.worker_connector.bind_connector_metadata(kv_connector_metadata)
+            self.worker_connector.start_load_kv(self._dummy_ctx)
+
+            if scheduler_output.total_num_scheduled_tokens > 0:
+                self.worker_connector.wait_for_save()
+
+            if complete_transfers:
+                self.offloading_spec.complete_transfers()
+
+            finished_sending, finished_recving = self.worker_connector.get_finished(
+                scheduler_output.finished_req_ids
+            )
+
+            self.worker_connector.clear_connector_metadata()
+
+            model_runner_output = create_model_runner_output(
+                reqs=self.scheduler.running,
+                finished_sending=finished_sending,
+                finished_recving=finished_recving,
+                token_id=token_id or 0,
+            )
+
+            prev_token_id = token_id
+            if self.scheduler.running:
+                token_id = next(tokens_iter, None)
+
+            if self.async_scheduling:
+                # in async scheduling we update the output of the previous step
+                if prev_model_runner_output is not None:
+                    self.scheduler.update_from_output(
+                        prev_scheduler_output, prev_model_runner_output
+                    )
+                prev_scheduler_output = scheduler_output
+                prev_model_runner_output = model_runner_output
+            else:
+                self.scheduler.update_from_output(scheduler_output, model_runner_output)
+
+            if (
+                prev_token_id == EOS_TOKEN_ID
+                and prev_token_id != token_id
+                and self.scheduler.requests
+            ):
+                # continue for one more step to allow offloading to kick off
+                continue
+
+            if token_id is None:
+                if self.async_scheduling:
+                    # sample last token
+                    self.scheduler.update_from_output(
+                        prev_scheduler_output, prev_model_runner_output
+                    )
+                break
+
+        self._parse_transfers()
+
+        # run one more step to update finished stored
+        if EOS_TOKEN_ID in decoded_tokens:
+            assert not self.scheduler.running
+
+            while self.scheduler.requests:
+                scheduler_output = self.scheduler.schedule()
+
+                finished_sending, finished_recving = self.worker_connector.get_finished(
+                    scheduler_output.finished_req_ids
+                )
+
+                assert not finished_recving
+
+                model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+                model_runner_output.kv_connector_output = KVConnectorOutput(
+                    finished_sending=finished_sending
+                )
+
+                self.scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    def run(
+        self,
+        decoded_tokens: list[int],
+        complete_transfers: bool = True,
+        expected_stored_gpu_block_indexes: tuple[int, ...] = (),
+        expected_loaded_gpu_block_indexes: tuple[int, ...] = (),
+        expected_flushed_gpu_block_indexes: tuple[int, ...] = (),
+    ):
+        """
+        Runs multiple engine (scheduler + worker) steps.
+        Assumes a single request is running.
+
+        Args:
+            decoded_tokens: the tokens to yield at each step.
+            complete_transfers: complete transfers immediately
+            expected_stored_gpu_block_indexes: GPU block indexes
+                that are expected to be written during the run.
+            expected_loaded_gpu_block_indexes: GPU block indexes
+                that are expected to be loaded during the run.
+            expected_flushed_gpu_block_indexes: GPU block indexes
+                that are expected to be flushed during the run.
+        """
+
+        self.manager.reset_mock()
+        self._run(decoded_tokens, complete_transfers)
+
+        loaded_gpu_block_indexes: set[int] = set()
+        for transfer in self.completed_loads:
+            for gpu_block_idx, offloaded_address in zip(
+                transfer.gpu_block_indices, transfer.offload_addresses
+            ):
+                loaded_gpu_block_indexes.add(gpu_block_idx)
+                assert gpu_block_idx == self.offloaded[offloaded_address]
+
+        assert set(expected_loaded_gpu_block_indexes) == loaded_gpu_block_indexes
+        self.completed_loads.clear()
+
+        stored_gpu_block_indexes: set[int] = set()
+        for transfer in self.completed_stores:
+            for gpu_block_idx, offloaded_address in zip(
+                transfer.gpu_block_indices, transfer.offload_addresses
+            ):
+                stored_gpu_block_indexes.add(gpu_block_idx)
+                self.offloaded[offloaded_address] = gpu_block_idx
+
+        assert set(expected_stored_gpu_block_indexes) == stored_gpu_block_indexes
+        self.completed_stores.clear()
+
+        assert set(expected_flushed_gpu_block_indexes) == self.flushed_gpu_block_indexes
+        self.flushed_gpu_block_indexes.clear()
+
+
+@pytest.fixture
+def request_runner():
+    runners = []
+
+    def runner_factory(
+        offloaded_block_size, gpu_block_size, num_gpu_blocks, async_scheduling
+    ):
+        runner = RequestRunner(
+            offloaded_block_size=offloaded_block_size,
+            gpu_block_size=gpu_block_size,
+            num_gpu_blocks=num_gpu_blocks,
+            async_scheduling=async_scheduling,
+        )
+        runners.append(runner)
+        return runner
+
+    yield runner_factory  # pass factory to the test
+
+
+def generate_store_output(block_hashes: Iterable[BlockHash]):
+    block_hashes = list(block_hashes)
+    return PrepareStoreOutput(
+        block_hashes_to_store=list(block_hashes),
+        store_spec=MockLoadStoreSpec(block_hashes),
+        block_hashes_evicted=[],
+    )
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_offloading_connector(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+    block_size_factor = offloaded_block_size // gpu_block_size
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    # 3 blocks, store just the middle block (skip first and last)
+    # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
+    runner.new_request(token_ids=[0] * offloaded_block_size * 3)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
+    )
+    runner.run(decoded_tokens=[0])
+
+    # add block missing 1 token -> no offload
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size - 1),
+        expected_stored_gpu_block_indexes=(3, 4, 5),
+    )
+    runner.manager.prepare_store.assert_not_called()
+
+    # +1 token -> single block, fail prepare_store
+    runner.manager.prepare_store.side_effect = lambda block_hashes: None
+    runner.run(decoded_tokens=[0])
+    runner.manager.prepare_store.assert_called()
+
+    # 1 more block (+ token for async scheduling)
+    # now set block_hashes_to_store = []
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
+
+    # 1 more block (+ token for kicking off offloading)
+    # now check touch was called with all 6 blocks
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size + 1),
+        expected_stored_gpu_block_indexes=(15, 16, 17),
+    )
+    runner.manager.touch.assert_called()
+    block_hashes1 = list(runner.manager.touch.call_args.args[0])
+    assert len(block_hashes1) == 6
+
+    # terminate request
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    # create a new request differing only on the last token
+    runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
+    runner.run(decoded_tokens=[0])
+    runner.manager.touch.assert_called()
+    block_hashes2 = list(runner.manager.touch.call_args.args[0])
+    assert len(block_hashes2) == 6
+
+    # verify hashes are the same, except for the last block
+    assert block_hashes1[:5] == block_hashes2[:5]
+    assert block_hashes1[5] != block_hashes2[5]
+
+    # terminate request
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
+    )
+
+    # full_block_tokens - num_computed_tokens < offloaded_block_size
+    runner.new_request(
+        token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
+    )
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+    runner.manager.lookup.assert_not_called()
+
+    # single block lookup with no hits
+    runner.new_request(token_ids=[1] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+    runner.manager.lookup.assert_called()
+    assert len(list(runner.manager.lookup.call_args.args[0])) == 1
+
+    # single block lookup with a hit
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
+    )
+
+    # single block lookup with a hit in a middle block
+    runner.new_request(
+        token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
+    )
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
+    )
+
+    # test take_events
+    def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
+        return [BlockHash(str(i).encode()) for i in int_hashes]
+
+    def take_events() -> Iterable[OffloadingEvent]:
+        yield OffloadingEvent(
+            block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
+        )
+        yield OffloadingEvent(
+            block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
+        )
+
+    runner.manager.take_events.side_effect = take_events
+    events = list(runner.scheduler_connector.take_events())
+    assert len(events) == 2
+    event = events[0]
+    assert isinstance(event, BlockStored)
+    assert event.block_hashes == to_hashes([1, 2, 3])
+    assert event.block_size == 16
+    assert event.medium == "A"
+    assert event.token_ids == []
+    assert event.parent_block_hash is None
+    assert event.lora_id is None
+    assert event.lora_name is None
+    event = events[1]
+    assert isinstance(event, BlockRemoved)
+    assert event.block_hashes == to_hashes([4, 5, 6])
+    assert event.medium == "B"
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_request_preemption(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
+    num_free_blocks_empty = free_block_queue.num_free_blocks
+
+    # 2 blocks, store all, without flushing
+    # blocks = [0, 1, 2], [3, 4, 5]
+    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0],
+        complete_transfers=False,
+    )
+
+    # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
+        complete_transfers=False,
+    )
+
+    # simulate KV cache running out of space
+    free_block_queue.num_free_blocks = 0
+
+    # request should be preempted now
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+        expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
+        expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
+    )
+
+    # restore KV cache space and reset GPU prefix cache
+    free_block_queue.num_free_blocks = num_free_blocks_empty
+    runner.scheduler.reset_prefix_cache()
+
+    # request should now return from preemption
+    # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
+    runner.manager.lookup.return_value = 3
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0] * gpu_block_size,
+        expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
+    )
+
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=(9, 10, 11),
+    )
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    # store 1 blocks
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # start a request to load the first block, but don't complete
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+    )
+
+    # request triggered a load
+    transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
+    assert transfer_jobs
+
+    # start a new request to load the same first block
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+    )
+
+    # request did not trigger a load
+    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
+
+    # complete transfers
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_loaded_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # second request will use the GPU prefix cache
+    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_abort_loading_requests(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    # store 1 blocks
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # start a request to load the first block, but don't complete
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+    )
+
+    # request triggered a load
+    transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
+    assert transfer_jobs
+
+    # abort request
+    req_id = str(runner.req_id)
+    runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
+
+    # verify request is not deleted
+    assert req_id in runner.scheduler.requests
+
+    # complete loading request
+    runner.run(
+        decoded_tokens=[],
+        expected_loaded_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # assert request is deleted
+    assert req_id not in runner.scheduler.requests
+
+
+class TestOffloadingConnectorStats:
+    """Tests for OffloadingConnector stats reconstruction and operations."""
+
+    def test_build_kv_connector_stats_with_none(self):
+        """Test that build_kv_connector_stats returns empty stats when given None."""
+        stats = OffloadingConnector.build_kv_connector_stats(data=None)
+
+        assert stats is not None
+        assert isinstance(stats, OffloadingConnectorStats)
+        assert len(stats.data) == 0
+        assert stats.is_empty()
+
+    def test_build_kv_connector_stats_with_empty_dict(self):
+        """Test that build_kv_connector_stats returns empty stats with empty dict."""
+        stats = OffloadingConnector.build_kv_connector_stats(data={})
+
+        assert stats is not None
+        assert isinstance(stats, OffloadingConnectorStats)
+        assert len(stats.data) == 0
+        assert stats.is_empty()
+
+    def test_build_kv_connector_stats_reconstructs_offload_stats(self):
+        """Test that OffloadingConnector stats are properly reconstructed with
+        correct data."""
+        serialized_data = {
+            "CPU_to_GPU": [
+                {"op_size": 16, "op_time": 1.0},
+                {"op_size": 8, "op_time": 0.5},
+            ],
+            "GPU_to_CPU": [
+                {"op_size": 1, "op_time": 0.1},
+                {"op_size": 2, "op_time": 0.2},
+            ],
+        }
+
+        stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
+
+        offload_connector_stats = stats
+        assert isinstance(offload_connector_stats, OffloadingConnectorStats)
+        assert offload_connector_stats.data["CPU_to_GPU"] == [
+            {"op_size": 16, "op_time": 1.0},
+            {"op_size": 8, "op_time": 0.5},
+        ]
+        assert offload_connector_stats.data["GPU_to_CPU"] == [
+            {"op_size": 1, "op_time": 0.1},
+            {"op_size": 2, "op_time": 0.2},
+        ]
+
+    def test_aggregate_same_connector(self):
+        """Test aggregating stats from the same connector type."""
+        stats1 = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 16, "op_time": 1.0},
+                    {"op_size": 8, "op_time": 0.5},
+                ],
+                "GPU_to_CPU": [
+                    {"op_size": 1, "op_time": 0.1},
+                    {"op_size": 2, "op_time": 0.2},
+                ],
+            }
+        )
+
+        stats2 = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 3, "op_time": 0.2},
+                    {"op_size": 7, "op_time": 0.9},
+                ],
+                "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
+            }
+        )
+
+        result = stats1.aggregate(stats2)
+
+        assert result is stats1  # Should return self
+        offload_connector_stats = result
+        assert offload_connector_stats.data["CPU_to_GPU"] == [
+            {"op_size": 16, "op_time": 1.0},
+            {"op_size": 8, "op_time": 0.5},
+            {"op_size": 3, "op_time": 0.2},
+            {"op_size": 7, "op_time": 0.9},
+        ]
+        assert offload_connector_stats.data["GPU_to_CPU"] == [
+            {"op_size": 1, "op_time": 0.1},
+            {"op_size": 2, "op_time": 0.2},
+            {"op_size": 16, "op_time": 2},
+        ]
+
+    def test_reduce(self):
+        """Test that reduce() correctly reduces all nested connector stats."""
+        stats = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 16, "op_time": 1.0},
+                    {"op_size": 8, "op_time": 0.5},
+                    {"op_size": 3, "op_time": 0.2},
+                    {"op_size": 7, "op_time": 0.9},
+                ],
+                "GPU_to_CPU": [
+                    {"op_size": 1, "op_time": 0.1},
+                    {"op_size": 2, "op_time": 0.2},
+                    {"op_size": 16, "op_time": 2},
+                ],
+            }
+        )
+
+        reduced = stats.reduce()
+
+        assert isinstance(reduced, dict)
+        # Check that the stats were reduced (should have aggregated values)
+        assert "CPU_to_GPU_total_bytes" in reduced
+        assert "CPU_to_GPU_total_time" in reduced
+        assert "GPU_to_CPU_total_bytes" in reduced
+        assert "GPU_to_CPU_total_time" in reduced
+        assert reduced["CPU_to_GPU_total_bytes"] == 34
+        assert reduced["CPU_to_GPU_total_time"] == 2.6
+        assert reduced["GPU_to_CPU_total_time"] == 2.3
+        assert reduced["GPU_to_CPU_total_bytes"] == 19
+
+    def test_reset(self):
+        """Test that reset() resets all nested connector stats."""
+        offload_connector_stats = OffloadingConnectorStats(
+            data={
+                "CPU_to_GPU": [
+                    {"op_size": 3, "op_time": 0.2},
+                    {"op_size": 7, "op_time": 0.9},
+                ],
+                "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
+            }
+        )
+
+        assert not offload_connector_stats.is_empty()
+
+        offload_connector_stats.reset()
+
+        # After reset, stats should be empty
+        assert offload_connector_stats.is_empty()
+        assert len(offload_connector_stats.data) == 0

third_party/vllm/tests/v1/kv_connector/unit/test_output_aggregator.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
+from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
+
+pytestmark = pytest.mark.cpu_test
+
+
+class DummyModelRunnerOutput(ModelRunnerOutput):
+    def __init__(
+        self,
+        finished_sending: set[str] | None = None,
+        finished_recving: set[str] | None = None,
+        invalid_block_ids: set[int] | None = None,
+        expected_finished_count: int = 0,
+    ):
+        self.kv_connector_output = KVConnectorOutput(
+            finished_sending=finished_sending,
+            finished_recving=finished_recving,
+            invalid_block_ids=invalid_block_ids or set(),
+            expected_finished_count=expected_finished_count,
+        )
+
+    def __repr__(self):
+        return (
+            f"DummyModelRunnerOutput("
+            f"finished_sending={self.kv_connector_output.finished_sending},"
+            f"finished_recving={self.kv_connector_output.finished_recving})"
+            f"invalid_block_ids={self.kv_connector_output.invalid_block_ids})"
+        )
+
+
+def test_aggregate_workers_output():
+    aggregator = KVOutputAggregator(expected_finished_count=2)
+
+    output1 = DummyModelRunnerOutput()
+    output2 = DummyModelRunnerOutput()
+
+    aggregated = aggregator.aggregate([output1, output2])
+
+    assert aggregated is output1
+    aggregated = aggregated.kv_connector_output
+    assert aggregated.finished_sending is None
+    assert aggregated.finished_recving is None
+    assert not aggregated.invalid_block_ids
+
+    output1 = DummyModelRunnerOutput(
+        finished_sending={"req1"}, finished_recving={"req2"}
+    )
+    output2 = DummyModelRunnerOutput(invalid_block_ids={1})
+
+    aggregated = aggregator.aggregate([output1, output2])
+
+    assert aggregated is output1
+    aggregated = aggregated.kv_connector_output
+    assert aggregated.finished_sending is None
+    assert aggregated.finished_recving is None
+    assert aggregated.invalid_block_ids == {1}
+
+    output1 = DummyModelRunnerOutput(invalid_block_ids={2})
+    output2 = DummyModelRunnerOutput(finished_sending={"req1"})
+
+    aggregated = aggregator.aggregate([output1, output2])
+
+    assert aggregated is output1
+    aggregated = aggregated.kv_connector_output
+    assert aggregated.finished_sending == {"req1"}
+    assert aggregated.finished_recving is None
+    assert aggregated.invalid_block_ids == {2}
+
+    output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})
+    output2 = DummyModelRunnerOutput(
+        finished_recving={"req2"}, invalid_block_ids={4, 5}
+    )
+
+    aggregated = aggregator.aggregate([output1, output2])
+
+    assert aggregated is output1
+    aggregated = aggregated.kv_connector_output
+    assert aggregated.finished_sending is None
+    assert aggregated.finished_recving == {"req2"}
+    assert aggregated.invalid_block_ids == {3, 4, 5}
+
+
+def test_aggregate_workers_output_with_expected_finished_count():
+    # We create the aggregator expecting to collect from 4 workers
+    aggregator = KVOutputAggregator(expected_finished_count=4)
+    assert aggregator._expected_finished_count == 4
+    # Some request with default expected finished requests
+    output1 = DummyModelRunnerOutput(finished_sending={"req1"})
+    aggregated = aggregator.aggregate([output1])
+    # still expecting to collect from 4 workers
+    assert aggregator._send_remaining_count["req1"] == 3
+    assert not aggregated.kv_connector_output.finished_sending
+    assert not aggregated.kv_connector_output.finished_recving
+
+    # Workers discover and find that in this setup they only need to
+    # collect from 2
+    output1 = DummyModelRunnerOutput(
+        finished_sending={"req1"}, expected_finished_count=2
+    )
+    output2 = DummyModelRunnerOutput(
+        finished_recving={"req2"}, expected_finished_count=2
+    )
+    output3 = DummyModelRunnerOutput(finished_recving={"req2"})
+    # Req2 only needs 2 acks
+    aggregated = aggregator.aggregate([output1, output2, output3])
+    assert aggregated.kv_connector_output.expected_finished_count == 2
+
+    assert not aggregated.kv_connector_output.finished_sending
+
+    # Req2 is finished
+    assert "req2" not in aggregator._recv_remaining_count
+    assert aggregated.kv_connector_output.finished_recving == {"req2"}
+
+    # Req1 is still waiting for 2 more acks (expected_finished_count has no effect)
+    # NOTE: This is to showcase dynamic update. Workers are responsible for
+    # ensuring "req1" termination in this case
+    assert aggregator._send_remaining_count["req1"] == 2

third_party/vllm/tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py

@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+
+import pytest
+
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
+from vllm.v1.request import FinishReason, RequestStatus
+
+from .utils import (
+    assert_scheduler_empty,
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def test_basic_lifecycle():
+    """Test lifecycle of a Remote Decode request."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        max_tokens=1,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=True,
+    )
+
+    scheduler.add_request(request)
+    request_id = request.request_id
+
+    # STEP (1): Prefill.
+    # (1a): schedule()
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
+    assert len(scheduler.running) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 1
+
+    # (1b): execute_model()
+    model_runner_output = create_model_runner_output(reqs=[request])
+
+    # (1c): update_from_output()
+    engine_core_outputs = scheduler.update_from_output(
+        scheduler_output, model_runner_output
+    )
+
+    # Ensure the request is finished after 1 token.
+    assert request.is_finished()
+    assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
+    output = engine_core_outputs[0].outputs[0]
+    assert output.finish_reason == FinishReason.LENGTH
+    assert output.kv_transfer_params is not None
+
+    # Request freed in Scheduler and in Persistent Batch ...
+    assert request_id in scheduler.finished_req_ids
+    assert len(scheduler.running) == 0
+    assert len(scheduler.waiting) == 0
+
+    # ... but blocks should not be freed.
+    assert len(scheduler.requests) == 1
+    blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
+        0
+    ].req_to_blocks[request_id]
+    for block in blocks:
+        assert block.ref_cnt == 1
+
+    # STEP (2): Send Finished to PB.
+    # (2a): schedule() - pass finished request to PB.
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
+    assert len(scheduler.running) == 0
+    assert len(scheduler_output.finished_req_ids) == 1
+    assert request_id in scheduler_output.finished_req_ids
+    assert len(scheduler_output.scheduled_new_reqs) == 0
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
+    assert len(scheduler.finished_req_ids) == 0
+
+    # (2b): execute_model()
+    model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
+
+    # (2c): update_from_output()
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # STEP (3): Finished sending.
+    # (3a): schedule() - pass finished request to PB.
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
+    assert len(scheduler.running) == 0
+    assert len(scheduler_output.finished_req_ids) == 0
+    assert len(scheduler_output.scheduled_new_reqs) == 0
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
+    assert len(scheduler.finished_req_ids) == 0
+
+    # (3b): execute_model()
+    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    model_runner_output.kv_connector_output = KVConnectorOutput(
+        finished_sending={request_id}
+    )
+
+    # (3c): update_from_output()
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # Confirm we do not have any memory leaks after req lifecycle.
+    assert_scheduler_empty(scheduler)
+
+
+def test_short_prompt_lifecycle():
+    """Test lifecycle of a Remote Decode request with short prompt."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # Not enough tokens for full block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_TOKENS = BLOCK_SIZE // 2
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        max_tokens=1,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=True,
+    )
+
+    scheduler.add_request(request)
+
+    # STEP (1): Prefill.
+    # (1a): schedule()
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.requests) == 1
+    assert len(scheduler.running) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 1
+
+    # (1b): execute_model()
+    model_runner_output = create_model_runner_output(reqs=[request])
+
+    # (1c): update_from_output()
+    # Even though tokens < block_size, there will be kv xfer for partial block.
+    eco = scheduler.update_from_output(scheduler_output, model_runner_output)
+    kv_transfer_params = eco[0].outputs[0].kv_transfer_params
+
+    assert len(kv_transfer_params["remote_block_ids"]) == 1
+
+    # Confirm we do not have any memory leaks after req lifecycle.
+    # We need to mark sending finish to clear data for persistent batch.
+    scheduler_output = scheduler.schedule()
+    # Use create_model_runner_output to pass kv_connector_output along
+    model_runner_output = create_model_runner_output(
+        reqs=[request], finished_sending={request.request_id}
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert_scheduler_empty(scheduler)
+
+
+def test_prefix_cache_lifecycle():
+    """Test that remote decode params still work with a prefix cache hit."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # Prime the KVCache.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 3
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request_normal = create_request(
+        request_id=1, block_size=BLOCK_SIZE, num_tokens=NUM_TOKENS
+    )
+
+    scheduler.add_request(request_normal)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[request_normal], use_eos=True
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    scheduler_output = scheduler.schedule()
+    scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
+
+    #####################
+    # Actual Test: confirm we send all blocks.
+
+    # Step (1): Send the KV Transfer.
+    NUM_EXTERNAL_FULL_BLOCKS -= 1
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request_remote = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=True,
+    )
+
+    scheduler.add_request(request_remote)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_remote])
+    eco = scheduler.update_from_output(scheduler_output, model_runner_output)
+    kv_transfer_params = eco[0].outputs[0].kv_transfer_params
+
+    # Ensure we send all block ids, including the partial blocks,
+    # even if there is a cache hit.
+    # remote_block_ids is BlockIds (tuple of lists); sum block counts across groups.
+    num_remote_blocks = sum(len(g) for g in kv_transfer_params["remote_block_ids"])
+    assert num_remote_blocks == (NUM_EXTERNAL_FULL_BLOCKS + 1)
+
+    # STEP (2): Ensure it is freed.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    model_runner_output.kv_connector_output = KVConnectorOutput(
+        finished_sending={request_remote.request_id}
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert_scheduler_empty(scheduler)
+
+
+def test_abort_during_kv_transfer():
+    """Test aborting request does not release blocks for remote decode."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # Prime the KVCache.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_decode=True,
+    )
+
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    scheduler_output = scheduler.schedule()
+    scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
+
+    # Request removed from PB but blocks should not be freed.
+    assert len(scheduler.requests) == 1
+
+    # Abort the request, and check the blocks are still not freed
+    scheduler.finish_requests([request.request_id], RequestStatus.FINISHED_ABORTED)
+    assert len(scheduler.requests) == 1
+
+    # Simulate a finished sending notification
+    scheduler_output = scheduler.schedule()
+    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    model_runner_output.kv_connector_output = KVConnectorOutput(
+        finished_sending=[request.request_id]
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert_scheduler_empty(scheduler)

third_party/vllm/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

@@ -0,0 +1,581 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+
+import pytest
+
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
+from vllm.v1.request import FinishReason, RequestStatus
+
+from .utils import (
+    assert_scheduler_empty,
+    create_model_runner_output,
+    create_request,
+    create_scheduler,
+    create_vllm_config,
+)
+
+pytestmark = pytest.mark.cpu_test
+
+
+def _num_waiting_requests(scheduler) -> int:
+    return len(scheduler.waiting) + len(scheduler.skipped_waiting)
+
+
+def test_basic_lifecycle():
+    """Test lifecycle of a remote prefill."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+    START_FREE_BLOCK_QUEUE_SIZE = (
+        scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks
+    )
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_prefill=True,
+    )
+
+    scheduler.add_request(request)
+    request_id = request.request_id
+
+    # STEP (1):
+    # (1a): schedule()
+    scheduler_output = scheduler.schedule()
+
+    # Nothing running and empty scheduler output.
+    assert len(scheduler.running) == 0
+    assert len(scheduler_output.scheduled_new_reqs) == 0
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
+    assert len(scheduler_output.num_scheduled_tokens) == 0
+    assert scheduler_output.total_num_scheduled_tokens == 0
+
+    # Req waiting for KVs with no computed/scheduled toks ...
+    assert _num_waiting_requests(scheduler) == 1
+    assert request in scheduler.skipped_waiting
+    assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+    assert request.num_computed_tokens == NUM_TOKENS
+
+    # ... but should have (uncached) blocks allocated to it.
+    block_pool = scheduler.kv_cache_manager.block_pool
+    assert block_pool.free_block_queue.num_free_blocks < START_FREE_BLOCK_QUEUE_SIZE
+    assert len(block_pool.cached_block_hash_to_block) == 0
+    blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
+        0
+    ].req_to_blocks[request_id]
+    for block in blocks:
+        assert block._block_hash is None
+
+    # (1b): forward()
+    model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
+
+    # (1c): update_from_output()
+    engine_core_outputs = scheduler.update_from_output(
+        scheduler_output, model_runner_output
+    )
+    assert not engine_core_outputs or not engine_core_outputs[0].outputs
+
+    # STEP (2):
+    # (2a): schedule(): nothing happens!
+    scheduler_output = scheduler.schedule()
+    assert _num_waiting_requests(scheduler) == 1
+    assert len(scheduler.running) == 0
+
+    # (2b): forward(): request finishes recv.
+    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    model_runner_output.kv_connector_output = KVConnectorOutput(
+        finished_recving={request_id}
+    )
+
+    # (2c): update_from_output():
+    engine_core_outputs = scheduler.update_from_output(
+        scheduler_output, model_runner_output
+    )
+    assert _num_waiting_requests(scheduler) == 1
+    assert request_id in scheduler.finished_recving_kv_req_ids
+
+    # STEP (3):
+    # (3a): schedule(): this should actually schedule.
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.running) == 1
+
+    # Confirm the block are actually allocated.
+    num_hashed_blocks = 0
+    blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
+        0
+    ].req_to_blocks[request_id]
+    for block in blocks:
+        assert block.ref_cnt == 1
+        num_hashed_blocks += 1 if block._block_hash is not None else 0
+    assert num_hashed_blocks == NUM_EXTERNAL_FULL_BLOCKS
+
+    # Confirm the rest of the prompt is scheduled in this step.
+    scheduled_req = scheduler_output.scheduled_new_reqs[0]
+    num_scheduled_tokens = scheduler_output.num_scheduled_tokens[request_id]
+    num_computed_tokens = scheduled_req.num_computed_tokens
+    total_prompt_tokens = len(scheduled_req.prompt_token_ids)
+    assert num_scheduled_tokens == total_prompt_tokens - num_computed_tokens
+
+    # (3b): execute_model()
+    model_runner_output = create_model_runner_output([request])
+    # (3c): update_from_output()
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # Step (4): Hit EOS.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output([request], use_eos=True)
+    engine_core_outputs = scheduler.update_from_output(
+        scheduler_output, model_runner_output
+    )
+    scheduler.schedule()
+
+    outputs = engine_core_outputs[0].outputs
+    assert len(outputs) == 1
+    output = outputs[0]
+    assert output.finish_reason == FinishReason.STOP
+    assert_scheduler_empty(scheduler)
+
+
+def test_interleaved_lifecycle():
+    """Test Remote Prefills Work Well With Other Requests."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request_remote = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_prefill=True,
+    )
+    request_local_a = create_request(
+        request_id=2,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+    )
+    request_local_b = create_request(
+        request_id=3,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+    )
+
+    # STEP 1: Regular request is running.
+    scheduler.add_request(request_local_a)
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.running) == 1
+
+    model_runner_output = create_model_runner_output([request_local_a])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # STEP 2: Add a local and remote request.
+    scheduler.add_request(request_local_b)
+    scheduler.add_request(request_remote)
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.running) == 2
+    assert _num_waiting_requests(scheduler) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 1
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 1
+
+    model_runner_output = create_model_runner_output([request_local_a, request_local_b])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # STEP 3: continue running, KVs not arrived yet.
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.running) == 2
+    assert _num_waiting_requests(scheduler) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 0
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
+
+    model_runner_output = create_model_runner_output(
+        reqs=[request_local_a, request_local_b]
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 2
+    assert _num_waiting_requests(scheduler) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 0
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
+
+    # STEP 4: KVs arrive.
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.running) == 2
+    assert _num_waiting_requests(scheduler) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 0
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
+
+    model_runner_output = create_model_runner_output(
+        [request_local_a, request_local_b], finished_recving={request_remote.request_id}
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # STEP 5: RECVed KVs are sent to ModelRunner.
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.running) == 3
+    assert _num_waiting_requests(scheduler) == 0
+    assert len(scheduler_output.scheduled_new_reqs) == 1
+    assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
+
+    model_runner_output = create_model_runner_output(
+        [request_local_a, request_local_b, request_remote]
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # STEP 6: Hit EOS and free.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        [request_local_a, request_local_b, request_remote],
+        use_eos=True,
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    scheduler.schedule()
+    assert_scheduler_empty(scheduler)
+
+
+def test_no_spurious_prefix_caching():
+    """
+    With P/D, blocks can be allocated but uncomputed for
+    multiple engine steps. This test confirms that we do
+    not accidentally have cache hits against uncomputed
+    blocks.
+    """
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 and a half full external blocks.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    # Both of these requests have prompts like [1,1,1,1,1, ...]
+    request_remote = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        common_prefix_len=NUM_TOKENS,
+        do_remote_prefill=True,
+    )
+
+    request_local = create_request(
+        request_id=2,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        common_prefix_len=NUM_TOKENS,
+        do_remote_prefill=False,
+    )
+
+    # Schedule the remote prefill request. This should not
+    # cause any blocks to be cached.
+    scheduler.add_request(request_remote)
+    scheduler_output = scheduler.schedule()
+    scheduler.update_from_output(scheduler_output, EMPTY_MODEL_RUNNER_OUTPUT)
+    assert _num_waiting_requests(scheduler) == 1
+
+    # Schedule the local prefill request. This should
+    # cause blocks to be cached, but separately from
+    scheduler.add_request(request_local)
+    scheduler_output = scheduler.schedule()
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 1
+
+    local_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
+        0
+    ].req_to_blocks[request_local.request_id]
+    remote_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
+        0
+    ].req_to_blocks[request_remote.request_id]
+
+    # Local should have cached blocks (but not all due to preallocate).
+    num_hashed_blocks = 0
+    for block in local_blocks:
+        assert block.ref_cnt == 1
+        num_hashed_blocks += 1 if block._block_hash is not None else 0
+    assert num_hashed_blocks > 0
+
+    # Remote blocks should not be cached.
+    for block in remote_blocks:
+        assert block.ref_cnt == 1
+        assert block._block_hash is None
+
+
+def test_full_block_prompt():
+    """Test that we handle a prompt that is the full block size."""
+
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * NUM_EXTERNAL_FULL_BLOCKS)
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_prefill=True,
+    )
+
+    scheduler.add_request(request)
+    request_id = request.request_id
+
+    # STEP (1): Initialize a recv.
+    scheduler_output = scheduler.schedule()
+    # All blocks should be allocated.
+    num_blocks = len(
+        scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
+            request_id
+        ]
+    )
+    assert num_blocks == NUM_EXTERNAL_FULL_BLOCKS
+    model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # # STEP (2): Recv.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = copy.deepcopy(EMPTY_MODEL_RUNNER_OUTPUT)
+    model_runner_output.kv_connector_output = KVConnectorOutput(
+        finished_recving={request_id}
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert _num_waiting_requests(scheduler) == 1
+    assert request_id in scheduler.finished_recving_kv_req_ids
+
+    # # STEP (3): Run as usual.
+    scheduler_output = scheduler.schedule()
+
+    # We need to recompute the final token of the prompt to generate
+    # the first new token, so we should not have a new block.
+    num_blocks = len(
+        scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
+            request_id
+        ]
+    )
+    assert num_blocks == NUM_EXTERNAL_FULL_BLOCKS
+    assert scheduler_output.scheduled_new_reqs[0].num_computed_tokens == NUM_TOKENS - 1
+    assert scheduler_output.num_scheduled_tokens[request_id] == 1
+
+    model_runner_output = create_model_runner_output([request])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+
+    # # Step (4): Hit EOS.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output([request], use_eos=True)
+    engine_core_outputs = scheduler.update_from_output(
+        scheduler_output, model_runner_output
+    )
+    scheduler.schedule()
+
+    outputs = engine_core_outputs[0].outputs
+    assert len(outputs) == 1
+    output = outputs[0]
+    assert output.finish_reason == FinishReason.STOP
+    assert_scheduler_empty(scheduler)
+
+
+def test_cannot_schedule_after_recv():
+    """
+    Test that we can handle no schedule after recv due to not
+    enough remaining KV blocks.
+    """
+
+    # NOTE: the KVCacheManager will use 1 null block.
+    # So there are 5 total working blocks.
+    TOTAL_NUM_BLOCKS = 6
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config, num_blocks=TOTAL_NUM_BLOCKS)
+
+    # Prime the KVCache.
+    NUM_PROMPT_BLOCKS = 2
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    # Prompt will use 2 blocks + 1 block after we schedule.
+    NUM_TOKENS_LOCAL = int(BLOCK_SIZE * NUM_PROMPT_BLOCKS)
+    NUM_TOKENS_REMOTE = int(BLOCK_SIZE * NUM_PROMPT_BLOCKS)
+
+    request_normal = create_request(
+        request_id=1, block_size=BLOCK_SIZE, num_tokens=NUM_TOKENS_LOCAL
+    )
+    request_remote = create_request(
+        request_id=2,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS_REMOTE,
+        do_remote_prefill=True,
+    )
+
+    # STEP 1: 3 blocks are in use (2 for prompt, 1 for decode).
+    scheduler.add_request(request_normal)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_normal])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 0
+
+    # Step 2: 5 blocks are in use (2 new for remote blocks).
+    scheduler.add_request(request_remote)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_normal])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 1
+
+    # Step 3: finish recving (5 blocks in use)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[request_normal], finished_recving={request_remote.request_id}
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 1
+
+    # Step 4: try to schedule, remote request is put to running list
+    # because the transfer is completed.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[request_normal, request_remote]
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 2
+    assert _num_waiting_requests(scheduler) == 0
+
+    # Step 5: Remote request will be put back to waiting list
+    # because it needs new block to hold generated token.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_normal])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 1
+
+    # Step 6: finish the request, free it.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[request_normal], use_eos=True
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 0
+    assert _num_waiting_requests(scheduler) == 1
+
+    # Step 7: now we can schedule (with 2 blocks computed),
+    # request is retrieved from preempted list.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_remote])
+    assert (
+        scheduler_output.scheduled_cached_reqs.num_computed_tokens[0]
+        == NUM_PROMPT_BLOCKS * BLOCK_SIZE
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 0
+
+    # Step 8: free everything.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[request_remote], use_eos=True
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    _ = scheduler.schedule()
+    assert_scheduler_empty(scheduler)
+
+
+def test_cannot_recv():
+    """
+    Test that we can handle no schedule KV block transfer due to not
+    enough remaining KV blocks.
+    """
+
+    # NOTE: the KVCacheManager will use 1 null block.
+    # So there are 5 total working blocks.
+    TOTAL_NUM_BLOCKS = 6
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config, num_blocks=TOTAL_NUM_BLOCKS)
+
+    # Prime the KVCache.
+    NUM_PROMPT_BLOCKS = 2
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    # Prompt will use 2 blocks + 1 block after we schedule.
+    NUM_TOKENS_LOCAL = int(BLOCK_SIZE * NUM_PROMPT_BLOCKS)
+    NUM_TOKENS_REMOTE = int(BLOCK_SIZE * (NUM_PROMPT_BLOCKS + 0.5))
+
+    request_normal = create_request(
+        request_id=1, block_size=BLOCK_SIZE, num_tokens=NUM_TOKENS_LOCAL
+    )
+    request_remote = create_request(
+        request_id=2,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS_REMOTE,
+        do_remote_prefill=True,
+    )
+
+    # STEP 1: 3 blocks are in use (2 for prompt, 1 for decode).
+    scheduler.add_request(request_normal)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_normal])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 0
+
+    # Step 2: 3 blocks are in use,
+    # need 3 new for remote blocks but only 2 are available.
+    scheduler.add_request(request_remote)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_normal])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 1
+    # Should not have KV transfer in progress.
+    assert request_remote.status != RequestStatus.WAITING_FOR_REMOTE_KVS
+
+    # Step 3: finish the request, free it.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[request_normal], use_eos=True
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 0
+    assert _num_waiting_requests(scheduler) == 1
+
+    # Step 4: now we can initiate KV transfer (with 2 blocks computed).
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 0
+    assert _num_waiting_requests(scheduler) == 1
+    assert request_remote.status == RequestStatus.WAITING_FOR_REMOTE_KVS
+
+    # Step 5: finish recving (5 blocks in use)
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[], finished_recving={request_remote.request_id}
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 0
+    assert _num_waiting_requests(scheduler) == 1
+
+    # Step 6: schedule remote request
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(reqs=[request_remote])
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    assert len(scheduler.running) == 1
+    assert _num_waiting_requests(scheduler) == 0
+
+    # Step 7: free everything.
+    scheduler_output = scheduler.schedule()
+    model_runner_output = create_model_runner_output(
+        reqs=[request_remote], use_eos=True
+    )
+    scheduler.update_from_output(scheduler_output, model_runner_output)
+    _ = scheduler.schedule()
+    assert_scheduler_empty(scheduler)

third_party/vllm/tests/v1/kv_connector/unit/utils.py

@@ -0,0 +1,456 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import tempfile
+from collections import defaultdict
+from collections.abc import Callable
+from dataclasses import dataclass
+from itertools import chain, count
+from typing import Any, Literal
+
+import torch
+
+from vllm import SamplingParams
+from vllm.config import (
+    AttentionConfig,
+    CacheConfig,
+    DeviceConfig,
+    KVTransferConfig,
+    ModelConfig,
+    SchedulerConfig,
+    VllmConfig,
+)
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (  # noqa
+    ExampleConnector,
+)
+from vllm.utils.hashing import sha256
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
+from vllm.v1.core.sched.scheduler import Scheduler, SchedulerOutput
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    SlidingWindowSpec,
+)
+from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
+from vllm.v1.request import Request
+from vllm.v1.structured_output import StructuredOutputManager
+
+EOS_TOKEN_ID = 50256
+
+
+def assert_scheduler_empty(scheduler: Scheduler):
+    """Confirm the scheduler is "empty" - i.e. no leaks."""
+    # Scheduler Metadata.
+    assert len(scheduler.requests) == 0
+    assert len(scheduler.waiting) == 0
+    assert len(scheduler.running) == 0
+    assert len(scheduler.finished_req_ids) == 0
+    assert len(scheduler.finished_recving_kv_req_ids) == 0
+
+    # EncoderCacheManager.
+    assert len(scheduler.encoder_cache_manager.freed) == 0
+    assert len(scheduler.encoder_cache_manager.cached) == 0
+
+    # KVCache Manager.
+    assert (
+        len(
+            scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks
+        )
+        == 0
+    )
+    assert (
+        len(
+            scheduler.kv_cache_manager.coordinator.single_type_managers[
+                0
+            ].num_cached_block
+        )
+        == 0
+    )
+    num_free_blocks = (
+        scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks
+    )
+    assert num_free_blocks == (scheduler.kv_cache_manager.block_pool.num_gpu_blocks - 1)
+
+    # NOTE(rob): just the ref count on blocks will be 0. The hash
+    # value, etc will remain since we lazily evict for prefix cache.
+    for block in scheduler.kv_cache_manager.block_pool.blocks:
+        assert block.ref_cnt == 0
+
+
+def create_vllm_config(
+    model: str = "facebook/opt-125m",
+    max_num_seqs: int = 16,
+    max_num_batched_tokens: int = 64,
+    block_size: int = 16,
+    max_model_len: int = 10000,
+    enable_chunked_prefill: bool = True,
+    enable_permute_local_kv: bool = False,
+    kv_connector_extra_config: dict[str, Any] | None = None,
+    dtype: str = "float16",
+    cache_dtype: str = "auto",
+    hf_overrides: dict[str, Any] | None = None,
+    attention_backend: str | None = None,
+    kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
+) -> VllmConfig:
+    """Initialize VllmConfig For Testing."""
+    model_config = ModelConfig(
+        model=model,
+        trust_remote_code=True,
+        dtype=dtype,
+        seed=42,
+        hf_overrides=hf_overrides or {},
+    )
+    scheduler_config = SchedulerConfig(
+        max_num_seqs=max_num_seqs,
+        max_num_batched_tokens=max_num_batched_tokens,
+        max_model_len=max_model_len,
+        enable_chunked_prefill=enable_chunked_prefill,
+        is_encoder_decoder=model_config.is_encoder_decoder,
+    )
+    # Cache config, optionally force APC
+    cache_config = CacheConfig(
+        block_size=block_size,
+        gpu_memory_utilization=0.9,
+        cache_dtype=cache_dtype,
+        enable_prefix_caching=True,
+    )
+    kv_transfer_config = KVTransferConfig(
+        kv_connector="NixlConnector",
+        kv_role="kv_both",
+        enable_permute_local_kv=enable_permute_local_kv,
+        kv_connector_extra_config=kv_connector_extra_config or {},
+        kv_load_failure_policy=kv_load_failure_policy,
+    )
+    attention_config = AttentionConfig(backend=attention_backend)
+    return VllmConfig(
+        scheduler_config=scheduler_config,
+        model_config=model_config,
+        cache_config=cache_config,
+        kv_transfer_config=kv_transfer_config,
+        device_config=DeviceConfig("cpu"),
+        attention_config=attention_config,
+    )
+
+
+def create_scheduler(
+    vllm_config: VllmConfig,
+    num_blocks: int = 10000,
+    kv_cache_config: KVCacheConfig | None = None,
+) -> Scheduler | AsyncScheduler:
+    """Initialize Scheduler For Testing."""
+    block_size = vllm_config.cache_config.block_size
+    if kv_cache_config is None:
+        kv_cache_config = KVCacheConfig(
+            num_blocks=num_blocks,  # A large number of blocks to hold all requests
+            kv_cache_tensors=[],
+            kv_cache_groups=[
+                KVCacheGroupSpec(
+                    ["layer"],
+                    FullAttentionSpec(
+                        block_size=block_size,
+                        num_kv_heads=1,
+                        head_size=1,
+                        dtype=torch.float32,
+                    ),
+                )
+            ],
+        )
+    vllm_config.cache_config.num_gpu_blocks = num_blocks
+
+    scheduler_cls = (
+        AsyncScheduler if vllm_config.scheduler_config.async_scheduling else Scheduler
+    )
+    return scheduler_cls(
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+        log_stats=True,
+        structured_output_manager=StructuredOutputManager(vllm_config),
+        block_size=block_size,
+    )
+
+
+_request_count = count(1)
+_none_hash_initialized = False
+
+
+def create_request(
+    request_id: int | None = None,
+    num_tokens: int = 10,
+    common_prefix_len=0,
+    max_tokens: int = 16,
+    do_remote_decode: bool = False,
+    do_remote_prefill: bool = False,
+    num_remote_blocks: int = 3,
+    block_size: int = 16,
+    hash_fn: Callable = sha256,
+) -> Request:
+    """Make dummy request for testing."""
+    assert num_tokens >= common_prefix_len >= 0
+
+    if request_id is None:
+        request_id = next(_request_count)
+
+    global _none_hash_initialized
+    if not _none_hash_initialized:
+        init_none_hash(hash_fn)
+        _none_hash_initialized = True
+
+    kv_transfer_params: dict[str, Any] | None = None
+
+    if do_remote_decode:
+        assert not do_remote_prefill
+        kv_transfer_params = dict(do_remote_prefill=False, do_remote_decode=True)
+    elif do_remote_prefill:
+        kv_transfer_params = dict(
+            do_remote_prefill=True,
+            do_remote_decode=False,
+            remote_engine_id="my-engine-id",
+            remote_request_id=f"prefill-{request_id}",
+            remote_block_ids=list(range(num_remote_blocks)),
+            remote_host="my-host",
+            remote_port=1234,
+        )
+
+    max_tokens = 1 if do_remote_decode else max_tokens
+    sampling_params = SamplingParams(max_tokens=max_tokens)
+    sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
+
+    common_prefix = [1] * common_prefix_len if common_prefix_len > 0 else []
+    suffix = [i * request_id for i in range(num_tokens - common_prefix_len)]
+    prompt_token_ids = common_prefix + suffix
+
+    req = Request(
+        request_id=f"id-{request_id}",
+        prompt_token_ids=prompt_token_ids,
+        sampling_params=sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=get_request_block_hasher(block_size, hash_fn),
+    )
+    req.kv_transfer_params = kv_transfer_params
+    return req
+
+
+def create_model_runner_output(
+    reqs: list[Request],
+    finished_sending: set[str] | None = None,
+    finished_recving: set[str] | None = None,
+    invalid_block_ids: set[int] | None = None,
+    use_eos: bool = False,
+    token_id: int = 0,
+) -> ModelRunnerOutput:
+    """Make dummy model runner output for testing."""
+
+    # Make request data.
+    req_ids = [req.request_id for req in reqs]
+    req_id_to_index = {req_id: idx for idx, req_id in enumerate(req_ids)}
+
+    # Make sampled tokens.
+    sampled_token = EOS_TOKEN_ID if use_eos else token_id
+    sampled_token_ids = [[sampled_token] for _ in req_ids]
+
+    kv_connector_output = (
+        None
+        if (
+            finished_sending is None
+            and finished_recving is None
+            and invalid_block_ids is None
+        )
+        else KVConnectorOutput(
+            finished_sending=finished_sending,
+            finished_recving=finished_recving,
+            invalid_block_ids=invalid_block_ids or set(),
+        )
+    )
+
+    # Make output data structure.
+    return ModelRunnerOutput(
+        req_ids=req_ids,
+        req_id_to_index=req_id_to_index,
+        sampled_token_ids=sampled_token_ids,
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        kv_connector_output=kv_connector_output,
+    )
+
+
+class TestExampleConnector(ExampleConnector):
+    def __init__(self, config: VllmConfig, role, kv_cache_config):
+        self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
+        self._connector = ExampleConnector(config, role)
+        self.call_record: dict[str, int] = defaultdict(int)
+        # Use a unique temp file per connector
+        self._event_file = (
+            tempfile.gettempdir()
+            + f"/connector_{self.name}-{self.role.name}_events.log"
+        )
+        # Start with an empty file
+        with open(self._event_file, "w") as _:
+            pass
+
+    def __getattribute__(self, name):
+        if name in (
+            "_connector",
+            "call_record",
+            "name",
+            "_event_file",
+            "__class__",
+            "__dict__",
+            "__getattribute__",
+            "__init__",
+        ):  # avoid recursion
+            return object.__getattribute__(self, name)
+        if not hasattr(self._connector, name):
+            return object.__getattribute__(self, name)
+        attr = getattr(self._connector, name)
+
+        # Intercept calls to the connector interface and write an event
+        # for each one to a file, which can be read back in the main test proc.
+        if callable(attr):
+
+            def wrapper(*args, **kwargs):
+                self.call_record[name] += 1
+
+                # Include args that we're interested in
+                to_log = [name]
+                for arg in args:
+                    if isinstance(arg, int):
+                        to_log.append(str(arg))
+                    elif isinstance(arg, KVCacheBlocks):
+                        to_log.append(f"num_blocks={[len(b) for b in arg.blocks]}")
+
+                # Log the event as a line to the file
+                try:
+                    with open(self._event_file, "a") as f:
+                        f.write(" ".join(to_log) + "\n")
+                except Exception as e:
+                    print(f"[ERROR] Could not log event {name} for {self.name}: {e}")
+                return attr(*args, **kwargs)
+
+            return wrapper
+        return attr
+
+
+@dataclass(frozen=True)
+class MockKVConfig:
+    matched_tokens: int = 0
+    is_async: bool = False
+
+
+class MockKVConnectorMetadata(KVConnectorMetadata):
+    def __init__(self):
+        # Scheduler tests check metadata.requests
+        self.requests: list = []
+
+
+class MockKVConnector(KVConnectorBase_V1):
+    """Mock KV connector for scheduler tests, supporting both sync and async mode."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: KVCacheConfig | None = None,
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+        extra_config = self._kv_transfer_config.kv_connector_extra_config
+        self.config = MockKVConfig(
+            matched_tokens=extra_config["matched_tokens"],
+            is_async=extra_config["is_async"],
+        )
+
+    def get_num_new_matched_tokens(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        return (self.config.matched_tokens, self.config.is_async)
+
+    def update_state_after_alloc(
+        self,
+        request: Request,
+        blocks: KVCacheBlocks,
+        num_external_tokens: int,
+    ):
+        pass
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> KVConnectorMetadata:
+        metadata = MockKVConnectorMetadata()
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        for req_id in chain(
+            (req.req_id for req in scheduler_output.scheduled_new_reqs),
+            (
+                req_id
+                for req_id in cached_reqs.req_ids
+                if req_id in cached_reqs.resumed_req_ids
+            ),
+        ):
+            metadata.requests.append({"req_id": req_id})
+        return metadata
+
+    def start_load_kv(self, kv_caches, finished_req_ids):
+        pass
+
+    def wait_for_layer_load(self, layer_name):
+        pass
+
+    def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
+        pass
+
+    def wait_for_save(self):
+        pass
+
+
+KVConnectorFactory.register_connector(
+    "TestExampleConnector", __name__, TestExampleConnector.__name__
+)
+
+KVConnectorFactory.register_connector(
+    "MockKVConnector", __name__, MockKVConnector.__name__
+)
+
+
+def make_kv_cache_config(
+    block_size: int,
+    hma_enabled: bool = False,
+    sw_size: int = 128,
+    num_blocks: int = 100,
+) -> KVCacheConfig:
+    kv_cache_groups = [
+        KVCacheGroupSpec(
+            ["layer0", "layer2"],
+            FullAttentionSpec(
+                block_size=block_size,
+                num_kv_heads=4,
+                head_size=16,
+                dtype=torch.float16,
+            ),
+        )
+    ]
+    if hma_enabled:
+        kv_cache_groups.append(
+            KVCacheGroupSpec(
+                ["layer1", "layer3"],
+                SlidingWindowSpec(
+                    block_size=block_size,
+                    num_kv_heads=4,
+                    head_size=16,
+                    dtype=torch.float16,
+                    sliding_window=sw_size,
+                ),
+            )
+        )
+    return KVCacheConfig(
+        num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
+    )