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

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, 22:35 UTC

[Core] Move EngineCoreRequest to Request conversion out of EngineCore (#21627)

Signed-off-by: linzebing <linzebing1995@gmail.com>

Zebing Lin committed Jul 30, 2025 at 18:00 UTC ca9e2be3ed6320b51f52f536595cd24e254f8bb2

3 files changed +73 -48

tests/v1/engine/test_engine_core.py

+26 -18

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

65 """Test basic request lifecycle."""

67 # First request.

68 - engine_core.add_request(make_request())

68 + engine_core.add_request(

69 + *engine_core.preprocess_add_request(make_request()))

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

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

@@ -74,7 +75,8 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

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

77 # Second request.

77 - engine_core.add_request(make_request())

78 + engine_core.add_request(

79 + *engine_core.preprocess_add_request(make_request()))

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

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

@@ -83,8 +85,10 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

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

87 # Add two requests in a row.

86 - engine_core.add_request(make_request())

87 - engine_core.add_request(make_request())

88 + engine_core.add_request(

89 + *engine_core.preprocess_add_request(make_request()))

90 + engine_core.add_request(

91 + *engine_core.preprocess_add_request(make_request()))

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

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

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

108 req = make_request()

109 request_id = req.request_id

110

107 - engine_core.add_request(req)

111 + engine_core.add_request(*engine_core.preprocess_add_request(req))

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

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

114 assert engine_core.scheduler.has_unfinished_requests()

@@ -131,8 +135,8 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

135 req1 = make_request()

136 req2 = make_request()

137

134 - engine_core.add_request(req0)

135 - engine_core.add_request(req1)

138 + engine_core.add_request(*engine_core.preprocess_add_request(req0))

139 + engine_core.add_request(*engine_core.preprocess_add_request(req1))

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

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

142

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

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

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

146

143 - engine_core.add_request(req2)

147 + engine_core.add_request(*engine_core.preprocess_add_request(req2))

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

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

150

@@ -166,12 +170,12 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):

170 req0 = make_request()

171 req1 = make_request()

172 req0.request_id = req1.request_id = "test"

169 - engine_core.add_request(req0)

173 + engine_core.add_request(*engine_core.preprocess_add_request(req0))

174

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

176 pass

177

174 - engine_core.add_request(req1)

178 + engine_core.add_request(*engine_core.preprocess_add_request(req1))

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

180 pass

181

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

211 repetition_penalty=0.1,

212 stop_token_ids=[1001, 1002],

213 )

210 - engine_core.add_request(request)

214 + engine_core.add_request(*engine_core.preprocess_add_request(request))

215

216 def _check_engine_state():

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

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

230 top_p=0.99,

231 top_k=50,

232 )

229 - engine_core.add_request(request2)

233 + engine_core.add_request(*engine_core.preprocess_add_request(request2))

234 _check_engine_state()

235

236

@@ -298,9 +302,9 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):

302

303 # Add two requests in a row. Each request have 12 prompt tokens.

304 req0 = make_request_with_max_tokens("0", 5)

301 - engine_core.add_request(req0)

305 + engine_core.add_request(*engine_core.preprocess_add_request(req0))

306 req1 = make_request_with_max_tokens("1", 5)

303 - engine_core.add_request(req1)

307 + engine_core.add_request(*engine_core.preprocess_add_request(req1))

308

309 # Schedule Batch 1: (10, req0)

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

@@ -436,7 +440,8 @@ def test_engine_core_invalid_request_id_type(monkeypatch: pytest.MonkeyPatch):

440

441 with pytest.raises(TypeError,

442 match="request_id must be a string, got.*UUID"):

439 - engine_core.add_request(uuid_request)

443 + engine_core.add_request(

444 + *engine_core.preprocess_add_request(uuid_request))

445

446 # Test with integer

447 int_request = make_request()

@@ -444,7 +449,8 @@ def test_engine_core_invalid_request_id_type(monkeypatch: pytest.MonkeyPatch):

449

450 with pytest.raises(TypeError,

451 match="request_id must be a string, got.*int"):

447 - engine_core.add_request(int_request)

452 + engine_core.add_request(

453 + *engine_core.preprocess_add_request(int_request))

454

455 # Test with None

456 none_request = make_request()

@@ -452,10 +458,12 @@ def test_engine_core_invalid_request_id_type(monkeypatch: pytest.MonkeyPatch):

458

459 with pytest.raises(TypeError,

460 match="request_id must be a string, got.*NoneType"):

455 - engine_core.add_request(none_request)

461 + engine_core.add_request(

462 + *engine_core.preprocess_add_request(none_request))

463

464 # Verify engine is still functional after errors

465 valid_request = make_request()

459 - engine_core.add_request(valid_request)

466 + engine_core.add_request(

467 + *engine_core.preprocess_add_request(valid_request))

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

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

vllm/v1/engine/core.py

+45 -29

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

205 def get_supported_tasks(self) -> tuple[SupportedTask, ...]:

206 return self.model_executor.supported_tasks

207

208 - def add_request(self, request: EngineCoreRequest):

209 - """Add request to the scheduler."""

208 + def add_request(self, request: Request, request_wave: int = 0):

209 + """Add request to the scheduler.

210 +

211 + `request_wave`: indicate which wave of requests this is expected to

212 + belong to in DP case

213 + """

214 # Validate the request_id type.

215 if not isinstance(request.request_id, str):

216 raise TypeError(

@@ -222,27 +226,12 @@ class EngineCore:

226 raise ValueError(f"Unsupported task: {pooling_params.task!r} "

227 f"Supported tasks: {supported_pooling_tasks}")

228

225 - if request.mm_hashes is not None:

226 - # Here, if hash exists for a multimodal input, then it will be

227 - # fetched from the cache, else it will be added to the cache.

228 - # Note that the cache here is mirrored with the client cache, so

229 - # anything that has a hash must have a HIT cache entry here

230 - # as well.

231 - assert request.mm_inputs is not None

232 - request.mm_inputs = self.mm_input_cache_server.get_and_update_p1(

233 - request.mm_inputs, request.mm_hashes)

234 -

235 - req = Request.from_engine_core_request(request)

236 - if req.use_structured_output:

237 - # Start grammar compilation asynchronously

238 - self.structured_output_manager.grammar_init(req)

239 -

240 - if req.kv_transfer_params is not None and (

229 + if request.kv_transfer_params is not None and (

230 not self.scheduler.get_kv_connector()):

231 logger.warning("Got kv_transfer_params, but no KVConnector found. "

232 "Disabling KVTransfer for this request.")

233

245 - self.scheduler.add_request(req)

234 + self.scheduler.add_request(request)

235

236 def abort_requests(self, request_ids: list[str]):

237 """Abort requests from the scheduler."""

@@ -414,6 +403,31 @@ class EngineCore:

403 self.model_executor.save_tensorized_model(

404 tensorizer_config=tensorizer_config, )

405

406 + def preprocess_add_request(

407 + self, request: EngineCoreRequest) -> tuple[Request, int]:

408 + """Preprocess the request.

409 +

410 + This function could be directly used in input processing thread to allow

411 + request initialization running in parallel with Model forward

412 + """

413 + if request.mm_hashes is not None:

414 + assert request.mm_inputs is not None

415 + # Note on thread safety: no race condition.

416 + # `mm_input_cache_server` is reset at the end of LLMEngine init,

417 + # and will only accessed in the input processing thread afterwards.

418 + request.mm_inputs = self.mm_input_cache_server.get_and_update_p1(

419 + request.mm_inputs, request.mm_hashes)

420 +

421 + req = Request.from_engine_core_request(request)

422 + if req.use_structured_output:

423 + # Note on thread safety: no race condition.

424 + # `grammar_init` is only invoked in input processing thread. For

425 + # `structured_output_manager`, each request is independent and

426 + # grammar compilation is async. Scheduler always checks grammar

427 + # compilation status before scheduling request.

428 + self.structured_output_manager.grammar_init(req)

429 + return req, request.current_wave

430 +

431

432 class EngineCoreProc(EngineCore):

433 """ZMQ-wrapper for running EngineCore in background process."""

@@ -707,7 +721,8 @@ class EngineCoreProc(EngineCore):

721 """Dispatch request from client."""

722

723 if request_type == EngineCoreRequestType.ADD:

710 - self.add_request(request)

724 + req, request_wave = request

725 + self.add_request(req, request_wave)

726 elif request_type == EngineCoreRequestType.ABORT:

727 self.abort_requests(request)

728 elif request_type == EngineCoreRequestType.UTILITY:

@@ -806,10 +821,11 @@ class EngineCoreProc(EngineCore):

821 bytes(type_frame.buffer))

822

823 # Deserialize the request data.

809 - decoder = add_request_decoder if (

810 - request_type

811 - == EngineCoreRequestType.ADD) else generic_decoder

812 - request = decoder.decode(data_frames)

824 + if request_type == EngineCoreRequestType.ADD:

825 + request = add_request_decoder.decode(data_frames)

826 + request = self.preprocess_add_request(request)

827 + else:

828 + request = generic_decoder.decode(data_frames)

829

830 # Push to input queue for core busy loop.

831 self.input_queue.put_nowait((request_type, request))

@@ -939,17 +955,17 @@ class DPEngineCoreProc(EngineCoreProc):

955 if dp_group := getattr(self, "dp_group", None):

956 stateless_destroy_torch_distributed_process_group(dp_group)

957

942 - def add_request(self, request: EngineCoreRequest):

943 - if self.has_coordinator and request.current_wave != self.current_wave:

944 - if request.current_wave > self.current_wave:

945 - self.current_wave = request.current_wave

958 + def add_request(self, request: Request, request_wave: int = 0):

959 + if self.has_coordinator and request_wave != self.current_wave:

960 + if request_wave > self.current_wave:

961 + self.current_wave = request_wave

962 elif not self.engines_running:

963 # Request received for an already-completed wave, notify

964 # front-end that we need to start the next one.

965 self.output_queue.put_nowait(

966 (-1, EngineCoreOutputs(start_wave=self.current_wave)))

967

952 - super().add_request(request)

968 + super().add_request(request, request_wave)

969

970 def _handle_client_request(self, request_type: EngineCoreRequestType,

971 request: Any) -> None:

vllm/v1/engine/core_client.py

+2 -1

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

250 return self.engine_core.get_supported_tasks()

251

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

253 - self.engine_core.add_request(request)

253 + req, request_wave = self.engine_core.preprocess_add_request(request)

254 + self.engine_core.add_request(req, request_wave)

255

256 def abort_requests(self, request_ids: list[str]) -> None:

257 if len(request_ids) > 0: