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 10, 20:15 UTC
[Core] Deduplicate generate/encode logic in `AsyncLLM` (#31510)
Signed-off-by: njhill <nickhill123@gmail.com>
Nick Hill committed
Dec 29, 2025 at 18:42 UTC
e54ee3ea33bb8554eb6fa2eedb7007ae101197cf
1 file changed
+32
-48
vllm/v1/engine/async_llm.py
+32
-48
@@ -281,6 +281,25 @@ class AsyncLLM(EngineClient):
281
282
is_pooling = isinstance(params, PoolingParams)
283
284
+ if (
285
+ self.vllm_config.cache_config.kv_sharing_fast_prefill
286
+ and not is_pooling
287
+ and params.prompt_logprobs
288
+ ):
289
+ raise ValueError(
290
+ "--kv-sharing-fast-prefill produces incorrect logprobs for "
291
+ "prompt tokens, please disable it when the requests need "
292
+ "prompt logprobs"
293
+ )
294
+
295
+ if tokenization_kwargs is None:
296
+ tokenization_kwargs = {}
297
+ _validate_truncation_size(
298
+ self.model_config.max_model_len,
299
+ params.truncate_prompt_tokens,
300
+ tokenization_kwargs,
301
+ )
302
+
303
# Convert Input --> Request.
304
if isinstance(prompt, EngineCoreRequest):
305
request = prompt
@@ -291,7 +310,10 @@ class AsyncLLM(EngineClient):
310
"latter will be used, and the former will be ignored."
311
)
312
else:
294
- assert prompt_text is None
313
+ if prompt_text is not None:
314
+ raise ValueError(
315
+ "should only provide prompt_text with EngineCoreRequest"
316
+ )
317
request = self.input_processor.process_inputs(
318
request_id,
319
prompt,
@@ -310,6 +332,15 @@ class AsyncLLM(EngineClient):
332
333
self.input_processor.assign_request_id(request)
334
335
+ # We start the output_handler on the first call to add_request() so
336
+ # we can call __init__ before the event loop, which enables us
337
+ # to handle startup failure gracefully in the OpenAI server.
338
+ self._run_output_handler()
339
+
340
+ # Respect pause state before accepting new requests.
341
+ async with self._pause_cond:
342
+ await self._pause_cond.wait_for(lambda: not self._paused)
343
+
344
# Create a new output collector for the request.
345
queue = RequestOutputCollector(params.output_kind, request.request_id)
346
@@ -385,37 +416,8 @@ class AsyncLLM(EngineClient):
416
returning the RequestOutput back to the caller.
417
"""
418
388
- if (
389
- self.vllm_config.cache_config.kv_sharing_fast_prefill
390
- and sampling_params.prompt_logprobs
391
- ):
392
- raise ValueError(
393
- "--kv-sharing-fast-prefill produces incorrect logprobs for "
394
- "prompt tokens, please disable it when the requests need "
395
- "prompt logprobs"
396
- )
397
-
419
q: RequestOutputCollector | None = None
420
try:
400
- # We start the output_handler on the first call to generate() so
401
- # we can call __init__ before the event loop, which enables us
402
- # to handle startup failure gracefully in the OpenAI server.
403
- self._run_output_handler()
404
-
405
- # Wait until generation is resumed if the engine is paused.
406
- async with self._pause_cond:
407
- await self._pause_cond.wait_for(lambda: not self._paused)
408
-
409
- if tokenization_kwargs is None:
410
- tokenization_kwargs = {}
411
- truncate_prompt_tokens = sampling_params.truncate_prompt_tokens
412
-
413
- _validate_truncation_size(
414
- self.model_config.max_model_len,
415
- truncate_prompt_tokens,
416
- tokenization_kwargs,
417
- )
418
-
421
q = await self.add_request(
422
request_id,
423
prompt,
@@ -639,18 +641,6 @@ class AsyncLLM(EngineClient):
641
642
q: RequestOutputCollector | None = None
643
try:
642
- # We start the output_handler on the first call to generate() so
643
- # we can call __init__ before the event loop, which enables us
644
- # to handle startup failure gracefully in the OpenAI server.
645
- self._run_output_handler()
646
-
647
- # Respect pause state before accepting new requests.
648
- async with self._pause_cond:
649
- await self._pause_cond.wait_for(lambda: not self._paused)
650
-
651
- if tokenization_kwargs is None:
652
- tokenization_kwargs = {}
653
-
644
if truncate_prompt_tokens is not None:
645
warnings.warn(
646
"The `truncate_prompt_tokens` parameter in `AsyncLLM.encode()` "
@@ -660,12 +650,6 @@ class AsyncLLM(EngineClient):
650
stacklevel=2,
651
)
652
663
- _validate_truncation_size(
664
- self.model_config.max_model_len,
665
- pooling_params.truncate_prompt_tokens,
666
- tokenization_kwargs,
667
- )
668
-
653
q = await self.add_request(
654
request_id,
655
prompt,