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 6, 13:15 UTC
[Hardware][TPU] Add supports_async_scheduling() method to Executor interface so that it can be extended for Executor implementations. (#36924)
Signed-off-by: Guangxiang Du <gxd@google.com>
gxd3 committed
Mar 17, 2026 at 21:53 UTC
a0dd1995c76cba4905671522ade7733cfc1ca966
5 files changed
+44
-10
tests/v1/executor/test_executor.py
+23
@@ -14,12 +14,35 @@ from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
14
from vllm.sampling_params import SamplingParams
15
from vllm.v1.engine.async_llm import AsyncLLM
16
from vllm.v1.engine.llm_engine import LLMEngine
17
+from vllm.v1.executor.abstract import Executor
18
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
19
+from vllm.v1.executor.uniproc_executor import (
20
+ ExecutorWithExternalLauncher,
21
+ UniProcExecutor,
22
+)
23
24
25
class Mock: ...
26
27
28
+def test_supports_async_scheduling_base_executor():
29
+ assert Executor.supports_async_scheduling() is False
30
+
31
+
32
+def test_supports_async_scheduling_uniproc_executor():
33
+ assert UniProcExecutor.supports_async_scheduling() is True
34
+
35
+
36
+def test_supports_async_scheduling_executor_with_external_launcher():
37
+ # ExecutorWithExternalLauncher inherits from UniProcExecutor and does not
38
+ # override supports_async_scheduling, so it should return True.
39
+ assert ExecutorWithExternalLauncher.supports_async_scheduling() is True
40
+
41
+
42
+def test_supports_async_scheduling_multiproc_executor():
43
+ assert MultiprocExecutor.supports_async_scheduling() is True
44
+
45
+
46
class CustomMultiprocExecutor(MultiprocExecutor):
47
def collective_rpc(
48
self,
vllm/config/vllm.py
+6
-10
@@ -682,12 +682,11 @@ class VllmConfig:
682
self.model_config, self.load_config
683
)
684
685
+ from vllm.v1.executor.abstract import Executor
686
+
687
executor_backend = self.parallel_config.distributed_executor_backend
686
- executor_supports_async_sched = executor_backend in (
687
- "mp",
688
- "uni",
689
- "external_launcher",
690
- )
688
+ executor_class = Executor.get_class(self)
689
+ executor_supports_async_sched = executor_class.supports_async_scheduling()
690
691
if self.scheduler_config.async_scheduling:
692
# Async scheduling explicitly enabled, hard fail any incompatibilities.
@@ -711,9 +710,7 @@ class VllmConfig:
710
)
711
if not executor_supports_async_sched:
712
raise ValueError(
714
- "Currently, async scheduling only supports `mp`, `uni`, or "
715
- "`external_launcher` distributed executor backend, but you chose "
716
- f"`{executor_backend}`."
713
+ f"`{executor_backend}` does not support async scheduling yet."
714
)
715
elif self.scheduler_config.async_scheduling is None:
716
# Enable async scheduling unless there is an incompatible option.
@@ -742,8 +739,7 @@ class VllmConfig:
739
elif not executor_supports_async_sched:
740
logger.warning_once(
741
"Async scheduling will be disabled because it is not supported "
745
- "with the `%s` distributed executor backend (only `mp`, `uni`, and "
746
- "`external_launcher` are supported).",
742
+ "with the `%s` distributed executor backend. ",
743
executor_backend,
744
scope="local",
745
)
vllm/v1/executor/abstract.py
+7
@@ -353,6 +353,13 @@ class Executor(ABC):
353
) -> None:
354
raise NotImplementedError
355
356
+ @classmethod
357
+ def supports_async_scheduling(cls) -> bool:
358
+ """
359
+ Whether the executor supports async scheduling.
360
+ """
361
+ return False
362
+
363
364
from vllm.v1.executor.uniproc_executor import ( # noqa: E402
365
ExecutorWithExternalLauncher as _ExecutorWithExternalLauncher,
vllm/v1/executor/multiproc_executor.py
+4
@@ -487,6 +487,10 @@ class MultiprocExecutor(Executor):
487
* self.parallel_config.prefill_context_parallel_size
488
)
489
490
+ @classmethod
491
+ def supports_async_scheduling(cls) -> bool:
492
+ return True
493
+
494
495
@dataclass
496
class UnreadyWorkerProcHandle:
vllm/v1/executor/uniproc_executor.py
+4
@@ -134,6 +134,10 @@ class UniProcExecutor(Executor):
134
if worker := self.driver_worker:
135
worker.shutdown()
136
137
+ @classmethod
138
+ def supports_async_scheduling(cls) -> bool:
139
+ return True
140
+
141
142
class ExecutorWithExternalLauncher(UniProcExecutor):
143
"""An executor that uses external launchers to launch engines,