@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b712be98c7
Fork 0

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 15, 07:15 UTC

feat: add data parallel rank to KVEventBatch (#18925)

Yan Ru Pei committed Jun 3, 2025 at 17:14 UTC b712be98c790794479030313f2c2b9dae17ea7de
6 files changed +362 -86
.buildkite/test-pipeline.yaml
+2
@@ -145,6 +145,7 @@ steps:
145 - examples/offline_inference/rlhf_colocate.py
146 - tests/examples/offline_inference/data_parallel.py
147 - tests/v1/test_async_llm_dp.py
148 + - tests/v1/engine/test_engine_core_client.py
149 commands:
150 # test with tp=2 and external_dp=2
151 - VLLM_USE_V1=0 torchrun --nproc-per-node=4 distributed/test_torchrun_example.py
@@ -154,6 +155,7 @@ steps:
155 # test with internal dp
156 - python3 ../examples/offline_inference/data_parallel.py
157 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
158 + - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
159 - pytest -v -s distributed/test_utils.py
160 - pytest -v -s compile/test_basic_correctness.py
161 - pytest -v -s distributed/test_pynccl.py
tests/distributed/conftest.py
+66 -41
@@ -13,11 +13,13 @@ from vllm.distributed.kv_events import EventPublisherFactory
13
14 from .test_events import SampleBatch
15
16 +DP_RANK = 0
17 +
18
19 @pytest.fixture
20 def random_port():
21 """Generate a random port number for testing"""
20 - return random.randint(10000, 60000)
22 + return random.randint(10000, 59900)
23
24
25 @pytest.fixture
@@ -30,21 +32,23 @@ def publisher_config(random_port, request):
32 replay_endpoint = endpoint + "-replay"
33 else:
34 endpoint = f"tcp://*:{random_port}"
33 - replay_endpoint = f"tcp://*:{random_port + 1}"
35 + replay_endpoint = f"tcp://*:{random_port + 100}"
36
35 - return KVEventsConfig(enable_kv_cache_events=True,
36 - publisher="zmq",
37 - endpoint=endpoint,
38 - replay_endpoint=replay_endpoint,
39 - buffer_steps=100,
40 - hwm=1000,
41 - topic="test")
37 + return KVEventsConfig(
38 + enable_kv_cache_events=True,
39 + publisher="zmq",
40 + endpoint=endpoint,
41 + replay_endpoint=replay_endpoint,
42 + buffer_steps=100,
43 + hwm=1000,
44 + topic="test",
45 + )
46
47
48 @pytest.fixture
49 def publisher(publisher_config):
50 """Create and return a publisher instance"""
47 - pub = EventPublisherFactory.create(publisher_config)
51 + pub = EventPublisherFactory.create(publisher_config, DP_RANK)
52 yield pub
53 pub.shutdown()
54
@@ -60,7 +64,11 @@ def subscriber(publisher_config):
64 if replay_endpoint and replay_endpoint.startswith("tcp://*"):
65 replay_endpoint = replay_endpoint.replace("*", "127.0.0.1")
66
63 - sub = MockSubscriber(endpoint, replay_endpoint, publisher_config.topic)
67 + sub = MockSubscriber(
68 + [endpoint],
69 + [replay_endpoint] if replay_endpoint else None,
70 + publisher_config.topic,
71 + )
72 yield sub
73 sub.close()
74
@@ -68,26 +76,37 @@ def subscriber(publisher_config):
76 class MockSubscriber:
77 """Helper class to receive and verify published events"""
78
71 - def __init__(self,
72 - pub_endpoint: str,
73 - replay_endpoint: Optional[str] = None,
74 - topic: str = "",
75 - decode_type=SampleBatch):
79 + def __init__(
80 + self,
81 + pub_endpoints: Union[str, list[str]],
82 + replay_endpoints: Optional[Union[str, list[str]]] = None,
83 + topic: str = "",
84 + decode_type=SampleBatch,
85 + ):
86 self.ctx = zmq.Context.instance()
87
78 - # Set up subscriber socket
79 - self.sub = self.ctx.socket(zmq.SUB)
80 - self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode('utf-8'))
81 - self.sub.connect(pub_endpoint)
88 + # Convert single endpoint to list for consistency
89 + if isinstance(pub_endpoints, str):
90 + pub_endpoints = [pub_endpoints]
91 + if isinstance(replay_endpoints, str):
92 + replay_endpoints = [replay_endpoints]
93
83 - # Set up replay socket if provided
84 - self.replay = None
85 - if replay_endpoint:
86 - self.replay = self.ctx.socket(zmq.REQ)
87 - self.replay.connect(replay_endpoint)
94 + # Set up subscriber socket - connect to all endpoints
95 + self.sub = self.ctx.socket(zmq.SUB)
96 + self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode("utf-8"))
97 + for endpoint in pub_endpoints:
98 + self.sub.connect(endpoint)
99 +
100 + # Set up replay sockets if provided
101 + self.replay_sockets = []
102 + if replay_endpoints:
103 + for replay_endpoint in replay_endpoints:
104 + replay = self.ctx.socket(zmq.REQ)
105 + replay.connect(replay_endpoint)
106 + self.replay_sockets.append(replay)
107
108 self.topic = topic
90 - self.topic_bytes = topic.encode('utf-8')
109 + self.topic_bytes = topic.encode("utf-8")
110 self.received_msgs: list[tuple[int, SampleBatch]] = []
111 self.last_seq = -1
112 self.decoder = msgspec.msgpack.Decoder(type=decode_type)
@@ -107,25 +126,31 @@ class MockSubscriber:
126 self.received_msgs.append((seq, data))
127 return seq, data
128
110 - def request_replay(self, start_seq: int) -> None:
129 + def request_replay(self, start_seq: int, socket_idx: int = 0) -> None:
130 """Request replay of messages starting from start_seq"""
112 - if not self.replay:
113 - raise ValueError("Replay socket not initialized")
114 -
115 - self.replay.send(start_seq.to_bytes(8, "big"))
116 -
117 - def receive_replay(self) -> list[tuple[int, SampleBatch]]:
118 - """Receive replayed messages"""
119 - if not self.replay:
120 - raise ValueError("Replay socket not initialized")
121 -
131 + if not self.replay_sockets:
132 + raise ValueError("Replay sockets not initialized")
133 + if socket_idx >= len(self.replay_sockets):
134 + raise ValueError(f"Invalid socket index {socket_idx}")
135 +
136 + self.replay_sockets[socket_idx].send(start_seq.to_bytes(8, "big"))
137 +
138 + def receive_replay(self,
139 + socket_idx: int = 0) -> list[tuple[int, SampleBatch]]:
140 + """Receive replayed messages from a specific replay socket"""
141 + if not self.replay_sockets:
142 + raise ValueError("Replay sockets not initialized")
143 + if socket_idx >= len(self.replay_sockets):
144 + raise ValueError(f"Invalid socket index {socket_idx}")
145 +
146 + replay_socket = self.replay_sockets[socket_idx]
147 replayed: list[tuple[int, SampleBatch]] = []
148 while True:
149 try:
125 - if not self.replay.poll(1000):
150 + if not replay_socket.poll(1000):
151 break
152
128 - frames = self.replay.recv_multipart()
153 + frames = replay_socket.recv_multipart()
154 if not frames or not frames[-1]:
155 # End of replay marker
156 break
@@ -142,5 +167,5 @@ class MockSubscriber:
167 def close(self):
168 """Clean up resources"""
169 self.sub.close()
145 - if self.replay:
146 - self.replay.close()
170 + for replay in self.replay_sockets:
171 + replay.close()
tests/distributed/test_events.py
+67 -2
@@ -9,6 +9,8 @@ import pytest
9 from vllm.distributed.kv_events import (EventBatch, EventPublisherFactory,
10 NullEventPublisher)
11
12 +DP_RANK = 0
13 +
14
15 class EventSample(
16 msgspec.Struct,
@@ -121,7 +123,7 @@ def test_topic_filtering(publisher_config):
123 publisher_config.replay_endpoint = None
124
125 publisher_config.topic = "foo"
124 - pub = EventPublisherFactory.create(publisher_config)
126 + pub = EventPublisherFactory.create(publisher_config, DP_RANK)
127
128 from .conftest import MockSubscriber
129 sub_foo = MockSubscriber(publisher_config.endpoint, None, "foo")
@@ -185,9 +187,72 @@ def test_high_volume(publisher, subscriber):
187
188 def test_null_publisher():
189 """Test that NullEventPublisher can be used without errors"""
188 - publisher = NullEventPublisher()
190 + publisher = NullEventPublisher(DP_RANK)
191
192 # This should not raise any errors
193 batch = create_test_events(5)
194 publisher.publish(batch)
195 publisher.shutdown()
196 +
197 +
198 +def test_data_parallel_rank_tagging(publisher_config):
199 + """Test that events are properly tagged with their data parallel rank"""
200 +
201 + publisher_config.topic = "foo"
202 + pub_0 = EventPublisherFactory.create(publisher_config, DP_RANK)
203 + pub_1 = EventPublisherFactory.create(publisher_config, DP_RANK + 1)
204 +
205 + # Hardcode the expected endpoints based on port offsetting behavior
206 + # Both ranks get offsets according to _offset_endpoint_port function
207 + base_endpoint = publisher_config.endpoint
208 + if "tcp://" in base_endpoint:
209 + # For TCP endpoints: tcp://localhost:5557 -> tcp://localhost:5557, tcp://localhost:5558
210 + expected_endpoint_0 = base_endpoint # rank 0 gets port + 0 = same port
211 + expected_endpoint_1 = base_endpoint.replace(
212 + ":5557", ":5558") # rank 1 gets port + 1
213 + else:
214 + # For inproc endpoints: inproc://test -> inproc://test_dp0, inproc://test_dp1
215 + expected_endpoint_0 = base_endpoint # rank 0 gets base
216 + expected_endpoint_1 = base_endpoint + "_dp1" # rank 1 gets _dp1
217 +
218 + from .conftest import MockSubscriber
219 + sub_0 = MockSubscriber(expected_endpoint_0, None, publisher_config.topic)
220 + sub_1 = MockSubscriber(expected_endpoint_1, None, publisher_config.topic)
221 +
222 + try:
223 + time.sleep(0.1) # Let publishers start up
224 +
225 + # Publish events from different ranks
226 + batch_0 = create_test_events(2)
227 + batch_1 = create_test_events(3)
228 +
229 + pub_0.publish(batch_0)
230 + pub_1.publish(batch_1)
231 +
232 + # Receive events from rank 0
233 + result_0 = sub_0.receive_one(timeout=200)
234 + assert result_0 is not None, "No message received from rank 0"
235 + seq_0, received_0 = result_0
236 +
237 + # Receive events from rank 1
238 + result_1 = sub_1.receive_one(timeout=200)
239 + assert result_1 is not None, "No message received from rank 1"
240 + seq_1, received_1 = result_1
241 +
242 + # Verify DP rank tagging
243 + assert received_0.data_parallel_rank == 0, (
244 + f"Expected DP rank 0, got {received_0.data_parallel_rank}")
245 + assert received_1.data_parallel_rank == 1, (
246 + f"Expected DP rank 1, got {received_1.data_parallel_rank}")
247 +
248 + # Verify event content is correct
249 + assert len(
250 + received_0.events) == 2, "Wrong number of events from rank 0"
251 + assert len(
252 + received_1.events) == 3, "Wrong number of events from rank 1"
253 +
254 + finally:
255 + pub_0.shutdown()
256 + pub_1.shutdown()
257 + sub_0.close()
258 + sub_1.close()
tests/v1/engine/test_engine_core_client.py
+156 -33
@@ -12,8 +12,10 @@ from typing import Optional
12 import pytest
13 from transformers import AutoTokenizer
14
15 +from tests.utils import multi_gpu_test
16 from vllm import SamplingParams
16 -from vllm.distributed.kv_events import BlockStored, KVEventBatch
17 +from vllm.distributed.kv_events import (BlockStored, KVEventBatch,
18 + ZmqEventPublisher)
19 from vllm.engine.arg_utils import EngineArgs
20 from vllm.platforms import current_platform
21 from vllm.usage.usage_lib import UsageContext
@@ -37,10 +39,15 @@ PROMPT = "Hello my name is Robert and I love quantization kernels"
39 PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids
40
41
40 -def make_request(params: SamplingParams) -> EngineCoreRequest:
42 +def make_request(
43 + params: SamplingParams,
44 + prompt_tokens_ids: Optional[list[int]] = None) -> EngineCoreRequest:
45 + if not prompt_tokens_ids:
46 + prompt_tokens_ids = PROMPT_TOKENS
47 +
48 return EngineCoreRequest(
49 request_id=str(uuid.uuid4()),
43 - prompt_token_ids=PROMPT_TOKENS,
50 + prompt_token_ids=prompt_tokens_ids,
51 mm_inputs=None,
52 mm_hashes=None,
53 mm_placeholders=None,
@@ -88,6 +95,25 @@ async def loop_until_done_async(client: EngineCoreClient, outputs: dict):
95 break
96
97
98 +async def loop_until_fully_done_async(client: EngineCoreClient, outputs: dict):
99 +
100 + while True:
101 + engine_core_outputs = (await client.get_output_async()).outputs
102 +
103 + if len(engine_core_outputs) == 0:
104 + continue
105 +
106 + # Add outputs to the dict
107 + for out in engine_core_outputs:
108 + outputs[out.request_id].append(out)
109 +
110 + # Check if all request IDs in outputs have finished
111 + if all(outs and outs[-1].finished for outs in outputs.values()):
112 + break
113 +
114 + await asyncio.sleep(0.1)
115 +
116 +
117 # Dummy utility function to monkey-patch into engine core.
118 def echo(self, msg: str, err_msg: Optional[str] = None) -> str:
119 print(f"echo util function called: {msg}, {err_msg}")
@@ -273,10 +299,12 @@ def test_kv_cache_events(
299 block_size = 16
300 num_blocks = 2
301
276 - engine_args = EngineArgs(model=MODEL_NAME,
277 - enforce_eager=True,
278 - enable_prefix_caching=True,
279 - block_size=block_size)
302 + engine_args = EngineArgs(
303 + model=MODEL_NAME,
304 + enforce_eager=True,
305 + enable_prefix_caching=True,
306 + block_size=block_size,
307 + )
308 engine_args.kv_events_config = publisher_config
309
310 vllm_config = engine_args.create_engine_config(
@@ -297,19 +325,8 @@ def test_kv_cache_events(
325
326 try:
327 custom_tokens = list(range(num_blocks * block_size))
300 - request = EngineCoreRequest(
301 - request_id=str(uuid.uuid4()),
302 - prompt_token_ids=custom_tokens,
303 - mm_inputs=None,
304 - mm_hashes=None,
305 - mm_placeholders=None,
306 - sampling_params=SamplingParams(
307 - max_tokens=1), # Short completion for speed
308 - eos_token_id=None,
309 - arrival_time=time.time(),
310 - lora_request=None,
311 - cache_salt=None,
312 - )
328 + sampling_params = SamplingParams(max_tokens=1)
329 + request = make_request(sampling_params, custom_tokens)
330 client.add_request(request)
331
332 outputs: dict[str, list] = {request.request_id: []}
@@ -321,24 +338,130 @@ def test_kv_cache_events(
338 seq, received = result
339
340 assert seq == 0, "Sequence number mismatch"
324 - assert len(received.events) == 1, (
325 - "We should have exactly one BlockStored event")
341 + assert (len(received.events) == 1
342 + ), "We should have exactly one BlockStored event"
343 event = received.events[0]
344 assert isinstance(
328 - event, BlockStored), ("We should have a BlockStored event")
329 - assert len(event.block_hashes) == num_blocks, (
330 - "We should have a BlockStored event with 2 block_hashes")
331 - assert event.block_size == block_size, (
332 - "Block size should be the same as the block size")
333 - assert event.parent_block_hash is None, (
334 - "Parent block hash should be None")
345 + event, BlockStored), "We should have a BlockStored event"
346 + assert (len(event.block_hashes) == num_blocks
347 + ), "We should have a BlockStored event with 2 block_hashes"
348 + assert (event.block_size == block_size
349 + ), "Block size should be the same as the block size"
350 + assert (event.parent_block_hash
351 + is None), "Parent block hash should be None"
352 assert event.lora_id is None, "Lora id should be None"
336 - assert len(event.token_ids) == num_blocks * block_size, (
337 - "Token ids should be the same as the custom tokens")
338 - assert event.token_ids == custom_tokens, (
339 - "Token ids should be the same as the custom tokens")
353 + assert (len(event.token_ids) == num_blocks * block_size
354 + ), "Token ids should be the same as the custom tokens"
355 + assert (event.token_ids == custom_tokens
356 + ), "Token ids should be the same as the custom tokens"
357 + finally:
358 + client.shutdown()
359 + subscriber.close()
360 +
361 +
362 +@pytest.mark.asyncio
363 +@pytest.mark.parametrize(
364 + "multiprocessing_mode,publisher_config",
365 + [(True, "tcp")],
366 + indirect=["publisher_config"],
367 +)
368 +@multi_gpu_test(num_gpus=4)
369 +async def test_kv_cache_events_dp(
370 + monkeypatch: pytest.MonkeyPatch,
371 + multiprocessing_mode: bool,
372 + publisher_config,
373 +):
374 +
375 + with monkeypatch.context() as m:
376 + m.setenv("VLLM_USE_V1", "1")
377 + block_size = 16
378 + num_blocks = 2
379 + dp_size = 2
380 + tp_size = 2
381 +
382 + engine_args = EngineArgs(
383 + model=MODEL_NAME,
384 + enforce_eager=True,
385 + enable_prefix_caching=True,
386 + data_parallel_size=dp_size,
387 + tensor_parallel_size=tp_size,
388 + block_size=block_size,
389 + )
390 + engine_args.kv_events_config = publisher_config
391 +
392 + vllm_config = engine_args.create_engine_config(
393 + UsageContext.UNKNOWN_CONTEXT)
394 +
395 + executor_class = Executor.get_class(vllm_config)
396 + client = EngineCoreClient.make_client(
397 + multiprocess_mode=multiprocessing_mode,
398 + asyncio_mode=True,
399 + vllm_config=vllm_config,
400 + executor_class=executor_class,
401 + log_stats=False,
402 + )
403 + await asyncio.sleep(1)
404 +
405 + # Build endpoints for all DP ranks
406 + base_endpoint = publisher_config.endpoint.replace("*", "127.0.0.1")
407 + endpoints = []
408 + for i in range(dp_size):
409 + offset_endpoint = ZmqEventPublisher.offset_endpoint_port(
410 + base_endpoint, i)
411 + endpoints.append(offset_endpoint)
412 +
413 + subscriber = MockSubscriber(endpoints,
414 + topic=publisher_config.topic,
415 + decode_type=KVEventBatch)
416 +
417 + try:
418 + custom_tokens = list(range(num_blocks * block_size))
419 + sampling_params = SamplingParams(max_tokens=1)
420 + all_request_ids = []
421 +
422 + # Create and add 25 requests
423 + # NOTE: attempts to force routing to both dp groups but can be flaky
424 + for i in range(25):
425 + await asyncio.sleep(0.01)
426 + request = make_request(sampling_params, custom_tokens)
427 + await client.add_request_async(request)
428 + all_request_ids.append(request.request_id)
429 +
430 + await asyncio.sleep(0.1)
431 +
432 + # Initialize outputs dict for all requests
433 + outputs: dict[str, list] = {
434 + req_id: []
435 + for req_id in all_request_ids
436 + }
437 +
438 + print("processing requests...")
439 + await asyncio.wait_for(loop_until_fully_done_async(
440 + client, outputs),
441 + timeout=20.0)
442 +
443 + # Receive from subscriber until no more messages
444 + print("collecting results...")
445 + results = []
446 + while True:
447 + result = subscriber.receive_one(timeout=1)
448 + print(result)
449 + if result is None:
450 + break
451 + results.append(result)
452 +
453 + # Collect all events and data_parallel_ranks from all results
454 + all_dp_ranks = [
455 + received.data_parallel_rank for (_, received) in results
456 + ]
457 + unique_dps = set(all_dp_ranks)
458 + assert (
459 + len(unique_dps) == 2
460 + ), f"Expected 2 unique data_parallel_ranks, got {len(unique_dps)}"
461 +
462 finally:
463 client.shutdown()
464 + subscriber.close()
465
466
467 @pytest.mark.timeout(20)
vllm/distributed/kv_events.py
+68 -9
@@ -28,6 +28,7 @@ class EventBatch(
28 ):
29 ts: float
30 events: list[Any]
31 + data_parallel_rank: Optional[int] = None
32
33
34 class KVCacheEvent(
@@ -60,7 +61,22 @@ class KVEventBatch(EventBatch):
61
62
63 class EventPublisher(ABC):
63 - """Lightweight publisher for EventBatch batches."""
64 + """Lightweight publisher for EventBatch batches with data parallelism
65 + support.
66 +
67 + In data parallel setups, each DP rank runs its own EventPublisher instance
68 + to avoid duplicate events and ensure proper event attribution:
69 +
70 + - Each DP rank creates a separate publisher
71 + - Publishers automatically annotate events with their data_parallel_rank
72 + - This allows consumers to distinguish events from different DP ranks
73 +
74 + The publisher is responsible for adding DP metadata since the scheduler
75 + operates independently of DP topology and shouldn't need DP awareness.
76 + """
77 +
78 + def __init__(self, data_parallel_rank: int = 0) -> None:
79 + self._data_parallel_rank = data_parallel_rank
80
81 @abstractmethod
82 def publish(self, events: EventBatch) -> None:
@@ -113,6 +129,7 @@ class ZmqEventPublisher(EventPublisher):
129
130 def __init__(
131 self,
132 + data_parallel_rank: int,
133 endpoint: str = "tcp://*:5557",
134 replay_endpoint: Optional[str] = None,
135 buffer_steps: int = 10_000,
@@ -121,6 +138,7 @@ class ZmqEventPublisher(EventPublisher):
138 topic: str = "",
139 ) -> None:
140 # Storage
141 + super().__init__(data_parallel_rank)
142 self._event_queue = Queue[Optional[EventBatch]](maxsize=max_queue_size)
143 self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps)
144
@@ -128,8 +146,11 @@ class ZmqEventPublisher(EventPublisher):
146 self._ctx = zmq.Context.instance()
147 self._pub: Optional[zmq.Socket] = None
148 self._replay: Optional[zmq.Socket] = None
131 - self._endpoint = endpoint
132 - self._replay_endpoint = replay_endpoint
149 + self._dp_rank = data_parallel_rank
150 +
151 + self._endpoint = self.offset_endpoint_port(endpoint, self._dp_rank)
152 + self._replay_endpoint = self.offset_endpoint_port(
153 + replay_endpoint, self._dp_rank)
154 self._hwm = hwm
155 self._socket_setup()
156
@@ -149,6 +170,8 @@ class ZmqEventPublisher(EventPublisher):
170 def publish(self, events: EventBatch) -> None:
171 if not self._running:
172 raise RuntimeError("Publisher is closed")
173 + if events.data_parallel_rank is None:
174 + events.data_parallel_rank = self._data_parallel_rank
175 self._event_queue.put(events)
176
177 def shutdown(self) -> None:
@@ -191,11 +214,12 @@ class ZmqEventPublisher(EventPublisher):
214 self._pub.set_hwm(self._hwm)
215 # Heuristic: bind if wildcard / * present, else connect.
216 # bind stable, connect volatile convention
194 - if ("*" in self._endpoint or "::" in self._endpoint
195 - or self._endpoint.startswith("ipc://")
196 - or self._endpoint.startswith("inproc://")):
217 + if (self._endpoint is not None
218 + and ("*" in self._endpoint or "::" in self._endpoint
219 + or self._endpoint.startswith("ipc://")
220 + or self._endpoint.startswith("inproc://"))):
221 self._pub.bind(self._endpoint)
198 - else:
222 + elif self._endpoint is not None:
223 self._pub.connect(self._endpoint)
224
225 # Set up replay socket: use ROUTER
@@ -266,6 +290,38 @@ class ZmqEventPublisher(EventPublisher):
290 # receiving payload is (-1, b""")
291 self._replay.send_multipart((client_id, b"", self.END_SEQ, b""))
292
293 + @staticmethod
294 + def offset_endpoint_port(endpoint: Optional[str],
295 + data_parallel_rank: int) -> Optional[str]:
296 + """Helper function to offset the port in an endpoint by
297 + the data parallel rank.
298 +
299 + Args:
300 + endpoint: The endpoint string
301 + (e.g., "tcp://*:5557" or "inproc://cache")
302 + data_parallel_rank: The data parallel rank to offset by
303 +
304 + Returns:
305 + The endpoint with the port offset by data_parallel_rank
306 + or suffix appended
307 + """
308 + # Do nothing if input is None or data_parallel_rank is 0
309 + if not endpoint or data_parallel_rank == 0:
310 + return endpoint
311 +
312 + if "inproc" in endpoint:
313 + return f"{endpoint}_dp{data_parallel_rank}"
314 + if "tcp" in endpoint:
315 + if endpoint and ":" in endpoint:
316 + # Get everything after the last colon (the port)
317 + last_colon_idx = endpoint.rfind(":")
318 + base_addr = endpoint[:last_colon_idx]
319 + base_port = int(endpoint[last_colon_idx + 1:])
320 + new_port = base_port + data_parallel_rank
321 + return f"{base_addr}:{new_port}"
322 + return endpoint
323 + raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'")
324 +
325
326 class EventPublisherFactory:
327 _registry: dict[str, Callable[..., EventPublisher]] = {
@@ -281,7 +337,9 @@ class EventPublisherFactory:
337 cls._registry[name] = ctor
338
339 @classmethod
284 - def create(cls, config: Optional[KVEventsConfig]) -> EventPublisher:
340 + def create(cls,
341 + config: Optional[KVEventsConfig],
342 + data_parallel_rank: int = 0) -> EventPublisher:
343 """Create publisher from a config mapping."""
344 if not config:
345 return NullEventPublisher()
@@ -294,4 +352,5 @@ class EventPublisherFactory:
352 constructor = cls._registry[kind]
353 except KeyError as exc:
354 raise ValueError(f"Unknown event publisher '{kind}'") from exc
297 - return constructor(**config_dict)
355 + return constructor(data_parallel_rank=data_parallel_rank,
356 + **config_dict)
vllm/v1/core/sched/scheduler.py
+3 -1
@@ -80,7 +80,9 @@ class Scheduler(SchedulerInterface):
80 config=self.vllm_config, role=KVConnectorRole.SCHEDULER)
81
82 self.kv_event_publisher = EventPublisherFactory.create(
83 - self.kv_events_config)
83 + self.kv_events_config,
84 + vllm_config.parallel_config.data_parallel_rank,
85 + )
86
87 num_gpu_blocks = self.cache_config.num_gpu_blocks
88 assert num_gpu_blocks is not None and num_gpu_blocks > 0