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

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 14, 05:30 UTC

[BugFix] Async scheduling and PP compatibility with DP (#23770)

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

Nick Hill committed Aug 29, 2025 at 08:17 UTC d90d8eb674f3870b8c85515a468108d5f1bd609a

7 files changed +107 -100

tests/v1/engine/test_engine_core.py

+22 -37

@@ -306,17 +306,17 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

306

307 # Schedule Batch 1: (10, req0)

308 assert engine_core.step_with_batch_queue()[0] is None

309 - assert engine_core.batch_queue.qsize() == 1

310 - scheduler_output = engine_core.batch_queue.queue[-1][1]

309 + assert len(engine_core.batch_queue) == 1

310 + scheduler_output = engine_core.batch_queue[-1][1]

311 assert scheduler_output.num_scheduled_tokens["0"] == 10

312 # num_computed_tokens should have been updated immediately.

313 assert engine_core.scheduler.requests[

314 req0.request_id].num_computed_tokens == 10

315

316 # Schedule Batch 2: (2, req0), (8, req1)

317 - assert engine_core.step_with_batch_queue()[0] is None

318 - assert engine_core.batch_queue.qsize() == 2

319 - scheduler_output = engine_core.batch_queue.queue[-1][1]

317 + assert engine_core.step_with_batch_queue()[0] == {}

318 + assert len(engine_core.batch_queue) == 1

319 + scheduler_output = engine_core.batch_queue[-1][1]

320 assert scheduler_output.num_scheduled_tokens["0"] == 2

321 assert scheduler_output.num_scheduled_tokens["1"] == 8

322 # num_computed_tokens should have been updated immediately.

@@ -325,42 +325,32 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

325

326 assert engine_core.scheduler.get_num_unfinished_requests() == 2

327

328 - # Batch queue is full. Finish Batch 1.

329 - engine_core.step_with_batch_queue()

330 -

331 - # Schedule Batch 3: (4, req1). Note that req0 cannot be scheduled

328 + # Finish Batch 1 and schedule Batch 3: (4, req1).

329 + # Note that req0 cannot be scheduled

330 # because it is in the decoding stage now.

331 engine_core.step_with_batch_queue()

334 - assert engine_core.batch_queue.qsize() == 2

335 - scheduler_output = engine_core.batch_queue.queue[-1][1]

332 + assert len(engine_core.batch_queue) == 1

333 + scheduler_output = engine_core.batch_queue[-1][1]

334 assert scheduler_output.num_scheduled_tokens["1"] == 4

335

338 - # Batch queue is full. Finish Batch 2. Get first token of req0.

336 + # Finish Batch 2. Get first token of req0.

337 + # Schedule Batch 4: (1, req0).

338 output = engine_core.step_with_batch_queue()[0].get(0)

339 assert output is not None

340 assert len(output.outputs) == 1

341 assert engine_core.scheduler.requests[req0.request_id].num_tokens == 13

343 -

344 - # Schedule Batch 4: (1, req0).

345 - engine_core.step_with_batch_queue()

346 - assert engine_core.batch_queue.qsize() == 2

347 - scheduler_output = engine_core.batch_queue.queue[-1][1]

342 + scheduler_output = engine_core.batch_queue[-1][1]

343 assert scheduler_output.num_scheduled_tokens["0"] == 1

344

350 - # Batch queue is full. Finish Batch 3. Get first token of req1.

345 + # Finish Batch 3. Get first token of req1. Schedule Batch 5: (1, req1).

346 output = engine_core.step_with_batch_queue()[0].get(0)

347 assert output is not None

348 assert len(output.outputs) == 1

349 assert engine_core.scheduler.requests[req1.request_id].num_tokens == 13

355 -

356 - # Schedule Batch 5: (1, req1).

357 - engine_core.step_with_batch_queue()

358 - assert engine_core.batch_queue.qsize() == 2

359 - scheduler_output = engine_core.batch_queue.queue[-1][1]

350 + scheduler_output = engine_core.batch_queue[-1][1]

351 assert scheduler_output.num_scheduled_tokens["1"] == 1

352

353 # Loop until req0 is finished.

363 - step = 0

354 req_id = 0

355 expected_num_tokens = [

356 engine_core.scheduler.requests["0"].num_tokens + 1,

@@ -368,19 +358,14 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

358 ]

359 while engine_core.scheduler.get_num_unfinished_requests() == 2:

360 output = engine_core.step_with_batch_queue()[0]

371 - if step % 2 == 0:

372 - # Even steps consumes an output.

373 - assert output is not None

374 - assert len(output[0].outputs) == 1

375 - if req_id in engine_core.scheduler.requests:

376 - assert engine_core.scheduler.requests[

377 - req_id].num_tokens == expected_num_tokens[req_id]

378 - expected_num_tokens[req_id] += 1

379 - req_id = (req_id + 1) % 2

380 - else:

381 - # Odd steps schedules a new batch.

382 - assert output is None

383 - step += 1

361 + # Every step consumes an output.

362 + assert output is not None

363 + assert len(output[0].outputs) == 1

364 + if req_id in engine_core.scheduler.requests:

365 + assert engine_core.scheduler.requests[

366 + req_id].num_tokens == expected_num_tokens[req_id]

367 + expected_num_tokens[req_id] += 1

368 + req_id = (req_id + 1) % 2

369

370

371 @multi_gpu_test(num_gpus=2)

tests/v1/test_async_llm_dp.py

+4 -2

@@ -75,9 +75,10 @@ async def generate(

75 ],

76 )

77 @pytest.mark.parametrize("data_parallel_backend", ["mp", "ray"])

78 +@pytest.mark.parametrize("async_scheduling", [True, False])

79 @pytest.mark.asyncio

79 -async def test_load(output_kind: RequestOutputKind,

80 - data_parallel_backend: str):

80 +async def test_load(output_kind: RequestOutputKind, data_parallel_backend: str,

81 + async_scheduling: bool):

83 stats_loggers = {}

@@ -105,6 +106,7 @@ async def test_load(output_kind: RequestOutputKind,

106 prompt = "This is a test of data parallel"

107

108 engine_args.data_parallel_backend = data_parallel_backend

109 + engine_args.async_scheduling = async_scheduling

110 engine = AsyncLLM.from_engine_args(engine_args,

111 stat_loggers=[SimpleStatsLogger])

112 after.callback(engine.shutdown)

vllm/executor/ray_utils.py

@@ -10,6 +10,7 @@ import msgspec

11 import vllm.platforms

12 from vllm.config import ParallelConfig

13 +from vllm.distributed import get_pp_group

14 from vllm.executor.msgspec_utils import decode_hook, encode_hook

15 from vllm.logger import init_logger

16 from vllm.platforms import current_platform

@@ -136,6 +137,11 @@ try:

137 scheduler_output, intermediate_tensors)

138 if isinstance(output, IntermediateTensors):

139 output = scheduler_output, output

140 + elif not get_pp_group().is_last_rank:

141 + # Case where there are no scheduled requests

142 + # but may still be finished requests.

143 + assert not output or not output.req_ids

144 + output = scheduler_output, None

145 return output

146

147 def override_env_vars(self, vars: Dict[str, str]):

vllm/v1/engine/core.py

+36 -33

@@ -138,12 +138,12 @@ class EngineCore:

138 # schedule and execute batches, and is required by pipeline parallelism

139 # to eliminate pipeline bubbles.

140 self.batch_queue_size = self.model_executor.max_concurrent_batches

141 - self.batch_queue: Optional[queue.Queue[tuple[Future[ModelRunnerOutput],

142 - SchedulerOutput]]] = None

141 + self.batch_queue: Optional[deque[tuple[Future[ModelRunnerOutput],

142 + SchedulerOutput]]] = None

143 if self.batch_queue_size > 1:

144 logger.info("Batch queue is enabled with size %d",

145 self.batch_queue_size)

146 - self.batch_queue = queue.Queue(self.batch_queue_size)

146 + self.batch_queue = deque(maxlen=self.batch_queue_size)

147

148 self.request_block_hasher: Optional[Callable[[Request],

149 list[BlockHash]]] = None

@@ -319,41 +319,43 @@ class EngineCore:

319 batch in the job queue is finished.

320 3. Update the scheduler from the output.

321 """

322 - assert self.batch_queue is not None

322 + batch_queue = self.batch_queue

323 + assert batch_queue is not None

324

324 - engine_core_outputs = None

325 - scheduler_output = None

325 # Try to schedule a new batch if the batch queue is not full, but

326 # the scheduler may return an empty batch if all requests are scheduled.

327 # Note that this is not blocking.

329 - if not self.batch_queue.full():

330 - scheduler_output = self.scheduler.schedule()

331 - if scheduler_output.total_num_scheduled_tokens > 0:

332 - future = self.model_executor.execute_model(scheduler_output)

333 - self.batch_queue.put_nowait(

334 - (future, scheduler_output)) # type: ignore

335 -

336 - scheduled_batch = (scheduler_output is not None

337 - and scheduler_output.total_num_scheduled_tokens > 0)

338 -

339 - # If no more requests can be scheduled and the job queue is not empty,

340 - # block until the first batch in the job queue is finished.

341 - # TODO(comaniac): Ideally we should peek the first batch in the

342 - # job queue to check if it's finished before scheduling a new batch,

343 - # but peeking the first element in a queue is not thread-safe,

344 - # so we need more work.

345 - if not scheduled_batch and not self.batch_queue.empty():

346 - future, scheduler_output = self.batch_queue.get_nowait()

328 + assert len(batch_queue) < self.batch_queue_size

329

348 - # Blocking until the first result is available.

349 - model_output = self.execute_model_with_error_logging(

350 - lambda _: future.result(), scheduler_output)

330 + model_executed = False

331 + if self.scheduler.has_requests():

332 + scheduler_output = self.scheduler.schedule()

333 + future = self.model_executor.execute_model(scheduler_output)

334 + batch_queue.appendleft(

335 + (future, scheduler_output)) # type: ignore[arg-type]

336 +

337 + model_executed = scheduler_output.total_num_scheduled_tokens > 0

338 + if model_executed and len(batch_queue) < self.batch_queue_size \

339 + and not batch_queue[-1][0].done():

340 + # Don't block on next worker response unless the queue is full

341 + # or there are no more requests to schedule.

342 + return None, True

343 +

344 + elif not batch_queue:

345 + # Queue is empty. We should not reach here since this method should

346 + # only be called when the scheduler contains requests or the queue

347 + # is non-empty.

348 + return None, False

349 +

350 + # Block until the next result is available.

351 + future, scheduler_output = batch_queue.pop()

352 + model_output = self.execute_model_with_error_logging(

353 + lambda _: future.result(), scheduler_output)

354

352 - self.batch_queue.task_done()

353 - engine_core_outputs = (self.scheduler.update_from_output(

354 - scheduler_output, model_output))

355 + engine_core_outputs = self.scheduler.update_from_output(

356 + scheduler_output, model_output)

357

356 - return engine_core_outputs, scheduled_batch

358 + return engine_core_outputs, model_executed

359

360 def shutdown(self):

361 self.structured_output_manager.clear_backend()

@@ -388,7 +390,7 @@ class EngineCore:

390 return self.model_executor.is_sleeping

391

392 def execute_dummy_batch(self):

391 - self.model_executor.collective_rpc("execute_dummy_batch")

393 + self.model_executor.execute_dummy_batch()

394

395 def add_lora(self, lora_request: LoRARequest) -> bool:

396 return self.model_executor.add_lora(lora_request)

@@ -733,7 +735,8 @@ class EngineCoreProc(EngineCore):

735 """Exits when an engine step needs to be performed."""

736

737 waited = False

736 - while not self.engines_running and not self.scheduler.has_requests():

738 + while not self.engines_running and not self.scheduler.has_requests() \

739 + and not self.batch_queue:

740 if logger.isEnabledFor(DEBUG) and self.input_queue.empty():

741 logger.debug("EngineCore waiting for work.")

742 waited = True

vllm/v1/executor/abstract.py

+5 -4

@@ -81,12 +81,10 @@ class Executor(ExecutorBase):

81 pass

83 def determine_available_memory(self) -> list[int]: # in bytes

84 - output = self.collective_rpc("determine_available_memory")

85 - return output

84 + return self.collective_rpc("determine_available_memory")

86 def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:

88 - output = self.collective_rpc("get_kv_cache_spec")

89 - return output

87 + return self.collective_rpc("get_kv_cache_spec")

89 def execute_model(

90 self,

@@ -96,6 +94,9 @@ class Executor(ExecutorBase):

94 args=(scheduler_output, ))

95 return output[0]

97 + def execute_dummy_batch(self) -> None:

98 + self.collective_rpc("execute_dummy_batch")

99 +

100 def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

101 output = self.collective_rpc("take_draft_token_ids")

102 return output[0]

vllm/v1/executor/multiproc_executor.py

+12 -3

@@ -191,6 +191,10 @@ class MultiprocExecutor(Executor):

191 outputs, self.output_rank)

192 return self.kv_output_aggregator.aggregate(outputs, self.output_rank)

193

194 + def execute_dummy_batch(self) -> None:

195 + self.collective_rpc("execute_dummy_batch",

196 + unique_reply_rank=self.output_rank)

197 +

198 def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

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

200 outputs = self.collective_rpc("take_draft_token_ids",

@@ -242,12 +246,17 @@ class MultiprocExecutor(Executor):

246 dequeue_timeout = None if deadline is None else (

247 deadline - time.monotonic())

248

245 - if non_block:

249 + if self.io_thread_pool is not None:

250 + # We must consume worker_response_mq from a single thread.

251 result = self.io_thread_pool.submit( # type: ignore

252 get_response, w, dequeue_timeout, self.shutdown_event)

248 - else:

253 + if not non_block:

254 + result = result.result()

255 + elif not non_block:

256 result = get_response(w, dequeue_timeout)

250 -

257 + else:

258 + raise RuntimeError("non_block can only be used when"

259 + " max_concurrent_batches > 1")

260 responses.append(result)

261

262 return responses

vllm/v1/worker/gpu_worker.py

+22 -21

@@ -354,36 +354,37 @@ class Worker(WorkerBase):

354 scheduler_output: "SchedulerOutput",

355 ) -> Optional[ModelRunnerOutput]:

356 intermediate_tensors = None

357 - if not get_pp_group().is_first_rank:

357 + forward_pass = scheduler_output.total_num_scheduled_tokens > 0

358 + if forward_pass and not get_pp_group().is_first_rank:

359 intermediate_tensors = IntermediateTensors(

360 get_pp_group().recv_tensor_dict(

361 all_gather_group=get_tp_group()))

362

363 output = self.model_runner.execute_model(scheduler_output,

364 intermediate_tensors)

365 + if isinstance(output, ModelRunnerOutput):

366 + return output

367

368 + assert isinstance(output, IntermediateTensors)

369 parallel_config = self.vllm_config.parallel_config

366 - if parallel_config.distributed_executor_backend != "external_launcher" \

367 - and not get_pp_group().is_last_rank:

368 - assert isinstance(output, IntermediateTensors)

369 - get_pp_group().send_tensor_dict(output.tensors,

370 - all_gather_group=get_tp_group())

371 -

372 - kv_connector_output = output.kv_connector_output

373 - if not kv_connector_output:

374 - return None

375 -

376 - # In case of PP with kv transfer, we need to pass through the

377 - # kv_connector_output

378 - if (not kv_connector_output.finished_sending

379 - and not kv_connector_output.finished_recving):

380 - return EMPTY_MODEL_RUNNER_OUTPUT

381 -

382 - output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

383 - output.kv_connector_output = kv_connector_output

384 - return output

370 + assert parallel_config.distributed_executor_backend != (

371 + "external_launcher") and not get_pp_group().is_last_rank

372 +

373 + get_pp_group().send_tensor_dict(output.tensors,

374 + all_gather_group=get_tp_group())

375 +

376 + kv_connector_output = output.kv_connector_output

377 + if not kv_connector_output:

378 + return None

379 +

380 + # In case of PP with kv transfer, we need to pass through the

381 + # kv_connector_output

382 + if (not kv_connector_output.finished_sending

383 + and not kv_connector_output.finished_recving):

384 + return EMPTY_MODEL_RUNNER_OUTPUT

385

386 - assert isinstance(output, ModelRunnerOutput)

386 + output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

387 + output.kv_connector_output = kv_connector_output

388 return output

389

390 def take_draft_token_ids(self) -> Optional[DraftTokenIds]: