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 14, 05:30 UTC
[Core] Support async scheduling with uniproc executor (#24219)
Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Ronald1995 <ronaldautomobile@163.com> Co-authored-by: Ronald1995 <ronaldautomobile@163.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Nick Hill committed
Sep 12, 2025 at 16:34 UTC
4fdd6f5cbf877de7c4de33086fe41bb0ac1d3cf3
9 files changed
+103
-55
tests/basic_correctness/test_basic_correctness.py
+19
-5
@@ -62,6 +62,8 @@ def _fix_prompt_embed_outputs(
62
@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
63
@pytest.mark.parametrize("max_tokens", [5])
64
@pytest.mark.parametrize("enforce_eager", [False])
65
+@pytest.mark.parametrize("async_scheduling", [True, False])
66
+@pytest.mark.parametrize("model_executor", ["uni", "mp"])
67
@pytest.mark.parametrize("enable_prompt_embeds", [True, False])
68
def test_models(
69
monkeypatch: pytest.MonkeyPatch,
@@ -70,6 +72,8 @@ def test_models(
72
backend: str,
73
max_tokens: int,
74
enforce_eager: bool,
75
+ async_scheduling: bool,
76
+ model_executor: str,
77
enable_prompt_embeds: bool,
78
) -> None:
79
@@ -77,6 +81,12 @@ def test_models(
81
"VLLM_USE_V1") and envs.VLLM_USE_V1:
82
pytest.skip("enable_prompt_embeds is not supported in v1.")
83
84
+ if not envs.VLLM_USE_V1:
85
+ if async_scheduling:
86
+ pytest.skip("async_scheduling only supported in v1.")
87
+ if model_executor != "uni":
88
+ pytest.skip("only test uniproc executor for v0.")
89
+
90
if backend == "XFORMERS" and model == "google/gemma-2-2b-it":
91
pytest.skip(
92
f"{backend} does not support gemma2 with full context length.")
@@ -98,11 +108,15 @@ def test_models(
108
prompt_embeds = hf_model.get_prompt_embeddings(
109
example_prompts)
110
101
- with VllmRunner(model,
102
- max_model_len=8192,
103
- enforce_eager=enforce_eager,
104
- enable_prompt_embeds=enable_prompt_embeds,
105
- gpu_memory_utilization=0.7) as vllm_model:
111
+ with VllmRunner(
112
+ model,
113
+ max_model_len=8192,
114
+ enforce_eager=enforce_eager,
115
+ enable_prompt_embeds=enable_prompt_embeds,
116
+ gpu_memory_utilization=0.7,
117
+ async_scheduling=async_scheduling,
118
+ distributed_executor_backend=model_executor,
119
+ ) as vllm_model:
120
if enable_prompt_embeds:
121
vllm_outputs = vllm_model.generate_greedy(
122
prompt_embeds, max_tokens)
tests/v1/engine/test_engine_core.py
+4
@@ -257,9 +257,13 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):
257
def execute_model(
258
self,
259
scheduler_output,
260
+ non_block=False,
261
) -> Future[ModelRunnerOutput]:
262
"""Make execute_model non-blocking."""
263
264
+ # DummyExecutor used only for testing async case.
265
+ assert non_block
266
+
267
def _execute():
268
output = self.collective_rpc("execute_model",
269
args=(scheduler_output, ))
vllm/engine/arg_utils.py
+2
-5
@@ -1296,11 +1296,8 @@ class EngineArgs:
1296
# Async scheduling does not work with the uniprocess backend.
1297
if self.distributed_executor_backend is None:
1298
self.distributed_executor_backend = "mp"
1299
- logger.info("Using mp-based distributed executor backend "
1300
- "for async scheduling.")
1301
- if self.distributed_executor_backend == "uni":
1302
- raise ValueError("Async scheduling is not supported with "
1303
- "uni-process backend.")
1299
+ logger.info("Defaulting to mp-based distributed executor "
1300
+ "backend for async scheduling.")
1301
if self.pipeline_parallel_size > 1:
1302
raise ValueError("Async scheduling is not supported with "
1303
"pipeline-parallel-size > 1.")
vllm/executor/uniproc_executor.py
+46
-28
@@ -1,7 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
import os
4
+from concurrent.futures import Future, ThreadPoolExecutor
5
+from functools import cached_property
6
from multiprocessing import Lock
7
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
8
@@ -17,6 +18,7 @@ from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
18
run_method)
19
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
20
from vllm.v1.executor.utils import get_and_update_mm_cache
21
+from vllm.v1.outputs import AsyncModelRunnerOutput
22
from vllm.worker.worker_base import WorkerWrapperBase
23
24
logger = init_logger(__name__)
@@ -31,15 +33,7 @@ class UniProcExecutor(ExecutorBase):
33
"""
34
self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,
35
rpc_rank=0)
34
- distributed_init_method = get_distributed_init_method(
35
- get_ip(), get_open_port())
36
- local_rank = 0
37
- # set local rank as the device index if specified
38
- device_info = self.vllm_config.device_config.device.__str__().split(
39
- ":")
40
- if len(device_info) > 1:
41
- local_rank = int(device_info[1])
42
- rank = 0
36
+ distributed_init_method, rank, local_rank = self._distributed_args()
37
is_driver_worker = True
38
kwargs = dict(
39
vllm_config=self.vllm_config,
@@ -50,21 +44,56 @@ class UniProcExecutor(ExecutorBase):
44
)
45
self.mm_receiver_cache = worker_receiver_cache_from_config(
46
self.vllm_config, MULTIMODAL_REGISTRY, Lock())
47
+
48
+ self.async_output_thread: Optional[ThreadPoolExecutor] = None
49
+ if self.max_concurrent_batches > 1:
50
+ self.async_output_thread = ThreadPoolExecutor(
51
+ max_workers=1, thread_name_prefix="WorkerAsyncOutput")
52
+
53
self.collective_rpc("init_worker", args=([kwargs], ))
54
self.collective_rpc("init_device")
55
self.collective_rpc("load_model")
56
57
+ def _distributed_args(self) -> tuple[str, int, int]:
58
+ """Return (distributed_init_method, rank, local_rank)."""
59
+ distributed_init_method = get_distributed_init_method(
60
+ get_ip(), get_open_port())
61
+ # set local rank as the device index if specified
62
+ device_info = self.vllm_config.device_config.device.__str__().split(
63
+ ":")
64
+ local_rank = int(device_info[1]) if len(device_info) > 1 else 0
65
+ return distributed_init_method, 0, local_rank
66
+
67
+ @cached_property
68
+ def max_concurrent_batches(self) -> int:
69
+ return 2 if self.scheduler_config.async_scheduling else 1
70
+
71
def collective_rpc(self,
72
method: Union[str, Callable],
73
timeout: Optional[float] = None,
74
args: Tuple = (),
61
- kwargs: Optional[Dict] = None) -> List[Any]:
75
+ kwargs: Optional[Dict] = None,
76
+ non_block: bool = False) -> List[Any]:
77
if kwargs is None:
78
kwargs = {}
79
if self.mm_receiver_cache is not None and method == "execute_model":
80
get_and_update_mm_cache(self.mm_receiver_cache, args)
66
- answer = run_method(self.driver_worker, method, args, kwargs)
67
- return [answer]
81
+
82
+ if not non_block:
83
+ return [run_method(self.driver_worker, method, args, kwargs)]
84
+
85
+ try:
86
+ result = run_method(self.driver_worker, method, args, kwargs)
87
+ if isinstance(result, AsyncModelRunnerOutput):
88
+ if (async_thread := self.async_output_thread) is not None:
89
+ return [async_thread.submit(result.get_output)]
90
+ result = result.get_output()
91
+ future = Future[Any]()
92
+ future.set_result(result)
93
+ except Exception as e:
94
+ future = Future[Any]()
95
+ future.set_exception(e)
96
+ return [future]
97
98
def check_health(self) -> None:
99
# UniProcExecutor will always be healthy as long as
@@ -116,8 +145,9 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
145
assert not envs.VLLM_ENABLE_V1_MULTIPROCESSING, \
146
("To get deterministic execution in V1, "
147
"please set VLLM_ENABLE_V1_MULTIPROCESSING=0")
119
- self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,
120
- rpc_rank=0)
148
+ super()._init_executor()
149
+
150
+ def _distributed_args(self) -> tuple[str, int, int]:
151
# engines are launched in torchrun-compatible launchers
152
# so we can use the env:// method.
153
# required env vars:
@@ -128,19 +158,7 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
158
distributed_init_method = "env://"
159
rank = int(os.environ["RANK"])
160
local_rank = int(os.environ["LOCAL_RANK"])
131
- is_driver_worker = True
132
- kwargs = dict(
133
- vllm_config=self.vllm_config,
134
- local_rank=local_rank,
135
- rank=rank,
136
- distributed_init_method=distributed_init_method,
137
- is_driver_worker=is_driver_worker,
138
- )
139
- self.mm_receiver_cache = worker_receiver_cache_from_config(
140
- self.vllm_config, MULTIMODAL_REGISTRY, Lock())
141
- self.collective_rpc("init_worker", args=([kwargs], ))
142
- self.collective_rpc("init_device")
143
- self.collective_rpc("load_model")
161
+ return distributed_init_method, rank, local_rank
162
163
def determine_num_available_blocks(self) -> Tuple[int, int]:
164
"""
vllm/v1/engine/core.py
+5
-4
@@ -159,6 +159,9 @@ class EngineCore:
159
self.request_block_hasher = get_request_block_hasher(
160
block_size, caching_hash_fn)
161
162
+ self.step_fn = (self.step if self.batch_queue is None else
163
+ self.step_with_batch_queue)
164
+
165
def _initialize_kv_caches(
166
self, vllm_config: VllmConfig) -> tuple[int, int, KVCacheConfig]:
167
start = time.time()
@@ -331,7 +334,8 @@ class EngineCore:
334
model_executed = False
335
if self.scheduler.has_requests():
336
scheduler_output = self.scheduler.schedule()
334
- future = self.model_executor.execute_model(scheduler_output)
337
+ future = self.model_executor.execute_model(scheduler_output,
338
+ non_block=True)
339
batch_queue.appendleft(
340
(future, scheduler_output)) # type: ignore[arg-type]
341
@@ -534,9 +538,6 @@ class EngineCoreProc(EngineCore):
538
assert addresses.coordinator_input is not None
539
logger.info("Waiting for READY message from DP Coordinator...")
540
537
- self.step_fn = (self.step if self.batch_queue is None else
538
- self.step_with_batch_queue)
539
-
541
# Mark the startup heap as static so that it's ignored by GC.
542
# Reduces pause times of oldest generation collections.
543
gc.collect()
vllm/v1/engine/core_client.py
+2
-2
@@ -245,8 +245,8 @@ class InprocClient(EngineCoreClient):
245
self.engine_core = EngineCore(*args, **kwargs)
246
247
def get_output(self) -> EngineCoreOutputs:
248
- outputs, _ = self.engine_core.step()
249
- return outputs.get(0) or EngineCoreOutputs()
248
+ outputs, _ = self.engine_core.step_fn()
249
+ return outputs and outputs.get(0) or EngineCoreOutputs()
250
251
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
252
return self.engine_core.get_supported_tasks()
vllm/v1/executor/abstract.py
+14
-3
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from concurrent.futures import Future
5
-from typing import Callable, Optional, Union
5
+from typing import Any, Callable, Optional, Union
6
7
import torch
8
import torch.distributed as dist
@@ -14,6 +14,7 @@ from vllm.executor.uniproc_executor import ( # noqa
14
from vllm.executor.uniproc_executor import ( # noqa
15
UniProcExecutor as UniProcExecutorV0)
16
from vllm.utils import resolve_obj_by_qualname
17
+from vllm.v1.core.sched.output import SchedulerOutput
18
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
19
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
20
@@ -86,12 +87,22 @@ class Executor(ExecutorBase):
87
def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
88
return self.collective_rpc("get_kv_cache_spec")
89
90
+ def collective_rpc(self,
91
+ method: Union[str, Callable],
92
+ timeout: Optional[float] = None,
93
+ args: tuple = (),
94
+ kwargs: Optional[dict] = None,
95
+ non_block: bool = False) -> list[Any]:
96
+ raise NotImplementedError
97
+
98
def execute_model(
99
self,
91
- scheduler_output,
100
+ scheduler_output: SchedulerOutput,
101
+ non_block: bool = False,
102
) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:
103
output = self.collective_rpc("execute_model",
94
- args=(scheduler_output, ))
104
+ args=(scheduler_output, ),
105
+ non_block=non_block)
106
return output[0]
107
108
def execute_dummy_batch(self) -> None:
vllm/v1/executor/multiproc_executor.py
+7
-5
@@ -11,7 +11,7 @@ import weakref
11
from concurrent.futures import Future, ThreadPoolExecutor
12
from dataclasses import dataclass
13
from enum import Enum, auto
14
-from functools import partial
14
+from functools import cached_property, partial
15
from multiprocessing.connection import Connection
16
from multiprocessing.process import BaseProcess
17
from multiprocessing.synchronize import Lock as LockType
@@ -37,6 +37,7 @@ from vllm.multimodal.cache import worker_receiver_cache_from_config
37
from vllm.utils import (decorate_logs, get_distributed_init_method,
38
get_loopback_ip, get_mp_context, get_open_port,
39
set_process_title)
40
+from vllm.v1.core.sched.output import SchedulerOutput
41
from vllm.v1.executor.abstract import Executor, FailureCallback
42
from vllm.v1.executor.utils import get_and_update_mm_cache
43
from vllm.v1.outputs import (AsyncModelRunnerOutput, DraftTokenIds,
@@ -174,9 +175,9 @@ class MultiprocExecutor(Executor):
175
176
def execute_model(
177
self,
177
- scheduler_output,
178
+ scheduler_output: SchedulerOutput,
179
+ non_block: bool = False,
180
) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:
179
- non_block = self.max_concurrent_batches > 1
181
182
if not self.has_connector:
183
# get output only from a single worker (output_rank)
@@ -328,7 +329,7 @@ class MultiprocExecutor(Executor):
329
self.collective_rpc("check_health", timeout=10)
330
return
331
331
- @property
332
+ @cached_property
333
def max_concurrent_batches(self) -> int:
334
if self.scheduler_config.async_scheduling:
335
return 2
@@ -632,7 +633,8 @@ class WorkerProc:
633
result = (WorkerProc.ResponseStatus.FAILURE, str(output))
634
else:
635
result = (WorkerProc.ResponseStatus.SUCCESS, output)
635
- self.worker_response_mq.enqueue(result)
636
+ if (response_mq := self.worker_response_mq) is not None:
637
+ response_mq.enqueue(result)
638
639
def handle_output(self, output: Any):
640
"""Handles output from the worker. If async scheduling is enabled,
vllm/v1/executor/ray_distributed_executor.py
+4
-3
@@ -66,11 +66,13 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
66
def execute_model(
67
self,
68
scheduler_output: SchedulerOutput,
69
+ non_block: bool = False,
70
) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:
71
"""Execute the model on the Ray workers.
72
73
Args:
74
scheduler_output: The scheduler output to execute.
75
+ non_block: If True, the method will return a Future.
76
77
Returns:
78
The model runner output.
@@ -84,7 +86,7 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
86
if not self.has_connector:
87
# Get output only from a single worker (output_rank)
88
# When PP is not used, we block here until the result is available.
87
- if self.max_concurrent_batches == 1:
89
+ if not non_block:
90
return refs[0].get()
91
92
# When PP is used, we return a FutureWrapper immediately so that
@@ -92,7 +94,7 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
94
return FutureWrapper(refs)
95
96
# Get output from all workers when connector is present
95
- if self.max_concurrent_batches == 1:
97
+ if not non_block:
98
# Block and get results from all workers
99
outputs = [ref.get() for ref in refs]
100
return self.kv_output_aggregator.aggregate(outputs)
@@ -106,4 +108,3 @@ class RayDistributedExecutor(RayDistributedExecutorV0, Executor):
108
if reconfig_request.new_data_parallel_rank == \
109
ReconfigureRankType.SHUTDOWN_CURRENT_RANK:
110
self.shutdown()
109
- return
\ No newline at end of file