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, 01:00 UTC
[BugFix] Make PD work with Ray (#21072)
Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>
kourosh hakhamaneshi committed
Jul 19, 2025 at 08:46 UTC
9f414a12adb991d04d2adf0b80f1f115d6281fad
11 files changed
+329
-221
tests/v1/kv_connector/unit/test_nixl_connector.py
+42
-75
@@ -1,13 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import os
5
+import tempfile
6
+import textwrap
7
import time
5
-import uuid
6
-from collections import defaultdict
7
-from typing import Optional
8
from unittest.mock import patch
9
10
import pytest
11
+import ray
12
13
from vllm import LLM
14
from vllm.config import KVTransferConfig
@@ -15,11 +16,32 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
16
KVConnectorRole, NixlAgentMetadata, NixlConnector, NixlConnectorMetadata,
17
NixlConnectorWorker)
18
from vllm.forward_context import ForwardContext
19
+from vllm.mocks.mock_nixl_connector import FakeNixlWrapper
20
from vllm.sampling_params import SamplingParams
21
22
from .utils import create_request, create_scheduler, create_vllm_config
23
24
25
+def _make_stub_pkg() -> str:
26
+ """Return a directory that makes
27
+ `from nixl._api import nixl_agent` resolve to our FakeNixlWrapper."""
28
+ td = tempfile.mkdtemp()
29
+ pkg_root = os.path.join(td, "nixl", "_api")
30
+ os.makedirs(pkg_root, exist_ok=True)
31
+
32
+ stub = textwrap.dedent("""\
33
+ # Forward the real FakeNixlWrapper that the driver already defined.
34
+ print("In fake package")
35
+ from vllm.mocks.mock_nixl_connector import FakeNixlWrapper as nixl_agent
36
+ """)
37
+ with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
38
+ f.write(stub)
39
+
40
+ # touch parent package
41
+ open(os.path.join(td, "nixl", "__init__.py"), "w").close()
42
+ return td
43
+
44
+
45
def test_basic_interface():
46
"""Unit test for basic NixlConnector interface functionality."""
47
@@ -87,77 +109,6 @@ def test_prompt_less_than_block_size():
109
assert len(scheduler_output.scheduled_new_reqs) == 1
110
111
90
-class FakeNixlWrapper:
91
- """Mock implementation of NixlWrapper for testing.
92
-
93
- We don't inherit from nixl._api.nixl_agent because nixl may not be
94
- installed.
95
- """
96
-
97
- AGENT_METADATA = b"fake_agent_metadata"
98
- REMOTE_AGENT_NAME = "remote_agent"
99
-
100
- def __init__(self, agent_name: str, *args, **kwargs):
101
- self._cycles_before_xfer_done = 0
102
- self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(
103
- lambda: 0)
104
-
105
- def get_reg_descs(self, caches_data, memory_type: str) -> list:
106
- return [str(uuid.uuid4()) for _ in caches_data]
107
-
108
- def register_memory(self, descs) -> None:
109
- pass
110
-
111
- def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
112
- return [str(uuid.uuid4()) for _ in blocks_data]
113
-
114
- def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
115
- return uuid.uuid4().int
116
-
117
- def get_agent_metadata(self) -> bytes:
118
- return self.AGENT_METADATA
119
-
120
- def add_remote_agent(self, agent_metadata: bytes) -> str:
121
- return self.REMOTE_AGENT_NAME
122
-
123
- def get_new_notifs(self) -> dict[str, list[bytes]]:
124
- # Used to collect done_sending, which we don't test yet.
125
- return {}
126
-
127
- def check_xfer_state(self, handle: int) -> str:
128
- if self._check_xfer_state_cycles[
129
- handle] >= self._cycles_before_xfer_done:
130
- return "DONE"
131
- self._check_xfer_state_cycles[handle] += 1
132
- return "PROC"
133
-
134
- def release_xfer_handle(self, handle: int) -> None:
135
- pass
136
-
137
- def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
138
- pass
139
-
140
- def make_prepped_xfer(self,
141
- xfer_type: str,
142
- local_xfer_side_handle: int,
143
- local_block_descs_ids: list[int],
144
- remote_xfer_side_handle: int,
145
- remote_block_descs_ids: list[int],
146
- notif_msg: Optional[bytes] = None) -> int:
147
- return uuid.uuid4().int
148
-
149
- def transfer(self, handle: int) -> str:
150
- return "PROC"
151
-
152
- ############################################################
153
- # Follow are for changing the behavior during testing.
154
- ############################################################
155
-
156
- def set_cycles_before_xfer_done(self, cycles: int):
157
- """Set the number of cycles before a transfer is considered done."""
158
- self._cycles_before_xfer_done = cycles
159
-
160
-
112
class FakeNixlConnectorWorker(NixlConnectorWorker):
113
114
REMOTE_ENGINE_ID = "remote_engine"
@@ -378,10 +329,14 @@ class TestNixlHandshake:
329
raise TimeoutError("Took too long to complete async handshake.")
330
331
332
+# NOTE: resource cleanup in mp backend is a bit finicky, so the order in which
333
+# we put here is important. First run ray, it will clean up the resources, then
334
+# the rest of the tests.
335
+@pytest.mark.parametrize("distributed_executor_backend", ["ray", None])
336
@patch(
337
"vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
338
FakeNixlWrapper)
384
-def test_abort_timeout_on_prefiller(monkeypatch):
339
+def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
340
"""
341
Test lifecycle of an aborted Remote Prefill request hitting the timeout.
342
-----> P
@@ -399,11 +354,23 @@ def test_abort_timeout_on_prefiller(monkeypatch):
354
timeout = 6
355
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
356
monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout))
357
+
358
+ # Build runtime_env only if we’re using Ray
359
+ if distributed_executor_backend == "ray":
360
+ runtime_env = {
361
+ "working_dir": _make_stub_pkg(), # ship stub package
362
+ "env_vars": {
363
+ "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),
364
+ },
365
+ }
366
+ ray.init(runtime_env=runtime_env)
367
+
368
llm = LLM(
369
model=model_name,
370
enforce_eager=True,
371
gpu_memory_utilization=0.5,
372
kv_transfer_config=kv_transfer_config,
373
+ distributed_executor_backend=distributed_executor_backend,
374
)
375
remote_prefill_opts = {
376
"do_remote_decode": True,
tests/v1/kv_connector/unit/test_output_aggreagator.py
renamed
+9
-28
@@ -1,28 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import threading
4
-from collections import defaultdict
3
from concurrent.futures import Future
4
from typing import Optional
5
8
-from vllm.v1.executor.multiproc_executor import MultiprocExecutor
6
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
7
from vllm.v1.outputs import ModelRunnerOutput
8
9
12
-class DummyMultiprocExecutor(MultiprocExecutor):
13
-
14
- def __init__(self, output_rank, world_size):
15
- # Manually initialize minimal required fields
16
- self.output_rank = output_rank
17
- self.world_size = world_size
18
- self._send_remaining_count = defaultdict[str,
19
- int](lambda: self.world_size)
20
- self._recv_remaining_count = defaultdict[str,
21
- int](lambda: self.world_size)
22
- self.io_thread_pool = None
23
- self.shutdown_event = threading.Event()
24
-
25
-
10
class DummyModelRunnerOutput(ModelRunnerOutput):
11
12
def __init__(self,
@@ -33,14 +17,14 @@ class DummyModelRunnerOutput(ModelRunnerOutput):
17
18
19
def test_aggregate_workers_output():
36
- executor = DummyMultiprocExecutor(output_rank=0, world_size=2)
20
+ aggregator = KVOutputAggregator(world_size=2)
21
22
output1 = DummyModelRunnerOutput(finished_sending={'req1'},
23
finished_recving={'req2'})
24
output2 = DummyModelRunnerOutput(finished_sending=None,
25
finished_recving=None)
26
43
- aggregated = executor._aggregate_workers_output([output1, output2])
27
+ aggregated = aggregator.aggregate([output1, output2])
28
29
assert aggregated is output1
30
assert aggregated.finished_sending is None
@@ -51,7 +35,7 @@ def test_aggregate_workers_output():
35
output2 = DummyModelRunnerOutput(finished_sending={'req1'},
36
finished_recving=None)
37
54
- aggregated = executor._aggregate_workers_output([output1, output2])
38
+ aggregated = aggregator.aggregate([output1, output2])
39
40
assert aggregated is output1
41
assert aggregated.finished_sending == {'req1'}
@@ -62,7 +46,7 @@ def test_aggregate_workers_output():
46
output2 = DummyModelRunnerOutput(finished_sending={'req1'},
47
finished_recving={'req2'})
48
65
- aggregated = executor._aggregate_workers_output([output1, output2])
49
+ aggregated = aggregator.aggregate([output1, output2])
50
51
assert aggregated is output1
52
assert aggregated.finished_sending is None
@@ -70,12 +54,11 @@ def test_aggregate_workers_output():
54
55
56
def test_async_aggregate_workers_output():
73
- executor = DummyMultiprocExecutor(output_rank=0, world_size=2)
57
+ aggregator = KVOutputAggregator(world_size=2)
58
59
future1: Future[DummyModelRunnerOutput] = Future()
60
future2: Future[DummyModelRunnerOutput] = Future()
77
- result_future = executor._async_aggregate_workers_output(
78
- [future1, future2])
61
+ result_future = aggregator.async_aggregate([future1, future2])
62
63
output1 = DummyModelRunnerOutput(finished_sending={'req1'},
64
finished_recving={'req2'})
@@ -92,8 +75,7 @@ def test_async_aggregate_workers_output():
75
76
future1 = Future()
77
future2 = Future()
95
- result_future = executor._async_aggregate_workers_output(
96
- [future1, future2])
78
+ result_future = aggregator.async_aggregate([future1, future2])
79
80
output1 = DummyModelRunnerOutput(finished_sending=None,
81
finished_recving=None)
@@ -110,8 +92,7 @@ def test_async_aggregate_workers_output():
92
93
future1 = Future()
94
future2 = Future()
113
- result_future = executor._async_aggregate_workers_output(
114
- [future1, future2])
95
+ result_future = aggregator.async_aggregate([future1, future2])
96
97
output1 = DummyModelRunnerOutput(finished_sending=None,
98
finished_recving=None)
vllm/distributed/kv_transfer/kv_connector/utils.py
+90
@@ -3,12 +3,18 @@
3
"""
4
KV cache helper for store.
5
"""
6
+from collections import defaultdict
7
+from collections.abc import Sequence
8
+from concurrent.futures import CancelledError, Future
9
+from typing import Optional, cast
10
+
11
import torch
12
13
import vllm.envs as envs
14
from vllm import _custom_ops as ops
15
from vllm.config import VllmConfig, get_current_vllm_config
16
from vllm.logger import init_logger
17
+from vllm.v1.outputs import ModelRunnerOutput
18
19
logger = init_logger(__name__)
20
@@ -107,3 +113,87 @@ def get_kv_connector_cache_layout():
113
"layout to HND for better xfer performance.")
114
return "HND"
115
return "NHD"
116
+
117
+
118
+class KVOutputAggregator:
119
+ """Utility class to aggregate the output of all workers into a single
120
+ output corresponding to Rank 0 for scheduler."""
121
+
122
+ def __init__(self, world_size: int):
123
+ # Complete transfer tracker. Used by to track finished requests
124
+ # [req_id -> n_finished_workers]
125
+ self._recv_remaining_count = defaultdict[str, int](lambda: world_size)
126
+ self._send_remaining_count = defaultdict[str, int](lambda: world_size)
127
+
128
+ def aggregate(self,
129
+ outputs: list[ModelRunnerOutput],
130
+ output_rank: int = 0) -> ModelRunnerOutput:
131
+ # aggregate finished_sending, finished_recving from all workers
132
+
133
+ def update_finished_set(req_ids: Optional[set[str]],
134
+ remaining_count_dict: dict[str, int],
135
+ finished_set: set[str]) -> None:
136
+ for req_id in req_ids or ():
137
+ new_count = remaining_count_dict[req_id] - 1
138
+ if new_count == 0:
139
+ finished_set.add(req_id)
140
+ del remaining_count_dict[req_id]
141
+ else:
142
+ remaining_count_dict[req_id] = new_count
143
+
144
+ finished_sending = set[str]()
145
+ finished_recving = set[str]()
146
+ for output in outputs:
147
+ update_finished_set(output.finished_sending,
148
+ self._send_remaining_count, finished_sending)
149
+ update_finished_set(output.finished_recving,
150
+ self._recv_remaining_count, finished_recving)
151
+
152
+ # select output of the worker specified by output_rank
153
+ output = outputs[output_rank]
154
+
155
+ # set the aggregated finished_sending / finished_recving
156
+ # if output.finished_sending/recving is not empty, but the other ranks
157
+ # still have unfinished send/recv, we want to set the aggregated
158
+ # finished_sending/recving to None until all ranks have finished
159
+ # send/recv
160
+ output.finished_sending = finished_sending if finished_sending else None
161
+ output.finished_recving = finished_recving if finished_recving else None
162
+
163
+ return output
164
+
165
+ def async_aggregate(self,
166
+ output_futures: Sequence[Future[ModelRunnerOutput]],
167
+ output_rank: int = 0) -> Future[ModelRunnerOutput]:
168
+ """Takes a list of futures and returns a single future which resolves
169
+ to the respective list of outputs."""
170
+ result_future: Future[ModelRunnerOutput] = Future()
171
+
172
+ outputs: list[Optional[ModelRunnerOutput]] = [None
173
+ ] * len(output_futures)
174
+
175
+ def make_callback(idx):
176
+
177
+ def callback(fut):
178
+ if result_future.done():
179
+ return
180
+
181
+ try:
182
+ outputs[idx] = fut.result()
183
+ except CancelledError:
184
+ result_future.cancel()
185
+ except Exception as e:
186
+ result_future.set_exception(e)
187
+
188
+ # this check assumes io_thread_pool uses a single thread
189
+ if all(outputs):
190
+ result_future.set_result(
191
+ self.aggregate(cast(list[ModelRunnerOutput], outputs),
192
+ output_rank))
193
+
194
+ return callback
195
+
196
+ for i, output_future in enumerate(output_futures):
197
+ output_future.add_done_callback(make_callback(i))
198
+
199
+ return result_future
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+1
-1
@@ -194,7 +194,7 @@ class KVConnectorBase_V1(ABC):
194
"""
195
Notifies worker-side connector ids of requests that have
196
finished generating tokens on the worker.
197
- The scheduler process (via the MultiprocExecutor) will use this output
197
+ The scheduler process (via the Executors) will use this output
198
to track which workers are done.
199
200
Returns:
vllm/mocks/__init__.py
vllm/mocks/mock_nixl_connector.py
new
+76
@@ -0,0 +1,76 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import uuid
4
+from collections import defaultdict
5
+from typing import Optional
6
+
7
+
8
+class FakeNixlWrapper:
9
+ """Mock implementation of NixlWrapper for testing.
10
+
11
+ We don't inherit from nixl._api.nixl_agent because nixl may not be
12
+ installed.
13
+ """
14
+
15
+ AGENT_METADATA = b"fake_agent_metadata"
16
+ REMOTE_AGENT_NAME = "remote_agent"
17
+
18
+ def __init__(self, agent_name: str, *args, **kwargs):
19
+ self._cycles_before_xfer_done = 0
20
+ self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(
21
+ lambda: 0)
22
+
23
+ def get_reg_descs(self, caches_data, memory_type: str) -> list:
24
+ return [str(uuid.uuid4()) for _ in caches_data]
25
+
26
+ def register_memory(self, descs) -> None:
27
+ pass
28
+
29
+ def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
30
+ return [str(uuid.uuid4()) for _ in blocks_data]
31
+
32
+ def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
33
+ return uuid.uuid4().int
34
+
35
+ def get_agent_metadata(self) -> bytes:
36
+ return self.AGENT_METADATA
37
+
38
+ def add_remote_agent(self, agent_metadata: bytes) -> str:
39
+ return self.REMOTE_AGENT_NAME
40
+
41
+ def get_new_notifs(self) -> dict[str, list[bytes]]:
42
+ # Used to collect done_sending, which we don't test yet.
43
+ return {}
44
+
45
+ def check_xfer_state(self, handle: int) -> str:
46
+ if self._check_xfer_state_cycles[
47
+ handle] >= self._cycles_before_xfer_done:
48
+ return "DONE"
49
+ self._check_xfer_state_cycles[handle] += 1
50
+ return "PROC"
51
+
52
+ def release_xfer_handle(self, handle: int) -> None:
53
+ pass
54
+
55
+ def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
56
+ pass
57
+
58
+ def make_prepped_xfer(self,
59
+ xfer_type: str,
60
+ local_xfer_side_handle: int,
61
+ local_block_descs_ids: list[int],
62
+ remote_xfer_side_handle: int,
63
+ remote_block_descs_ids: list[int],
64
+ notif_msg: Optional[bytes] = None) -> int:
65
+ return uuid.uuid4().int
66
+
67
+ def transfer(self, handle: int) -> str:
68
+ return "PROC"
69
+
70
+ ############################################################
71
+ # Follow are for changing the behavior during testing.
72
+ ############################################################
73
+
74
+ def set_cycles_before_xfer_done(self, cycles: int):
75
+ """Set the number of cycles before a transfer is considered done."""
76
+ self._cycles_before_xfer_done = cycles
vllm/sequence.py
+6
@@ -1188,9 +1188,15 @@ class IntermediateTensors:
1188
"""For all pipeline stages except the last, we need to return the hidden
1189
states and residuals to be sent to the next stage. This data structure
1190
contains the hidden states and residuals for a request.
1191
+
1192
+ Each stage also needs to handle its own finished_sending and
1193
+ finished_recving in case of kv transfer.
1194
"""
1195
1196
tensors: dict[str, torch.Tensor]
1197
+ # [req_ids]
1198
+ finished_sending: Optional[set[str]] = None
1199
+ finished_recving: Optional[set[str]] = None
1200
1201
def __init__(self, tensors):
1202
# manually define this function, so that
vllm/v1/executor/multiproc_executor.py
+7
-79
@@ -9,8 +9,7 @@ import threading
9
import time
10
import traceback
11
import weakref
12
-from collections import defaultdict
13
-from concurrent.futures import CancelledError, Future, ThreadPoolExecutor
12
+from concurrent.futures import Future, ThreadPoolExecutor
13
from dataclasses import dataclass
14
from enum import Enum, auto
15
from functools import partial
@@ -27,6 +26,7 @@ from vllm.distributed import (destroy_distributed_environment,
26
destroy_model_parallel)
27
from vllm.distributed.device_communicators.shm_broadcast import (Handle,
28
MessageQueue)
29
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
30
from vllm.executor.multiproc_worker_utils import (
31
_add_prefix, set_multiprocessing_worker_envs)
32
from vllm.logger import init_logger
@@ -118,13 +118,8 @@ class MultiprocExecutor(Executor):
118
119
self.output_rank = self._get_output_rank()
120
self.has_connector = self.vllm_config.kv_transfer_config is not None
121
-
122
- # Complete transfer tracker. Used by to track finished requests
123
- # [req_id -> n_finished_workers]
124
- self._recv_remaining_count = defaultdict[str,
125
- int](lambda: self.world_size)
126
- self._send_remaining_count = defaultdict[str,
127
- int](lambda: self.world_size)
121
+ self.kv_output_aggregator = KVOutputAggregator(
122
+ self.parallel_config.world_size)
123
124
def start_worker_monitor(self):
125
workers = self.workers
@@ -186,8 +181,9 @@ class MultiprocExecutor(Executor):
181
182
# aggregate all workers output to a single output
183
if non_block:
189
- return self._async_aggregate_workers_output(outputs)
190
- return self._aggregate_workers_output(outputs)
184
+ return self.kv_output_aggregator.async_aggregate(
185
+ outputs, self.output_rank)
186
+ return self.kv_output_aggregator.aggregate(outputs, self.output_rank)
187
188
def collective_rpc(self,
189
method: Union[str, Callable],
@@ -246,74 +242,6 @@ class MultiprocExecutor(Executor):
242
except TimeoutError as e:
243
raise TimeoutError(f"RPC call to {method} timed out.") from e
244
249
- def _aggregate_workers_output(
250
- self, outputs: list[ModelRunnerOutput]) -> ModelRunnerOutput:
251
- # aggregate finished_sending, finished_recving from all workers
252
-
253
- def update_finished_set(req_ids: Optional[set[str]],
254
- remaining_count_dict: dict[str, int],
255
- finished_set: set[str]) -> None:
256
- for req_id in req_ids or ():
257
- new_count = remaining_count_dict[req_id] - 1
258
- if new_count == 0:
259
- finished_set.add(req_id)
260
- del remaining_count_dict[req_id]
261
- else:
262
- remaining_count_dict[req_id] = new_count
263
-
264
- finished_sending = set[str]()
265
- finished_recving = set[str]()
266
- for output in outputs:
267
- update_finished_set(output.finished_sending,
268
- self._send_remaining_count, finished_sending)
269
- update_finished_set(output.finished_recving,
270
- self._recv_remaining_count, finished_recving)
271
-
272
- # select output of the worker specified by output_rank
273
- output = outputs[self.output_rank]
274
-
275
- # set the aggregated finished_sending / finished_recving
276
- output.finished_sending = finished_sending if finished_sending else None
277
- output.finished_recving = finished_recving if finished_recving else None
278
-
279
- return output
280
-
281
- def _async_aggregate_workers_output(
282
- self, output_futures: list[Future[ModelRunnerOutput]]
283
- ) -> (Future[ModelRunnerOutput]):
284
- """Takes a list of futures and returns a single future which resolves
285
- to the respective list of outputs."""
286
- result_future: Future[ModelRunnerOutput] = Future()
287
-
288
- outputs: list[Optional[ModelRunnerOutput]] = [None
289
- ] * len(output_futures)
290
-
291
- def make_callback(idx):
292
-
293
- def callback(fut):
294
- if result_future.done():
295
- return
296
-
297
- try:
298
- outputs[idx] = fut.result()
299
- except CancelledError:
300
- result_future.cancel()
301
- except Exception as e:
302
- result_future.set_exception(e)
303
-
304
- # this check assumes io_thread_pool uses a single thread
305
- if all(outputs):
306
- result_future.set_result(
307
- self._aggregate_workers_output(
308
- cast(list[ModelRunnerOutput], outputs)))
309
-
310
- return callback
311
-
312
- for i, output_future in enumerate(output_futures):
313
- output_future.add_done_callback(make_callback(i))
314
-
315
- return result_future
316
-
245
@staticmethod
246
def _ensure_worker_termination(worker_procs: list[BaseProcess]):
247
"""Ensure that all worker processes are terminated. Assumes workers have
vllm/v1/executor/ray_distributed_executor.py
+45
-12
@@ -2,33 +2,55 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from concurrent.futures import Future
5
-from typing import Union
5
+from typing import Optional, Union
6
7
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
8
from vllm.executor.ray_distributed_executor import ( # noqa
9
RayDistributedExecutor as RayDistributedExecutorV0)
10
+from vllm.logger import init_logger
11
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
12
from vllm.v1.executor.abstract import Executor
13
from vllm.v1.outputs import ModelRunnerOutput
14
15
+logger = init_logger(__name__)
16
+
17
18
class FutureWrapper(Future):
15
- """A wrapper around a Ray output reference to meet the interface
16
- of .execute_model().
19
+ """A wrapper around Ray output reference to meet the interface
20
+ of .execute_model(): The top level (core busy loop) expects .result() api
21
+ to block and return a single output.
22
+
23
+ If aggregator is provided, the outputs from all workers are aggregated upon
24
+ the result() call. If not only the first worker's output is returned.
25
"""
26
19
- def __init__(self, ref):
27
+ def __init__(self, refs, aggregator: Optional[KVOutputAggregator] = None):
28
super().__init__()
21
- self.ref = ref
29
+ self.refs = refs
30
+ self.aggregator = aggregator
31
32
def result(self, timeout=None):
33
if timeout is not None:
34
raise NotImplementedError("timeout is not supported")
26
- return self.ref.get()
35
+
36
+ if self.aggregator is None:
37
+ return self.refs[0].get()
38
+
39
+ outputs = [ref.get() for ref in self.refs]
40
+ return self.aggregator.aggregate(outputs, output_rank=0)
41
42
43
class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
44
"""Ray distributed executor using Ray Compiled Graphs."""
45
46
+ def _init_executor(self) -> None:
47
+ super()._init_executor()
48
+
49
+ # KV connector setup
50
+ self.has_connector = self.vllm_config.kv_transfer_config is not None
51
+ self.kv_output_aggregator = KVOutputAggregator(
52
+ self.parallel_config.world_size)
53
+
54
@property
55
def max_concurrent_batches(self) -> int:
56
"""Ray distributed executor supports pipeline parallelism,
@@ -56,13 +78,24 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
78
79
refs = self.forward_dag.execute(scheduler_output) # type: ignore
80
59
- # When PP is not used, we block here until the result is available.
81
+ if not self.has_connector:
82
+ # Get output only from a single worker (output_rank)
83
+ # When PP is not used, we block here until the result is available.
84
+ if self.max_concurrent_batches == 1:
85
+ return refs[0].get()
86
+
87
+ # When PP is used, we return a FutureWrapper immediately so that
88
+ # the scheduler can yield to the next batch.
89
+ return FutureWrapper(refs)
90
+
91
+ # Get output from all workers when connector is present
92
if self.max_concurrent_batches == 1:
61
- return refs[0].get()
93
+ # Block and get results from all workers
94
+ outputs = [ref.get() for ref in refs]
95
+ return self.kv_output_aggregator.aggregate(outputs)
96
63
- # When PP is used, we return a FutureWrapper immediately so that
64
- # the scheduler can yield to the next batch.
65
- return FutureWrapper(refs[0])
97
+ # Return a future that will aggregate outputs from all workers
98
+ return FutureWrapper(refs, self.kv_output_aggregator)
99
100
def reinitialize_distributed(
101
self, reconfig_request: ReconfigureDistributedRequest) -> None:
@@ -70,4 +103,4 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
103
if reconfig_request.new_data_parallel_rank == \
104
ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
105
self.shutdown()
73
- return
106
+ return
\ No newline at end of file
vllm/v1/worker/gpu_model_runner.py
+43
-6
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import copy
5
import gc
6
import time
7
from contextlib import contextmanager
@@ -1270,6 +1271,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1271
hidden_states: torch.Tensor,
1272
num_scheduled_tokens: int,
1273
num_scheduled_tokens_np: np.ndarray,
1274
+ finished_sending: Optional[set[str]],
1275
+ finished_recving: Optional[set[str]],
1276
) -> ModelRunnerOutput:
1277
assert self.input_batch.num_reqs ==\
1278
len(self.input_batch.pooling_params), \
@@ -1304,6 +1307,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1307
logprobs=None,
1308
prompt_logprobs_dict={},
1309
pooler_output=pooler_output,
1310
+ finished_sending=finished_sending,
1311
+ finished_recving=finished_recving,
1312
)
1313
1314
@torch.inference_mode()
@@ -1314,12 +1319,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1319
) -> Union[ModelRunnerOutput, IntermediateTensors]:
1320
self._update_states(scheduler_output)
1321
if not scheduler_output.total_num_scheduled_tokens:
1317
- if has_kv_transfer_group():
1318
- with set_forward_context(None, self.vllm_config):
1319
- self.maybe_setup_kv_connector(scheduler_output)
1322
+ if not has_kv_transfer_group():
1323
+ # Return empty ModelRunnerOutput if there's no work to do.
1324
+ return EMPTY_MODEL_RUNNER_OUTPUT
1325
1321
- # Return empty ModelRunnerOutput if there's no work to do.
1322
- return EMPTY_MODEL_RUNNER_OUTPUT
1326
+ return self.kv_connector_no_forward(scheduler_output)
1327
1328
# Prepare the decoder inputs.
1329
(attn_metadata, attention_cuda_graphs, logits_indices,
@@ -1412,6 +1416,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1416
)
1417
1418
self.maybe_wait_for_kv_save()
1419
+ finished_sending, finished_recving = (
1420
+ self.get_finished_kv_transfers(scheduler_output))
1421
1422
if self.use_aux_hidden_state_outputs:
1423
hidden_states, aux_hidden_states = model_output
@@ -1429,6 +1435,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1435
if not get_pp_group().is_last_rank:
1436
# For mid-pipeline stages, return the hidden states.
1437
if not broadcast_pp_output:
1438
+ if finished_sending or finished_recving:
1439
+ hidden_states.finished_sending = finished_sending
1440
+ hidden_states.finished_recving = finished_recving
1441
return hidden_states
1442
assert isinstance(hidden_states, IntermediateTensors)
1443
get_pp_group().send_tensor_dict(hidden_states.tensors,
@@ -1437,7 +1446,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1446
else:
1447
if self.input_batch.pooling_params:
1448
return self._pool(hidden_states, num_scheduled_tokens,
1440
- num_scheduled_tokens_np)
1449
+ num_scheduled_tokens_np, finished_sending,
1450
+ finished_recving)
1451
1452
sample_hidden_states = hidden_states[logits_indices]
1453
logits = self.model.compute_logits(sample_hidden_states, None)
@@ -1587,6 +1597,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1597
logprobs=logprobs_lists,
1598
prompt_logprobs_dict=prompt_logprobs_dict,
1599
pooler_output=[],
1600
+ finished_sending=finished_sending,
1601
+ finished_recving=finished_recving,
1602
num_nans_in_logits=num_nans_in_logits,
1603
)
1604
@@ -1711,6 +1723,31 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1723
if has_kv_transfer_group():
1724
get_kv_transfer_group().wait_for_save()
1725
1726
+ @staticmethod
1727
+ def get_finished_kv_transfers(
1728
+ scheduler_output: "SchedulerOutput",
1729
+ ) -> tuple[Optional[set[str]], Optional[set[str]]]:
1730
+ if has_kv_transfer_group():
1731
+ return get_kv_transfer_group().get_finished(
1732
+ scheduler_output.finished_req_ids)
1733
+ return None, None
1734
+
1735
+ def kv_connector_no_forward(
1736
+ self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
1737
+ # KV send/recv even if no work to do.
1738
+ with set_forward_context(None, self.vllm_config):
1739
+ self.maybe_setup_kv_connector(scheduler_output)
1740
+ finished_sending, finished_recving = (
1741
+ self.get_finished_kv_transfers(scheduler_output))
1742
+
1743
+ if not finished_sending and not finished_recving:
1744
+ return EMPTY_MODEL_RUNNER_OUTPUT
1745
+
1746
+ output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
1747
+ output.finished_sending = finished_sending
1748
+ output.finished_recving = finished_recving
1749
+ return output
1750
+
1751
def propose_ngram_draft_token_ids(
1752
self,
1753
sampled_token_ids: list[list[int]],
vllm/v1/worker/gpu_worker.py
+10
-20
@@ -15,9 +15,7 @@ from vllm.config import VllmConfig
15
from vllm.distributed import (ensure_model_parallel_initialized,
16
init_distributed_environment,
17
set_custom_all_reduce)
18
-from vllm.distributed.kv_transfer import (ensure_kv_transfer_initialized,
19
- get_kv_transfer_group,
20
- has_kv_transfer_group)
18
+from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
19
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
20
from vllm.logger import init_logger
21
from vllm.lora.request import LoRARequest
@@ -335,25 +333,17 @@ class Worker(WorkerBase):
333
assert isinstance(output, IntermediateTensors)
334
get_pp_group().send_tensor_dict(output.tensors,
335
all_gather_group=get_tp_group())
338
- output = EMPTY_MODEL_RUNNER_OUTPUT
336
340
- assert isinstance(output, ModelRunnerOutput)
341
- if has_kv_transfer_group():
342
- finished_sending, finished_recving = (
343
- get_kv_transfer_group().get_finished(
344
- scheduler_output.finished_req_ids))
345
- if finished_sending or finished_recving:
346
- if output is EMPTY_MODEL_RUNNER_OUTPUT:
347
- output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
348
- output.finished_sending = finished_sending
349
- output.finished_recving = finished_recving
350
-
351
- # Clear KVConnector state for this step.
352
- get_kv_transfer_group().clear_connector_metadata()
353
-
354
- # with a connector, the scheduler expects output from all workers
355
- return output
337
+ # In case of PP with kv transfer, we need to pass through the
338
+ # finished_sending and finished_recving buffers.
339
+ empty_output = EMPTY_MODEL_RUNNER_OUTPUT
340
+ if output.finished_sending or output.finished_recving:
341
+ empty_output = copy.copy(empty_output)
342
+ empty_output.finished_sending = output.finished_sending
343
+ empty_output.finished_recving = output.finished_recving
344
+ output = empty_output
345
346
+ assert isinstance(output, ModelRunnerOutput)
347
# return output only from the driver worker
348
return output if self.is_driver_worker else None
349