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 14, 05:30 UTC
[Core] Scheduler: Publish connector events after output (#25875)
Signed-off-by: Or Ozeri <oro@il.ibm.com>
Or Ozeri committed
Oct 28, 2025 at 23:01 UTC
111faf11185a75435403e3843e9a21a57cf84ccb
2 files changed
+135
-43
tests/v1/kv_offload/test_cpu_offloading.py
+118
-26
@@ -1,15 +1,68 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import socket
4
import time
5
6
+import msgspec
7
+import msgspec.msgpack
8
import pytest
9
+import zmq
10
+from tqdm import tqdm
11
7
-from vllm import LLM, SamplingParams
8
-from vllm.config import KVTransferConfig
12
+from vllm import LLM, SamplingParams, TokensPrompt
13
+from vllm.config import KVEventsConfig, KVTransferConfig
14
+from vllm.distributed.kv_events import BlockStored, KVEventBatch
15
16
CPU_BLOCK_SIZES = [16, 48]
17
18
19
+class MockSubscriber:
20
+ """Helper class to receive and verify published events"""
21
+
22
+ def __init__(
23
+ self,
24
+ endpoint: str,
25
+ topic: str,
26
+ ):
27
+ self.ctx = zmq.Context.instance()
28
+ self.topic_bytes = topic.encode("utf-8")
29
+
30
+ # Set up subscriber socket
31
+ self.sub = self.ctx.socket(zmq.SUB)
32
+ self.sub.setsockopt(zmq.SUBSCRIBE, self.topic_bytes)
33
+ self.sub.connect(endpoint)
34
+
35
+ self.decoder = msgspec.msgpack.Decoder(type=KVEventBatch)
36
+
37
+ def get_new_cpu_stored_events(self) -> list[BlockStored]:
38
+ cpu_stored_events: list[BlockStored] = []
39
+
40
+ poller = zmq.Poller()
41
+ poller.register(self.sub, zmq.POLLIN)
42
+
43
+ timeout = 1000 # 1 second
44
+ while True:
45
+ events = dict(poller.poll(timeout))
46
+
47
+ if events.get(self.sub) != zmq.POLLIN:
48
+ return cpu_stored_events
49
+
50
+ topic_bytes, _, payload = self.sub.recv_multipart()
51
+
52
+ assert topic_bytes == self.topic_bytes
53
+
54
+ event_batch = self.decoder.decode(payload)
55
+ assert isinstance(event_batch, KVEventBatch)
56
+ for event in event_batch.events:
57
+ if isinstance(event, BlockStored) and event.medium == "CPU":
58
+ cpu_stored_events.append(event)
59
+ timeout = 100
60
+
61
+ def close(self):
62
+ """Clean up resources"""
63
+ self.sub.close()
64
+
65
+
66
@pytest.mark.parametrize("cpu_block_size", CPU_BLOCK_SIZES)
67
def test_cpu_offloading(cpu_block_size: int) -> None:
68
"""
@@ -20,41 +73,80 @@ def test_cpu_offloading(cpu_block_size: int) -> None:
73
kv_transfer_config = KVTransferConfig(
74
kv_connector="OffloadingConnector",
75
kv_role="kv_both",
23
- kv_connector_extra_config={"num_cpu_blocks": 100, "block_size": cpu_block_size},
76
+ kv_connector_extra_config={
77
+ "num_cpu_blocks": 1000,
78
+ "block_size": cpu_block_size,
79
+ },
80
+ )
81
+
82
+ port: int
83
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
84
+ s.bind(("0.0.0.0", 0))
85
+ port = s.getsockname()[1]
86
+
87
+ events_endpoint = f"tcp://*:{port}"
88
+ kv_events_config = KVEventsConfig(
89
+ enable_kv_cache_events=True,
90
+ publisher="zmq",
91
+ endpoint=events_endpoint,
92
+ topic="test",
93
)
94
95
llm = LLM(
96
model="meta-llama/Llama-3.2-1B-Instruct",
97
gpu_memory_utilization=0.5,
98
+ kv_events_config=kv_events_config,
99
kv_transfer_config=kv_transfer_config,
30
- disable_hybrid_kv_cache_manager=True,
100
)
101
33
- prompts = ["Hi " * 100]
34
- sampling_params = SamplingParams(temperature=0, max_tokens=20)
102
+ sampling_params = SamplingParams(temperature=0, max_tokens=1)
103
+
104
+ events_endpoint = events_endpoint.replace("*", "127.0.0.1")
105
+ subscriber = MockSubscriber(events_endpoint, topic=kv_events_config.topic)
106
+
107
+ try:
108
+ num_times_cpu_better_than_cold = 0
109
+ num_tests = 10
110
+ total_cold_time = 0.0
111
+ total_gpu_hit_time = 0.0
112
+ total_cpu_hit_time = 0.0
113
+ prompt_token_ids = [0] * 10001
114
+ for i in tqdm(range(num_tests), desc="Running tests"):
115
+ prompt_token_ids[0] = i
116
+ prompts = [TokensPrompt(prompt_token_ids=prompt_token_ids)]
117
+
118
+ # run generation - this should trigger saving KV cache
119
+ start_time = time.time()
120
+ llm.generate(prompts, sampling_params, use_tqdm=False)
121
+ cold_time = time.time() - start_time
122
+ total_cold_time += cold_time
123
+
124
+ # run generation again - should hit the GPU prefix cache
125
+ start_time = time.time()
126
+ llm.generate(prompts, sampling_params, use_tqdm=False)
127
+ gpu_hit_time = time.time() - start_time
128
+ total_gpu_hit_time += gpu_hit_time
129
36
- # run generation - this should trigger saving KV cache
37
- start_time = time.time()
38
- llm.generate(prompts, sampling_params, use_tqdm=False)
39
- cold_time = time.time() - start_time
130
+ # reset prefix cache to avoid GPU hit.
131
+ llm.reset_prefix_cache()
132
41
- # run generation again - should hit the GPU prefix cache
42
- start_time = time.time()
43
- llm.generate(prompts, sampling_params, use_tqdm=False)
44
- gpu_hit_time = time.time() - start_time
133
+ assert subscriber.get_new_cpu_stored_events()
134
46
- # reset prefix cache to avoid GPU hit.
47
- llm.reset_prefix_cache()
135
+ # run generation again - this should trigger loading from CPU
136
+ start_time = time.time()
137
+ llm.generate(prompts, sampling_params, use_tqdm=False)
138
+ cpu_hit_time = time.time() - start_time
139
+ total_cpu_hit_time += cpu_hit_time
140
49
- # sleep for a sec to make sure CPU finished storing
50
- time.sleep(1)
141
+ if cpu_hit_time < cold_time:
142
+ num_times_cpu_better_than_cold += 1
143
52
- # run generation again - this should trigger loading from CPU
53
- start_time = time.time()
54
- llm.generate(prompts, sampling_params, use_tqdm=False)
55
- cpu_hit_time = time.time() - start_time
144
+ print("Average times:")
145
+ print(f" Cold: {total_cold_time * 1000 / num_tests:.2f}ms")
146
+ print(f" GPU hit: {total_gpu_hit_time * 1000 / num_tests:.2f}ms")
147
+ print(f" CPU hit: {total_cpu_hit_time * 1000 / num_tests:.2f}ms")
148
57
- print("Generation times:")
58
- print(f" Cold: {cold_time * 1000:.2f}ms")
59
- print(f" GPU hit: {gpu_hit_time * 1000:.2f}ms")
60
- print(f" CPU hit: {cpu_hit_time * 1000:.2f}ms")
149
+ assert num_times_cpu_better_than_cold >= 0.8 * num_tests
150
+ finally:
151
+ subscriber.close()
152
+ del llm
vllm/v1/core/sched/scheduler.py
+17
-17
@@ -646,23 +646,6 @@ class Scheduler(SchedulerInterface):
646
meta = self.connector.build_connector_meta(scheduler_output)
647
scheduler_output.kv_connector_metadata = meta
648
649
- # collect KV cache events from KV cache manager
650
- events = self.kv_cache_manager.take_events()
651
-
652
- # collect KV cache events from connector
653
- if self.connector is not None:
654
- connector_events = self.connector.take_events()
655
- if connector_events:
656
- if events is None:
657
- events = list(connector_events)
658
- else:
659
- events.extend(connector_events)
660
-
661
- # publish collected KV cache events
662
- if events:
663
- batch = KVEventBatch(ts=time.time(), events=events)
664
- self.kv_event_publisher.publish(batch)
665
-
649
self._update_after_schedule(scheduler_output)
650
return scheduler_output
651
@@ -1057,6 +1040,23 @@ class Scheduler(SchedulerInterface):
1040
if kv_connector_output:
1041
self._update_from_kv_xfer_finished(kv_connector_output)
1042
1043
+ # collect KV cache events from KV cache manager
1044
+ events = self.kv_cache_manager.take_events()
1045
+
1046
+ # collect KV cache events from connector
1047
+ if self.connector is not None:
1048
+ connector_events = self.connector.take_events()
1049
+ if connector_events:
1050
+ if events is None:
1051
+ events = list(connector_events)
1052
+ else:
1053
+ events.extend(connector_events)
1054
+
1055
+ # publish collected KV cache events
1056
+ if events:
1057
+ batch = KVEventBatch(ts=time.time(), events=events)
1058
+ self.kv_event_publisher.publish(batch)
1059
+
1060
# Create EngineCoreOutputs for all clients that have requests with
1061
# outputs in this step.
1062
engine_core_outputs = {