@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 67a2da890e

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 13, 14:30 UTC

[PerfFix] Avoid separate thread for MP executor shm spin (take 2) (#28319)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Nov 7, 2025 at 14:11 UTC 67a2da890eef2a6fd40384aa5ae80e03beb39490

9 files changed +153 -128

tests/v1/executor/test_executor.py

+2 -1

@@ -4,6 +4,7 @@

4 import asyncio

5 import os

6 from collections.abc import Callable

7 +from concurrent.futures import Future

8 from typing import Any

10 import pytest

@@ -27,7 +28,7 @@ class CustomMultiprocExecutor(MultiprocExecutor):

28 kwargs: dict | None = None,

29 non_block: bool = False,

30 unique_reply_rank: int | None = None,

30 - ) -> list[Any]:

31 + ) -> Any | list[Any] | Future[Any | list[Any]]:

32 # Drop marker to show that this was run

33 with open(".marker", "w"):

34 ...

tests/v1/kv_connector/unit/test_output_aggregator.py

+12 -20

@@ -89,14 +89,12 @@ def test_aggregate_workers_output():

89 def test_async_aggregate_workers_output():

90 aggregator = KVOutputAggregator(expected_finished_count=2)

92 - future1: Future[DummyModelRunnerOutput] = Future()

93 - future2: Future[DummyModelRunnerOutput] = Future()

94 - result_future = aggregator.async_aggregate([future1, future2])

92 + future: Future[list[DummyModelRunnerOutput]] = Future()

93 + result_future = aggregator.async_aggregate(future)

95 output1 = DummyModelRunnerOutput()

96 output2 = DummyModelRunnerOutput()

98 - future1.set_result(output1)

99 - future2.set_result(output2)

97 + future.set_result([output1, output2])

99 assert result_future.done()

100 aggregated = result_future.result()

@@ -106,16 +104,14 @@ def test_async_aggregate_workers_output():

104 assert aggregated.finished_recving is None

105 assert not aggregated.invalid_block_ids

106

109 - future1 = Future()

110 - future2 = Future()

111 - result_future = aggregator.async_aggregate([future1, future2])

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})

117 - future1.set_result(output1)

118 - future2.set_result(output2)

114 + future.set_result([output1, output2])

115

116 assert result_future.done()

117 aggregated = result_future.result()

@@ -125,14 +121,12 @@ def test_async_aggregate_workers_output():

121 assert aggregated.finished_recving is None

122 assert aggregated.invalid_block_ids == {1}

123

128 - future1 = Future()

129 - future2 = Future()

130 - result_future = aggregator.async_aggregate([future1, future2])

124 + future = Future()

125 + result_future = aggregator.async_aggregate(future)

126

127 output1 = DummyModelRunnerOutput(invalid_block_ids={2})

128 output2 = DummyModelRunnerOutput(finished_sending={"req1"})

134 - future1.set_result(output1)

135 - future2.set_result(output2)

129 + future.set_result([output1, output2])

130

131 assert result_future.done()

132 aggregated = result_future.result()

@@ -142,16 +136,14 @@ def test_async_aggregate_workers_output():

136 assert aggregated.finished_recving is None

137 assert aggregated.invalid_block_ids == {2}

138

145 - future1 = Future()

146 - future2 = Future()

147 - result_future = aggregator.async_aggregate([future1, future2])

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 )

153 - future1.set_result(output1)

154 - future2.set_result(output2)

146 + future.set_result([output1, output2])

147

148 assert result_future.done()

149 aggregated = result_future.result()

vllm/distributed/kv_transfer/kv_connector/utils.py

+24 -23

@@ -4,6 +4,7 @@

4 KV cache helper for store.

5 """

7 +import contextlib

8 from collections.abc import Sequence

9 from concurrent.futures import CancelledError, Future

10 from typing import TYPE_CHECKING, Literal

@@ -221,38 +222,38 @@ class KVOutputAggregator:

222

223 def async_aggregate(

224 self,

224 - output_futures: Sequence[Future[ModelRunnerOutput | None]],

225 + output_future: Future[Sequence[ModelRunnerOutput | None]],

226 output_rank: int = 0,

227 ) -> Future[ModelRunnerOutput | None]:

227 - """Takes a list of futures and returns a single future which resolves

228 - to the respective list of outputs."""

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

231 - outputs: list[ModelRunnerOutput | None] = [None] * len(output_futures)

232 - remaining = len(output_futures)

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

234 - def make_callback(idx):

235 - def callback(fut):

236 - if result_future.done():

237 - return

242 + output_future.add_done_callback(callback)

243

239 - try:

240 - outputs[idx] = fut.result()

241 - except CancelledError:

242 - result_future.cancel()

243 - except Exception as e:

244 - result_future.set_exception(e)

244 + from vllm.v1.executor.multiproc_executor import FutureWrapper

245

246 - # this check assumes io_thread_pool uses a single thread

247 - nonlocal remaining

248 - remaining -= 1

249 - if not remaining:

250 - result_future.set_result(self.aggregate(outputs, output_rank))

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

252 - return callback

251 + def result(timeout=None):

252 + with contextlib.suppress(Exception):

253 + output_future.result(timeout=timeout)

254 + return delegate_result()

255

254 - for i, output_future in enumerate(output_futures):

255 - output_future.add_done_callback(make_callback(i))

256 + result_future.result = result # type: ignore[method-assign]

257

258 return result_future

259

vllm/v1/executor/abstract.py

+2 -2

@@ -171,7 +171,7 @@ class Executor(ABC):

171 args: tuple = (),

172 kwargs: dict | None = None,

173 non_block: Literal[True] = True,

174 - ) -> list[Future[_R]]:

174 + ) -> Future[list[_R]]:

175 pass

176

177 @abstractmethod

@@ -219,7 +219,7 @@ class Executor(ABC):

219

220 def sample_tokens(

221 self, grammar_output: GrammarOutput | None, non_block: bool = False

222 - ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:

222 + ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:

223 output = self.collective_rpc( # type: ignore[call-overload]

224 "sample_tokens", args=(grammar_output,), non_block=non_block

225 )

vllm/v1/executor/multiproc_executor.py

+67 -61

@@ -9,8 +9,10 @@ import threading

9 import time

10 import traceback

11 import weakref

12 +from collections import deque

13 from collections.abc import Callable

13 -from concurrent.futures import Future, ThreadPoolExecutor

14 +from concurrent.futures import Future, InvalidStateError

15 +from contextlib import suppress

16 from dataclasses import dataclass

17 from enum import Enum, auto

18 from functools import cached_property, partial

@@ -54,6 +56,30 @@ from vllm.v1.worker.worker_base import WorkerWrapperBase

56 logger = init_logger(__name__)

59 +class FutureWrapper(Future):

60 + def __init__(self, futures_queue: deque[tuple["FutureWrapper", Callable]]):

61 + self.futures_queue = futures_queue

62 + super().__init__()

63 +

64 + def result(self, timeout=None):

65 + if timeout is not None:

66 + raise RuntimeError("timeout not implemented")

67 + # Drain any futures ahead of us in the queue.

68 + while not self.done():

69 + future, get_response = self.futures_queue.pop()

70 + future.wait_for_response(get_response)

71 + return super().result()

72 +

73 + def wait_for_response(self, get_response: Callable):

74 + try:

75 + response = get_response()

76 + with suppress(InvalidStateError):

77 + self.set_result(response)

78 + except Exception as e:

79 + with suppress(InvalidStateError):

80 + self.set_exception(e)

81 +

82 +

83 class MultiprocExecutor(Executor):

84 supports_pp: bool = True

@@ -64,7 +90,6 @@ class MultiprocExecutor(Executor):

90 self.is_failed = False

91 self.shutdown_event = threading.Event()

92 self.failure_callback: FailureCallback | None = None

67 - self.io_thread_pool: ThreadPoolExecutor | None = None

94 self.world_size = self.parallel_config.world_size

95 tensor_parallel_size = self.parallel_config.tensor_parallel_size

@@ -132,12 +157,7 @@ class MultiprocExecutor(Executor):

157 uw.death_writer.close()

158 self._ensure_worker_termination([uw.proc for uw in unready_workers])

159

135 - # Note: must use only 1 IO thread to keep dequeue sequence

136 - # from the response queue.

137 - # _async_aggregate_workers_output also assumes a single IO thread.

138 - self.io_thread_pool = ThreadPoolExecutor(

139 - max_workers=1, thread_name_prefix="mp_exec_io"

140 - )

160 + self.futures_queue = deque[tuple[FutureWrapper, Callable]]()

161

162 self.output_rank = self._get_output_rank()

163 self.has_connector = self.vllm_config.kv_transfer_config is not None

@@ -195,14 +215,13 @@ class MultiprocExecutor(Executor):

215 ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:

216 if not self.has_connector:

217 # get output only from a single worker (output_rank)

198 - (output,) = self.collective_rpc(

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 )

205 - return output

225

226 # get output from all workers

227 outputs = self.collective_rpc(

@@ -223,12 +242,11 @@ class MultiprocExecutor(Executor):

242

243 def take_draft_token_ids(self) -> DraftTokenIds | None:

244 # OPTIMIZATION: Get output only from a single worker (output_rank)

226 - outputs = self.collective_rpc(

245 + return self.collective_rpc(

246 "take_draft_token_ids", unique_reply_rank=self.output_rank

247 )

229 - return outputs[0]

248

231 - def collective_rpc(

249 + def collective_rpc( # type: ignore[override]

250 self,

251 method: str | Callable,

252 timeout: float | None = None,

@@ -236,7 +254,9 @@ class MultiprocExecutor(Executor):

254 kwargs: dict | None = None,

255 non_block: bool = False,

256 unique_reply_rank: int | None = None,

239 - ) -> list[Any]:

257 + ) -> Any | list[Any] | Future[Any | list[Any]]:

258 + """Returns single result if unique_reply_rank is provided, otherwise list."""

259 +

260 if self.is_failed:

261 raise RuntimeError("Executor failed.")

262

@@ -246,63 +266,52 @@ class MultiprocExecutor(Executor):

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.

249 - try:

250 - if isinstance(method, str):

251 - send_method = method

252 - else:

253 - send_method = cloudpickle.dumps(

254 - method, protocol=pickle.HIGHEST_PROTOCOL

255 - )

256 - self.rpc_broadcast_mq.enqueue(

257 - (send_method, args, kwargs, unique_reply_rank)

258 - )

269

260 - workers = (

261 - (self.workers[unique_reply_rank],)

262 - if unique_reply_rank is not None

263 - else self.workers

264 - )

265 - responses = []

270 + if isinstance(method, str):

271 + send_method = method

272 + else:

273 + send_method = cloudpickle.dumps(method, protocol=pickle.HIGHEST_PROTOCOL)

274 + self.rpc_broadcast_mq.enqueue((send_method, args, kwargs, unique_reply_rank))

275

267 - def get_response(

268 - w: WorkerProcHandle,

269 - dequeue_timeout: float | None = None,

270 - cancel_event: threading.Event | None = None,

271 - ):

272 - status, result = w.worker_response_mq.dequeue(

273 - timeout=dequeue_timeout, cancel=cancel_event

274 - )

276 + workers = (

277 + (self.workers[unique_reply_rank],)

278 + if unique_reply_rank is not None

279 + else self.workers

280 + )

281

276 - if status != WorkerProc.ResponseStatus.SUCCESS:

277 - raise RuntimeError(

278 - f"Worker failed with error '{result}', please check the"

279 - " stack trace above for the root cause"

280 - )

281 - return result

282 + shutdown_event = self.shutdown_event

283

284 + def get_response():

285 + responses = []

286 for w in workers:

287 dequeue_timeout = (

288 None if deadline is None else (deadline - time.monotonic())

289 )

287 -

288 - if self.io_thread_pool is not None:

289 - # We must consume worker_response_mq from a single thread.

290 - result = self.io_thread_pool.submit( # type: ignore

291 - get_response, w, dequeue_timeout, self.shutdown_event

290 + try:

291 + status, result = w.worker_response_mq.dequeue(

292 + timeout=dequeue_timeout, cancel=shutdown_event

293 )

293 - if not non_block:

294 - result = result.result()

295 - elif not non_block:

296 - result = get_response(w, dequeue_timeout, self.shutdown_event)

297 - else:

294 + except TimeoutError as e:

295 + raise TimeoutError(f"RPC call to {method} timed out.") from e

296 + if status != WorkerProc.ResponseStatus.SUCCESS:

297 raise RuntimeError(

299 - "non_block can only be used when max_concurrent_batches > 1"

298 + f"Worker failed with error '{result}', please check the"

299 + " stack trace above for the root cause"

300 )

301 responses.append(result)

302 + return responses[0] if unique_reply_rank is not None else responses

303 +

304 + if non_block:

305 + future = FutureWrapper(self.futures_queue)

306 + self.futures_queue.appendleft((future, get_response))

307 + return future

308 +

309 + # First drain any pending futures in the queue.

310 + while self.futures_queue:

311 + future, get_fut_response = self.futures_queue.pop()

312 + future.wait_for_response(get_fut_response)

313

303 - return responses

304 - except TimeoutError as e:

305 - raise TimeoutError(f"RPC call to {method} timed out.") from e

314 + return get_response()

315

316 @staticmethod

317 def _ensure_worker_termination(worker_procs: list[BaseProcess]):

@@ -348,9 +357,6 @@ class MultiprocExecutor(Executor):

357 self._ensure_worker_termination([w.proc for w in workers])

358

359 self.shutdown_event.set()

351 - if self.io_thread_pool is not None:

352 - self.io_thread_pool.shutdown(wait=False, cancel_futures=True)

353 - del self.io_thread_pool

360

361 self.rpc_broadcast_mq = None

362

vllm/v1/executor/ray_executor.py

+5 -6

@@ -435,26 +435,25 @@ class RayDistributedExecutor(Executor):

435

436 # When PP is used, we return a FutureWrapper immediately so that

437 # the scheduler can yield to the next batch.

438 - return FutureWrapper(refs)

438 + return FutureWrapper(refs[0])

439

440 # Get output from all workers when connector is present

441 assert self.kv_output_aggregator is not None

442 if not non_block:

443 # Block and get results from all workers

444 - outputs = [ref.get() for ref in refs]

445 - return self.kv_output_aggregator.aggregate(outputs)

444 + return self.kv_output_aggregator.aggregate(ray.get(refs))

445

446 # Return a future that will aggregate outputs from all workers

447 return FutureWrapper(refs, self.kv_output_aggregator)

448

450 - def collective_rpc(

449 + def collective_rpc( # type: ignore[override]

450 self,

451 method: str | Callable,

452 timeout: float | None = None,

453 args: tuple = (),

454 kwargs: dict[str, Any] | None = None,

455 non_block: bool = False,

457 - ) -> list[Any]:

456 + ) -> list[Any] | Future[list[Any]]:

457 """Runs the given method on all workers."""

458 sent_method = method if isinstance(method, str) else cloudpickle.dumps(method)

459 del method

@@ -470,7 +469,7 @@ class RayDistributedExecutor(Executor):

469

470 # Get the results of the ray workers.

471 if non_block:

473 - return [FutureWrapper((output,)) for output in ray_worker_outputs]

472 + return FutureWrapper(ray_worker_outputs)

473

474 return ray.get(ray_worker_outputs, timeout=timeout)

475

vllm/v1/executor/ray_utils.py

+4 -7

@@ -141,19 +141,16 @@ class FutureWrapper(Future):

141 the result() call. If not only the first worker's output is returned.

142 """

143

144 - def __init__(self, refs, aggregator: KVOutputAggregator | None = None):

144 + def __init__(self, ref_or_refs, aggregator: KVOutputAggregator | None = None):

145 super().__init__()

146 - self.refs = refs

146 + self.ref_or_refs = ref_or_refs

147 self.aggregator = aggregator

148

149 def result(self, timeout=None):

150 - if timeout is not None:

151 - raise NotImplementedError("timeout is not supported")

152 -

150 + outputs = ray.get(self.ref_or_refs, timeout=timeout)

151 if self.aggregator is None:

154 - return self.refs[0].get()

152 + return outputs

153

156 - outputs = [ref.get() for ref in self.refs]

154 return self.aggregator.aggregate(outputs, output_rank=0)

155

156

vllm/v1/executor/uniproc_executor.py

+36 -7

@@ -13,9 +13,10 @@ import torch.distributed as dist

13 import vllm.envs as envs

14 from vllm.logger import init_logger

15 from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port

16 +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

17 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

18 from vllm.v1.executor.abstract import Executor

18 -from vllm.v1.outputs import AsyncModelRunnerOutput

19 +from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput

20 from vllm.v1.serial_utils import run_method

21 from vllm.v1.worker.worker_base import WorkerWrapperBase

@@ -58,32 +59,60 @@ class UniProcExecutor(Executor):

59 def max_concurrent_batches(self) -> int:

60 return 2 if self.scheduler_config.async_scheduling else 1

61 - def collective_rpc(

62 + def collective_rpc( # type: ignore[override]

63 self,

64 method: str | Callable,

65 timeout: float | None = None,

66 args: tuple = (),

67 kwargs: dict | None = None,

68 non_block: bool = False,

68 - ) -> list[Any]:

69 + single_value: bool = False,

70 + ) -> Any | list[Any] | Future[Any | list[Any]]:

71 if kwargs is None:

72 kwargs = {}

74 if not non_block:

73 - return [run_method(self.driver_worker, method, args, kwargs)]

75 + result = run_method(self.driver_worker, method, args, kwargs)

76 + return result if single_value else [result]

78 try:

79 result = run_method(self.driver_worker, method, args, kwargs)

80 if isinstance(result, AsyncModelRunnerOutput):

81 if (async_thread := self.async_output_thread) is not None:

79 - return [async_thread.submit(result.get_output)]

82 + get_output = result.get_output

83 + if not single_value:

84 + get_output = lambda go=result.get_output: [go()]

85 + return async_thread.submit(get_output)

86 result = result.get_output()

87 future = Future[Any]()

82 - future.set_result(result)

88 + future.set_result(result if single_value else [result])

89 except Exception as e:

90 future = Future[Any]()

91 future.set_exception(e)

86 - return [future]

92 + return future

93 +

94 + def execute_model( # type: ignore[override]

95 + self, scheduler_output: SchedulerOutput, non_block: bool = False

96 + ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:

97 + return self.collective_rpc(

98 + "execute_model",

99 + args=(scheduler_output,),

100 + non_block=non_block,

101 + single_value=True,

102 + )

103 +

104 + def sample_tokens( # type: ignore[override]

105 + self, grammar_output: GrammarOutput | None, non_block: bool = False

106 + ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:

107 + return self.collective_rpc(

108 + "sample_tokens",

109 + args=(grammar_output,),

110 + non_block=non_block,

111 + single_value=True,

112 + )

113 +

114 + def take_draft_token_ids(self) -> DraftTokenIds | None:

115 + return self.collective_rpc("take_draft_token_ids", single_value=True)

116

117 def check_health(self) -> None:

118 # UniProcExecutor will always be healthy as long as

vllm/v1/worker/gpu_worker.py

+1 -1

@@ -524,7 +524,7 @@ class Worker(WorkerBase):

524

525 @torch.inference_mode()

526 def sample_tokens(

527 - self, grammar_output: "GrammarOutput"

527 + self, grammar_output: "GrammarOutput | None"

528 ) -> ModelRunnerOutput | AsyncModelRunnerOutput:

529 return self.model_runner.sample_tokens(grammar_output)

530