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 8, 01:15 UTC
[Refactor] Implement output type check in LLM (#34794)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Feb 20, 2026 at 11:57 UTC
ac900c89bba77f69ed42e8a19a5006bd215eeb80
2 files changed
+58
-38
vllm/entrypoints/llm.py
+58
-34
@@ -10,7 +10,7 @@ import cloudpickle
10
import torch.nn as nn
11
from pydantic import ValidationError
12
from tqdm.auto import tqdm
13
-from typing_extensions import TypeVar
13
+from typing_extensions import TypeVar, overload
14
15
from vllm.beam_search import (
16
BeamSearchInstance,
@@ -94,6 +94,11 @@ if TYPE_CHECKING:
94
95
logger = init_logger(__name__)
96
97
+_O = TypeVar(
98
+ "_O",
99
+ bound=RequestOutput | PoolingRequestOutput,
100
+ default=RequestOutput | PoolingRequestOutput,
101
+)
102
_P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
103
_R = TypeVar("_R", default=Any)
104
@@ -447,17 +452,16 @@ class LLM:
452
if sampling_params is None:
453
sampling_params = self.get_default_sampling_params()
454
450
- outputs = self._run_completion(
455
+ return self._run_completion(
456
prompts=prompts,
457
params=sampling_params,
458
+ output_type=RequestOutput,
459
use_tqdm=use_tqdm,
460
lora_request=lora_request,
461
tokenization_kwargs=tokenization_kwargs,
462
priority=priority,
463
)
464
459
- return self.engine_class.validate_outputs(outputs, RequestOutput)
460
-
465
def enqueue(
466
self,
467
prompts: PromptType | Sequence[PromptType],
@@ -524,23 +528,43 @@ class LLM:
528
529
return request_ids
530
531
+ @overload
532
def wait_for_completion(
533
self,
534
+ *,
535
use_tqdm: bool | Callable[..., tqdm] = True,
530
- ) -> list[RequestOutput]:
536
+ ) -> list[RequestOutput | PoolingRequestOutput]: ...
537
+
538
+ @overload
539
+ def wait_for_completion(
540
+ self,
541
+ output_type: type[_O] | tuple[type[_O], ...],
542
+ *,
543
+ use_tqdm: bool | Callable[..., tqdm] = True,
544
+ ) -> list[_O]: ...
545
+
546
+ def wait_for_completion(
547
+ self,
548
+ output_type: type[Any] | tuple[type[Any], ...] | None = None,
549
+ *,
550
+ use_tqdm: bool | Callable[..., tqdm] = True,
551
+ ) -> list[Any]:
552
"""Wait for all enqueued requests to complete and return results.
553
554
This method processes all requests currently in the engine queue
555
and returns their outputs. Use after enqueue() to get results.
556
557
Args:
558
+ output_type: The expected output type, defaults to RequestOutput.
559
use_tqdm: If True, shows a tqdm progress bar.
560
561
Returns:
540
- A list of RequestOutput objects for all completed requests.
562
+ A list of output objects for all completed requests.
563
"""
542
- outputs = self._run_engine(use_tqdm=use_tqdm)
543
- return self.engine_class.validate_outputs(outputs, RequestOutput)
564
+ if output_type is None:
565
+ output_type = (RequestOutput, PoolingRequestOutput)
566
+
567
+ return self._run_engine(output_type, use_tqdm=use_tqdm)
568
569
def _resolve_mm_lora(
570
self,
@@ -744,13 +768,13 @@ class LLM:
768
769
# only runs for one step
770
# we don't need to use tqdm here
747
- raw_output = self._render_and_run_requests(
771
+ output = self._render_and_run_requests(
772
prompts=(beam.get_prompt() for beam in all_beams),
773
params=self._params_to_seq(sampling_params, len(all_beams)),
774
+ output_type=RequestOutput,
775
lora_requests=[beam.lora_request for beam in all_beams],
776
use_tqdm=False,
777
)
753
- output = self.engine_class.validate_outputs(raw_output, RequestOutput)
778
779
for (start, end), instance in zip(
780
instance_start_and_end, instances_batch
@@ -987,9 +1011,10 @@ class LLM:
1011
if sampling_params is None:
1012
sampling_params = self.get_default_sampling_params()
1013
990
- outputs = self._run_chat(
1014
+ return self._run_chat(
1015
messages=messages,
1016
params=sampling_params,
1017
+ output_type=RequestOutput,
1018
use_tqdm=use_tqdm,
1019
lora_request=lora_request,
1020
chat_template=chat_template,
@@ -1002,8 +1027,6 @@ class LLM:
1027
mm_processor_kwargs=mm_processor_kwargs,
1028
)
1029
1005
- return self.engine_class.validate_outputs(outputs, RequestOutput)
1006
-
1030
def encode(
1031
self,
1032
prompts: PromptType | Sequence[PromptType] | DataPrompt,
@@ -1135,19 +1158,16 @@ class LLM:
1158
outputs = self._run_completion(
1159
prompts=prompts_seq,
1160
params=params_seq,
1161
+ output_type=PoolingRequestOutput,
1162
use_tqdm=use_tqdm,
1163
lora_request=lora_request,
1164
tokenization_kwargs=tokenization_kwargs,
1165
)
1166
1143
- model_outputs = self.engine_class.validate_outputs(
1144
- outputs, PoolingRequestOutput
1145
- )
1146
-
1167
if use_io_processor:
1168
# get the post-processed model outputs
1169
assert self.io_processor is not None
1150
- processed_outputs = self.io_processor.post_process(model_outputs)
1170
+ processed_outputs = self.io_processor.post_process(outputs)
1171
1172
return [
1173
PoolingRequestOutput[Any](
@@ -1160,8 +1180,8 @@ class LLM:
1180
finished=True,
1181
)
1182
]
1163
- else:
1164
- return model_outputs
1183
+
1184
+ return outputs
1185
1186
def embed(
1187
self,
@@ -1353,8 +1373,7 @@ class LLM:
1373
embed_2=encoded_output_2,
1374
)
1375
1356
- items = self.engine_class.validate_outputs(scores, PoolingRequestOutput)
1357
- return [ScoringRequestOutput.from_base(item) for item in items]
1376
+ return [ScoringRequestOutput.from_base(item) for item in scores]
1377
1378
def _late_interaction_score(
1379
self,
@@ -1393,7 +1412,7 @@ class LLM:
1412
)
1413
text_2.append(text)
1414
1396
- encoded_output: list[PoolingRequestOutput] = self.encode(
1415
+ encoded_output = self.encode(
1416
text_1 + text_2,
1417
use_tqdm=use_tqdm,
1418
lora_request=lora_request,
@@ -1402,8 +1421,8 @@ class LLM:
1421
tokenization_kwargs=tokenization_kwargs,
1422
)
1423
1405
- encoded_output_1: list[PoolingRequestOutput] = encoded_output[0 : len(text_1)]
1406
- encoded_output_2: list[PoolingRequestOutput] = encoded_output[len(text_1) :]
1424
+ encoded_output_1 = encoded_output[0 : len(text_1)]
1425
+ encoded_output_2 = encoded_output[len(text_1) :]
1426
1427
if len(encoded_output_1) == 1:
1428
encoded_output_1 = encoded_output_1 * len(encoded_output_2)
@@ -1434,8 +1453,7 @@ class LLM:
1453
)
1454
)
1455
1437
- items = self.engine_class.validate_outputs(scores, PoolingRequestOutput)
1438
- return [ScoringRequestOutput.from_base(item) for item in items]
1456
+ return [ScoringRequestOutput.from_base(item) for item in scores]
1457
1458
def _cross_encoding_score(
1459
self,
@@ -1491,13 +1509,12 @@ class LLM:
1509
outputs = self._run_completion(
1510
prompts=prompts,
1511
params=pooling_params_list,
1512
+ output_type=PoolingRequestOutput,
1513
use_tqdm=use_tqdm,
1514
lora_request=lora_request,
1515
)
1516
1498
- items = self.engine_class.validate_outputs(outputs, PoolingRequestOutput)
1499
-
1500
- return [ScoringRequestOutput.from_base(item) for item in items]
1517
+ return [ScoringRequestOutput.from_base(item) for item in outputs]
1518
1519
def score(
1520
self,
@@ -1759,6 +1776,7 @@ class LLM:
1776
params: SamplingParams
1777
| PoolingParams
1778
| Sequence[SamplingParams | PoolingParams],
1779
+ output_type: type[_O],
1780
*,
1781
use_tqdm: bool | Callable[..., tqdm] = True,
1782
lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
@@ -1790,6 +1808,7 @@ class LLM:
1808
)
1809
),
1810
params=seq_params,
1811
+ output_type=output_type,
1812
use_tqdm=use_tqdm,
1813
lora_requests=seq_lora_requests,
1814
priorities=seq_priority,
@@ -1802,6 +1821,7 @@ class LLM:
1821
params: SamplingParams
1822
| PoolingParams
1823
| Sequence[SamplingParams | PoolingParams],
1824
+ output_type: type[_O],
1825
*,
1826
use_tqdm: bool | Callable[..., tqdm] = True,
1827
lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
@@ -1848,6 +1868,7 @@ class LLM:
1868
)
1869
),
1870
params=seq_params,
1871
+ output_type=output_type,
1872
lora_requests=seq_lora_requests,
1873
use_tqdm=use_tqdm,
1874
)
@@ -1856,6 +1877,7 @@ class LLM:
1877
self,
1878
prompts: Iterable[ProcessorInputs],
1879
params: Sequence[SamplingParams | PoolingParams],
1880
+ output_type: type[_O],
1881
*,
1882
lora_requests: Sequence[LoRARequest | None] | None = None,
1883
priorities: Sequence[int] | None = None,
@@ -1878,7 +1900,7 @@ class LLM:
1900
priorities=priorities,
1901
)
1902
1881
- return self._run_engine(use_tqdm=use_tqdm)
1903
+ return self._run_engine(output_type, use_tqdm=use_tqdm)
1904
1905
def _render_and_add_requests(
1906
self,
@@ -1932,9 +1954,10 @@ class LLM:
1954
1955
def _run_engine(
1956
self,
1957
+ output_type: type[_O] | tuple[type[_O], ...],
1958
*,
1959
use_tqdm: bool | Callable[..., tqdm] = True,
1937
- ) -> list[RequestOutput | PoolingRequestOutput]:
1960
+ ) -> list[_O]:
1961
# Initialize tqdm.
1962
if use_tqdm:
1963
num_requests = self.llm_engine.get_num_unfinished_requests()
@@ -1947,14 +1970,15 @@ class LLM:
1970
)
1971
1972
# Run the engine.
1950
- outputs: list[RequestOutput | PoolingRequestOutput] = []
1973
+ outputs: list[_O] = []
1974
total_in_toks = 0
1975
total_out_toks = 0
1976
while self.llm_engine.has_unfinished_requests():
1977
step_outputs = self.llm_engine.step()
1978
for output in step_outputs:
1979
+ assert isinstance(output, output_type)
1980
if output.finished:
1957
- outputs.append(output)
1981
+ outputs.append(output) # type: ignore[arg-type]
1982
if use_tqdm:
1983
if isinstance(output, RequestOutput):
1984
# Calculate tokens only for RequestOutput
vllm/v1/engine/llm_engine.py
-4
@@ -199,10 +199,6 @@ class LLMEngine:
199
self.should_execute_dummy_batch = True
200
return aggregated_has_unfinished
201
202
- @classmethod
203
- def validate_outputs(cls, outputs, output_type):
204
- return outputs
205
-
202
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
203
if not hasattr(self, "_supported_tasks"):
204
# Cache the result