@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 221cfc2fea
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, 09:45 UTC

Feature/vllm/input embedding completion api (#17590)

Signed-off-by: Andrew Sansom <andrew@protopia.ai> Signed-off-by: Nan2018 <nan@protopia.ai> Co-authored-by: 临景 <linjing.yx@alibaba-inc.com> Co-authored-by: Bryce1010 <bryceyx@gmail.com> Co-authored-by: Andrew Sansom <andrew@protopia.ai> Co-authored-by: Andrew Sansom <qthequartermasterman@gmail.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

Nan Qin committed May 18, 2025 at 22:18 UTC 221cfc2feaf6ad60b59882249de80eeba66446a6
10 files changed +637 -40
docs/source/index.md
+1
@@ -119,6 +119,7 @@ serving/offline_inference
119 serving/openai_compatible_server
120 serving/serve_args
121 serving/multimodal_inputs
122 +serving/prompt_embeds
123 serving/distributed_serving
124 serving/metrics
125 serving/engine_args
docs/source/serving/prompt_embeds.md new
+142
@@ -0,0 +1,142 @@
1 +# Prompt Embedding Inputs
2 +
3 +This page teaches you how to pass prompt embedding inputs to vLLM.
4 +
5 +## What are prompt embeddings?
6 +
7 +The traditional flow of text data for a Large Language Model goes from text to token ids (via a tokenizer) then from token ids to prompt embeddings. For a traditional decoder-only model (such as meta-llama/Llama-3.1-8B-Instruct), this step of converting token ids to prompt embeddings happens via a look-up from a learned embedding matrix, but the model is not limited to processing only the embeddings corresponding to its token vocabulary.
8 +
9 +:::{note}
10 +Prompt embeddings are currently only supported in the v0 engine.
11 +:::
12 +
13 +## Offline Inference
14 +
15 +To input multi-modal data, follow this schema in {class}`vllm.inputs.EmbedsPrompt`:
16 +
17 +- `prompt_embeds`: A torch tensor representing a sequence of prompt/token embeddings. This has the shape (sequence_length, hidden_size), where sequence length is the number of tokens embeddings and hidden_size is the hidden size (embedding size) of the model.
18 +
19 +### Hugging Face Transformers Inputs
20 +
21 +You can pass prompt embeddings from Hugging Face Transformers models to the `'prompt_embeds'` field of the prompt embedding dictionary, as shown in the following examples:
22 +
23 +```python
24 +from vllm import LLM
25 +import transformers
26 +
27 +model_name = "meta-llama/Llama-3.2-1B-Instruct"
28 +
29 +# Transformers
30 +tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
31 +transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
32 +
33 +llm = LLM(model=model_name, enable_prompt_embeds=True)
34 +
35 +# Refer to the HuggingFace repo for the correct format to use
36 +chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
37 +token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt')
38 +
39 +prompt_embeds = embedding_layer(token_ids).squeeze(0)
40 +
41 +# Single prompt inference
42 +outputs = llm.generate({
43 + "prompt_embeds": prompt_embeds,
44 +})
45 +
46 +for o in outputs:
47 + generated_text = o.outputs[0].text
48 + print(generated_text)
49 +
50 +# Batch inference
51 +
52 +chats = [
53 + [{"role": "user", "content": "Please tell me about the capital of France."}],
54 + [{"role": "user", "content": "When is the day longest during the year?"}],
55 + [{"role": "user", "content": "Where is bigger, the moon or the sun?"}]
56 +]
57 +
58 +token_ids_list = [
59 + tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt') for chat in chats
60 +]
61 +prompt_embeds_list = [embedding_layer(token_ids).squeeze(0) for token_ids in token_ids_list]
62 +
63 +outputs = llm.generate(
64 + [
65 + {
66 + "prompt_embeds": prompt_embeds,
67 + } for prompt_embeds in prompt_embeds_list
68 + ]
69 +)
70 +
71 +for o in outputs:
72 + generated_text = o.outputs[0].text
73 + print(generated_text)
74 +```
75 +
76 +## Online Serving
77 +
78 +Our OpenAI-compatible server accepts prompt embeddings inputs via the [Completions API](https://platform.openai.com/docs/api-reference/completions). Prompt embeddings inputs are added via a new `'prompt_embeds'` key in the JSON package.
79 +
80 +When a mixture of `'prompt_embeds'` and `'prompt'` inputs are provided in a single request, the prompt embeds are always returned first.
81 +
82 +Prompt embeddings are passed in as base64 encoded torch tensors.
83 +
84 +### Transformers Inputs via OpenAI Client
85 +
86 +First, launch the OpenAI-compatible server:
87 +
88 +```bash
89 +vllm serve meta-llama/Llama-3.2-1B-Instruct --task generate \
90 + --max-model-len 4096 --enable-prompt-embeds
91 +```
92 +
93 +Then, you can use the OpenAI client as follows:
94 +
95 +```python
96 +from openai import OpenAI
97 +import transformers
98 +import torch
99 +
100 +openai_api_key = "EMPTY"
101 +openai_api_base = "http://localhost:8000/v1"
102 +
103 +client = OpenAI(
104 + api_key=openai_api_key,
105 + base_url=openai_api_base,
106 +)
107 +
108 +model_name = "meta-llama/Llama-3.2-1B-Instruct"
109 +
110 +# Transformers
111 +tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
112 +transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
113 +
114 +
115 +# Refer to the HuggingFace repo for the correct format to use
116 +chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
117 +token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors='pt')
118 +
119 +prompt_embeds = embedding_layer(token_ids).squeeze(0)
120 +
121 +# Prompt embeddings
122 +buffer = io.BytesIO()
123 +torch.save(prompt_embeds, buffer)
124 +buffer.seek(0)
125 +binary_data = buffer.read()
126 +encoded_embeds = base64.b64encode(binary_data).decode('utf-8')
127 +
128 +
129 +completion = client_with_prompt_embeds.completions.create(
130 + model=model_name,
131 + # NOTE: The OpenAI client does not allow `None` as an input to
132 + # `prompt`. Use an empty string if you have no text prompts.
133 + prompt="",
134 + max_tokens=5,
135 + temperature=0.0,
136 + # NOTE: The OpenAI client allows passing in extra JSON body via the
137 + # `extra_body` argument.
138 + extra_body={"prompt_embeds": encoded_embeds}
139 +)
140 +
141 +print(completion.choices[0].text)
142 +```
tests/entrypoints/openai/test_completion_with_prompt_embeds.py new
+257
@@ -0,0 +1,257 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import base64
4 +import io
5 +import shutil
6 +from tempfile import TemporaryDirectory
7 +
8 +import openai # use the official client for correctness check
9 +import pytest
10 +import pytest_asyncio
11 +import torch
12 +# downloading lora to test lora requests
13 +from huggingface_hub import snapshot_download
14 +from openai import BadRequestError
15 +from transformers import AutoConfig, AutoTokenizer
16 +
17 +from ...utils import RemoteOpenAIServer
18 +
19 +# any model with a chat template should work here
20 +MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
21 +LORA_NAME = "typeof/zephyr-7b-beta-lora"
22 +
23 +CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
24 +
25 +
26 +@pytest.fixture(scope="module")
27 +def zephyr_lora_files():
28 + return snapshot_download(repo_id=LORA_NAME)
29 +
30 +
31 +@pytest.fixture(scope="module")
32 +def zephyr_lora_added_tokens_files(zephyr_lora_files):
33 + tmp_dir = TemporaryDirectory()
34 + tmp_model_dir = f"{tmp_dir.name}/zephyr"
35 + shutil.copytree(zephyr_lora_files, tmp_model_dir)
36 + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
37 + # Copy tokenizer to adapter and add some unique tokens
38 + # 32000, 32001, 32002
39 + added = tokenizer.add_tokens(["vllm1", "vllm2", "vllm3"],
40 + special_tokens=True)
41 + assert added == 3
42 + tokenizer.save_pretrained(tmp_model_dir)
43 + yield tmp_model_dir
44 + tmp_dir.cleanup()
45 +
46 +
47 +@pytest.fixture(scope="module")
48 +def default_server_args(
49 + zephyr_lora_files,
50 + zephyr_lora_added_tokens_files,
51 +) -> list[str]:
52 + return [
53 + # use half precision for speed and memory savings in CI environment
54 + "--dtype",
55 + "bfloat16",
56 + "--max-model-len",
57 + "8192",
58 + "--max-num-seqs",
59 + "128",
60 + "--enforce-eager",
61 + # Prompt Embeds server args
62 + "--enable-prompt-embeds",
63 + "--no-enable-chunked-prefill",
64 + ]
65 +
66 +
67 +@pytest.fixture(scope="module",
68 + params=["", "--disable-frontend-multiprocessing"])
69 +def server_with_prompt_embeds(default_server_args, request):
70 + if request.param:
71 + default_server_args.append(request.param)
72 +
73 + with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
74 + yield remote_server
75 +
76 +
77 +@pytest_asyncio.fixture
78 +async def client_with_prompt_embeds(server_with_prompt_embeds):
79 + async with server_with_prompt_embeds.get_async_client() as async_client:
80 + yield async_client
81 +
82 +
83 +def create_dummy_embeds(num_tokens: int = 5) -> str:
84 + """Create dummy embeddings and return them as base64 encoded string."""
85 + dummy_embeds = torch.randn(num_tokens, CONFIG.hidden_size)
86 + buffer = io.BytesIO()
87 + torch.save(dummy_embeds, buffer)
88 + return base64.b64encode(buffer.getvalue()).decode('utf-8')
89 +
90 +
91 +@pytest.mark.asyncio
92 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
93 +async def test_completions_with_prompt_embeds(
94 + client_with_prompt_embeds: openai.AsyncOpenAI, model_name: str):
95 + # Test case: Single prompt embeds input
96 + encoded_embeds = create_dummy_embeds()
97 + completion = await client_with_prompt_embeds.completions.create(
98 + model=model_name,
99 + prompt="", # Add empty prompt as required parameter
100 + max_tokens=5,
101 + temperature=0.0,
102 + extra_body={"prompt_embeds": encoded_embeds})
103 + assert len(completion.choices[0].text) >= 1
104 + assert completion.choices[0].prompt_logprobs is None
105 +
106 + # Test case: batch completion with prompt_embeds
107 + encoded_embeds2 = create_dummy_embeds()
108 + completion = await client_with_prompt_embeds.completions.create(
109 + model=model_name,
110 + prompt="", # Add empty prompt as required parameter
111 + max_tokens=5,
112 + temperature=0.0,
113 + extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]})
114 + assert len(completion.choices) == 2
115 + assert len(completion.choices[0].text) >= 1
116 + assert len(completion.choices[1].text) >= 1
117 +
118 + # Test case: streaming with prompt_embeds
119 + encoded_embeds = create_dummy_embeds()
120 + single_completion = await client_with_prompt_embeds.completions.create(
121 + model=model_name,
122 + prompt="", # Add empty prompt as required parameter
123 + max_tokens=5,
124 + temperature=0.0,
125 + extra_body={"prompt_embeds": encoded_embeds})
126 + single_output = single_completion.choices[0].text
127 +
128 + stream = await client_with_prompt_embeds.completions.create(
129 + model=model_name,
130 + prompt="", # Add empty prompt as required parameter
131 + max_tokens=5,
132 + temperature=0.0,
133 + stream=True,
134 + extra_body={"prompt_embeds": encoded_embeds})
135 + chunks = []
136 + finish_reason_count = 0
137 + async for chunk in stream:
138 + chunks.append(chunk.choices[0].text)
139 + if chunk.choices[0].finish_reason is not None:
140 + finish_reason_count += 1
141 + assert finish_reason_count == 1
142 + assert chunk.choices[0].finish_reason == "length"
143 + assert chunk.choices[0].text
144 + assert "".join(chunks) == single_output
145 +
146 + # Test case: batch streaming with prompt_embeds
147 + encoded_embeds2 = create_dummy_embeds()
148 + stream = await client_with_prompt_embeds.completions.create(
149 + model=model_name,
150 + prompt="", # Add empty prompt as required parameter
151 + max_tokens=5,
152 + temperature=0.0,
153 + stream=True,
154 + extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]})
155 + chunks_stream_embeds: list[list[str]] = [[], []]
156 + finish_reason_count = 0
157 + async for chunk in stream:
158 + chunks_stream_embeds[chunk.choices[0].index].append(
159 + chunk.choices[0].text)
160 + if chunk.choices[0].finish_reason is not None:
161 + finish_reason_count += 1
162 + assert finish_reason_count == 2
163 + assert chunk.choices[0].finish_reason == "length"
164 + assert chunk.choices[0].text
165 + assert len(chunks_stream_embeds[0]) > 0
166 + assert len(chunks_stream_embeds[1]) > 0
167 +
168 + # Test case: mixed text and prompt_embeds
169 + encoded_embeds = create_dummy_embeds()
170 + completion_mixed = await client_with_prompt_embeds.completions.create(
171 + model=model_name,
172 + prompt="This is a prompt",
173 + max_tokens=5,
174 + temperature=0.0,
175 + extra_body={"prompt_embeds": encoded_embeds})
176 + assert len(completion.choices) == 2
177 + completion_text_only = await client_with_prompt_embeds.completions.create(
178 + model=model_name,
179 + prompt="This is a prompt",
180 + max_tokens=5,
181 + temperature=0.0,
182 + )
183 + completion_embeds_only = await client_with_prompt_embeds.completions.create(
184 + model=model_name,
185 + prompt="",
186 + max_tokens=5,
187 + temperature=0.0,
188 + extra_body={"prompt_embeds": encoded_embeds})
189 + # Embeddings responses should be handled first
190 + assert completion_mixed.choices[0].text == completion_embeds_only.choices[
191 + 0].text
192 + assert completion_mixed.choices[1].text == completion_text_only.choices[
193 + 0].text
194 +
195 +
196 +@pytest.mark.asyncio
197 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
198 +async def test_completions_errors_with_prompt_embeds(
199 + client_with_prompt_embeds: openai.AsyncOpenAI, model_name: str):
200 + # Test error case: invalid prompt_embeds
201 + with pytest.raises(BadRequestError):
202 + await client_with_prompt_embeds.completions.create(
203 + prompt="",
204 + model=model_name,
205 + max_tokens=5,
206 + temperature=0.0,
207 + extra_body={"prompt_embeds": "invalid_base64"})
208 +
209 +
210 +@pytest.mark.asyncio
211 +@pytest.mark.parametrize("logprobs_arg", [1, 0])
212 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
213 +async def test_completions_with_logprobs_and_prompt_embeds(
214 + client_with_prompt_embeds: openai.AsyncOpenAI, logprobs_arg: int,
215 + model_name: str):
216 + # Test case: Logprobs using prompt_embeds
217 + encoded_embeds = create_dummy_embeds()
218 + completion = await client_with_prompt_embeds.completions.create(
219 + model=model_name,
220 + prompt="", # Add empty prompt as required parameter
221 + max_tokens=5,
222 + temperature=0.0,
223 + echo=False,
224 + logprobs=logprobs_arg,
225 + extra_body={"prompt_embeds": encoded_embeds})
226 +
227 + logprobs = completion.choices[0].logprobs
228 + assert logprobs is not None
229 + assert len(logprobs.text_offset) == 5
230 + assert len(logprobs.token_logprobs) == 5
231 + assert len(logprobs.top_logprobs) == 5
232 + for top_logprobs in logprobs.top_logprobs[1:]:
233 + assert max(logprobs_arg, 1) <= len(top_logprobs) <= logprobs_arg + 1
234 + assert len(logprobs.tokens) == 5
235 +
236 + # Test case: Log probs with batch completion and prompt_embeds
237 + encoded_embeds2 = create_dummy_embeds()
238 + completion = await client_with_prompt_embeds.completions.create(
239 + model=model_name,
240 + prompt="", # Add empty prompt as required parameter
241 + max_tokens=5,
242 + temperature=0.0,
243 + echo=False,
244 + logprobs=logprobs_arg,
245 + extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]})
246 +
247 + assert len(completion.choices) == 2
248 + for choice in completion.choices:
249 + logprobs = choice.logprobs
250 + assert logprobs is not None
251 + assert len(logprobs.text_offset) == 5
252 + assert len(logprobs.token_logprobs) == 5
253 + assert len(logprobs.top_logprobs) == 5
254 + for top_logprobs in logprobs.top_logprobs[1:]:
255 + assert max(logprobs_arg,
256 + 1) <= len(top_logprobs) <= logprobs_arg + 1
257 + assert len(logprobs.tokens) == 5
vllm/entrypoints/logger.py
+7 -2
@@ -2,6 +2,8 @@
2
3 from typing import Optional, Union
4
5 +import torch
6 +
7 from vllm.logger import init_logger
8 from vllm.lora.request import LoRARequest
9 from vllm.pooling_params import PoolingParams
@@ -23,6 +25,7 @@ class RequestLogger:
25 request_id: str,
26 prompt: Optional[str],
27 prompt_token_ids: Optional[list[int]],
28 + prompt_embeds: Optional[torch.Tensor],
29 params: Optional[Union[SamplingParams, PoolingParams,
30 BeamSearchParams]],
31 lora_request: Optional[LoRARequest],
@@ -39,6 +42,8 @@ class RequestLogger:
42 logger.info(
43 "Received request %s: prompt: %r, "
44 "params: %s, prompt_token_ids: %s, "
45 + "prompt_embeds shape: %s, "
46 "lora_request: %s, prompt_adapter_request: %s.", request_id,
43 - prompt, params, prompt_token_ids, lora_request,
44 - prompt_adapter_request)
47 + prompt, params, prompt_token_ids,
48 + prompt_embeds.shape if prompt_embeds is not None else None,
49 + lora_request, prompt_adapter_request)
vllm/entrypoints/openai/cli_args.py
+3
@@ -286,6 +286,9 @@ def validate_parsed_serve_args(args: argparse.Namespace):
286 if args.enable_auto_tool_choice and not args.tool_call_parser:
287 raise TypeError("Error: --enable-auto-tool-choice requires "
288 "--tool-call-parser")
289 + if args.enable_prompt_embeds and args.enable_prompt_adapter:
290 + raise ValueError(
291 + "Cannot use prompt embeds and prompt adapter at the same time.")
292
293
294 def log_non_default_args(args: argparse.Namespace):
vllm/entrypoints/openai/protocol.py
+10 -1
@@ -745,7 +745,8 @@ class CompletionRequest(OpenAIBaseModel):
745 # Ordered by official OpenAI API documentation
746 # https://platform.openai.com/docs/api-reference/completions/create
747 model: Optional[str] = None
748 - prompt: Union[list[int], list[list[int]], str, list[str]]
748 + prompt: Optional[Union[list[int], list[list[int]], str, list[str]]] = None
749 + prompt_embeds: Optional[Union[bytes, list[bytes]]] = None
750 best_of: Optional[int] = None
751 echo: Optional[bool] = False
752 frequency_penalty: Optional[float] = 0.0
@@ -1025,6 +1026,14 @@ class CompletionRequest(OpenAIBaseModel):
1026
1027 return data
1028
1029 + @model_validator(mode="before")
1030 + @classmethod
1031 + def validate_prompt_and_prompt_embeds(cls, data):
1032 + if data.get("prompt") is None and data.get("prompt_embeds") is None:
1033 + raise ValueError(
1034 + "At least one of `prompt` or `prompt_embeds` must be set.")
1035 + return data
1036 +
1037
1038 class EmbeddingCompletionRequest(OpenAIBaseModel):
1039 # Ordered by official OpenAI API documentation
vllm/entrypoints/openai/serving_completion.py
+39 -6
@@ -8,6 +8,7 @@ from typing import Optional, Union, cast
8
9 import jinja2
10 from fastapi import Request
11 +from typing_extensions import assert_never
12
13 from vllm.config import ModelConfig
14 from vllm.engine.protocol import EngineClient
@@ -25,8 +26,11 @@ from vllm.entrypoints.openai.protocol import (CompletionLogProbs,
26 UsageInfo)
27 # yapf: enable
28 from vllm.entrypoints.openai.serving_engine import (OpenAIServing,
28 - clamp_prompt_logprobs)
29 + clamp_prompt_logprobs,
30 + is_text_tokens_prompt)
31 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
32 +from vllm.inputs.data import (EmbedsPrompt, TokensPrompt, is_embeds_prompt,
33 + is_tokens_prompt)
34 from vllm.logger import init_logger
35 from vllm.outputs import RequestOutput
36 from vllm.sampling_params import BeamSearchParams, SamplingParams
@@ -90,6 +94,10 @@ class OpenAIServingCompletion(OpenAIServing):
94 return self.create_error_response(
95 "suffix is not currently supported")
96
97 + if request.echo and request.prompt_embeds is not None:
98 + return self.create_error_response(
99 + "Echo is unsupported with prompt embeds.")
100 +
101 request_id = f"cmpl-{self._base_request_id(raw_request)}"
102 created_time = int(time.time())
103
@@ -130,8 +138,24 @@ class OpenAIServingCompletion(OpenAIServing):
138 try:
139 for i, engine_prompt in enumerate(engine_prompts):
140 sampling_params: Union[SamplingParams, BeamSearchParams]
133 - default_max_tokens = self.max_model_len - len(
134 - engine_prompt["prompt_token_ids"])
141 + # Mypy does not infer that engine_prompt will have only one of
142 + # "prompt_token_ids" or "prompt_embeds" defined, and both of
143 + # these as Union[object, the expected type], where it infers
144 + # object if engine_prompt is a subclass of one of the
145 + # typeddicts that defines both keys. Worse, because of
146 + # https://github.com/python/mypy/issues/8586, mypy does not
147 + # infer the type of engine_prompt correctly because of the
148 + # enumerate. So we need an unnecessary cast here.
149 + engine_prompt = cast(Union[EmbedsPrompt, TokensPrompt],
150 + engine_prompt)
151 + if is_embeds_prompt(engine_prompt):
152 + input_length = len(engine_prompt["prompt_embeds"])
153 + elif is_tokens_prompt(engine_prompt):
154 + input_length = len(engine_prompt["prompt_token_ids"])
155 + else:
156 + assert_never(engine_prompt)
157 + default_max_tokens = self.max_model_len - input_length
158 +
159 if request.use_beam_search:
160 sampling_params = request.to_beam_search_params(
161 default_max_tokens, self.default_sampling_params)
@@ -152,6 +176,11 @@ class OpenAIServingCompletion(OpenAIServing):
176 trace_headers = (None if raw_request is None else await
177 self._get_trace_headers(raw_request.headers))
178
179 + # Mypy inconsistently requires this second cast in different
180 + # environments. It shouldn't be necessary (redundant from above)
181 + # but pre-commit in CI fails without it.
182 + engine_prompt = cast(Union[EmbedsPrompt, TokensPrompt],
183 + engine_prompt)
184 if isinstance(sampling_params, BeamSearchParams):
185 generator = self.engine_client.beam_search(
186 prompt=engine_prompt,
@@ -211,7 +240,11 @@ class OpenAIServingCompletion(OpenAIServing):
240 # We did not pass it into vLLM engine to avoid being redundant
241 # with the inputs token IDs
242 if final_res.prompt is None:
214 - final_res.prompt = request_prompts[i]["prompt"]
243 + request_prompt = request_prompts[i]
244 + if is_text_tokens_prompt(request_prompt):
245 + final_res.prompt = request_prompt["prompt"]
246 + else:
247 + final_res.prompt = None
248
249 final_res_batch_checked = cast(list[RequestOutput],
250 final_res_batch)
@@ -276,8 +309,8 @@ class OpenAIServingCompletion(OpenAIServing):
309 prompt_text = res.prompt
310
311 # Prompt details are excluded from later streamed outputs
279 - if res.prompt_token_ids is not None:
280 - num_prompt_tokens[prompt_idx] = len(res.prompt_token_ids)
312 + if prompt_token_ids is not None:
313 + num_prompt_tokens[prompt_idx] = len(prompt_token_ids)
314
315 delta_token_ids: GenericSequence[int]
316 out_logprobs: Optional[GenericSequence[Optional[dict[
vllm/entrypoints/openai/serving_engine.py
+163 -28
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -
2 +import base64
3 +import io
4 import json
5 import sys
6 import time
@@ -8,11 +9,18 @@ from collections.abc import (AsyncGenerator, Iterable, Iterator, Mapping,
9 from concurrent.futures.thread import ThreadPoolExecutor
10 from http import HTTPStatus
11 from typing import (Annotated, Any, Callable, ClassVar, Generic, Optional,
11 - TypeVar, Union)
12 + TypeVar, Union, cast, overload)
13
14 +import torch
15 from fastapi import Request
16 from pydantic import BaseModel, ConfigDict, Field
17 from starlette.datastructures import Headers
18 +from typing_extensions import TypeIs
19 +
20 +if sys.version_info >= (3, 12):
21 + from typing import TypedDict
22 +else:
23 + from typing_extensions import TypedDict
24
25 if sys.version_info >= (3, 12):
26 from typing import TypedDict
@@ -53,7 +61,8 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
61 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
62 from vllm.entrypoints.openai.tool_parsers import ToolParser
63 # yapf: enable
56 -from vllm.inputs import TokensPrompt
64 +from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt
65 +from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
66 from vllm.inputs.parse import parse_and_batch_prompt
67 from vllm.logger import init_logger
68 from vllm.lora.request import LoRARequest
@@ -100,7 +109,22 @@ class TextTokensPrompt(TypedDict):
109 prompt_token_ids: list[int]
110
111
103 -RequestPrompt = Union[list[int], str, TextTokensPrompt]
112 +class EmbedsPrompt(TypedDict):
113 + prompt_embeds: torch.Tensor
114 +
115 +
116 +RequestPrompt = Union[list[int], str, TextTokensPrompt, EmbedsPrompt]
117 +
118 +
119 +def is_text_tokens_prompt(prompt: RequestPrompt) -> TypeIs[TextTokensPrompt]:
120 + return (isinstance(prompt, dict) and "prompt_token_ids" in prompt
121 + and "prompt_embeds" not in prompt)
122 +
123 +
124 +def is_embeds_prompt(prompt: RequestPrompt) -> TypeIs[EmbedsPrompt]:
125 + return (isinstance(prompt, dict) and "prompt_token_ids" not in prompt
126 + and "prompt_embeds" in prompt)
127 +
128
129 RequestT = TypeVar("RequestT", bound=AnyRequest)
130
@@ -112,8 +136,9 @@ class RequestProcessingMixin(BaseModel):
136 """
137 request_prompts: Optional[Sequence[RequestPrompt]] = \
138 Field(default_factory=list)
115 - engine_prompts: Optional[list[TokensPrompt]] = \
116 - Field(default_factory=list)
139 + engine_prompts: Optional[Union[list[EngineTokensPrompt],
140 + list[EngineEmbedsPrompt]]] = Field(
141 + default_factory=list)
142
143 model_config = ConfigDict(arbitrary_types_allowed=True)
144
@@ -311,6 +336,12 @@ class OpenAIServing:
336 lora_request=ctx.lora_request,
337 prompt_adapter_request=ctx.prompt_adapter_request)
338
339 + # Mypy has an existing bug related to inferring the variance of
340 + # TypedDicts with `builtins.enumerate`:
341 + # https://github.com/python/mypy/issues/8586#issuecomment-2867698435
342 + engine_prompt = cast(
343 + Union[EngineTokensPrompt, EngineEmbedsPrompt],
344 + engine_prompt)
345 generator = self.engine_client.encode(
346 engine_prompt,
347 pooling_params,
@@ -596,10 +627,11 @@ class OpenAIServing:
627 self,
628 request: AnyRequest,
629 tokenizer: AnyTokenizer,
599 - input_or_inputs: Union[str, list[str], list[int], list[list[int]]],
630 + input_or_inputs: Optional[Union[str, list[str], list[int],
631 + list[list[int]]]],
632 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
633 add_special_tokens: bool = True,
602 - ) -> list[TextTokensPrompt]:
634 + ) -> tuple[list[TextTokensPrompt], list[EmbedsPrompt]]:
635 """
636 Tokenize/detokenize depending on the input format.
637
@@ -607,11 +639,25 @@ class OpenAIServing:
639 , each input can be a string or array of tokens. Note that each request
640 can pass one or more inputs.
641 """
642 + inputs_embeds = list[EmbedsPrompt]()
643 + inputs_text = list[TextTokensPrompt]()
644 +
645 + if (isinstance(request, CompletionRequest)
646 + and request.prompt_embeds is not None):
647 + inputs_embeds.extend(
648 + self._load_prompt_embeds(request.prompt_embeds,
649 + truncate_prompt_tokens))
650 +
651 + # Empty prompts are okay as long as there are prompt embeddings
652 + if input_or_inputs is None or (inputs_embeds
653 + and input_or_inputs == ""):
654 + return [], inputs_embeds
655 +
656 # Although our type checking is based on mypy,
657 # VSCode Pyright extension should still work properly
612 - # "is True" is required for Pyright to perform type narrowing
658 + # "is False" is required for Pyright to perform type narrowing
659 # See: https://github.com/microsoft/pyright/issues/7672
614 - return [
660 + inputs_text.extend([
661 self._normalize_prompt_text_to_input(
662 request,
663 tokenizer,
@@ -625,29 +671,88 @@ class OpenAIServing:
671 prompt_ids=prompt_input["content"],
672 truncate_prompt_tokens=truncate_prompt_tokens)
673 for prompt_input in parse_and_batch_prompt(input_or_inputs)
628 - ]
674 + ])
675 +
676 + return inputs_text, inputs_embeds
677
678 + @overload
679 async def _preprocess_completion(
680 self,
632 - request: CompletionLikeRequest,
681 + request: Union[DetokenizeRequest, EmbeddingCompletionRequest,
682 + RerankRequest, ClassificationRequest, ScoreRequest,
683 + TokenizeCompletionRequest],
684 tokenizer: AnyTokenizer,
685 input_or_inputs: Union[str, list[str], list[int], list[list[int]]],
686 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = ...,
687 + add_special_tokens: bool = ...,
688 + ) -> tuple[list[TextTokensPrompt], list[EngineTokensPrompt]]:
689 + ...
690 +
691 + @overload
692 + async def _preprocess_completion(
693 + self,
694 + request: CompletionRequest,
695 + tokenizer: AnyTokenizer,
696 + input_or_inputs: Optional[Union[str, list[str], list[int],
697 + list[list[int]]]],
698 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = ...,
699 + add_special_tokens: bool = ...,
700 + ) -> tuple[list[Union[TextTokensPrompt, EmbedsPrompt]], list[Union[
701 + EngineTokensPrompt, EngineEmbedsPrompt]]]:
702 + ...
703 +
704 + async def _preprocess_completion(
705 + self,
706 + request: CompletionLikeRequest,
707 + tokenizer: AnyTokenizer,
708 + input_or_inputs: Optional[Union[str, list[str], list[int],
709 + list[list[int]]]],
710 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
711 add_special_tokens: bool = True,
637 - ) -> tuple[list[TextTokensPrompt], list[TokensPrompt]]:
638 - request_prompts = await self._tokenize_prompt_input_or_inputs_async(
639 - request,
640 - tokenizer,
641 - input_or_inputs,
642 - truncate_prompt_tokens=truncate_prompt_tokens,
643 - add_special_tokens=add_special_tokens,
644 - )
712 + ) -> tuple[Union[list[TextTokensPrompt], list[Union[
713 + TextTokensPrompt, EmbedsPrompt]]], Union[
714 + list[EngineTokensPrompt], list[Union[EngineTokensPrompt,
715 + EngineEmbedsPrompt]]]]:
716 + if not isinstance(request,
717 + CompletionRequest) and input_or_inputs is None:
718 + raise ValueError(
719 + "Prompt embeds with non-completion requests is not"
720 + " currently supported.")
721 +
722 + (request_prompts_text, request_prompts_embeds
723 + ) = await self._tokenize_prompt_input_or_inputs_async(
724 + request,
725 + tokenizer,
726 + input_or_inputs,
727 + truncate_prompt_tokens=truncate_prompt_tokens,
728 + add_special_tokens=add_special_tokens,
729 + )
730 +
731 + engine_prompts_text = [
732 + EngineTokensPrompt(
733 + prompt_token_ids=request_prompt_text["prompt_token_ids"])
734 + for request_prompt_text in request_prompts_text
735 + ]
736
646 - engine_prompts = [
647 - TokensPrompt(prompt_token_ids=request_prompt["prompt_token_ids"])
648 - for request_prompt in request_prompts
737 + # This check is equivalent to simply checking if
738 + # `request_prompts_embeds` is empty, but it's difficult to propagate
739 + # overloads to the private helper functions to enable this check.
740 + # This overload is needed because only TextPrompts are allowed for
741 + # non-completion requests and if we don't add the overload here,
742 + # everywhere this function is used outside of serving_completion will
743 + # need logic asserting that only text prompts are in the request.
744 + if not isinstance(request,
745 + CompletionRequest) and input_or_inputs is not None:
746 + return request_prompts_text, engine_prompts_text
747 +
748 + engine_prompts_embeds = [
749 + EngineEmbedsPrompt(
750 + prompt_embeds=request_prompt_embeds["prompt_embeds"])
751 + for request_prompt_embeds in request_prompts_embeds
752 ]
753
754 + request_prompts = request_prompts_embeds + request_prompts_text
755 + engine_prompts = engine_prompts_embeds + engine_prompts_text
756 return request_prompts, engine_prompts
757
758 async def _preprocess_chat(
@@ -666,7 +771,7 @@ class OpenAIServing:
771 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None,
772 add_special_tokens: bool = False,
773 ) -> tuple[list[ConversationMessage], Sequence[RequestPrompt],
669 - list[TokensPrompt]]:
774 + list[EngineTokensPrompt]]:
775 model_config = self.model_config
776
777 resolved_content_format = resolve_chat_template_content_format(
@@ -739,7 +844,7 @@ class OpenAIServing:
844 prompt=tokenizer.decode(request_prompt),
845 prompt_token_ids=request_prompt)
846
742 - engine_prompt = TokensPrompt(
847 + engine_prompt = EngineTokensPrompt(
848 prompt_token_ids=prompt_inputs["prompt_token_ids"])
849 if mm_data is not None:
850 engine_prompt["multi_modal_data"] = mm_data
@@ -751,6 +856,35 @@ class OpenAIServing:
856
857 return conversation, [request_prompt], [engine_prompt]
858
859 + def _load_prompt_embeds(
860 + self,
861 + prompt_embeds: Optional[Union[bytes, list[bytes]]],
862 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
863 + ) -> list[EmbedsPrompt]:
864 +
865 + def _load_and_validate_embed(embed: bytes) -> EmbedsPrompt:
866 + tensor = torch.load(io.BytesIO(base64.b64decode(embed)),
867 + weights_only=True)
868 + assert isinstance(
869 + tensor,
870 + (torch.FloatTensor, torch.BFloat16Tensor, torch.HalfTensor))
871 + if tensor.dim() > 2:
872 + tensor = tensor.squeeze(0)
873 + assert tensor.dim() == 2
874 + if truncate_prompt_tokens is not None:
875 + tensor = tensor[-truncate_prompt_tokens:]
876 + return {"prompt_embeds": tensor}
877 +
878 + if prompt_embeds:
879 + if isinstance(prompt_embeds, list):
880 + return [
881 + _load_and_validate_embed(embed) for embed in prompt_embeds
882 + ]
883 + else:
884 + return [_load_and_validate_embed(prompt_embeds)]
885 + else:
886 + return []
887 +
888 def _log_inputs(
889 self,
890 request_id: str,
@@ -762,13 +896,13 @@ class OpenAIServing:
896 ) -> None:
897 if self.request_logger is None:
898 return
765 -
899 + prompt, prompt_token_ids, prompt_embeds = None, None, None
900 if isinstance(inputs, str):
901 prompt = inputs
768 - prompt_token_ids = None
902 elif isinstance(inputs, list):
770 - prompt = None
903 prompt_token_ids = inputs
904 + elif 'prompt_embeds' in inputs:
905 + prompt_embeds = inputs.get("prompt_embeds")
906 else:
907 prompt = inputs["prompt"]
908 prompt_token_ids = inputs["prompt_token_ids"]
@@ -777,6 +911,7 @@ class OpenAIServing:
911 request_id,
912 prompt,
913 prompt_token_ids,
914 + prompt_embeds,
915 params=params,
916 lora_request=lora_request,
917 prompt_adapter_request=prompt_adapter_request,
vllm/entrypoints/openai/serving_tokenization.py
+3 -2
@@ -106,8 +106,9 @@ class OpenAIServingTokenization(OpenAIServing):
106
107 # Silently ignore prompt adapter since it does not affect
108 # tokenization (Unlike in Embeddings API where an error is raised)
109 -
110 - input_ids.extend(engine_prompt["prompt_token_ids"])
109 + if isinstance(engine_prompt,
110 + dict) and "prompt_token_ids" in engine_prompt:
111 + input_ids.extend(engine_prompt["prompt_token_ids"])
112
113 return TokenizeResponse(tokens=input_ids,
114 count=len(input_ids),
vllm/inputs/data.py
+12 -1
@@ -3,7 +3,7 @@ from collections.abc import Iterable
3 from typing import TYPE_CHECKING, Any, Generic, Literal, Optional, Union, cast
4
5 import torch
6 -from typing_extensions import NotRequired, TypedDict, TypeVar
6 +from typing_extensions import NotRequired, TypedDict, TypeIs, TypeVar
7
8 if TYPE_CHECKING:
9 from vllm.multimodal.inputs import MultiModalDataDict, MultiModalInputs
@@ -98,6 +98,17 @@ where the decoder-prompt is not specified explicitly, or
98 more than one prompt, i.e. {class}`ExplicitEncoderDecoderPrompt`
99 """
100
101 +
102 +def is_tokens_prompt(prompt: SingletonPrompt) -> TypeIs[TokensPrompt]:
103 + return (isinstance(prompt, dict) and "prompt_token_ids" in prompt
104 + and "prompt_embeds" not in prompt)
105 +
106 +
107 +def is_embeds_prompt(prompt: SingletonPrompt) -> TypeIs[EmbedsPrompt]:
108 + return (isinstance(prompt, dict) and "prompt_token_ids" not in prompt
109 + and "prompt_embeds" in prompt)
110 +
111 +
112 _T1_co = TypeVar("_T1_co",
113 bound=SingletonPrompt,
114 default=SingletonPrompt,