A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 13:15 UTC
Make KV connector metadata build overridable via plugin (#37336)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Mar 17, 2026 at 14:29 UTC
de35c06c6667ed9e1853b8a7c0d97765cb81c457
3 files changed
+142
-4
tests/v1/core/utils.py
+6
-1
@@ -47,7 +47,7 @@ def create_scheduler(
47
enable_prefix_caching: bool = False,
48
long_prefill_token_threshold: int = 0,
49
disable_chunked_mm_input: bool = False,
50
- use_kv_connector: None | bool | MockKVConfig = None,
50
+ use_kv_connector: None | bool | str | MockKVConfig = None,
51
num_blocks: int = 10000,
52
block_size: int = 16,
53
max_model_len: int | None = None,
@@ -107,6 +107,11 @@ def create_scheduler(
107
"is_async": use_kv_connector.is_async,
108
},
109
)
110
+ elif isinstance(use_kv_connector, str):
111
+ kv_transfer_config = KVTransferConfig(
112
+ kv_connector=use_kv_connector,
113
+ kv_role="kv_both",
114
+ )
115
elif use_kv_connector:
116
kv_transfer_config = KVTransferConfig(
117
kv_connector="ExampleConnector",
tests/v1/kv_connector/unit/test_scheduler_kv_connector_override.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from unittest.mock import MagicMock, patch
4
+
5
+import pytest
6
+
7
+import vllm.plugins as plugins_module
8
+from tests.v1.core.utils import create_requests, create_scheduler
9
+from vllm.distributed.kv_transfer.kv_connector.factory import (
10
+ KVConnectorFactory,
11
+)
12
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
13
+ KVConnectorBase_V1,
14
+ KVConnectorMetadata,
15
+)
16
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
17
+from vllm.v1.core.kv_cache_utils import BlockHash
18
+from vllm.v1.core.sched.output import SchedulerOutput
19
+from vllm.v1.core.sched.scheduler import Scheduler
20
+from vllm.v1.request import Request
21
+
22
+
23
+class DummyConnectorMetadata(KVConnectorMetadata):
24
+ def __init__(self, block_hashes_by_req: dict[str, list[BlockHash]]):
25
+ self.block_hashes_by_req = block_hashes_by_req
26
+
27
+
28
+class DummyKVConnector(KVConnectorBase_V1):
29
+ def __init__(self, vllm_config, role, kv_cache_config=None):
30
+ super().__init__(vllm_config, role, kv_cache_config)
31
+
32
+ def get_num_new_matched_tokens(
33
+ self, request: Request, num_computed_tokens: int
34
+ ) -> tuple[int | None, bool]:
35
+ return (0, False)
36
+
37
+ def update_state_after_alloc(
38
+ self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
39
+ ):
40
+ pass
41
+
42
+ def build_connector_meta(
43
+ self, scheduler_output: SchedulerOutput
44
+ ) -> KVConnectorMetadata:
45
+ block_hashes_by_req = getattr(scheduler_output, "block_hashes_by_req", None)
46
+ assert block_hashes_by_req is not None, (
47
+ "DummyKVConnector expected 'block_hashes_by_req' on scheduler_output"
48
+ )
49
+ return DummyConnectorMetadata(
50
+ block_hashes_by_req=block_hashes_by_req,
51
+ )
52
+
53
+ def start_load_kv(self, kv_caches, finished_req_ids):
54
+ pass
55
+
56
+ def wait_for_layer_load(self, layer_name):
57
+ pass
58
+
59
+ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
60
+ pass
61
+
62
+ def wait_for_save(self):
63
+ pass
64
+
65
+
66
+def _my_plugin():
67
+ """Registers the dummy KV connector and overrides _build_kv_connector_meta"""
68
+ KVConnectorFactory.register_connector(
69
+ "DummyKVConnector",
70
+ __name__,
71
+ DummyKVConnector.__name__,
72
+ )
73
+
74
+ def _custom_build_kv_connector_meta(
75
+ self, connector: KVConnectorBase_V1, scheduler_output: SchedulerOutput
76
+ ) -> KVConnectorMetadata:
77
+ block_hashes_by_req: dict[str, list[BlockHash]] = {}
78
+ for req_id in scheduler_output.num_scheduled_tokens:
79
+ request = self.requests[req_id]
80
+ block_hashes_by_req[req_id] = request.block_hashes
81
+
82
+ scheduler_output.block_hashes_by_req = block_hashes_by_req # type: ignore[attr-defined]
83
+ return connector.build_connector_meta(scheduler_output)
84
+
85
+ Scheduler._build_kv_connector_meta = _custom_build_kv_connector_meta
86
+
87
+
88
+@pytest.fixture
89
+def _load_plugin():
90
+ """Load the fake plugin through the real load_general_plugins() path."""
91
+ ep = MagicMock()
92
+ ep.name = "dummy_kv_connector_plugin"
93
+ ep.value = f"{__name__}:_my_plugin"
94
+ ep.load.return_value = _my_plugin
95
+
96
+ # Reset the global guard so load_general_plugins() actually runs.
97
+ plugins_module.plugins_loaded = False
98
+ with patch("importlib.metadata.entry_points", return_value=[ep]):
99
+ plugins_module.load_general_plugins()
100
+ yield
101
+ # Reset again so other tests are not affected.
102
+ plugins_module.plugins_loaded = False
103
+
104
+
105
+def test_connector_receives_block_hashes(_load_plugin):
106
+ block_size = 16
107
+ num_tokens = 48 # 3 full blocks worth of tokens
108
+ scheduler = create_scheduler(
109
+ use_kv_connector="DummyKVConnector", block_size=block_size
110
+ )
111
+ requests = create_requests(
112
+ num_requests=3, num_tokens=num_tokens, block_size=block_size
113
+ )
114
+ for req in requests:
115
+ scheduler.add_request(req)
116
+
117
+ output = scheduler.schedule()
118
+
119
+ # Verify the connector metadata was built with block hashes.
120
+ meta = output.kv_connector_metadata
121
+ assert isinstance(meta, DummyConnectorMetadata)
122
+ assert len(meta.block_hashes_by_req) == 3
123
+
124
+ for req in requests:
125
+ assert req.request_id in meta.block_hashes_by_req
126
+ # Each request has num_tokens / block_size = 3 full block hashes.
127
+ assert len(meta.block_hashes_by_req[req.request_id]) == (
128
+ num_tokens // block_size
129
+ )
130
+ assert meta.block_hashes_by_req[req.request_id] == req.block_hashes
vllm/v1/core/sched/scheduler.py
+6
-3
@@ -910,9 +910,7 @@ class Scheduler(SchedulerInterface):
910
# 2. Wrap up all the KV cache load / save ops into an opaque object
911
# 3. Clear the internal states of the connector
912
if self.connector is not None:
913
- meta: KVConnectorMetadata = self.connector.build_connector_meta(
914
- scheduler_output
915
- )
913
+ meta = self._build_kv_connector_meta(self.connector, scheduler_output)
914
scheduler_output.kv_connector_metadata = meta
915
916
# Build the connector meta for ECConnector
@@ -926,6 +924,11 @@ class Scheduler(SchedulerInterface):
924
self._update_after_schedule(scheduler_output)
925
return scheduler_output
926
927
+ def _build_kv_connector_meta(
928
+ self, connector: KVConnectorBase_V1, scheduler_output: SchedulerOutput
929
+ ) -> KVConnectorMetadata:
930
+ return connector.build_connector_meta(scheduler_output)
931
+
932
def _preempt_request(self, request: Request, timestamp: float) -> None:
933
"""Preempt a request and put it back to the waiting queue.
934