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

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

[Core] Rework handling of async scheduling config (#28250)

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

Nick Hill committed Nov 7, 2025 at 12:01 UTC da786e339e985cbf32f43474e7b88b8da65c76e5

6 files changed +121 -71

tests/v1/engine/test_engine_core.py

+17 -15

@@ -66,7 +66,7 @@ def test_engine_core():

66 assert len(engine_core.scheduler.waiting) == 1

67 assert len(engine_core.scheduler.running) == 0

69 - _ = engine_core.step()

69 + _ = engine_core.step_fn()

70 assert len(engine_core.scheduler.waiting) == 0

71 assert len(engine_core.scheduler.running) == 1

@@ -75,7 +75,7 @@ def test_engine_core():

75 assert len(engine_core.scheduler.waiting) == 1

76 assert len(engine_core.scheduler.running) == 1

78 - _ = engine_core.step()

78 + _ = engine_core.step_fn()

79 assert len(engine_core.scheduler.waiting) == 0

80 assert len(engine_core.scheduler.running) == 2

@@ -85,12 +85,12 @@ def test_engine_core():

85 assert len(engine_core.scheduler.waiting) == 2

86 assert len(engine_core.scheduler.running) == 2

88 - _ = engine_core.step()

88 + _ = engine_core.step_fn()

89 assert len(engine_core.scheduler.waiting) == 0

90 assert len(engine_core.scheduler.running) == 4

92 # Loop through until they are all done.

93 - while (outs := engine_core.step()[0].get(0)) and outs.outputs:

93 + while (outs := engine_core.step_fn()[0].get(0)) and outs.outputs:

94 pass

96 assert len(engine_core.scheduler.waiting) == 0

@@ -107,7 +107,7 @@ def test_engine_core():

107 assert engine_core.scheduler.has_unfinished_requests()

108 assert not engine_core.scheduler.has_finished_requests()

109

110 - _ = engine_core.step()

110 + _ = engine_core.step_fn()

111 assert len(engine_core.scheduler.waiting) == 0

112 assert len(engine_core.scheduler.running) == 1

113 assert engine_core.scheduler.has_unfinished_requests()

@@ -119,7 +119,7 @@ def test_engine_core():

119 assert not engine_core.scheduler.has_unfinished_requests()

120 assert engine_core.scheduler.has_finished_requests()

121

122 - _ = engine_core.step()

122 + _ = engine_core.step_fn()

123 assert not engine_core.scheduler.has_unfinished_requests()

124 assert not engine_core.scheduler.has_finished_requests()

125

@@ -133,7 +133,7 @@ def test_engine_core():

133 assert len(engine_core.scheduler.waiting) == 2

134 assert len(engine_core.scheduler.running) == 0

135

136 - _ = engine_core.step()

136 + _ = engine_core.step_fn()

137 assert len(engine_core.scheduler.waiting) == 0

138 assert len(engine_core.scheduler.running) == 2

139

@@ -141,7 +141,7 @@ def test_engine_core():

141 assert len(engine_core.scheduler.waiting) == 1

142 assert len(engine_core.scheduler.running) == 2

143

144 - _ = engine_core.step()

144 + _ = engine_core.step_fn()

145 assert len(engine_core.scheduler.waiting) == 0

146 assert len(engine_core.scheduler.running) == 3

147

@@ -150,7 +150,7 @@ def test_engine_core():

150 assert len(engine_core.scheduler.waiting) == 0

151 assert len(engine_core.scheduler.running) == 2

152

153 - _ = engine_core.step()

153 + _ = engine_core.step_fn()

154 assert len(engine_core.scheduler.waiting) == 0

155 assert len(engine_core.scheduler.running) == 2

156

@@ -165,12 +165,12 @@ def test_engine_core():

165 req0.request_id = req1.request_id = "test"

166 engine_core.add_request(*engine_core.preprocess_add_request(req0))

167

168 - while (outs := engine_core.step()[0].get(0)) and outs.outputs:

169 - pass

168 + while engine_core.scheduler.has_requests():

169 + engine_core.step_fn()

170

171 engine_core.add_request(*engine_core.preprocess_add_request(req1))

172 - while (outs := engine_core.step()[0].get(0)) and outs.outputs:

173 - pass

172 + while engine_core.scheduler.has_requests():

173 + engine_core.step_fn()

174

175 assert len(engine_core.scheduler.waiting) == 0

176 assert len(engine_core.scheduler.running) == 0

@@ -208,8 +208,8 @@ def test_engine_core_advanced_sampling():

208 assert len(engine_core.scheduler.waiting) == 1

209 assert len(engine_core.scheduler.running) == 0

210 # Loop through until they are all done.

211 - while (outs := engine_core.step()[0].get(0)) and outs.outputs:

212 - pass

211 + while engine_core.scheduler.has_requests():

212 + engine_core.step_fn()

213 assert len(engine_core.scheduler.waiting) == 0

214 assert len(engine_core.scheduler.running) == 0

215

@@ -297,6 +297,8 @@ def test_engine_core_concurrent_batches():

297 max_num_batched_tokens=10,

298 # Reduce startup time.

299 enforce_eager=True,

300 + # Test concurrent batch behaviour independently of async scheduling.

301 + async_scheduling=False,

302 )

303 vllm_config = engine_args.create_engine_config()

304 with set_default_torch_num_threads(1):

vllm/config/scheduler.py

+34 -9

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

4 import hashlib

5 from collections.abc import Callable

6 from dataclasses import InitVar

7 -from typing import Any, Literal

7 +from typing import TYPE_CHECKING, Any, Literal, cast

9 from pydantic import Field, field_validator, model_validator

10 from pydantic.dataclasses import dataclass

@@ -17,6 +17,10 @@ from vllm.utils import (

17 MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,

18 POOLING_MODEL_MAX_NUM_BATCHED_TOKENS,

19 )

20 +from vllm.utils.import_utils import resolve_obj_by_qualname

21 +

22 +if TYPE_CHECKING:

23 + from vllm.v1.core.sched.interface import SchedulerInterface

25 logger = init_logger(__name__)

@@ -120,7 +124,7 @@ class SchedulerConfig:

124

125 # scheduler class or path. "vllm.v1.core.sched.scheduler.Scheduler"

126 # (default) or "mod.custom_class".

123 - scheduler_cls: str | type[object] = "vllm.v1.core.sched.scheduler.Scheduler"

127 + scheduler_cls: str | type[object] = Field(default=None)

128 """The scheduler class to use. "vllm.v1.core.sched.scheduler.Scheduler" is

129 the default scheduler. Can be a class directly or the path to a class of

130 form "mod.custom_class"."""

@@ -132,12 +136,34 @@ class SchedulerConfig:

136 """

137

138 async_scheduling: bool = False

135 - """EXPERIMENTAL: If set to True, perform async scheduling. This may help

136 - reduce the CPU overheads, leading to better latency and throughput. However,

137 - async scheduling is currently not supported with some features such as

138 - structured outputs, speculative decoding, and pipeline parallelism.

139 + """If set to True, perform async scheduling. This helps to avoid gaps in

140 + GPU utilization, leading to better latency and throughput.

141 + Async scheduling is currently not supported with some features such as

142 + speculative decoding and pipeline parallelism.

143 """

144

145 + def get_scheduler_cls(self) -> type["SchedulerInterface"]:

146 + if self.scheduler_cls is None:

147 + if self.async_scheduling:

148 + from vllm.v1.core.sched.async_scheduler import AsyncScheduler

149 +

150 + return AsyncScheduler

151 + from vllm.v1.core.sched.scheduler import Scheduler

152 +

153 + return Scheduler

154 +

155 + # This warning can be removed once the Scheduler interface is

156 + # finalized and we can maintain support for scheduler classes that

157 + # implement it

158 + logger.warning_once(

159 + "Using custom scheduler class %s. This scheduler interface is "

160 + "not public and compatibility may not be maintained.",

161 + self.scheduler_cls,

162 + )

163 + if not isinstance(self.scheduler_cls, str):

164 + return cast(type["SchedulerInterface"], self.scheduler_cls)

165 + return resolve_obj_by_qualname(self.scheduler_cls)

166 +

167 def compute_hash(self) -> str:

168 """

169 WARNING: Whenever a new field is added to this config,

@@ -161,6 +187,8 @@ class SchedulerConfig:

187 "max_num_seqs",

188 "max_model_len",

189 "enable_chunked_prefill",

190 + "scheduler_cls",

191 + "async_scheduling",

192 mode="wrap",

193 )

194 @classmethod

@@ -242,9 +270,6 @@ class SchedulerConfig:

270 self.long_prefill_token_threshold,

271 )

272

245 - if self.async_scheduling:

246 - self.scheduler_cls = "vllm.v1.core.sched.async_scheduler.AsyncScheduler"

247 -

273 @model_validator(mode="after")

274 def _verify_args(self) -> Self:

275 if (

vllm/config/vllm.py

+49 -2

@@ -353,6 +353,53 @@ class VllmConfig:

353 self.model_config, self.load_config

354 )

355

356 + executor_backend = self.parallel_config.distributed_executor_backend

357 + executor_supports_async_sched = executor_backend in (

358 + "mp",

359 + "uni",

360 + "external_launcher",

361 + )

362 +

363 + if self.scheduler_config.async_scheduling:

364 + # Async scheduling explicitly enabled, hard fail any incompatibilities.

365 + if self.parallel_config.pipeline_parallel_size > 1:

366 + raise ValueError(

367 + "Async scheduling is not yet compatible with "

368 + "pipeline_parallel_size > 1."

369 + )

370 + if self.speculative_config is not None:

371 + raise ValueError(

372 + "Async scheduling is not yet compatible with speculative decoding."

373 + )

374 + if not executor_supports_async_sched:

375 + raise ValueError(

376 + "Currently, async scheduling only supports `mp`, `uni`, or "

377 + "`external_launcher` distributed executor backend, but you chose "

378 + f"`{executor_backend}`."

379 + )

380 + elif self.scheduler_config.async_scheduling is None:

381 + # Enable async scheduling unless there is an incompatible option.

382 + # NOTE: we won't reach here until async scheduling is enabled by default.

383 + if (

384 + self.parallel_config.pipeline_parallel_size > 1

385 + or self.speculative_config is not None

386 + ):

387 + logger.warning(

388 + "Async scheduling is not yet supported with speculative decoding "

389 + " or pipeline_parallel_size > 1 and will be disabled."

390 + )

391 + self.scheduler_config.async_scheduling = False

392 + elif not executor_supports_async_sched:

393 + logger.warning(

394 + "Async scheduling will be disabled because it is not supported "

395 + "with the `%s` distributed executor backend (only `mp`, `uni`, and "

396 + "`external_launcher` are supported).",

397 + executor_backend,

398 + )

399 + self.scheduler_config.async_scheduling = False

400 + else:

401 + self.scheduler_config.async_scheduling = True

402 +

403 from vllm.platforms import current_platform

404

405 if (

@@ -467,7 +514,7 @@ class VllmConfig:

514 self.speculative_config is not None

515 and self.speculative_config.use_eagle()

516 ):

470 - raise NotImplementedError(

517 + raise ValueError(

518 "Fast prefill optimization for KV sharing is not "

519 "compatible with EAGLE as EAGLE requires correct logits "

520 "for all tokens while fast prefill gives incorrect logits "

@@ -491,7 +538,7 @@ class VllmConfig:

538 )

539 if not getattr(self.model_config.hf_config, "is_causal", True):

540 disable_chunked_prefill_reasons.append(

494 - "Only models using causal attention supports chunked "

541 + "Only models using causal attention support chunked "

542 "prefill and prefix caching; disabling both."

543 )

544 elif self.model_config.is_encoder_decoder:

vllm/engine/arg_utils.py

+2 -26

@@ -513,7 +513,7 @@ class EngineArgs:

513 ObservabilityConfig.collect_detailed_traces

514 )

515 scheduling_policy: SchedulerPolicy = SchedulerConfig.policy

516 - scheduler_cls: str | type[object] = SchedulerConfig.scheduler_cls

516 + scheduler_cls: str | type[object] | None = SchedulerConfig.scheduler_cls

517

518 pooler_config: PoolerConfig | None = ModelConfig.pooler_config

519 override_pooler_config: dict | PoolerConfig | None = (

@@ -552,7 +552,7 @@ class EngineArgs:

552 )

553 """Custom logitproc types"""

554

555 - async_scheduling: bool = SchedulerConfig.async_scheduling

555 + async_scheduling: bool | None = SchedulerConfig.async_scheduling

556

557 kv_sharing_fast_prefill: bool = CacheConfig.kv_sharing_fast_prefill

558

@@ -1479,20 +1479,6 @@ class EngineArgs:

1479 else ParallelConfig.data_parallel_rpc_port

1480 )

1481

1482 - if self.async_scheduling:

1483 - if self.pipeline_parallel_size > 1:

1484 - raise ValueError(

1485 - "Async scheduling is not supported with pipeline-parallel-size > 1."

1486 - )

1487 -

1488 - # Currently, async scheduling does not support speculative decoding.

1489 - # TODO(woosuk): Support it.

1490 - if self.speculative_config is not None:

1491 - raise ValueError(

1492 - "Currently, speculative decoding is not supported with "

1493 - "async scheduling."

1494 - )

1495 -

1482 # Forward the deprecated CLI args to the EPLB config.

1483 if self.num_redundant_experts is not None:

1484 self.eplb_config.num_redundant_experts = self.num_redundant_experts

@@ -1536,16 +1522,6 @@ class EngineArgs:

1522 _api_process_rank=self._api_process_rank,

1523 )

1524

1539 - if self.async_scheduling and (

1540 - parallel_config.distributed_executor_backend

1541 - not in ("mp", "uni", "external_launcher")

1542 - ):

1543 - raise ValueError(

1544 - "Currently, async scheduling only supports `mp`, `uni` or "

1545 - "`external_launcher` distributed executor backend, but you choose "

1546 - f"`{parallel_config.distributed_executor_backend}`."

1547 - )

1548 -

1525 speculative_config = self.create_speculative_config(

1526 target_model_config=model_config,

1527 target_parallel_config=parallel_config,

vllm/v1/core/sched/interface.py

+18

@@ -4,16 +4,34 @@ from abc import ABC, abstractmethod

4 from collections.abc import Iterable

5 from typing import TYPE_CHECKING, Optional

7 +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

8 +

9 if TYPE_CHECKING:

10 + from vllm.config import VllmConfig

11 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1

12 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput

13 from vllm.v1.engine import EngineCoreOutputs

14 + from vllm.v1.kv_cache_interface import KVCacheConfig

15 from vllm.v1.metrics.stats import SchedulerStats

16 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

17 from vllm.v1.request import Request, RequestStatus

18 + from vllm.v1.structured_output import StructuredOutputManager

21 class SchedulerInterface(ABC):

22 + @abstractmethod

23 + def __init__(

24 + self,

25 + vllm_config: "VllmConfig",

26 + kv_cache_config: "KVCacheConfig",

27 + structured_output_manager: "StructuredOutputManager",

28 + block_size: int,

29 + mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,

30 + include_finished_set: bool = False,

31 + log_stats: bool = False,

32 + ) -> None:

33 + raise NotImplementedError

34 +

35 @abstractmethod

36 def schedule(self) -> "SchedulerOutput":

37 """Schedule the requests to process in this scheduling step.

vllm/v1/engine/core.py

+1 -19

@@ -29,7 +29,6 @@ from vllm.tasks import POOLING_TASKS, SupportedTask

29 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value

30 from vllm.utils.gc_utils import maybe_attach_gc_debug_callback

31 from vllm.utils.hashing import get_hash_fn_by_name

32 -from vllm.utils.import_utils import resolve_obj_by_qualname

32 from vllm.utils.network_utils import make_zmq_socket

33 from vllm.utils.system_utils import decorate_logs, set_process_title

34 from vllm.v1.core.kv_cache_utils import (

@@ -41,7 +40,6 @@ from vllm.v1.core.kv_cache_utils import (

40 )

41 from vllm.v1.core.sched.interface import SchedulerInterface

42 from vllm.v1.core.sched.output import SchedulerOutput

44 -from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler

43 from vllm.v1.engine import (

44 EngineCoreOutputs,

45 EngineCoreRequest,

@@ -117,23 +115,7 @@ class EngineCore:

115 self.structured_output_manager = StructuredOutputManager(vllm_config)

116

117 # Setup scheduler.

120 - if isinstance(vllm_config.scheduler_config.scheduler_cls, str):

121 - Scheduler = resolve_obj_by_qualname(

122 - vllm_config.scheduler_config.scheduler_cls

123 - )

124 - else:

125 - Scheduler = vllm_config.scheduler_config.scheduler_cls

126 -

127 - # This warning can be removed once the V1 Scheduler interface is

128 - # finalized and we can maintain support for scheduler classes that

129 - # implement it

130 - if Scheduler is not V1Scheduler:

131 - logger.warning(

132 - "Using configured V1 scheduler class %s. "

133 - "This scheduler interface is not public and "

134 - "compatibility may not be maintained.",

135 - vllm_config.scheduler_config.scheduler_cls,

136 - )

118 + Scheduler = vllm_config.scheduler_config.get_scheduler_cls()

119

120 if len(kv_cache_config.kv_cache_groups) == 0:

121 # Encoder models without KV cache don't support