| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | from unittest.mock import Mock |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | import vllm.v1.worker.gpu.kv_connector as kv_connector_module |
| 11 | from vllm.config import KVTransferConfig |
| 12 | from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorTransferResults |
| 13 | from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector |
| 14 | |
| 15 | |
| 16 | def _make_connector( |
| 17 | monkeypatch: pytest.MonkeyPatch, |
| 18 | events: list[str], |
| 19 | ) -> ActiveKVConnector: |
| 20 | backend = Mock() |
| 21 | backend.handle_preemptions.side_effect = lambda _: events.append("handle") |
| 22 | backend.bind_connector_metadata.side_effect = lambda _: events.append("bind") |
| 23 | backend.start_load_kv.side_effect = lambda *_args, **_kwargs: events.append("start") |
| 24 | backend.wait_for_save.side_effect = lambda: events.append("wait") |
| 25 | backend.get_transfer_results.return_value = KVConnectorTransferResults() |
| 26 | backend.get_block_ids_with_load_errors.return_value = set() |
| 27 | backend.get_kv_connector_stats.return_value = None |
| 28 | backend.get_kv_connector_kv_cache_events.return_value = None |
| 29 | backend.build_connector_worker_meta.return_value = None |
| 30 | backend.clear_connector_metadata.side_effect = lambda: events.append("clear") |
| 31 | monkeypatch.setattr(kv_connector_module, "get_kv_transfer_group", lambda: backend) |
| 32 | monkeypatch.setattr( |
| 33 | kv_connector_module, "is_forward_context_available", lambda: True |
| 34 | ) |
| 35 | monkeypatch.setattr(kv_connector_module, "get_forward_context", object) |
| 36 | |
| 37 | kv_config = KVTransferConfig( |
| 38 | kv_connector="NixlConnector", |
| 39 | kv_role="kv_consumer", |
| 40 | kv_buffer_device="cpu", |
| 41 | ) |
| 42 | connector = ActiveKVConnector(SimpleNamespace(kv_transfer_config=kv_config), {}) |
| 43 | events.clear() |
| 44 | return connector |
| 45 | |
| 46 | |
| 47 | def _scheduler_output(has_sync_kv_loads: bool) -> SimpleNamespace: |
| 48 | return SimpleNamespace( |
| 49 | kv_connector_metadata=object(), |
| 50 | finished_req_ids=set(), |
| 51 | has_sync_kv_loads=has_sync_kv_loads, |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | @pytest.mark.parametrize("has_sync_kv_loads", [False, True]) |
| 56 | def test_load_start_phase( |
| 57 | monkeypatch: pytest.MonkeyPatch, |
| 58 | has_sync_kv_loads: bool, |
| 59 | ): |
| 60 | events: list[str] = [] |
| 61 | connector = _make_connector(monkeypatch, events) |
| 62 | output = _scheduler_output(has_sync_kv_loads) |
| 63 | |
| 64 | request_indices = torch.tensor([3, 1]) |
| 65 | request_ids = ["first", "second"] |
| 66 | attn_metadata = {"layer": object()} |
| 67 | connector.pre_forward( |
| 68 | output, |
| 69 | request_state_indices=request_indices, |
| 70 | request_ids=request_ids, |
| 71 | attn_metadata=attn_metadata, |
| 72 | ) |
| 73 | assert events == ( |
| 74 | ["handle", "bind", "start"] if has_sync_kv_loads else ["handle", "bind"] |
| 75 | ) |
| 76 | |
| 77 | connector.post_forward(set()) |
| 78 | assert events == ["handle", "bind", "start", "wait", "clear"] |
| 79 | |
| 80 | kwargs = connector.kv_connector.start_load_kv.call_args.kwargs |
| 81 | assert kwargs["request_state_indices"] is request_indices |
| 82 | assert kwargs["request_ids"] is request_ids |
| 83 | assert kwargs["attn_metadata"] is attn_metadata |
| 84 | |
| 85 | # A subsequent step without a forward must not reuse the prior batch. |
| 86 | connector.no_forward(_scheduler_output(False)) |
| 87 | assert connector.kv_connector.start_load_kv.call_count == 2 |
| 88 | assert connector.kv_connector.start_load_kv.call_args.kwargs == {} |
| 89 | |
| 90 | |
| 91 | def test_no_forward_starts_deferred_load_once(monkeypatch: pytest.MonkeyPatch): |
| 92 | events: list[str] = [] |
| 93 | connector = _make_connector(monkeypatch, events) |
| 94 | |
| 95 | connector.no_forward(_scheduler_output(False)) |
| 96 | |
| 97 | assert events == ["handle", "bind", "start", "wait", "clear"] |