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, 14:30 UTC
[Core] Simplify async KV output aggregation (#28327)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Nov 9, 2025 at 09:44 UTC
289eb6c5373f74bc643890d6e14445d504c2d62d
4 files changed
+45
-153
tests/v1/executor/test_executor.py
+9
-1
@@ -9,6 +9,7 @@ from typing import Any
9
10
import pytest
11
12
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
13
from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
14
from vllm.sampling_params import SamplingParams
15
from vllm.v1.engine.async_llm import AsyncLLM
@@ -28,12 +29,19 @@ class CustomMultiprocExecutor(MultiprocExecutor):
29
kwargs: dict | None = None,
30
non_block: bool = False,
31
unique_reply_rank: int | None = None,
32
+ kv_output_aggregator: KVOutputAggregator = None,
33
) -> Any | list[Any] | Future[Any | list[Any]]:
34
# Drop marker to show that this was run
35
with open(".marker", "w"):
36
...
37
return super().collective_rpc(
36
- method, timeout, args, kwargs, non_block, unique_reply_rank
38
+ method,
39
+ timeout,
40
+ args,
41
+ kwargs,
42
+ non_block,
43
+ unique_reply_rank,
44
+ kv_output_aggregator,
45
)
46
47
tests/v1/kv_connector/unit/test_output_aggregator.py
-69
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from concurrent.futures import Future
3
4
import pytest
5
@@ -86,74 +85,6 @@ def test_aggregate_workers_output():
85
assert aggregated.invalid_block_ids == {3, 4, 5}
86
87
89
-def test_async_aggregate_workers_output():
90
- aggregator = KVOutputAggregator(expected_finished_count=2)
91
-
92
- future: Future[list[DummyModelRunnerOutput]] = Future()
93
- result_future = aggregator.async_aggregate(future)
94
-
95
- output1 = DummyModelRunnerOutput()
96
- output2 = DummyModelRunnerOutput()
97
- future.set_result([output1, output2])
98
-
99
- assert result_future.done()
100
- aggregated = result_future.result()
101
- assert aggregated is output1
102
- aggregated = aggregated.kv_connector_output
103
- assert aggregated.finished_sending is None
104
- assert aggregated.finished_recving is None
105
- assert not aggregated.invalid_block_ids
106
-
107
- future = Future()
108
- result_future = aggregator.async_aggregate(future)
109
-
110
- output1 = DummyModelRunnerOutput(
111
- finished_sending={"req1"}, finished_recving={"req2"}
112
- )
113
- output2 = DummyModelRunnerOutput(invalid_block_ids={1})
114
- future.set_result([output1, output2])
115
-
116
- assert result_future.done()
117
- aggregated = result_future.result()
118
- assert aggregated is output1
119
- aggregated = aggregated.kv_connector_output
120
- assert aggregated.finished_sending is None
121
- assert aggregated.finished_recving is None
122
- assert aggregated.invalid_block_ids == {1}
123
-
124
- future = Future()
125
- result_future = aggregator.async_aggregate(future)
126
-
127
- output1 = DummyModelRunnerOutput(invalid_block_ids={2})
128
- output2 = DummyModelRunnerOutput(finished_sending={"req1"})
129
- future.set_result([output1, output2])
130
-
131
- assert result_future.done()
132
- aggregated = result_future.result()
133
- assert aggregated is output1
134
- aggregated = aggregated.kv_connector_output
135
- assert aggregated.finished_sending == {"req1"}
136
- assert aggregated.finished_recving is None
137
- assert aggregated.invalid_block_ids == {2}
138
-
139
- future = Future()
140
- result_future = aggregator.async_aggregate(future)
141
-
142
- output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})
143
- output2 = DummyModelRunnerOutput(
144
- finished_recving={"req2"}, invalid_block_ids={4, 5}
145
- )
146
- future.set_result([output1, output2])
147
-
148
- assert result_future.done()
149
- aggregated = result_future.result()
150
- assert aggregated is output1
151
- aggregated = aggregated.kv_connector_output
152
- assert aggregated.finished_sending is None
153
- assert aggregated.finished_recving == {"req2"}
154
- assert aggregated.invalid_block_ids == {3, 4, 5}
155
-
156
-
88
def test_aggregate_workers_output_with_expected_finished_count():
89
# We create the aggregator expecting to collect from 4 workers
90
aggregator = KVOutputAggregator(expected_finished_count=4)
vllm/distributed/kv_transfer/kv_connector/utils.py
-40
@@ -4,9 +4,6 @@
4
KV cache helper for store.
5
"""
6
7
-import contextlib
8
-from collections.abc import Sequence
9
-from concurrent.futures import CancelledError, Future
7
from typing import TYPE_CHECKING, Literal
8
9
import torch
@@ -220,43 +217,6 @@ class KVOutputAggregator:
217
218
return output
219
223
- def async_aggregate(
224
- self,
225
- output_future: Future[Sequence[ModelRunnerOutput | None]],
226
- output_rank: int = 0,
227
- ) -> Future[ModelRunnerOutput | None]:
228
- """Takes a future that resolves to a list of outputs and returns a future
229
- which resolves to a single aggregated output."""
230
- result_future: Future[ModelRunnerOutput | None] = Future()
231
-
232
- def callback(fut):
233
- if result_future.done():
234
- return
235
- try:
236
- result_future.set_result(self.aggregate(fut.result(), output_rank))
237
- except CancelledError:
238
- result_future.cancel()
239
- except Exception as e:
240
- result_future.set_exception(e)
241
-
242
- output_future.add_done_callback(callback)
243
-
244
- from vllm.v1.executor.multiproc_executor import FutureWrapper
245
-
246
- if isinstance(output_future, FutureWrapper):
247
- # Due to the threadless implementation of multiproc FutureWrapper,
248
- # we must block on the delegate future's result() method.
249
- delegate_result = result_future.result
250
-
251
- def result(timeout=None):
252
- with contextlib.suppress(Exception):
253
- output_future.result(timeout=timeout)
254
- return delegate_result()
255
-
256
- result_future.result = result # type: ignore[method-assign]
257
-
258
- return result_future
259
-
220
221
def _make_src_and_dst_indices(
222
src_block_ids: list[int],
vllm/v1/executor/multiproc_executor.py
+36
-43
@@ -29,6 +29,7 @@ import vllm.envs as envs
29
from vllm.config import VllmConfig
30
from vllm.distributed import destroy_distributed_environment, destroy_model_parallel
31
from vllm.distributed.device_communicators.shm_broadcast import Handle, MessageQueue
32
+from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
33
from vllm.distributed.parallel_state import (
34
get_dp_group,
35
get_ep_group,
@@ -57,8 +58,13 @@ logger = init_logger(__name__)
58
59
60
class FutureWrapper(Future):
60
- def __init__(self, futures_queue: deque[tuple["FutureWrapper", Callable]]):
61
+ def __init__(
62
+ self,
63
+ futures_queue: deque[tuple["FutureWrapper", Callable]],
64
+ aggregate: Callable = lambda x: x,
65
+ ):
66
self.futures_queue = futures_queue
67
+ self.aggregate = aggregate
68
super().__init__()
69
70
def result(self, timeout=None):
@@ -72,7 +78,7 @@ class FutureWrapper(Future):
78
79
def wait_for_response(self, get_response: Callable):
80
try:
75
- response = get_response()
81
+ response = self.aggregate(get_response())
82
with suppress(InvalidStateError):
83
self.set_result(response)
84
except Exception as e:
@@ -160,7 +166,6 @@ class MultiprocExecutor(Executor):
166
self.futures_queue = deque[tuple[FutureWrapper, Callable]]()
167
168
self.output_rank = self._get_output_rank()
163
- self.has_connector = self.vllm_config.kv_transfer_config is not None
169
170
def start_worker_monitor(self):
171
workers = self.workers
@@ -199,44 +204,27 @@ class MultiprocExecutor(Executor):
204
def execute_model( # type: ignore[override]
205
self, scheduler_output: SchedulerOutput, non_block: bool = False
206
) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
202
- return self._execute_with_aggregation(
203
- "execute_model", scheduler_output, non_block=non_block
207
+ return self.collective_rpc(
208
+ "execute_model",
209
+ args=(scheduler_output,),
210
+ unique_reply_rank=self.output_rank,
211
+ non_block=non_block,
212
+ timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
213
+ kv_output_aggregator=self.kv_output_aggregator,
214
)
215
216
def sample_tokens( # type: ignore[override]
217
self, grammar_output: GrammarOutput | None, non_block: bool = False
218
) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
209
- return self._execute_with_aggregation( # type: ignore[return-value]
210
- "sample_tokens", grammar_output, non_block=non_block
211
- )
212
-
213
- def _execute_with_aggregation(
214
- self, method: str, *args, non_block: bool = False
215
- ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
216
- if not self.has_connector:
217
- # get output only from a single worker (output_rank)
218
- return self.collective_rpc(
219
- method,
220
- args=args,
221
- unique_reply_rank=self.output_rank,
222
- non_block=non_block,
223
- timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
224
- )
225
-
226
- # get output from all workers
227
- outputs = self.collective_rpc(
228
- method,
229
- args=args,
219
+ return self.collective_rpc(
220
+ "sample_tokens",
221
+ args=(grammar_output,),
222
+ unique_reply_rank=self.output_rank,
223
non_block=non_block,
224
timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
225
+ kv_output_aggregator=self.kv_output_aggregator,
226
)
227
234
- # aggregate all workers output to a single output
235
- assert self.kv_output_aggregator is not None
236
- if non_block:
237
- return self.kv_output_aggregator.async_aggregate(outputs, self.output_rank)
238
- return self.kv_output_aggregator.aggregate(outputs, self.output_rank)
239
-
228
def execute_dummy_batch(self) -> None:
229
self.collective_rpc("execute_dummy_batch", unique_reply_rank=self.output_rank)
230
@@ -254,8 +242,10 @@ class MultiprocExecutor(Executor):
242
kwargs: dict | None = None,
243
non_block: bool = False,
244
unique_reply_rank: int | None = None,
245
+ kv_output_aggregator: KVOutputAggregator = None,
246
) -> Any | list[Any] | Future[Any | list[Any]]:
258
- """Returns single result if unique_reply_rank is provided, otherwise list."""
247
+ """Returns single result if unique_reply_rank and/or kv_output_aggregator
248
+ is provided, otherwise list."""
249
250
if self.is_failed:
251
raise RuntimeError("Executor failed.")
@@ -263,20 +253,23 @@ class MultiprocExecutor(Executor):
253
deadline = None if timeout is None else time.monotonic() + timeout
254
kwargs = kwargs or {}
255
266
- # NOTE: If the args are heterogeneous, then we pack them into a list,
267
- # and unpack them in the method of every worker, because every worker
268
- # knows their own rank.
256
+ if kv_output_aggregator is not None:
257
+ output_rank = None
258
+ aggregate: Callable[[Any], Any] = partial(
259
+ kv_output_aggregator.aggregate, output_rank=unique_reply_rank or 0
260
+ )
261
+ else:
262
+ output_rank = unique_reply_rank
263
+ aggregate = lambda x: x
264
265
if isinstance(method, str):
266
send_method = method
267
else:
268
send_method = cloudpickle.dumps(method, protocol=pickle.HIGHEST_PROTOCOL)
274
- self.rpc_broadcast_mq.enqueue((send_method, args, kwargs, unique_reply_rank))
269
+ self.rpc_broadcast_mq.enqueue((send_method, args, kwargs, output_rank))
270
271
workers = (
277
- (self.workers[unique_reply_rank],)
278
- if unique_reply_rank is not None
279
- else self.workers
272
+ (self.workers[output_rank],) if output_rank is not None else self.workers
273
)
274
275
shutdown_event = self.shutdown_event
@@ -299,10 +292,10 @@ class MultiprocExecutor(Executor):
292
" stack trace above for the root cause"
293
)
294
responses.append(result)
302
- return responses[0] if unique_reply_rank is not None else responses
295
+ return responses[0] if output_rank is not None else responses
296
297
if non_block:
305
- future = FutureWrapper(self.futures_queue)
298
+ future = FutureWrapper(self.futures_queue, aggregate=aggregate)
299
self.futures_queue.appendleft((future, get_response))
300
return future
301
@@ -311,7 +304,7 @@ class MultiprocExecutor(Executor):
304
future, get_fut_response = self.futures_queue.pop()
305
future.wait_for_response(get_fut_response)
306
314
- return get_response()
307
+ return aggregate(get_response())
308
309
@staticmethod
310
def _ensure_worker_termination(worker_procs: list[BaseProcess]):