| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | from typing import Any |
| 6 | |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.distributed.kv_transfer.kv_connector.v1.base import ( |
| 10 | KVConnectorBase_V1, |
| 11 | KVConnectorHandshakeMetadata, |
| 12 | ) |
| 13 | from vllm.v1.engine import core as engine_core_module |
| 14 | |
| 15 | pytestmark = pytest.mark.cpu_test |
| 16 | |
| 17 | |
| 18 | class _Metadata(KVConnectorHandshakeMetadata): |
| 19 | pass |
| 20 | |
| 21 | |
| 22 | class _FakeExecutor: |
| 23 | handshake_metadata_src: ( |
| 24 | list[dict[tuple[int, int], KVConnectorHandshakeMetadata] | None] | None |
| 25 | ) |
| 26 | last_instance: "_FakeExecutor | None" = None |
| 27 | |
| 28 | def __init__( |
| 29 | self, |
| 30 | vllm_config: Any, |
| 31 | ) -> None: |
| 32 | del vllm_config |
| 33 | self.handshake_metadata = self.handshake_metadata_src |
| 34 | self.handshake_calls = 0 |
| 35 | _FakeExecutor.last_instance = self |
| 36 | |
| 37 | def get_kv_connector_handshake_metadata( |
| 38 | self, |
| 39 | ) -> list[dict[tuple[int, int], KVConnectorHandshakeMetadata] | None] | None: |
| 40 | self.handshake_calls += 1 |
| 41 | return self.handshake_metadata |
| 42 | |
| 43 | def init_kv_output_aggregator(self, connector: KVConnectorBase_V1) -> None: |
| 44 | pass |
| 45 | |
| 46 | |
| 47 | def _run_engine_core_handshake( |
| 48 | monkeypatch: pytest.MonkeyPatch, |
| 49 | connector: KVConnectorBase_V1, |
| 50 | *, |
| 51 | handshake_metadata: ( |
| 52 | list[dict[tuple[int, int], KVConnectorHandshakeMetadata] | None] | None |
| 53 | ), |
| 54 | ) -> _FakeExecutor: |
| 55 | class _FakeScheduler: |
| 56 | def __init__(self, **kwargs: Any) -> None: |
| 57 | self.connector = connector |
| 58 | self.ec_connector = None |
| 59 | |
| 60 | def get_kv_connector(self) -> KVConnectorBase_V1: |
| 61 | return connector |
| 62 | |
| 63 | _FakeExecutor.handshake_metadata_src = handshake_metadata |
| 64 | _FakeExecutor.last_instance = None |
| 65 | |
| 66 | monkeypatch.setattr("vllm.plugins.load_general_plugins", lambda: None) |
| 67 | monkeypatch.setattr( |
| 68 | engine_core_module.EngineCore, |
| 69 | "_initialize_kv_caches", |
| 70 | lambda self, vllm_config: SimpleNamespace(kv_cache_groups=[object()]), |
| 71 | ) |
| 72 | monkeypatch.setattr( |
| 73 | engine_core_module, |
| 74 | "StructuredOutputManager", |
| 75 | lambda vllm_config: object(), |
| 76 | ) |
| 77 | monkeypatch.setattr( |
| 78 | engine_core_module, |
| 79 | "resolve_kv_cache_block_sizes", |
| 80 | lambda kv_cache_config, vllm_config: (16, 16), |
| 81 | ) |
| 82 | monkeypatch.setattr(engine_core_module, "freeze_gc_heap", lambda: None) |
| 83 | monkeypatch.setattr( |
| 84 | engine_core_module, "maybe_attach_gc_debug_callback", lambda: None |
| 85 | ) |
| 86 | monkeypatch.setattr(engine_core_module, "enable_envs_cache", lambda: None) |
| 87 | monkeypatch.setattr(engine_core_module, "get_hash_fn_by_name", lambda name: None) |
| 88 | monkeypatch.setattr(engine_core_module, "init_none_hash", lambda hash_fn: None) |
| 89 | monkeypatch.setattr( |
| 90 | engine_core_module, "get_request_block_hasher", lambda *args: None |
| 91 | ) |
| 92 | |
| 93 | vllm_config = SimpleNamespace( |
| 94 | parallel_config=SimpleNamespace(data_parallel_rank_local=0), |
| 95 | scheduler_config=SimpleNamespace( |
| 96 | get_scheduler_cls=lambda: _FakeScheduler, |
| 97 | enable_chunked_prefill=False, |
| 98 | async_scheduling=False, |
| 99 | ), |
| 100 | speculative_config=None, |
| 101 | ec_transfer_config=None, |
| 102 | is_mm_encoder_only=False, |
| 103 | max_concurrent_batches=1, |
| 104 | model_config=SimpleNamespace( |
| 105 | runner_type="generate", |
| 106 | is_diffusion=False, |
| 107 | supports_multimodal_inputs=False, |
| 108 | ), |
| 109 | cache_config=SimpleNamespace( |
| 110 | enable_prefix_caching=False, |
| 111 | prefix_caching_hash_algo="builtin", |
| 112 | ), |
| 113 | ) |
| 114 | |
| 115 | engine_core_module.EngineCore(vllm_config, _FakeExecutor, log_stats=False) |
| 116 | assert _FakeExecutor.last_instance is not None |
| 117 | return _FakeExecutor.last_instance |
| 118 | |
| 119 | |
| 120 | class _LegacyConnector(KVConnectorBase_V1): |
| 121 | def __init__(self) -> None: |
| 122 | self.legacy_metadata: dict[int, KVConnectorHandshakeMetadata] | None = None |
| 123 | |
| 124 | def start_load_kv(self, forward_context: Any, **kwargs: Any) -> None: |
| 125 | pass |
| 126 | |
| 127 | def wait_for_layer_load(self, layer_name: str) -> None: |
| 128 | pass |
| 129 | |
| 130 | def save_kv_layer( |
| 131 | self, |
| 132 | layer_name: str, |
| 133 | kv_layer: Any, |
| 134 | attn_metadata: Any, |
| 135 | **kwargs: Any, |
| 136 | ) -> None: |
| 137 | pass |
| 138 | |
| 139 | def wait_for_save(self) -> None: |
| 140 | pass |
| 141 | |
| 142 | def get_num_new_matched_tokens( |
| 143 | self, request: Any, num_computed_tokens: int |
| 144 | ) -> tuple[int | None, bool]: |
| 145 | return 0, False |
| 146 | |
| 147 | def update_state_after_alloc( |
| 148 | self, request: Any, blocks: Any, num_external_tokens: int |
| 149 | ) -> None: |
| 150 | pass |
| 151 | |
| 152 | def build_connector_meta(self, scheduler_output: Any) -> Any: |
| 153 | raise NotImplementedError |
| 154 | |
| 155 | def set_xfer_handshake_metadata( |
| 156 | self, metadata: dict[int, KVConnectorHandshakeMetadata] |
| 157 | ) -> None: |
| 158 | self.legacy_metadata = metadata |
| 159 | |
| 160 | |
| 161 | class _PPAwareConnector(_LegacyConnector): |
| 162 | def __init__(self) -> None: |
| 163 | super().__init__() |
| 164 | self.pp_aware_metadata: ( |
| 165 | dict[tuple[int, int], KVConnectorHandshakeMetadata] | None |
| 166 | ) = None |
| 167 | |
| 168 | def set_xfer_handshake_metadata_pp_aware( |
| 169 | self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata] |
| 170 | ) -> None: |
| 171 | self.pp_aware_metadata = metadata |
| 172 | |
| 173 | |
| 174 | def test_engine_unwraps_handshake_metadata_for_legacy_connector( |
| 175 | monkeypatch: pytest.MonkeyPatch, |
| 176 | ) -> None: |
| 177 | """Engine core always asks workers for `(pp_rank, tp_rank)`-keyed metadata, |
| 178 | then unwraps to `{tp_rank: metadata}` for a connector that has not opted |
| 179 | into PP-aware handshake (single-PP producer, all `pp_rank == 0`).""" |
| 180 | metadata_0 = _Metadata() |
| 181 | metadata_1 = _Metadata() |
| 182 | connector = _LegacyConnector() |
| 183 | |
| 184 | executor = _run_engine_core_handshake( |
| 185 | monkeypatch, |
| 186 | connector, |
| 187 | handshake_metadata=[ |
| 188 | {(0, 0): metadata_0}, |
| 189 | None, |
| 190 | {(0, 1): metadata_1}, |
| 191 | ], |
| 192 | ) |
| 193 | |
| 194 | assert executor.handshake_calls == 1 |
| 195 | assert connector.legacy_metadata == {0: metadata_0, 1: metadata_1} |
| 196 | |
| 197 | |
| 198 | def test_engine_rejects_pp_producer_for_legacy_connector( |
| 199 | monkeypatch: pytest.MonkeyPatch, |
| 200 | ) -> None: |
| 201 | """A connector that has not opted into PP-aware handshake must not silently |
| 202 | drop metadata from `pp_rank > 0`; engine core init raises instead.""" |
| 203 | connector = _LegacyConnector() |
| 204 | |
| 205 | with pytest.raises(ValueError, match="does not support PP-disaggregated"): |
| 206 | _run_engine_core_handshake( |
| 207 | monkeypatch, |
| 208 | connector, |
| 209 | handshake_metadata=[{(0, 0): _Metadata()}, {(1, 0): _Metadata()}], |
| 210 | ) |
| 211 | |
| 212 | |
| 213 | def test_engine_passes_handshake_metadata_through_for_pp_aware_connector( |
| 214 | monkeypatch: pytest.MonkeyPatch, |
| 215 | ) -> None: |
| 216 | """A PP-aware connector receives the full `(pp_rank, tp_rank)`-keyed dict |
| 217 | unchanged.""" |
| 218 | metadata_0 = _Metadata() |
| 219 | metadata_1 = _Metadata() |
| 220 | connector = _PPAwareConnector() |
| 221 | |
| 222 | executor = _run_engine_core_handshake( |
| 223 | monkeypatch, |
| 224 | connector, |
| 225 | handshake_metadata=[{(0, 0): metadata_0}, {(1, 0): metadata_1}], |
| 226 | ) |
| 227 | |
| 228 | assert executor.handshake_calls == 1 |
| 229 | assert connector.legacy_metadata is None |
| 230 | assert connector.pp_aware_metadata == { |
| 231 | (0, 0): metadata_0, |
| 232 | (1, 0): metadata_1, |
| 233 | } |