@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 68a72a5cc1

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

Revert "[PerfFix] Avoid separate thread for MP executor shm spin (#28012)" (#28289)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Nov 7, 2025 at 16:07 UTC 68a72a5cc1e29198730d1b2471e23675d9b964dd

9 files changed +131 -143

tests/v1/executor/test_executor.py

+1 -2

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

4 import asyncio

5 import os

6 from collections.abc import Callable

7 -from concurrent.futures import Future

7 from typing import Any

9 import pytest

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

27 kwargs: dict | None = None,

28 non_block: bool = False,

29 unique_reply_rank: int | None = None,

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

30 + ) -> list[Any]:

31 # Drop marker to show that this was run

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

33 ...

tests/v1/kv_connector/unit/test_output_aggregator.py

+20 -12

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

89 def test_async_aggregate_workers_output():

90 aggregator = KVOutputAggregator(expected_finished_count=2)

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

93 - result_future = aggregator.async_aggregate(future)

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

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

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

96 output1 = DummyModelRunnerOutput()

97 output2 = DummyModelRunnerOutput()

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

98 + future1.set_result(output1)

99 + future2.set_result(output2)

100

101 assert result_future.done()

102 aggregated = result_future.result()

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

106 assert aggregated.finished_recving is None

107 assert not aggregated.invalid_block_ids

108

107 - future = Future()

108 - result_future = aggregator.async_aggregate(future)

109 + future1 = Future()

110 + future2 = Future()

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

112

113 output1 = DummyModelRunnerOutput(

114 finished_sending={"req1"}, finished_recving={"req2"}

115 )

116 output2 = DummyModelRunnerOutput(invalid_block_ids={1})

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

117 + future1.set_result(output1)

118 + future2.set_result(output2)

119

120 assert result_future.done()

121 aggregated = result_future.result()

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

125 assert aggregated.finished_recving is None

126 assert aggregated.invalid_block_ids == {1}

127

124 - future = Future()

125 - result_future = aggregator.async_aggregate(future)

128 + future1 = Future()

129 + future2 = Future()

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

131

132 output1 = DummyModelRunnerOutput(invalid_block_ids={2})

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

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

134 + future1.set_result(output1)

135 + future2.set_result(output2)

136

137 assert result_future.done()

138 aggregated = result_future.result()

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

142 assert aggregated.finished_recving is None

143 assert aggregated.invalid_block_ids == {2}

144

139 - future = Future()

140 - result_future = aggregator.async_aggregate(future)

145 + future1 = Future()

146 + future2 = Future()

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

148

149 output1 = DummyModelRunnerOutput(invalid_block_ids={3, 4})

150 output2 = DummyModelRunnerOutput(

151 finished_recving={"req2"}, invalid_block_ids={4, 5}

152 )

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

153 + future1.set_result(output1)

154 + future2.set_result(output2)

155

156 assert result_future.done()

157 aggregated = result_future.result()

vllm/distributed/kv_transfer/kv_connector/utils.py

+29 -14

@@ -221,24 +221,39 @@ class KVOutputAggregator:

221

222 def async_aggregate(

223 self,

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

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

225 output_rank: int = 0,

226 ) -> Future[ModelRunnerOutput | None]:

227 - """Takes a future that resolves to a list of outputs and returns a future

228 - which resolves to a single aggregated output."""

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

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

229 result_future: Future[ModelRunnerOutput | None] = Future()

230

231 - def callback(fut):

232 - if result_future.done():

233 - return

234 - try:

235 - result_future.set_result(self.aggregate(fut.result(), output_rank))

236 - except CancelledError:

237 - result_future.cancel()

238 - except Exception as e:

239 - result_future.set_exception(e)

240 -

241 - output_future.add_done_callback(callback)

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

232 + remaining = len(output_futures)

233 +

234 + def make_callback(idx):

235 + def callback(fut):

236 + if result_future.done():

237 + return

238 +

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)

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

251 +

252 + return callback

253 +

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

255 + output_future.add_done_callback(make_callback(i))

256 +

257 return result_future

258

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 - ) -> Future[list[_R]]:

174 + ) -> list[Future[_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 | None | Future[ModelRunnerOutput | None]:

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

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

+61 -67

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

9 import time

10 import traceback

11 import weakref

12 -from collections import deque

12 from collections.abc import Callable

14 -from concurrent.futures import Future, InvalidStateError

15 -from contextlib import suppress

13 +from concurrent.futures import Future, ThreadPoolExecutor

14 from dataclasses import dataclass

15 from enum import Enum, auto

16 from functools import cached_property, partial

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

54 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 -

57 class MultiprocExecutor(Executor):

58 supports_pp: bool = True

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

64 self.is_failed = False

65 self.shutdown_event = threading.Event()

66 self.failure_callback: FailureCallback | None = None

67 + self.io_thread_pool: ThreadPoolExecutor | None = None

69 self.world_size = self.parallel_config.world_size

70 tensor_parallel_size = self.parallel_config.tensor_parallel_size

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

132 uw.death_writer.close()

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

134

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

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

141

142 self.output_rank = self._get_output_rank()

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

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

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

196 if not self.has_connector:

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

218 - return self.collective_rpc(

198 + (output,) = self.collective_rpc(

199 method,

200 args=args,

201 unique_reply_rank=self.output_rank,

202 non_block=non_block,

203 timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,

204 )

205 + return output

206

207 # get output from all workers

208 outputs = self.collective_rpc(

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

223

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

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

245 - return self.collective_rpc(

226 + outputs = self.collective_rpc(

227 "take_draft_token_ids", unique_reply_rank=self.output_rank

228 )

229 + return outputs[0]

230

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

231 + def collective_rpc(

232 self,

233 method: str | Callable,

234 timeout: float | None = None,

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

236 kwargs: dict | None = None,

237 non_block: bool = False,

238 unique_reply_rank: int | None = None,

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

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

259 -

239 + ) -> list[Any]:

240 if self.is_failed:

241 raise RuntimeError("Executor failed.")

242

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

246 # NOTE: If the args are heterogeneous, then we pack them into a list,

247 # and unpack them in the method of every worker, because every worker

248 # 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 + )

259

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

260 + workers = (

261 + (self.workers[unique_reply_rank],)

262 + if unique_reply_rank is not None

263 + else self.workers

264 + )

265 + responses = []

266

276 - workers = (

277 - (self.workers[unique_reply_rank],)

278 - if unique_reply_rank is not None

279 - else self.workers

280 - )

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

275

282 - shutdown_event = self.shutdown_event

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

284 - def get_response():

285 - responses = []

283 for w in workers:

284 dequeue_timeout = (

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

286 )

290 - try:

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

292 - timeout=dequeue_timeout, cancel=shutdown_event

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

292 )

294 - except TimeoutError as e:

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

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

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:

298 raise RuntimeError(

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

299 - " stack trace above for the root cause"

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

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)

302

314 - return get_response()

303 + return responses

304 + except TimeoutError as e:

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

306

307 @staticmethod

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

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

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

349

350 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

354

355 self.rpc_broadcast_mq = None

356

vllm/v1/executor/ray_executor.py

+6 -5

@@ -435,25 +435,26 @@ 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[0])

438 + return FutureWrapper(refs)

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 - return self.kv_output_aggregator.aggregate(ray.get(refs))

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

445 + return self.kv_output_aggregator.aggregate(outputs)

446

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

448 return FutureWrapper(refs, self.kv_output_aggregator)

449

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

450 + def collective_rpc(

451 self,

452 method: str | Callable,

453 timeout: float | None = None,

454 args: tuple = (),

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

456 non_block: bool = False,

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

457 + ) -> list[Any]:

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

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

460 del method

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

470

471 # Get the results of the ray workers.

472 if non_block:

472 - return FutureWrapper(ray_worker_outputs)

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

474

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

476

vllm/v1/executor/ray_utils.py

+4 -4

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

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

142 """

143

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

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

145 super().__init__()

146 - self.ref_or_refs = ref_or_refs

146 + self.refs = 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

153 - outputs = ray.get(self.ref_or_refs, timeout=timeout)

153 if self.aggregator is None:

155 - return outputs

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

155

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

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

158

159

vllm/v1/executor/uniproc_executor.py

+7 -36

@@ -13,10 +13,9 @@ 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

16 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

17 from vllm.v1.executor.abstract import Executor

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

18 +from vllm.v1.outputs import AsyncModelRunnerOutput

19 from vllm.v1.serial_utils import run_method

20 from vllm.v1.worker.worker_base import WorkerWrapperBase

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

58 def max_concurrent_batches(self) -> int:

59 return 2 if self.scheduler_config.async_scheduling else 1

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

61 + def collective_rpc(

62 self,

63 method: str | Callable,

64 timeout: float | None = None,

65 args: tuple = (),

66 kwargs: dict | None = None,

67 non_block: bool = False,

69 - single_value: bool = False,

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

68 + ) -> list[Any]:

69 if kwargs is None:

70 kwargs = {}

72 if not non_block:

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

76 - return result if single_value else [result]

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

75 try:

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

77 if isinstance(result, AsyncModelRunnerOutput):

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

82 - get_output = result.get_output

83 - if not single_value:

84 - get_output = lambda: [get_output()]

85 - return async_thread.submit(get_output)

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

80 result = result.get_output()

81 future = Future[Any]()

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

82 + future.set_result(result)

83 except Exception as e:

84 future = Future[Any]()

85 future.set_exception(e)

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)

86 + return [future]

88 def check_health(self) -> None:

89 # 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 | None"

527 + self, grammar_output: "GrammarOutput"

528 ) -> ModelRunnerOutput | AsyncModelRunnerOutput:

529 return self.model_runner.sample_tokens(grammar_output)

530