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

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

[Core] Remove prompt string from engine core data structures (#17214)

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

Nick Hill committed Apr 25, 2025 at 23:41 UTC df6f3ce883bbfb1884b040f178496a3da33209b9

21 files changed +40 -76

tests/tokenization/test_detokenize.py

+2 -2

@@ -60,8 +60,8 @@ def _run_incremental_decode(tokenizer,

60 skip_special_tokens=skip_special_tokens,

61 spaces_between_special_tokens=spaces_between_special_tokens,

62 )

63 - request = EngineCoreRequest("", "", prompt_token_ids, None, None, None,

64 - params, None, 0.0, None)

63 + request = EngineCoreRequest("", prompt_token_ids, None, None, None, params,

64 + None, 0.0, None)

66 if fast is None:

67 detokenizer = IncrementalDetokenizer.from_new_request(

tests/v1/core/test_kv_cache_utils.py

-1

@@ -37,7 +37,6 @@ def make_request(request_id,

38 return Request(

39 request_id=request_id,

40 - prompt=None,

40 prompt_token_ids=prompt_token_ids,

41 multi_modal_inputs=multi_modal_inputs,

42 multi_modal_hashes=mm_hashes,

tests/v1/core/test_prefix_caching.py

-1

@@ -29,7 +29,6 @@ def make_request(request_id,

30 return Request(

31 request_id=request_id,

32 - prompt=None,

32 prompt_token_ids=prompt_token_ids,

33 multi_modal_inputs=multi_modal_inputs,

34 multi_modal_hashes=mm_hashes,

tests/v1/core/test_scheduler.py

-1

@@ -132,7 +132,6 @@ def create_requests(num_requests: int,

132 mm_inputs = None

133 request = Request(

134 request_id=f"{i}",

135 - prompt=None,

135 prompt_token_ids=[i] * num_tokens,

136 sampling_params=sampling_params,

137 multi_modal_inputs=mm_inputs,

tests/v1/engine/test_engine_core.py

+1 -2

@@ -31,8 +31,7 @@ PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids

32 def make_request() -> EngineCoreRequest:

33 return EngineCoreRequest(

34 - request_id=uuid.uuid4(),

35 - prompt=PROMPT,

34 + request_id=str(uuid.uuid4()),

35 prompt_token_ids=PROMPT_TOKENS,

36 mm_inputs=None,

37 mm_hashes=None,

tests/v1/engine/test_engine_core_client.py

-1

@@ -35,7 +35,6 @@ PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids

35 def make_request(params: SamplingParams) -> EngineCoreRequest:

36 return EngineCoreRequest(

37 request_id=str(uuid.uuid4()),

38 - prompt=PROMPT,

38 prompt_token_ids=PROMPT_TOKENS,

39 mm_inputs=None,

40 mm_hashes=None,

tests/v1/engine/test_output_processor.py

+16 -26

@@ -50,7 +50,6 @@ def test_incremental_detokenization(request_output_kind: RequestOutputKind,

50 # Make N requests.

51 requests = [

52 EngineCoreRequest(request_id=f"request-{idx}",

53 - prompt=prompt,

53 prompt_token_ids=prompt_tokens,

54 arrival_time=0,

55 mm_inputs=None,

@@ -64,14 +63,13 @@ def test_incremental_detokenization(request_output_kind: RequestOutputKind,

63 output_kind=request_output_kind,

64 stop=[],

65 include_stop_str_in_output=False,

67 - )) for idx, (prompt, prompt_tokens) in enumerate(

68 - zip(dummy_test_vectors.prompt_strings,

69 - dummy_test_vectors.prompt_tokens))

66 + ))

67 + for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)

68 ]

70 # Add requests to the detokenizer.

73 - for request in requests:

74 - output_processor.add_request(request)

71 + for request, prompt in zip(requests, dummy_test_vectors.prompt_strings):

72 + output_processor.add_request(request, prompt)

74 gen_strings = {}

75 gen_tokens = {}

@@ -398,7 +396,6 @@ def test_logprobs_processor(request_output_kind: RequestOutputKind,

396 ]

397 requests = [

398 EngineCoreRequest(request_id=request_id_list[idx],

401 - prompt=prompt,

399 prompt_token_ids=prompt_tokens,

400 arrival_time=0,

401 mm_inputs=None,

@@ -414,14 +411,13 @@ def test_logprobs_processor(request_output_kind: RequestOutputKind,

411 include_stop_str_in_output=False,

412 logprobs=num_sample_logprobs,

413 prompt_logprobs=num_prompt_logprobs,

417 - )) for idx, (prompt, prompt_tokens) in enumerate(

418 - zip(dummy_test_vectors.prompt_strings,

419 - dummy_test_vectors.prompt_tokens))

414 + ))

415 + for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)

416 ]

417

418 # Add requests to the detokenizer.

423 - for request in requests:

424 - output_processor.add_request(request)

419 + for request, prompt in zip(requests, dummy_test_vectors.prompt_strings):

420 + output_processor.add_request(request, prompt)

421

422 gen_tokens = {}

423 gen_logprobs = {}

@@ -562,7 +558,6 @@ def test_stop_token(include_stop_str_in_output: bool,

558 request_id = "request-0"

559 request = EngineCoreRequest(

560 request_id=request_id,

565 - prompt=prompt_string,

561 prompt_token_ids=prompt_tokens,

562 arrival_time=0,

563 mm_inputs=None,

@@ -583,7 +578,7 @@ def test_stop_token(include_stop_str_in_output: bool,

578 ))

579

580 # Add request to the detokenizer.

586 - output_processor.add_request(request)

581 + output_processor.add_request(request, prompt_string)

582

583 # Loop over engine core steps; run output processor

584 gen_string = ""

@@ -659,7 +654,6 @@ def test_stop_string(include_stop_str_in_output: bool,

654 requests = [

655 EngineCoreRequest(

656 request_id=request_id_list[idx],

662 - prompt=prompt,

657 prompt_token_ids=prompt_tokens,

658 arrival_time=0,

659 mm_inputs=None,

@@ -675,14 +669,13 @@ def test_stop_string(include_stop_str_in_output: bool,

669 include_stop_str_in_output=include_stop_str_in_output,

670 logprobs=num_sample_logprobs,

671 prompt_logprobs=None,

678 - )) for idx, (prompt, prompt_tokens) in enumerate(

679 - zip(dummy_test_vectors.prompt_strings,

680 - dummy_test_vectors.prompt_tokens))

672 + ))

673 + for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)

674 ]

675

676 # Add requests to the detokenizer.

684 - for request in requests:

685 - output_processor.add_request(request)

677 + for request, prompt in zip(requests, dummy_test_vectors.prompt_strings):

678 + output_processor.add_request(request, prompt)

679

680 gen_strings = {}

681 gen_tokens = {}

@@ -774,7 +767,6 @@ def test_iteration_stats(dummy_test_vectors):

767 requests = [

768 EngineCoreRequest(

769 request_id=f"request-{idx}",

777 - prompt=prompt,

770 prompt_token_ids=prompt_tokens,

771 arrival_time=0,

772 mm_inputs=None,

@@ -783,15 +775,13 @@ def test_iteration_stats(dummy_test_vectors):

775 eos_token_id=None,

776 lora_request=None,

777 sampling_params=SamplingParams(),

786 - ) for idx, (prompt, prompt_tokens) in enumerate(

787 - zip(dummy_test_vectors.prompt_strings,

788 - dummy_test_vectors.prompt_tokens))

778 + ) for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)

779 ]

780

781 # Add all requests except one to the OutputProcessor.

782 num_active = len(dummy_test_vectors.generation_tokens) - 1

783 for request in requests[:num_active]:

794 - output_processor.add_request(request)

784 + output_processor.add_request(request, None)

785 inactive_request = requests[num_active]

786

787 # First iteration has 2 prefills.

@@ -817,7 +807,7 @@ def test_iteration_stats(dummy_test_vectors):

807 assert iteration_stats.num_generation_tokens == num_active

808

809 # Add a new request - prefill and 2 decodes in this step.

820 - output_processor.add_request(inactive_request)

810 + output_processor.add_request(inactive_request, None)

811 num_active += 1

812 outputs = engine_core.get_outputs()[:num_active]

813 iteration_stats = IterationStats()

tests/v1/tpu/worker/test_tpu_model_runner.py

-1

@@ -77,7 +77,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

77 NewRequestData(

78 req_id=req_id,

79 prompt_token_ids=[1, 2, 3],

80 - prompt="test",

80 mm_inputs=[],

81 mm_hashes=[],

82 mm_positions=[],

tests/v1/worker/test_gpu_input_batch.py

-1

@@ -195,7 +195,6 @@ def _construct_cached_request_state(req_id_suffix: int):

195 return CachedRequestState(

196 req_id=f"req_id_{req_id_suffix}",

197 prompt_token_ids=prompt_token_ids,

198 - prompt=None,

198 sampling_params=_create_sampling_params(),

199 mm_inputs=[],

200 mm_positions=[],

tests/v1/worker/test_gpu_model_runner.py

-1

@@ -50,7 +50,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

50 NewRequestData(

51 req_id=req_id,

52 prompt_token_ids=[1, 2, 3],

53 - prompt="test",

53 mm_inputs=[],

54 mm_hashes=[],

55 mm_positions=[],

vllm/v1/core/sched/output.py

-2

@@ -22,7 +22,6 @@ class NewRequestData:

23 req_id: str

24 prompt_token_ids: list[int]

25 - prompt: Optional[str]

25 mm_inputs: list[MultiModalKwargs]

26 mm_hashes: list[str]

27 mm_positions: list[PlaceholderRange]

@@ -40,7 +39,6 @@ class NewRequestData:

39 return cls(

40 req_id=request.request_id,

41 prompt_token_ids=request.prompt_token_ids,

43 - prompt=request.prompt,

42 mm_inputs=request.mm_inputs,

43 mm_hashes=request.mm_hashes,

44 mm_positions=request.mm_positions,

vllm/v1/engine/__init__.py

-3

@@ -49,9 +49,6 @@ class EngineCoreRequest(

49 # due to circular imports and typing we have in data.py

51 request_id: str

52 - # NOTE(ywang96): original text prompt is needed when a request is added to

53 - # Detokenizer, but set to None when it is added to EngineCoreClient.

54 - prompt: Optional[str]

52 prompt_token_ids: list[int]

53 mm_inputs: Optional[Sequence[Optional[MultiModalKwargs]]]

54 mm_hashes: Optional[list[str]]

vllm/v1/engine/async_llm.py

+9 -8

@@ -217,14 +217,12 @@ class AsyncLLM(EngineClient):

217 queue = RequestOutputCollector(output_kind=params.output_kind)

218

219 # Convert Input --> Request.

220 - request = self.processor.process_inputs(request_id, prompt, params,

221 - arrival_time, lora_request,

222 - trace_headers,

223 - prompt_adapter_request,

224 - priority)

220 + prompt_str, request = self.processor.process_inputs(

221 + request_id, prompt, params, arrival_time, lora_request,

222 + trace_headers, prompt_adapter_request, priority)

223

224 if params.n == 1:

227 - await self._add_request(request, None, 0, queue)

225 + await self._add_request(request, prompt_str, None, 0, queue)

226 return queue

227

228 # Fan out child requests (for n>1).

@@ -234,15 +232,18 @@ class AsyncLLM(EngineClient):

232 child_request = request if idx == params.n - 1 else copy(request)

233 child_request.request_id = request_id

234 child_request.sampling_params = params

237 - await self._add_request(child_request, parent_request, idx, queue)

235 + await self._add_request(child_request, prompt_str, parent_request,

236 + idx, queue)

237 return queue

238

239 async def _add_request(self, request: EngineCoreRequest,

240 + prompt: Optional[str],

241 parent_req: Optional[ParentRequest], index: int,

242 queue: RequestOutputCollector):

243

244 # Add the request to OutputProcessor (this process).

245 - self.output_processor.add_request(request, parent_req, index, queue)

245 + self.output_processor.add_request(request, prompt, parent_req, index,

246 + queue)

247

248 # Add the EngineCoreRequest to EngineCore (separate process).

249 await self.engine_core.add_request_async(request)

vllm/v1/engine/core_client.py

-9

@@ -583,9 +583,6 @@ class SyncMPClient(MPClient):

583 return future.result()

584

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

586 - # NOTE: text prompt is not needed in the core engine as it has been

587 - # tokenized.

588 - request.prompt = None

586 self._send_input(EngineCoreRequestType.ADD, request)

587

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

@@ -772,9 +769,6 @@ class AsyncMPClient(MPClient):

769 return await future

770

771 async def add_request_async(self, request: EngineCoreRequest) -> None:

775 - # NOTE: text prompt is not needed in the core engine as it has been

776 - # tokenized.

777 - request.prompt = None

772 await self._send_input(EngineCoreRequestType.ADD, request)

773 self._ensure_output_queue_task()

774

@@ -867,9 +861,6 @@ class DPAsyncMPClient(AsyncMPClient):

861 ]))[0]

862

863 async def add_request_async(self, request: EngineCoreRequest) -> None:

870 - # NOTE: text prompt is not needed in the core engine as it has been

871 - # tokenized.

872 - request.prompt = None

864 request.current_wave = self.current_wave

865

866 chosen_engine = self.get_core_engine_for_request()

vllm/v1/engine/llm_engine.py

+6 -7

@@ -180,17 +180,15 @@ class LLMEngine:

180 priority: int = 0,

181 ) -> None:

182 # Process raw inputs into the request.

183 - request = self.processor.process_inputs(request_id, prompt, params,

184 - arrival_time, lora_request,

185 - trace_headers,

186 - prompt_adapter_request,

187 - priority)

183 + prompt_str, request = self.processor.process_inputs(

184 + request_id, prompt, params, arrival_time, lora_request,

185 + trace_headers, prompt_adapter_request, priority)

186

187 n = params.n if isinstance(params, SamplingParams) else 1

188

189 if n == 1:

190 # Make a new RequestState and queue.

193 - self.output_processor.add_request(request, None, 0)

191 + self.output_processor.add_request(request, prompt_str, None, 0)

192 # Add the request to EngineCore.

193 self.engine_core.add_request(request)

194 return

@@ -204,7 +202,8 @@ class LLMEngine:

202 child_request.sampling_params = params

203

204 # Make a new RequestState and queue.

207 - self.output_processor.add_request(child_request, parent_req, idx)

205 + self.output_processor.add_request(child_request, prompt_str,

206 + parent_req, idx)

207 # Add the request to EngineCore.

208 self.engine_core.add_request(child_request)

209

vllm/v1/engine/output_processor.py

+4 -1

@@ -109,6 +109,7 @@ class RequestState:

109 cls,

110 tokenizer: AnyTokenizer,

111 request: EngineCoreRequest,

112 + prompt: Optional[str],

113 parent_req: Optional[ParentRequest],

114 request_index: int,

115 queue: Optional[RequestOutputCollector],

@@ -123,7 +124,7 @@ class RequestState:

124 lora_name=(request.lora_request.name

125 if request.lora_request is not None else None),

126 output_kind=request.sampling_params.output_kind,

126 - prompt=request.prompt,

127 + prompt=prompt,

128 prompt_token_ids=request.prompt_token_ids,

129 logprobs_processor=LogprobsProcessor.from_new_request(

130 tokenizer=tokenizer,

@@ -267,6 +268,7 @@ class OutputProcessor:

268 def add_request(

269 self,

270 request: EngineCoreRequest,

271 + prompt: Optional[str],

272 parent_req: Optional[ParentRequest] = None,

273 request_index: int = 0,

274 queue: Optional[RequestOutputCollector] = None,

@@ -278,6 +280,7 @@ class OutputProcessor:

280 req_state = RequestState.from_new_request(

281 tokenizer=self.tokenizer.get_lora_tokenizer(request.lora_request),

282 request=request,

283 + prompt=prompt,

284 parent_req=parent_req,

285 request_index=request_index,

286 queue=queue,

vllm/v1/engine/processor.py

+2 -3

@@ -202,7 +202,7 @@ class Processor:

202 trace_headers: Optional[Mapping[str, str]] = None,

203 prompt_adapter_request: Optional[PromptAdapterRequest] = None,

204 priority: int = 0,

205 - ) -> EngineCoreRequest:

205 + ) -> tuple[Optional[str], EngineCoreRequest]:

206

207 # TODO(woosuk): Support pooling models.

208 # TODO(woosuk): Support encoder-decoder models.

@@ -306,9 +306,8 @@ class Processor:

306 else:

307 sorted_mm_inputs = orig_sorted_mm_inputs

308

309 - return EngineCoreRequest(

309 + return decoder_inputs.get("prompt"), EngineCoreRequest(

310 request_id=request_id,

311 - prompt=decoder_inputs.get("prompt"),

311 prompt_token_ids=decoder_inputs["prompt_token_ids"],

312 mm_inputs=sorted_mm_inputs,

313 mm_hashes=sorted_mm_hashes,

vllm/v1/request.py

-3

@@ -20,7 +20,6 @@ class Request:

20 def __init__(

21 self,

22 request_id: str,

23 - prompt: Optional[str],

23 prompt_token_ids: list[int],

24 multi_modal_inputs: Optional[list[MultiModalKwargs]],

25 multi_modal_hashes: Optional[list[str]],

@@ -46,7 +45,6 @@ class Request:

45 assert sampling_params.max_tokens is not None

46 self.max_tokens = sampling_params.max_tokens

49 - self.prompt = prompt

48 self.prompt_token_ids = prompt_token_ids

49 self.num_prompt_tokens = len(self.prompt_token_ids)

50 self._output_token_ids: list[int] = []

@@ -81,7 +79,6 @@ class Request:

80 return cls(

81 request_id=request.request_id,

84 - prompt=request.prompt,

82 prompt_token_ids=request.prompt_token_ids,

83 multi_modal_inputs=request.mm_inputs,

84 multi_modal_hashes=request.mm_hashes,

vllm/v1/worker/gpu_input_batch.py

-1

@@ -24,7 +24,6 @@ class CachedRequestState:

25 req_id: str

26 prompt_token_ids: list[int]

27 - prompt: Optional[str]

27 mm_inputs: list[MultiModalKwargs]

28 mm_positions: list[PlaceholderRange]

29 sampling_params: SamplingParams

vllm/v1/worker/gpu_model_runner.py

-1

@@ -347,7 +347,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

347 self.requests[req_id] = CachedRequestState(

348 req_id=req_id,

349 prompt_token_ids=new_req_data.prompt_token_ids,

350 - prompt=new_req_data.prompt,

350 mm_inputs=new_req_data.mm_inputs,

351 mm_positions=new_req_data.mm_positions,

352 sampling_params=sampling_params,

vllm/v1/worker/tpu_model_runner.py

-1

@@ -356,7 +356,6 @@ class TPUModelRunner:

356 self.requests[req_id] = CachedRequestState(

357 req_id=req_id,

358 prompt_token_ids=new_req_data.prompt_token_ids,

359 - prompt=new_req_data.prompt,

359 mm_inputs=new_req_data.mm_inputs,

360 mm_positions=new_req_data.mm_positions,

361 sampling_params=sampling_params,