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 13, 07:00 UTC
[KV Connector] Test async mode in scheduler tests (#28550)
Signed-off-by: Mark McLoughlin <markmc@redhat.com>
Mark McLoughlin committed
Nov 13, 2025 at 23:30 UTC
6e25b1cddfd78eab307acdb5e3ec14475e465d90
3 files changed
+165
-45
tests/v1/core/test_scheduler.py
+63
-37
@@ -31,11 +31,11 @@ from vllm.v1.kv_cache_interface import (
31
KVCacheConfig,
32
KVCacheGroupSpec,
33
)
34
-from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
34
+from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
35
from vllm.v1.request import Request, RequestStatus
36
from vllm.v1.structured_output import StructuredOutputManager
37
38
-from .utils import EOS_TOKEN_ID, create_requests, create_scheduler
38
+from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
39
40
pytestmark = pytest.mark.cpu_test
41
@@ -888,27 +888,65 @@ def _step_until_done(
888
all_finished = all_done
889
890
891
-def test_kv_connector_basic():
891
+def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
892
+ """Cycle requests through a KV transfer cyle."""
893
+
894
+ # Requests should first transition to WAITING_FOR_REMOTE_KVS
895
+ output = scheduler.schedule()
896
+ assert len(scheduler.waiting) == len(req_ids)
897
+ assert len(scheduler.running) == 0
898
+ assert len(output.scheduled_new_reqs) == 0
899
+ for req in scheduler.requests.values():
900
+ assert req.status == RequestStatus.WAITING_FOR_REMOTE_KVS
901
+
902
+ # No model execution yet
903
+ EMPTY_OUTPUT = ModelRunnerOutput(
904
+ req_ids=[],
905
+ req_id_to_index={},
906
+ sampled_token_ids=[],
907
+ logprobs=None,
908
+ prompt_logprobs_dict={},
909
+ pooler_output=[],
910
+ )
911
+ scheduler.update_from_output(output, EMPTY_OUTPUT)
912
+
913
+ # Simulate KV transfer completion using KVConnectorOutput.finished_recving
914
+ output = scheduler.schedule()
915
+ assert len(scheduler.waiting) == len(req_ids)
916
+ assert len(scheduler.running) == 0
917
+
918
+ MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
919
+ req_ids=[],
920
+ req_id_to_index={},
921
+ sampled_token_ids=[],
922
+ logprobs=None,
923
+ prompt_logprobs_dict={},
924
+ pooler_output=[],
925
+ kv_connector_output=KVConnectorOutput(finished_recving=req_ids),
926
+ )
927
+ scheduler.update_from_output(output, MODEL_RUNNER_OUTPUT)
928
+ for req_id in req_ids:
929
+ assert req_id in scheduler.finished_recving_kv_req_ids
930
+
931
+
932
+@pytest.mark.parametrize("is_async", [False, True])
933
+def test_kv_connector_basic(is_async: bool):
934
"""
935
Test whether Scheduler with KVConnector schedules tokens, allocates
936
memory, and cleans up requests as expected under normal operation.
937
"""
938
939
# Setup Scheduler.
940
+ BLOCK_SIZE = 16
941
+ NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
942
scheduler = create_scheduler(
943
enable_prefix_caching=True,
900
- use_kv_connector=True,
944
+ use_kv_connector=mock_kv(
945
+ matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async
946
+ ),
947
+ block_size=BLOCK_SIZE,
948
)
949
NUM_TOTAL_BLOCKS = scheduler.kv_cache_manager.block_pool.get_num_free_blocks()
903
- BLOCK_SIZE = scheduler.cache_config.block_size
904
-
905
- # Mock External Cache Hit.
906
- NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
907
- scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
908
- scheduler.connector.get_num_new_matched_tokens.return_value = (
909
- NUM_MATCHED_NEW_TOKENS,
910
- False,
911
- )
950
951
######################################################
952
# FIRST SET OF REQUESTS - External Hit Only
@@ -928,6 +966,9 @@ def test_kv_connector_basic():
966
req_ids.append(request.request_id)
967
req_to_index[request.request_id] = i
968
969
+ if is_async:
970
+ _step_until_kv_transfer_finished(scheduler, req_ids)
971
+
972
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
973
req_ids=req_ids,
974
req_id_to_index=req_to_index,
@@ -978,6 +1019,9 @@ def test_kv_connector_basic():
1019
req_ids.append(request.request_id)
1020
req_to_index[request.request_id] = i
1021
1022
+ if is_async:
1023
+ _step_until_kv_transfer_finished(scheduler, req_ids)
1024
+
1025
MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
1026
req_ids=req_ids,
1027
req_id_to_index=req_to_index,
@@ -1020,17 +1064,10 @@ def test_external_prefix_cache_metrics():
1064
"""
1065
1066
# Setup Scheduler.
1067
+ NUM_MATCHED_NEW_TOKENS = 4
1068
scheduler = create_scheduler(
1069
enable_prefix_caching=False,
1025
- use_kv_connector=True,
1026
- )
1027
-
1028
- # Mock connector to simulate a partial external cache hit
1029
- NUM_MATCHED_NEW_TOKENS = 4
1030
- scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
1031
- scheduler.connector.get_num_new_matched_tokens.return_value = (
1032
- NUM_MATCHED_NEW_TOKENS,
1033
- False,
1070
+ use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=False),
1071
)
1072
1073
# --- Prepare simple requests ---
@@ -1085,21 +1122,16 @@ def test_kv_connector_unable_to_allocate(use_ec_connector, ec_role):
1122
# Setup Scheduler With Mock External Cache Hit.
1123
BLOCK_SIZE = 4
1124
NUM_BLOCKS = 10
1125
+ NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
1126
scheduler = create_scheduler(
1127
enable_prefix_caching=True,
1090
- use_kv_connector=True,
1128
+ use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=False),
1129
block_size=BLOCK_SIZE,
1130
num_blocks=NUM_BLOCKS,
1131
# encoder connector should not affect test results
1132
use_ec_connector=use_ec_connector,
1133
ec_role=ec_role,
1134
)
1097
- NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2
1098
- scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
1099
- scheduler.connector.get_num_new_matched_tokens.return_value = (
1100
- NUM_MATCHED_NEW_TOKENS,
1101
- False,
1102
- )
1135
1136
# Create two requests. The second request will not be able to
1137
# allocate slots because it will not have enough blocks.
@@ -1174,9 +1206,10 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):
1206
BLOCK_SIZE = 2
1207
# NOTE: there is 1 null block, so this is 6 blocks.
1208
NUM_BLOCKS = 7
1209
+ NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE
1210
scheduler = create_scheduler(
1211
enable_prefix_caching=True,
1179
- use_kv_connector=True,
1212
+ use_kv_connector=mock_kv(matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=False),
1213
block_size=BLOCK_SIZE,
1214
num_blocks=NUM_BLOCKS,
1215
# encoder connector should not affect test results
@@ -1184,13 +1217,6 @@ def test_kv_connector_handles_preemption(use_ec_connector, ec_role):
1217
ec_role=ec_role,
1218
)
1219
1187
- NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE
1188
- scheduler.connector.get_num_new_matched_tokens = Mock(name="method")
1189
- scheduler.connector.get_num_new_matched_tokens.return_value = (
1190
- NUM_MATCHED_NEW_TOKENS,
1191
- False,
1192
- )
1193
-
1220
# Create two requests.
1221
# Both can be scheduled at first, but the second request
1222
# will be preempted and re-scheduled.
tests/v1/core/utils.py
+18
-6
@@ -3,6 +3,7 @@
3
4
import torch
5
6
+from tests.v1.kv_connector.unit.utils import MockKVConfig
7
from vllm.config import (
8
CacheConfig,
9
ECTransferConfig,
@@ -33,6 +34,10 @@ from vllm.v1.structured_output import StructuredOutputManager
34
EOS_TOKEN_ID = 50256
35
36
37
+def mock_kv(matched_tokens: int, is_async: bool):
38
+ return MockKVConfig(matched_tokens=matched_tokens, is_async=is_async)
39
+
40
+
41
def create_scheduler(
42
model: str = "facebook/opt-125m",
43
max_num_seqs: int = 16,
@@ -40,7 +45,7 @@ def create_scheduler(
45
enable_prefix_caching: bool | None = None,
46
long_prefill_token_threshold: int = 0,
47
disable_chunked_mm_input: bool = False,
43
- use_kv_connector: bool = False,
48
+ use_kv_connector: None | bool | MockKVConfig = None,
49
num_blocks: int = 10000,
50
block_size: int = 16,
51
max_model_len: int | None = None,
@@ -94,15 +99,22 @@ def create_scheduler(
99
cache_dtype="auto",
100
**kwargs_cache,
101
)
97
- kv_transfer_config = (
98
- KVTransferConfig(
102
+ kv_transfer_config = None
103
+ if isinstance(use_kv_connector, MockKVConfig):
104
+ kv_transfer_config = KVTransferConfig(
105
+ kv_connector="MockKVConnector",
106
+ kv_role="kv_both",
107
+ kv_connector_extra_config={
108
+ "matched_tokens": use_kv_connector.matched_tokens,
109
+ "is_async": use_kv_connector.is_async,
110
+ },
111
+ )
112
+ elif use_kv_connector:
113
+ kv_transfer_config = KVTransferConfig(
114
kv_connector="SharedStorageConnector",
115
kv_role="kv_both",
116
kv_connector_extra_config={"shared_storage_path": "local_storage"},
117
)
103
- if use_kv_connector
104
- else None
105
- )
118
119
speculative_config: SpeculativeConfig | None = None
120
if num_speculative_tokens is not None:
tests/v1/kv_connector/unit/utils.py
+84
-2
@@ -3,7 +3,8 @@
3
import tempfile
4
from collections import defaultdict
5
from collections.abc import Callable
6
-from itertools import count
6
+from dataclasses import dataclass
7
+from itertools import chain, count
8
from typing import Any
9
10
import torch
@@ -18,13 +19,18 @@ from vllm.config import (
19
VllmConfig,
20
)
21
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
22
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
23
+ KVConnectorBase_V1,
24
+ KVConnectorMetadata,
25
+ KVConnectorRole,
26
+)
27
from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import ( # noqa
28
SharedStorageConnector,
29
)
30
from vllm.utils.hashing import sha256
31
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
32
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
27
-from vllm.v1.core.sched.scheduler import Scheduler
33
+from vllm.v1.core.sched.scheduler import Scheduler, SchedulerOutput
34
from vllm.v1.kv_cache_interface import (
35
FullAttentionSpec,
36
KVCacheConfig,
@@ -307,6 +313,82 @@ class TestSharedStorageConnector(SharedStorageConnector):
313
return attr
314
315
316
+@dataclass(frozen=True)
317
+class MockKVConfig:
318
+ matched_tokens: int = 0
319
+ is_async: bool = False
320
+
321
+
322
+class MockKVConnectorMetadata(KVConnectorMetadata):
323
+ def __init__(self):
324
+ # Scheduler tests check metadata.requests
325
+ self.requests: list = []
326
+
327
+
328
+class MockKVConnector(KVConnectorBase_V1):
329
+ """Mock KV connector for scheduler tests, supporting both sync and async mode."""
330
+
331
+ def __init__(
332
+ self,
333
+ vllm_config: VllmConfig,
334
+ role: KVConnectorRole,
335
+ kv_cache_config: KVCacheConfig | None = None,
336
+ ):
337
+ super().__init__(vllm_config, role, kv_cache_config)
338
+ extra_config = self._kv_transfer_config.kv_connector_extra_config
339
+ self.config = MockKVConfig(
340
+ matched_tokens=extra_config["matched_tokens"],
341
+ is_async=extra_config["is_async"],
342
+ )
343
+
344
+ def get_num_new_matched_tokens(
345
+ self,
346
+ request: Request,
347
+ num_computed_tokens: int,
348
+ ) -> tuple[int | None, bool]:
349
+ return (self.config.matched_tokens, self.config.is_async)
350
+
351
+ def update_state_after_alloc(
352
+ self,
353
+ request: Request,
354
+ blocks: KVCacheBlocks,
355
+ num_external_tokens: int,
356
+ ):
357
+ pass
358
+
359
+ def build_connector_meta(
360
+ self, scheduler_output: SchedulerOutput
361
+ ) -> KVConnectorMetadata:
362
+ metadata = MockKVConnectorMetadata()
363
+ cached_reqs = scheduler_output.scheduled_cached_reqs
364
+ for req_id in chain(
365
+ (req.req_id for req in scheduler_output.scheduled_new_reqs),
366
+ (
367
+ req_id
368
+ for req_id in cached_reqs.req_ids
369
+ if req_id in cached_reqs.resumed_req_ids
370
+ ),
371
+ ):
372
+ metadata.requests.append({"req_id": req_id})
373
+ return metadata
374
+
375
+ def start_load_kv(self, kv_caches, finished_req_ids):
376
+ pass
377
+
378
+ def wait_for_layer_load(self, layer_name):
379
+ pass
380
+
381
+ def save_kv_layer(self, layer_name, kv_layer, attn_metadata, **kwargs):
382
+ pass
383
+
384
+ def wait_for_save(self):
385
+ pass
386
+
387
+
388
KVConnectorFactory.register_connector(
389
"TestSharedStorageConnector", __name__, TestSharedStorageConnector.__name__
390
)
391
+
392
+KVConnectorFactory.register_connector(
393
+ "MockKVConnector", __name__, MockKVConnector.__name__
394
+)