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, 05:30 UTC
[Tests] conftest: Extending VllmRunner and HfRunner to accept token_ids as input (#26295)
Signed-off-by: Yannick Schnider <yannick.schnider1@ibm.com> Signed-off-by: Yannick Schnider <Yannick.Schnider1@ibm.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Yannick Schnider committed
Oct 6, 2025 at 19:19 UTC
6431be808f0c83cb131caeb3fa487df6582d9912
2 files changed
+65
-65
tests/conftest.py
+48
-31
@@ -57,7 +57,7 @@ from vllm.multimodal.utils import fetch_image
57
from vllm.outputs import RequestOutput
58
from vllm.sampling_params import BeamSearchParams
59
from vllm.transformers_utils.utils import maybe_model_redirect
60
-from vllm.utils import set_default_torch_num_threads
60
+from vllm.utils import is_list_of, set_default_torch_num_threads
61
62
logger = init_logger(__name__)
63
@@ -406,11 +406,11 @@ class HfRunner:
406
407
def get_inputs(
408
self,
409
- prompts: list[str],
409
+ prompts: Union[list[str], list[list[int]]],
410
images: Optional[PromptImageInput] = None,
411
videos: Optional[PromptVideoInput] = None,
412
audios: Optional[PromptAudioInput] = None,
413
- ) -> list[Union[BatchFeature, BatchEncoding]]:
413
+ ) -> list[Union[BatchFeature, BatchEncoding, dict[str, torch.Tensor]]]:
414
if images is not None:
415
assert len(prompts) == len(images)
416
@@ -420,31 +420,48 @@ class HfRunner:
420
if audios is not None:
421
assert len(prompts) == len(audios)
422
423
- all_inputs: list[Union[BatchFeature, BatchEncoding]] = []
423
+ all_inputs: list[
424
+ Union[BatchFeature, BatchEncoding, dict[str, torch.Tensor]]
425
+ ] = []
426
for i, prompt in enumerate(prompts):
425
- processor_kwargs: dict[str, Any] = {
426
- "text": prompt,
427
- "return_tensors": "pt",
428
- }
429
- if images is not None and (image := images[i]) is not None:
430
- processor_kwargs["images"] = image
431
- if videos is not None and (video := videos[i]) is not None:
432
- processor_kwargs["videos"] = video
433
- if audios is not None and (audio_inputs := audios[i]) is not None:
434
- # HACK - not all processors take sampling_rate; we should
435
- # clean this up in the future.
436
- if len(audio_inputs) == 2:
437
- audio, sr = audio_inputs
438
- processor_kwargs["audio"] = audio
439
- processor_kwargs["sampling_rate"] = sr
440
- else:
441
- processor_kwargs["audio"] = audio_inputs
442
-
443
- inputs = self.processor(**processor_kwargs)
444
- if isinstance(inputs, BatchFeature):
445
- inputs = inputs.to(dtype=self.dtype)
446
-
447
- all_inputs.append(inputs)
427
+ if isinstance(prompt, str):
428
+ processor_kwargs: dict[str, Any] = {
429
+ "text": prompt,
430
+ "return_tensors": "pt",
431
+ }
432
+ if images is not None and (image := images[i]) is not None:
433
+ processor_kwargs["images"] = image
434
+ if videos is not None and (video := videos[i]) is not None:
435
+ processor_kwargs["videos"] = video
436
+ if audios is not None and (audio_inputs := audios[i]) is not None:
437
+ # HACK - not all processors take sampling_rate; we should
438
+ # clean this up in the future.
439
+ if len(audio_inputs) == 2:
440
+ audio, sr = audio_inputs
441
+ processor_kwargs["audio"] = audio
442
+ processor_kwargs["sampling_rate"] = sr
443
+ else:
444
+ processor_kwargs["audio"] = audio_inputs
445
+
446
+ inputs = self.processor(**processor_kwargs)
447
+ if isinstance(inputs, BatchFeature):
448
+ inputs = inputs.to(dtype=self.dtype)
449
+ all_inputs.append(inputs)
450
+ else:
451
+ # check that prompt is (batched) list of integers (token ids)
452
+ if not is_list_of(prompt, typ=int, check="all"):
453
+ raise ValueError(
454
+ "Prompt must be a list of ints corresponding to the prompt token ids."
455
+ )
456
+ # check that no multimodal input is provided
457
+ if images or videos or audios:
458
+ raise ValueError(
459
+ "When providing prompt token ids multimodal inputs are not supported."
460
+ )
461
+ input_dict = {
462
+ "input_ids": torch.tensor(prompt, dtype=torch.long).unsqueeze(0),
463
+ }
464
+ all_inputs.append(input_dict)
465
466
return all_inputs
467
@@ -477,7 +494,7 @@ class HfRunner:
494
495
def generate(
496
self,
480
- prompts: list[str],
497
+ prompts: Union[list[str], list[list[int]]],
498
images: Optional[PromptImageInput] = None,
499
videos: Optional[PromptVideoInput] = None,
500
audios: Optional[PromptAudioInput] = None,
@@ -505,7 +522,7 @@ class HfRunner:
522
523
def generate_greedy(
524
self,
508
- prompts: list[str],
525
+ prompts: Union[list[str], list[list[int]]],
526
max_tokens: int,
527
images: Optional[PromptImageInput] = None,
528
videos: Optional[PromptVideoInput] = None,
@@ -807,7 +824,7 @@ class VllmRunner:
824
825
def generate(
826
self,
810
- prompts: Union[list[str], list[torch.Tensor]],
827
+ prompts: Union[list[str], list[torch.Tensor], list[list[int]]],
828
sampling_params: SamplingParams,
829
images: Optional[PromptImageInput] = None,
830
videos: Optional[PromptVideoInput] = None,
@@ -877,7 +894,7 @@ class VllmRunner:
894
895
def generate_greedy(
896
self,
880
- prompts: Union[list[str], list[torch.Tensor]],
897
+ prompts: Union[list[str], list[torch.Tensor], list[list[int]]],
898
max_tokens: int,
899
images: Optional[PromptImageInput] = None,
900
videos: Optional[PromptVideoInput] = None,
tests/v1/e2e/test_context_length.py
+17
-34
@@ -23,13 +23,10 @@ the 1st will be sampled after the prefill and the 2nd after the first decode
23
"""
24
25
import pytest
26
-import torch
27
-from transformers import AutoModelForCausalLM
26
27
+from tests.conftest import HfRunner, VllmRunner
28
from tests.models.utils import check_outputs_equal
29
from tests.utils import create_new_process_for_each_test
31
-from vllm import LLM, SamplingParams
32
-from vllm.inputs import TokensPrompt
30
31
32
@create_new_process_for_each_test()
@@ -43,6 +40,8 @@ from vllm.inputs import TokensPrompt
40
)
41
def test_max_context_length(
42
model: str,
43
+ vllm_runner: type[VllmRunner],
44
+ hf_runner: type[HfRunner],
45
prompt_len: int,
46
max_tokens: int,
47
) -> None:
@@ -57,42 +56,26 @@ def test_max_context_length(
56
# Construct a prompt of size prompt_len
57
prompt_ids = [[43] * prompt_len]
58
60
- # Generate max_tokens new tokens deterministically.
61
- sampling_params = [
62
- SamplingParams(max_tokens=max_tokens, temperature=0.0, ignore_eos=True)
63
- ]
64
-
59
# --- vLLM generation ---
66
- llm = LLM(
67
- model=model,
68
- tokenizer=model,
60
+ with vllm_runner(
61
+ model_name=model,
62
+ tokenizer_name=model,
63
max_model_len=2048,
64
max_num_seqs=1,
65
tensor_parallel_size=1,
72
- )
73
-
74
- vllm_token_prompts = [TokensPrompt(prompt_token_ids=prompt_ids[0])]
75
- vllm_results = llm.generate(vllm_token_prompts, sampling_params)
76
-
77
- vllm_output_ids = vllm_results[0].outputs[0].token_ids
78
-
79
- # --- HuggingFace generation ---
80
- with torch.no_grad():
81
- hf_model = AutoModelForCausalLM.from_pretrained(model)
82
-
83
- # HF expects a tensor of input ids shaped (batch, seq_len).
84
- hf_input_tokens = torch.tensor(prompt_ids[0]).unsqueeze(0)
85
-
66
+ ) as vllm_model:
67
# Generate max_tokens new tokens deterministically.
87
- hf_generated = hf_model.generate(
88
- hf_input_tokens,
89
- do_sample=False,
90
- min_new_tokens=max_tokens,
91
- max_new_tokens=max_tokens,
92
- )
68
+ vllm_outputs = vllm_model.generate_greedy(prompt_ids, max_tokens)
69
94
- # HF returns the prompt + generated tokens. Slice off the prompt.
95
- hf_output_ids = hf_generated.cpu().tolist()[0][len(prompt_ids[0]) :]
70
+ # --- HuggingFace generation ---
71
+ with hf_runner(
72
+ model_name=model,
73
+ ) as hf_model:
74
+ hf_outputs = hf_model.generate_greedy(prompt_ids, max_tokens)
75
+
76
+ # vLLM and HF runners return prompt + generated tokens. Slice off the prompt.
77
+ vllm_output_ids = vllm_outputs[0][0][prompt_len:]
78
+ hf_output_ids = hf_outputs[0][0][prompt_len:]
79
80
# check that exactly max_tokens tokens were generated with vLLM and HF
81
assert len(vllm_output_ids) == len(hf_output_ids) == max_tokens