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, 12:15 UTC
[Core] Gate `prompt_embeds` behind a feature flag (#17607)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
May 4, 2025 at 00:19 UTC
887d7af882d2bcb0ce4884d8369418e226add0c4
8 files changed
+84
-33
tests/engine/test_options.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from contextlib import nullcontext
3
+
4
+import pytest
5
+
6
+from vllm.entrypoints.llm import LLM
7
+from vllm.sampling_params import SamplingParams
8
+
9
+
10
+@pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
11
+def test_skip_tokenizer_initialization(model: str):
12
+ # This test checks if the flag skip_tokenizer_init skips the initialization
13
+ # of tokenizer and detokenizer. The generated output is expected to contain
14
+ # token ids.
15
+ llm = LLM(
16
+ model=model,
17
+ skip_tokenizer_init=True,
18
+ enforce_eager=True,
19
+ )
20
+ sampling_params = SamplingParams(prompt_logprobs=True, detokenize=True)
21
+
22
+ with pytest.raises(ValueError, match="cannot pass text prompts when"):
23
+ llm.generate("abc", sampling_params)
24
+
25
+ outputs = llm.generate({"prompt_token_ids": [1, 2, 3]},
26
+ sampling_params=sampling_params)
27
+ assert len(outputs) > 0
28
+ completions = outputs[0].outputs
29
+ assert len(completions) > 0
30
+ assert completions[0].text == ""
31
+ assert completions[0].token_ids
32
+
33
+
34
+@pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
35
+@pytest.mark.parametrize("enable_prompt_embeds", [True, False])
36
+def test_enable_prompt_embeds(hf_runner, model: str,
37
+ enable_prompt_embeds: bool):
38
+ prompt = "abc"
39
+
40
+ with hf_runner(model) as hf_model:
41
+ token_ids = hf_model.tokenizer(prompt, return_tensors="pt").input_ids
42
+ token_ids = token_ids.to(hf_model.model.device)
43
+
44
+ embed_layer = hf_model.model.get_input_embeddings()
45
+ prompt_embeds = embed_layer(token_ids).squeeze(0)
46
+
47
+ ctx = (nullcontext() if enable_prompt_embeds else pytest.raises(
48
+ ValueError, match="set `--enable-prompt-embeds`"))
49
+
50
+ # This test checks if the flag skip_tokenizer_init skips the initialization
51
+ # of tokenizer and detokenizer. The generated output is expected to contain
52
+ # token ids.
53
+ llm = LLM(
54
+ model=model,
55
+ enable_prompt_embeds=enable_prompt_embeds,
56
+ enforce_eager=True,
57
+ )
58
+
59
+ with ctx:
60
+ llm.generate({"prompt_embeds": prompt_embeds})
tests/engine/test_skip_tokenizer_init.py
deleted
-29
@@ -1,29 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-
3
-import pytest
4
-
5
-from vllm.entrypoints.llm import LLM
6
-from vllm.sampling_params import SamplingParams
7
-
8
-
9
-@pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
10
-def test_skip_tokenizer_initialization(model: str):
11
- # This test checks if the flag skip_tokenizer_init skips the initialization
12
- # of tokenizer and detokenizer. The generated output is expected to contain
13
- # token ids.
14
- llm = LLM(
15
- model=model,
16
- skip_tokenizer_init=True,
17
- )
18
- sampling_params = SamplingParams(prompt_logprobs=True, detokenize=True)
19
-
20
- with pytest.raises(ValueError, match="cannot pass text prompts when"):
21
- llm.generate("abc", sampling_params)
22
-
23
- outputs = llm.generate({"prompt_token_ids": [1, 2, 3]},
24
- sampling_params=sampling_params)
25
- assert len(outputs) > 0
26
- completions = outputs[0].outputs
27
- assert len(completions) > 0
28
- assert completions[0].text == ""
29
- assert completions[0].token_ids
tests/models/language/generation/test_common.py
+6
-2
@@ -109,12 +109,15 @@ def test_models(hf_runner, vllm_runner, example_prompts, model: str,
109
# in parts of the operators
110
pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
111
112
+ use_prompt_embeds = os.getenv("VLLM_USE_V1") == "0"
113
+
114
with hf_runner(model) as hf_model:
115
hf_outputs = hf_model.generate_greedy_logprobs_limit(
116
example_prompts, max_tokens, num_logprobs)
117
116
- prompt_embeds: Optional[list[torch.Tensor]] = [] if os.getenv(
117
- "VLLM_USE_V1") == "0" else None
118
+ prompt_embeds: Optional[list[torch.Tensor]] = ([] if use_prompt_embeds
119
+ else None)
120
+
121
prompt_token_ids = []
122
for prompt in example_prompts:
123
token_ids = hf_model.tokenizer(prompt,
@@ -131,6 +134,7 @@ def test_models(hf_runner, vllm_runner, example_prompts, model: str,
134
tokenizer_mode=model_info.tokenizer_mode,
135
trust_remote_code=model_info.trust_remote_code,
136
max_num_seqs=2,
137
+ enable_prompt_embeds=use_prompt_embeds,
138
) as vllm_model:
139
vllm_outputs = vllm_model.generate_greedy_logprobs(
140
example_prompts, max_tokens, num_logprobs)
tests/worker/test_model_runner.py
+3
@@ -43,6 +43,7 @@ def test_prepare_prompt(batch_size, use_prompt_embeds, monkeypatch):
43
max_num_batched_tokens=100000,
44
max_num_seqs=100000,
45
enable_chunked_prefill=False,
46
+ enable_prompt_embeds=True,
47
)
48
49
seq_lens: list[int] = []
@@ -179,6 +180,7 @@ def test_prepare_decode_cuda_graph(batch_size, use_prompt_embeds, monkeypatch):
180
max_num_batched_tokens=100000,
181
max_num_seqs=100000,
182
enable_chunked_prefill=False,
183
+ enable_prompt_embeds=True,
184
)
185
186
context_lens: list[int] = []
@@ -359,6 +361,7 @@ def test_hybrid_batches(batch_size, enforce_eager, use_prompt_embeds,
361
max_num_batched_tokens=100000,
362
max_num_seqs=100000,
363
enable_chunked_prefill=True,
364
+ enable_prompt_embeds=True,
365
)
366
367
# Add prefill requests.
vllm/config.py
+4
@@ -321,6 +321,10 @@ class ModelConfig:
321
"""Skip initialization of tokenizer and detokenizer. Expects valid
322
`prompt_token_ids` and `None` for prompt from the input. The generated
323
output will contain token ids."""
324
+ enable_prompt_embeds: bool = False
325
+ """If `True`, enables passing text embeddings as inputs via the
326
+ `prompt_embeds` key. Note that enabling this will double the time required
327
+ for graph compilation."""
328
served_model_name: Optional[Union[str, list[str]]] = None
329
"""The model name(s) used in the API. If multiple names are provided, the
330
server will respond to any of the provided names. The model name in the
vllm/engine/arg_utils.py
+4
@@ -234,6 +234,7 @@ class EngineArgs:
234
hf_config_path: Optional[str] = ModelConfig.hf_config_path
235
task: TaskOption = ModelConfig.task
236
skip_tokenizer_init: bool = ModelConfig.skip_tokenizer_init
237
+ enable_prompt_embeds: bool = ModelConfig.enable_prompt_embeds
238
tokenizer_mode: TokenizerMode = ModelConfig.tokenizer_mode
239
trust_remote_code: bool = ModelConfig.trust_remote_code
240
allowed_local_media_path: str = ModelConfig.allowed_local_media_path
@@ -445,6 +446,8 @@ class EngineArgs:
446
**model_kwargs["disable_cascade_attn"])
447
model_group.add_argument("--skip-tokenizer-init",
448
**model_kwargs["skip_tokenizer_init"])
449
+ model_group.add_argument("--enable-prompt-embeds",
450
+ **model_kwargs["enable_prompt_embeds"])
451
model_group.add_argument("--served-model-name",
452
**model_kwargs["served_model_name"])
453
# This one is a special case because it is the
@@ -874,6 +877,7 @@ class EngineArgs:
877
disable_sliding_window=self.disable_sliding_window,
878
disable_cascade_attn=self.disable_cascade_attn,
879
skip_tokenizer_init=self.skip_tokenizer_init,
880
+ enable_prompt_embeds=self.enable_prompt_embeds,
881
served_model_name=self.served_model_name,
882
limit_mm_per_prompt=self.limit_mm_per_prompt,
883
use_async_output_proc=not self.disable_async_output_proc,
vllm/inputs/preprocess.py
+4
-1
@@ -303,8 +303,11 @@ class InputPreprocessor:
303
self,
304
parsed_content: EmbedsPrompt,
305
) -> EmbedsInputs:
306
+ if not self.model_config.enable_prompt_embeds:
307
+ raise ValueError("You must set `--enable-prompt-embeds` to input "
308
+ "`prompt_embeds`.")
309
if envs.VLLM_USE_V1:
307
- raise ValueError("prompt_embeds is only available in V0.")
310
+ raise ValueError("`prompt_embeds` is only available in V0.")
311
312
prompt_embeds = parsed_content["prompt_embeds"]
313
vllm/worker/model_runner.py
+3
-1
@@ -1565,7 +1565,9 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1565
# product.
1566
cudagraph_capture_sizes = self.vllm_config.compilation_config\
1567
.cudagraph_capture_sizes
1568
- cudagraph_inputs_embeds = (True, False)
1568
+ cudagraph_inputs_embeds = ((
1569
+ True, False) if self.model_config.enable_prompt_embeds else
1570
+ (False, ))
1571
compilation_cases = itertools.product(
1572
cudagraph_capture_sizes,
1573
cudagraph_inputs_embeds,