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

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 15, 08:30 UTC

[BugFix] Make DP work with connector-delayed new requests (#18559)

Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Will Eaton <weaton@redhat.com>

Nick Hill committed May 29, 2025 at 11:04 UTC d1d61f33511d46fda02944d8000a5ab22465c142

4 files changed +37 -41

tests/v1/engine/test_engine_core.py

+9 -9

@@ -88,7 +88,7 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

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

90 # Loop through until they are all done.

91 - while len(engine_core.step().outputs) > 0:

91 + while len(engine_core.step()[0].outputs) > 0:

92 pass

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

@@ -163,11 +163,11 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

163 req0.request_id = req1.request_id = "test"

164 engine_core.add_request(req0)

165

166 - while len(engine_core.step().outputs) > 0:

166 + while len(engine_core.step()[0].outputs) > 0:

167 pass

168

169 engine_core.add_request(req1)

170 - while len(engine_core.step().outputs) > 0:

170 + while len(engine_core.step()[0].outputs) > 0:

171 pass

172

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

@@ -207,7 +207,7 @@ def test_engine_core_advanced_sampling(monkeypatch: pytest.MonkeyPatch):

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

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

209 # Loop through until they are all done.

210 - while len(engine_core.step().outputs) > 0:

210 + while len(engine_core.step()[0].outputs) > 0:

211 pass

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

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

@@ -296,7 +296,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

296 engine_core.add_request(req1)

297

298 # Schedule Batch 1: (10, req0)

299 - assert engine_core.step_with_batch_queue() is None

299 + assert engine_core.step_with_batch_queue()[0] is None

300 assert engine_core.batch_queue.qsize() == 1

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

302 assert scheduler_output.num_scheduled_tokens[0] == 10

@@ -305,7 +305,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

305 req0.request_id].num_computed_tokens == 10

306

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

308 - assert engine_core.step_with_batch_queue() is None

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

309 assert engine_core.batch_queue.qsize() == 2

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

311 assert scheduler_output.num_scheduled_tokens[0] == 2

@@ -327,7 +327,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

327 assert scheduler_output.num_scheduled_tokens[1] == 4

328

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

330 - output = engine_core.step_with_batch_queue()

330 + output = engine_core.step_with_batch_queue()[0]

331 assert output is not None

332 assert len(output.outputs) == 1

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

@@ -339,7 +339,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

339 assert scheduler_output.num_scheduled_tokens[0] == 1

340

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

342 - output = engine_core.step_with_batch_queue()

342 + output = engine_core.step_with_batch_queue()[0]

343 assert output is not None

344 assert len(output.outputs) == 1

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

@@ -358,7 +358,7 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

358 engine_core.scheduler.requests[1].num_tokens + 1,

359 ]

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

361 - output = engine_core.step_with_batch_queue()

361 + output = engine_core.step_with_batch_queue()[0]

362 if step % 2 == 0:

363 # Even steps consumes an output.

364 assert output is not None

vllm/forward_context.py

+4 -3

@@ -101,7 +101,7 @@ def get_forward_context() -> ForwardContext:

101 def set_forward_context(attn_metadata: Any,

102 vllm_config: VllmConfig,

103 virtual_engine: int = 0,

104 - num_tokens: int = 0):

104 + num_tokens: Optional[int] = None):

105 """A context manager that stores the current forward context,

106 can be attention metadata, etc.

107 Here we can inject common logic for every model forward pass.

@@ -111,9 +111,10 @@ def set_forward_context(attn_metadata: Any,

111 if need_to_track_batchsize:

112 forward_start_time = time.perf_counter()

113 dp_metadata: Optional[DPMetadata] = None

114 - if vllm_config.parallel_config.data_parallel_size > 1:

114 + if vllm_config.parallel_config.data_parallel_size > 1 and (

115 + attn_metadata is not None or num_tokens is not None):

116 dp_metadata = DPMetadata.make(vllm_config.parallel_config,

116 - attn_metadata, num_tokens)

117 + attn_metadata, num_tokens or 0)

118

119 global _forward_context

120 prev_context = _forward_context

vllm/v1/engine/core.py

+22 -28

@@ -211,8 +211,12 @@ class EngineCore:

211 # Re-raise exception

212 raise err

213

214 - def step(self) -> EngineCoreOutputs:

215 - """Schedule, execute, and make output."""

214 + def step(self) -> tuple[EngineCoreOutputs, bool]:

215 + """Schedule, execute, and make output.

216 +

217 + Returns tuple of outputs and a flag indicating whether the model

218 + was executed.

219 + """

220

221 # Check for any requests remaining in the scheduler - unfinished,

222 # or finished and not yet removed from the batch.

@@ -220,15 +224,17 @@ class EngineCore:

224 return EngineCoreOutputs(

225 outputs=[],

226 scheduler_stats=self.scheduler.make_stats(),

223 - )

227 + ), False

228 scheduler_output = self.scheduler.schedule()

229 model_output = self.execute_model(scheduler_output)

230 engine_core_outputs = self.scheduler.update_from_output(

231 scheduler_output, model_output) # type: ignore

232

229 - return engine_core_outputs

233 + return (engine_core_outputs,

234 + scheduler_output.total_num_scheduled_tokens > 0)

235

231 - def step_with_batch_queue(self) -> Optional[EngineCoreOutputs]:

236 + def step_with_batch_queue(

237 + self) -> tuple[Optional[EngineCoreOutputs], bool]:

238 """Schedule and execute batches with the batch queue.

239 Note that if nothing to output in this step, None is returned.

240

@@ -273,7 +279,7 @@ class EngineCore:

279 engine_core_outputs = self.scheduler.update_from_output(

280 scheduler_output, model_output)

281

276 - return engine_core_outputs

282 + return engine_core_outputs, scheduled_batch

283

284 def shutdown(self):

285 self.structured_output_manager.clear_backend()

@@ -537,15 +543,17 @@ class EngineCoreProc(EngineCore):

543 req = self.input_queue.get_nowait()

544 self._handle_client_request(*req)

545

540 - def _process_engine_step(self):

546 + def _process_engine_step(self) -> bool:

547 """Called only when there are unfinished local requests."""

548

549 # Step the engine core.

544 - outputs = self.step_fn()

550 + outputs, model_executed = self.step_fn()

551 # Put EngineCoreOutputs into the output queue.

552 if outputs is not None:

553 self.output_queue.put_nowait(outputs)

554

555 + return model_executed

556 +

557 def _handle_client_request(self, request_type: EngineCoreRequestType,

558 request: Any) -> None:

559 """Dispatch request from client."""

@@ -749,30 +757,16 @@ class DPEngineCoreProc(EngineCoreProc):

757 # 1) Poll the input queue until there is work to do.

758 self._process_input_queue()

759

760 + # 2) Step the engine core.

761 + executed = self._process_engine_step()

762 local_unfinished_reqs = self.scheduler.has_unfinished_requests()

753 -

754 - if local_unfinished_reqs:

755 - # 2) Step the engine core.

756 - self._process_engine_step()

757 -

758 - # Check if we have now finished all requests.

759 - local_unfinished_reqs = (

760 - self.scheduler.has_unfinished_requests())

761 - else:

762 - if self.scheduler.has_finished_requests():

763 - # There are no unfinished requests, but there are some

764 - # finished requests remaining to be removed from the

765 - # batch state. This engine step won't perform a forward

766 - # pass but will flush the finished requests to ensure

767 - # up-to-date state is returned in the engine outputs.

768 - self._process_engine_step()

769 -

770 - if not self.engines_running:

763 + if not executed:

764 + if not local_unfinished_reqs and not self.engines_running:

765 # All engines are idle.

766 continue

767

774 - # There must be unfinished requests in DP peers, run a

775 - # dummy forward pass.

768 + # We are in a running state and so must execute a dummy pass

769 + # if the model didn't execute any ready requests.

770 self.execute_dummy_batch()

771

772 # 3) All-reduce operation to determine global unfinished reqs.

vllm/v1/engine/core_client.py

+2 -1

@@ -206,7 +206,8 @@ class InprocClient(EngineCoreClient):

206 self.engine_core = EngineCore(*args, **kwargs)

207

208 def get_output(self) -> EngineCoreOutputs:

209 - return self.engine_core.step()

209 + outputs, _ = self.engine_core.step()

210 + return outputs

211

212 def add_request(self, request: EngineCoreRequest) -> None:

213 self.engine_core.add_request(request)