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 4, 03:45 UTC
[Core] Avoid using extra thread in `UniProcExecutor` (#40891)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
May 7, 2026 at 15:33 UTC
10ebb40d62e024116c1a4473f8e357a3e72761ed
1 file changed
+21
-16
vllm/v1/executor/uniproc_executor.py
+21
-16
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import os
4
from collections.abc import Callable
5
-from concurrent.futures import Future, ThreadPoolExecutor
5
+from concurrent.futures import Future
6
from functools import cached_property
7
from multiprocessing import Lock
8
from typing import Any
@@ -23,6 +23,25 @@ from vllm.v1.worker.worker_base import WorkerWrapperBase
23
logger = init_logger(__name__)
24
25
26
+class AsyncOutputFuture(Future):
27
+ def __init__(self, async_output: AsyncModelRunnerOutput, single_value: bool):
28
+ self.async_output = async_output
29
+ self.single_value = single_value
30
+ super().__init__()
31
+
32
+ def result(self, timeout=None):
33
+ if timeout is not None:
34
+ raise RuntimeError("timeout not implemented")
35
+
36
+ if not super().done():
37
+ try:
38
+ output = self.async_output.get_output()
39
+ self.set_result(output if self.single_value else [output])
40
+ except Exception as e:
41
+ self.set_exception(e)
42
+ return super().result()
43
+
44
+
45
class UniProcExecutor(Executor):
46
def _init_executor(self) -> None:
47
"""Initialize the worker and load the model."""
@@ -37,12 +56,6 @@ class UniProcExecutor(Executor):
56
shared_worker_lock=Lock(),
57
)
58
40
- self.async_output_thread: ThreadPoolExecutor | None = None
41
- if self.max_concurrent_batches > 1:
42
- self.async_output_thread = ThreadPoolExecutor(
43
- max_workers=1, thread_name_prefix="WorkerAsyncOutput"
44
- )
45
-
59
self.driver_worker.init_worker(all_kwargs=[kwargs])
60
self.driver_worker.init_device()
61
@@ -83,15 +96,7 @@ class UniProcExecutor(Executor):
96
try:
97
result = run_method(self.driver_worker, method, args, kwargs)
98
if isinstance(result, AsyncModelRunnerOutput):
86
- if (async_thread := self.async_output_thread) is not None:
87
- if single_value:
88
- return async_thread.submit(result.get_output)
89
-
90
- def get_output_list() -> list[Any]:
91
- return [result.get_output()]
92
-
93
- return async_thread.submit(get_output_list)
94
- result = result.get_output()
99
+ return AsyncOutputFuture(result, single_value)
100
future = Future[Any]()
101
future.set_result(result if single_value else [result])
102
except Exception as e: