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 4, 10:00 UTC
feat: Enable `prompt_embeds` Content Part Support in vLLM Chat Completions API (#40720)
Signed-off-by: Luis Robaina <luis@protopia.ai> Signed-off-by: Luis Robaina 🚀 <luisfabian1545@gmail.com> Signed-off-by: LuisRobaina <luis@protopia.ai> Co-authored-by: Andrew Sansom <qthequartermasterman@gmail.com>
Luis 🚀 committed
Apr 30, 2026 at 22:05 UTC
14043dfecd35dd2f12b4d51eb9fa166184a0ca0f
28 files changed
+2301
-82
docs/features/README.md
+2
-2
@@ -52,10 +52,10 @@ th:not(:first-child) {
52
| [mm](multimodal_inputs.md) | ✅ | ✅ | [🟠](https://github.com/vllm-project/vllm/pull/4194)<sup>^</sup> | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | | |
53
| best-of | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/7968) | ✅ | ✅ | | |
54
| beam-search | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/7968) | ❔ | ✅ | ✅ | |
55
-| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ❌ | ❔ | ❔ | ✅ |
55
+| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ✅ | ❔ | ❔ | ✅ |
56
57
\* Chunked prefill and prefix caching are only applicable to last-token or all pooling with causal attention.
58
-<sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
58
+<sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
59
60
### Feature x Hardware
61
docs/features/prompt_embeds.md
+36
-1
@@ -20,12 +20,47 @@ You can pass prompt embeddings from Hugging Face Transformers models to the `'p
20
21
## Online Serving
22
23
-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 and are enabled by the `--enable-prompt-embeds` flag in `vllm serve`.
23
+Our OpenAI-compatible server accepts prompt embeddings inputs via both the [Completions API](https://platform.openai.com/docs/api-reference/completions) and the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat). Both are enabled by the `--enable-prompt-embeds` flag in `vllm serve`.
24
+
25
+### Completions API
26
+
27
+Prompt embeddings inputs are added via a `'prompt_embeds'` key in the JSON request body.
28
29
When a mixture of `'prompt_embeds'` and `'prompt'` inputs are provided in a single request, the prompt embeds are always returned first.
30
31
Prompt embeddings are passed in as base64 encoded torch tensors.
32
33
+The Completions endpoint does **not** apply a chat template to `prompt_embeds`. If the model assumes some chat template, the caller is responsible for producing embeddings for the full, already-templated prompt: apply the chat template, then embed the resulting token IDs. Anything the model would normally need (system prompt, role markers, generation prompt, etc.) must already be baked into the embedded tokens.
34
+
35
+### Chat Completions API
36
+
37
+Prompt embeddings can be included as content parts in chat messages, interleaved with text:
38
+
39
+```json
40
+{
41
+ "messages": [
42
+ {
43
+ "role": "system",
44
+ "content": [
45
+ {"type": "text", "text": "You are a helpful assistant."},
46
+ {"type": "prompt_embeds", "data": "<base64_encoded_tensor>"}
47
+ ]
48
+ },
49
+ {
50
+ "role": "user",
51
+ "content": [
52
+ {"type": "prompt_embeds", "data": "<base64_encoded_tensor>"},
53
+ {"type": "text", "text": "Summarize the above."}
54
+ ]
55
+ }
56
+ ]
57
+}
58
+```
59
+
60
+Each `prompt_embeds` content part contains a `data` field with a base64-encoded `torch.Tensor` of shape `(num_tokens, hidden_size)`. Multiple `prompt_embeds` parts can appear in any message, in any position relative to text parts. The server expands each part into the correct number of placeholder tokens during chat template rendering, then splices the pre-computed embeddings into the model's input at the corresponding positions.
61
+
62
+Unlike the Completions API, a `prompt_embeds` content part should encode **only** the content, not a templated conversation. The server wraps the chat template around the embedded content at request time, the same way it would for a plain text `content` string. Embedding a full templated conversation here would double-apply the template and produce incorrect inputs to the model.
63
+
64
!!! warning
65
The vLLM engine may crash if incorrect shape of embeddings is passed.
66
Only enable this flag for trusted users!
examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py
+105
-24
@@ -1,12 +1,29 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-vLLM OpenAI-Compatible Client with Prompt Embeddings
3
+"""vLLM OpenAI-Compatible Client with Prompt Embeddings.
4
5
This script demonstrates how to:
7
-1. Generate prompt embeddings using Hugging Face Transformers
8
-2. Encode them in base64 format
9
-3. Send them to a vLLM server via the OpenAI-compatible Completions API
6
+1. Generate prompt embeddings using Hugging Face Transformers.
7
+2. Encode them in base64 format.
8
+3. Send them to a vLLM server for inference via both:
9
+ - OpenAI-compatible Chat Completions API
10
+ - OpenAI-compatible Completions API
11
+
12
+Important distinction between the two APIs:
13
+
14
+- Chat Completions API: `prompt_embeds` content parts should encode ONLY
15
+ the user-provided content, not a templated conversation. The server
16
+ renders the surrounding chat template around the embedded content at
17
+ request time, the same way it would for a plain text `content` string.
18
+ Embedding a full templated conversation here would double-apply the
19
+ template and likely produce undesirable results.
20
+
21
+- Completions API: the server does NOT apply a chat template to
22
+ `prompt_embeds`. The caller is responsible for producing embeddings for
23
+ the full, already-templated prompt (i.e. apply the chat template first,
24
+ then embed the resulting token IDs). Anything the model would normally
25
+ need (system prompt, role markers, generation prompt, etc.) must already
26
+ be baked into the embedded tokens.
27
28
Run the vLLM server first:
29
vllm serve meta-llama/Llama-3.2-1B-Instruct \
@@ -34,34 +51,68 @@ from openai import OpenAI
51
from vllm.utils.serial_utils import tensor2base64
52
53
37
-def main():
38
- client = OpenAI(
39
- api_key="EMPTY",
40
- base_url="http://localhost:8000/v1",
54
+def run_chat_completion_prompt_embeds(
55
+ client: OpenAI,
56
+ model_name: str,
57
+ tokenizer: transformers.PreTrainedTokenizerBase,
58
+ embedding_layer,
59
+ messages: list[dict],
60
+) -> None:
61
+ """Run a Chat Completions API request using prompt_embeds content parts.
62
+
63
+ This example embeds ONLY the user-provided content of the final user turn, the
64
+ vLLM server applies the chat template around it at request time.
65
+ """
66
+ user_content = messages[-1]["content"]
67
+ content_token_ids = tokenizer(
68
+ user_content, return_tensors="pt", add_special_tokens=False
69
+ ).input_ids
70
+ content_prompt_embeds = embedding_layer(content_token_ids).squeeze(0)
71
+ encoded_embeds = tensor2base64(content_prompt_embeds)
72
+
73
+ api_messages = [
74
+ *messages[:-1],
75
+ {
76
+ "role": messages[-1]["role"],
77
+ "content": [{"type": "prompt_embeds", "data": encoded_embeds}],
78
+ },
79
+ ]
80
+
81
+ chat_completion = client.chat.completions.create(
82
+ model=model_name,
83
+ max_tokens=6,
84
+ temperature=0.0,
85
+ messages=api_messages,
86
)
87
43
- model_name = "meta-llama/Llama-3.2-1B-Instruct"
88
+ print("-" * 30)
89
+ print("Chat Completions API")
90
+ print(chat_completion.choices[0].message.content)
91
+ print("-" * 30)
92
45
- # Transformers
46
- tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
47
- transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
93
49
- # Refer to the HuggingFace repo for the correct format to use
50
- chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
51
- token_ids = tokenizer.apply_chat_template(
52
- chat, add_generation_prompt=True, return_tensors="pt", return_dict=True
94
+def run_completion_prompt_embeds(
95
+ client: OpenAI,
96
+ model_name: str,
97
+ tokenizer: transformers.PreTrainedTokenizerBase,
98
+ embedding_layer,
99
+ messages: list[dict],
100
+) -> None:
101
+ """Run a Completions API request using prompt embeddings.
102
+
103
+ The Completions endpoint does not apply a chat template,
104
+ so the caller must apply it and embed the full templated prompt.
105
+ """
106
+ templated_token_ids = tokenizer.apply_chat_template(
107
+ messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
108
).input_ids
54
-
55
- embedding_layer = transformers_model.get_input_embeddings()
56
- prompt_embeds = embedding_layer(token_ids).squeeze(0)
57
-
58
- # Prompt embeddings
59
- encoded_embeds = tensor2base64(prompt_embeds)
109
+ templated_prompt_embeds = embedding_layer(templated_token_ids).squeeze(0)
110
+ encoded_embeds = tensor2base64(templated_prompt_embeds)
111
112
completion = client.completions.create(
113
model=model_name,
114
prompt=None,
64
- max_tokens=5,
115
+ max_tokens=6,
116
temperature=0.0,
117
# NOTE: The OpenAI client allows passing in extra JSON body via the
118
# `extra_body` argument.
@@ -69,9 +120,39 @@ def main():
120
)
121
122
print("-" * 30)
123
+ print("Completions API")
124
print(completion.choices[0].text)
125
print("-" * 30)
126
127
128
+def main() -> None:
129
+ client = OpenAI(
130
+ api_key="EMPTY",
131
+ base_url="http://localhost:8000/v1",
132
+ )
133
+
134
+ model_name = "meta-llama/Llama-3.2-1B-Instruct"
135
+
136
+ tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
137
+ transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
138
+ embedding_layer = transformers_model.get_input_embeddings()
139
+
140
+ messages = [
141
+ {"role": "user", "content": "Please tell me about the capital of France."}
142
+ ]
143
+
144
+ # Chat Completions API: embed ONLY the user content. The server wraps
145
+ # the embedding in the chat template when it renders the messages.
146
+ run_chat_completion_prompt_embeds(
147
+ client, model_name, tokenizer, embedding_layer, messages
148
+ )
149
+
150
+ # Completions API: embed the FULL templated prompt. The caller must
151
+ # apply the chat template up-front.
152
+ run_completion_prompt_embeds(
153
+ client, model_name, tokenizer, embedding_layer, messages
154
+ )
155
+
156
+
157
if __name__ == "__main__":
158
main()
tests/entrypoints/llm/test_mm_processor_kwargs.py
+3
-1
@@ -11,7 +11,9 @@ from vllm import LLM, SamplingParams
11
12
def _make_mock_llm() -> LLM:
13
llm = object.__new__(LLM)
14
- llm.model_config = SimpleNamespace(runner_type="generate")
14
+ llm.model_config = SimpleNamespace(
15
+ runner_type="generate", enable_prompt_embeds=False
16
+ )
17
return llm
18
19
tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py
new
+190
@@ -0,0 +1,190 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""E2E test for mixing `prompt_embeds` with `audio_embeds` in a single
5
+Chat Completions request."""
6
+
7
+import json
8
+
9
+import openai
10
+import pytest
11
+import pytest_asyncio
12
+import safetensors
13
+import torch
14
+import torch.nn as nn
15
+from huggingface_hub import hf_hub_download
16
+from transformers import AutoConfig, AutoTokenizer
17
+
18
+from tests.utils import RemoteOpenAIServer
19
+from vllm.utils.serial_utils import tensor2base64
20
+
21
+QWEN2AUDIO_MODEL = "Qwen/Qwen2-Audio-7B-Instruct"
22
+
23
+# Use the model's native dtype to avoid an implicit cast inside
24
+# `safe_load_prompt_embeds` (mismatched floating-point dtypes are cast to the
25
+# model's dtype automatically, matching here just skips the conversion).
26
+QWEN2AUDIO_DTYPE = torch.bfloat16
27
+
28
+
29
+@pytest.fixture(scope="module")
30
+def qwen2audio_server_args() -> list[str]:
31
+ return [
32
+ "--dtype",
33
+ "bfloat16",
34
+ "--max-model-len",
35
+ "2048",
36
+ "--max-num-seqs",
37
+ "4",
38
+ "--enforce-eager",
39
+ "--trust-remote-code",
40
+ "--gpu-memory-utilization",
41
+ "0.85",
42
+ "--limit-mm-per-prompt",
43
+ json.dumps({"audio": 1}),
44
+ "--enable-prompt-embeds",
45
+ "--enable-mm-embeds",
46
+ ]
47
+
48
+
49
+@pytest.fixture(scope="module")
50
+def qwen2audio_server(qwen2audio_server_args):
51
+ with RemoteOpenAIServer(
52
+ QWEN2AUDIO_MODEL,
53
+ qwen2audio_server_args,
54
+ max_wait_seconds=600,
55
+ ) as remote_server:
56
+ yield remote_server
57
+
58
+
59
+@pytest_asyncio.fixture
60
+async def qwen2audio_client(qwen2audio_server):
61
+ async with qwen2audio_server.get_async_client() as async_client:
62
+ yield async_client
63
+
64
+
65
+@pytest.fixture(scope="module")
66
+def qwen2audio_hidden_size() -> int:
67
+ config = AutoConfig.from_pretrained(QWEN2AUDIO_MODEL, trust_remote_code=True)
68
+ return config.text_config.hidden_size
69
+
70
+
71
+@pytest.fixture(scope="module")
72
+def qwen2audio_prompt_embeds_b64(qwen2audio_hidden_size: int) -> str:
73
+ tensor = torch.randn(4, qwen2audio_hidden_size, dtype=QWEN2AUDIO_DTYPE)
74
+ return tensor2base64(tensor)
75
+
76
+
77
+@pytest.fixture(scope="module")
78
+def qwen2audio_audio_embeds_b64(qwen2audio_hidden_size: int) -> str:
79
+ # Shape matches the `audio_embeds` unit-test fixture.
80
+ torch.manual_seed(0)
81
+ tensor = torch.randn(1, 128, qwen2audio_hidden_size, dtype=QWEN2AUDIO_DTYPE)
82
+ return tensor2base64(tensor)
83
+
84
+
85
+@pytest.mark.asyncio
86
+async def test_prompt_embeds_plus_audio_embeds(
87
+ qwen2audio_client: openai.AsyncOpenAI,
88
+ qwen2audio_prompt_embeds_b64: str,
89
+ qwen2audio_audio_embeds_b64: str,
90
+):
91
+ """Single user message carrying both prompt_embeds and audio_embeds parts."""
92
+ chat = await qwen2audio_client.chat.completions.create(
93
+ model=QWEN2AUDIO_MODEL,
94
+ max_tokens=5,
95
+ temperature=0.0,
96
+ messages=[
97
+ {
98
+ "role": "user",
99
+ "content": [
100
+ {
101
+ "type": "prompt_embeds",
102
+ "data": qwen2audio_prompt_embeds_b64,
103
+ },
104
+ {
105
+ "type": "audio_embeds",
106
+ "audio_embeds": qwen2audio_audio_embeds_b64,
107
+ },
108
+ {"type": "text", "text": "Continue."},
109
+ ],
110
+ }
111
+ ],
112
+ )
113
+ assert chat.choices[0].message.content is not None
114
+ assert len(chat.choices[0].message.content) > 0
115
+
116
+
117
+@pytest.fixture(scope="module")
118
+def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]:
119
+ """Return `(content, base64_embeds)` where the embeddings are the model's
120
+ embedding of `content` tokenized WITHOUT special tokens.
121
+
122
+ Loads only the `embed_tokens` shard from disk on CPU (~1.1 GB of host
123
+ RAM) instead of the full 7B model on GPU.
124
+ """
125
+ content = "Describe this audio."
126
+ tokenizer = AutoTokenizer.from_pretrained(QWEN2AUDIO_MODEL, trust_remote_code=True)
127
+
128
+ index_path = hf_hub_download(QWEN2AUDIO_MODEL, "model.safetensors.index.json")
129
+ with open(index_path) as f:
130
+ weight_map = json.load(f)["weight_map"]
131
+ embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight"))
132
+ shard_path = hf_hub_download(QWEN2AUDIO_MODEL, weight_map[embed_key])
133
+ with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
134
+ embed_weight = f.get_tensor(embed_key)
135
+ embed_layer = nn.Embedding.from_pretrained(embed_weight.to(QWEN2AUDIO_DTYPE))
136
+
137
+ ids = tokenizer(content, add_special_tokens=False, return_tensors="pt").input_ids
138
+ embeds = embed_layer(ids).squeeze(0)
139
+ return content, tensor2base64(embeds)
140
+
141
+
142
+@pytest.mark.asyncio
143
+@pytest.mark.parametrize(
144
+ "audio_first",
145
+ [True, False],
146
+ ids=["audio_embeds-then-text", "text-then-audio_embeds"],
147
+)
148
+async def test_text_content_and_prompt_embeds_match_with_audio_embeds(
149
+ qwen2audio_client: openai.AsyncOpenAI,
150
+ qwen2audio_audio_embeds_b64: str,
151
+ qwen2audio_aligned_content_and_embeds_b64: tuple[str, str],
152
+ audio_first: bool,
153
+):
154
+ """Same content as text vs `prompt_embeds` should yield identical Chat
155
+ Completions output when mixed with `audio_embeds` in the same message.
156
+ """
157
+ content, encoded_text_embeds = qwen2audio_aligned_content_and_embeds_b64
158
+
159
+ audio_part = {
160
+ "type": "audio_embeds",
161
+ "audio_embeds": qwen2audio_audio_embeds_b64,
162
+ }
163
+ text_part = {"type": "text", "text": content}
164
+ embeds_part = {"type": "prompt_embeds", "data": encoded_text_embeds}
165
+
166
+ if audio_first:
167
+ text_content = [audio_part, text_part]
168
+ embeds_content = [audio_part, embeds_part]
169
+ else:
170
+ text_content = [text_part, audio_part]
171
+ embeds_content = [embeds_part, audio_part]
172
+
173
+ text_resp = await qwen2audio_client.chat.completions.create(
174
+ model=QWEN2AUDIO_MODEL,
175
+ max_tokens=10,
176
+ temperature=0.0,
177
+ messages=[{"role": "user", "content": text_content}],
178
+ )
179
+ embeds_resp = await qwen2audio_client.chat.completions.create(
180
+ model=QWEN2AUDIO_MODEL,
181
+ max_tokens=10,
182
+ temperature=0.0,
183
+ messages=[{"role": "user", "content": embeds_content}],
184
+ )
185
+
186
+ text_out = text_resp.choices[0].message.content
187
+ embeds_out = embeds_resp.choices[0].message.content
188
+ assert text_out is not None and len(text_out) > 0
189
+ assert embeds_out is not None and len(embeds_out) > 0
190
+ assert text_out == embeds_out
tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py
new
+212
@@ -0,0 +1,212 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""E2E tests for mixing `prompt_embeds` with image content parts in a single
5
+Chat Completions request.
6
+"""
7
+
8
+import json
9
+
10
+import openai
11
+import pytest
12
+import pytest_asyncio
13
+import safetensors
14
+import torch
15
+import torch.nn as nn
16
+from huggingface_hub import hf_hub_download
17
+from transformers import AutoTokenizer
18
+
19
+from tests.utils import RemoteOpenAIServer
20
+from vllm.assets.image import ImageAsset
21
+from vllm.multimodal.utils import encode_image_url
22
+from vllm.utils.serial_utils import tensor2base64
23
+
24
+MODEL_NAME = "Qwen/Qwen2-VL-2B-Instruct"
25
+
26
+# Use the model's native dtype to skip the implicit cast inside
27
+# `safe_load_prompt_embeds` (mismatched floating-point dtypes are cast to the
28
+# model's dtype automatically).
29
+MODEL_DTYPE = torch.bfloat16
30
+
31
+
32
+@pytest.fixture(scope="module")
33
+def server_args() -> list[str]:
34
+ return [
35
+ "--dtype",
36
+ "bfloat16",
37
+ "--max-model-len",
38
+ "2048",
39
+ "--max-num-seqs",
40
+ "4",
41
+ "--enforce-eager",
42
+ "--gpu-memory-utilization",
43
+ "0.4",
44
+ "--limit-mm-per-prompt",
45
+ json.dumps({"image": 1}),
46
+ "--enable-prompt-embeds",
47
+ "--enable-mm-embeds",
48
+ ]
49
+
50
+
51
+@pytest.fixture(scope="module")
52
+def server(server_args):
53
+ with RemoteOpenAIServer(
54
+ MODEL_NAME,
55
+ server_args,
56
+ max_wait_seconds=600,
57
+ ) as remote_server:
58
+ yield remote_server
59
+
60
+
61
+@pytest_asyncio.fixture
62
+async def client(server):
63
+ async with server.get_async_client() as async_client:
64
+ yield async_client
65
+
66
+
67
+@pytest.fixture(scope="module")
68
+def image_url() -> str:
69
+ """Stable real image as a data URL, kept identical across both the
70
+ text and prompt_embeds requests so any output difference must come from
71
+ how the text content is delivered."""
72
+ return encode_image_url(ImageAsset("stop_sign").pil_image)
73
+
74
+
75
+@pytest.fixture(scope="module")
76
+def aligned_content_and_embeds_b64() -> tuple[str, str]:
77
+ """`(content, base64_embeds)` where the embeddings are the model's
78
+ embedding of `content` tokenized WITHOUT special tokens.
79
+
80
+ Loads only the `embed_tokens` shard from disk on CPU instead of the full
81
+ model on GPU, so the fixture has zero VRAM footprint and won't contend
82
+ with the running vLLM server.
83
+ """
84
+ content = "Describe this image."
85
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
86
+
87
+ index_path = hf_hub_download(MODEL_NAME, "model.safetensors.index.json")
88
+ with open(index_path) as f:
89
+ weight_map = json.load(f)["weight_map"]
90
+ embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight"))
91
+ shard_path = hf_hub_download(MODEL_NAME, weight_map[embed_key])
92
+ with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
93
+ embed_weight = f.get_tensor(embed_key)
94
+ embed_layer = nn.Embedding.from_pretrained(embed_weight.to(MODEL_DTYPE))
95
+
96
+ ids = tokenizer(content, add_special_tokens=False, return_tensors="pt").input_ids
97
+ embeds = embed_layer(ids).squeeze(0)
98
+ return content, tensor2base64(embeds)
99
+
100
+
101
+@pytest.mark.asyncio
102
+@pytest.mark.parametrize(
103
+ "image_first",
104
+ [True, False],
105
+ ids=["image_url-then-text", "text-then-image_url"],
106
+)
107
+async def test_text_content_and_prompt_embeds_match_with_image_url(
108
+ client: openai.AsyncOpenAI,
109
+ image_url: str,
110
+ aligned_content_and_embeds_b64: tuple[str, str],
111
+ image_first: bool,
112
+):
113
+ """Same content as text vs `prompt_embeds` should yield identical Chat
114
+ Completions output when mixed with an `image_url` part in the same
115
+ message under greedy decoding.
116
+ """
117
+ content, encoded_text_embeds = aligned_content_and_embeds_b64
118
+
119
+ image_part = {"type": "image_url", "image_url": {"url": image_url}}
120
+ text_part = {"type": "text", "text": content}
121
+ embeds_part = {"type": "prompt_embeds", "data": encoded_text_embeds}
122
+
123
+ if image_first:
124
+ text_content = [image_part, text_part]
125
+ embeds_content = [image_part, embeds_part]
126
+ else:
127
+ text_content = [text_part, image_part]
128
+ embeds_content = [embeds_part, image_part]
129
+
130
+ text_resp = await client.chat.completions.create(
131
+ model=MODEL_NAME,
132
+ max_tokens=10,
133
+ temperature=0.0,
134
+ messages=[{"role": "user", "content": text_content}],
135
+ )
136
+ embeds_resp = await client.chat.completions.create(
137
+ model=MODEL_NAME,
138
+ max_tokens=10,
139
+ temperature=0.0,
140
+ messages=[{"role": "user", "content": embeds_content}],
141
+ )
142
+
143
+ text_out = text_resp.choices[0].message.content
144
+ embeds_out = embeds_resp.choices[0].message.content
145
+ assert text_out is not None and len(text_out) > 0
146
+ assert embeds_out is not None and len(embeds_out) > 0
147
+ assert text_out == embeds_out
148
+
149
+
150
+@pytest.fixture(scope="module")
151
+def image_embeds_b64() -> dict[str, str]:
152
+ """Synthetic but stable `image_embeds` for Qwen2-VL."""
153
+ grid = (1, 4, 4)
154
+ spatial_merge_size = 2
155
+ num_patches = (grid[1] // spatial_merge_size) * (grid[2] // spatial_merge_size)
156
+ text_hidden_size = 1536 # Qwen2-VL-2B
157
+ torch.manual_seed(0)
158
+ return {
159
+ "image_embeds": tensor2base64(
160
+ torch.randn(num_patches, text_hidden_size, dtype=MODEL_DTYPE)
161
+ ),
162
+ "image_grid_thw": tensor2base64(torch.tensor(grid)),
163
+ }
164
+
165
+
166
+@pytest.mark.asyncio
167
+@pytest.mark.parametrize(
168
+ "image_first",
169
+ [True, False],
170
+ ids=["image_embeds-then-text", "text-then-image_embeds"],
171
+)
172
+async def test_text_content_and_prompt_embeds_match_with_image_embeds(
173
+ client: openai.AsyncOpenAI,
174
+ image_embeds_b64: dict[str, str],
175
+ aligned_content_and_embeds_b64: tuple[str, str],
176
+ image_first: bool,
177
+):
178
+ """Same content as text vs `prompt_embeds` should yield identical Chat
179
+ Completions output when mixed with a precomputed `image_embeds` part in
180
+ the same message under greedy decoding.
181
+ """
182
+ content, encoded_text_embeds = aligned_content_and_embeds_b64
183
+
184
+ image_part = {"type": "image_embeds", "image_embeds": image_embeds_b64}
185
+ text_part = {"type": "text", "text": content}
186
+ embeds_part = {"type": "prompt_embeds", "data": encoded_text_embeds}
187
+
188
+ if image_first:
189
+ text_content = [image_part, text_part]
190
+ embeds_content = [image_part, embeds_part]
191
+ else:
192
+ text_content = [text_part, image_part]
193
+ embeds_content = [embeds_part, image_part]
194
+
195
+ text_resp = await client.chat.completions.create(
196
+ model=MODEL_NAME,
197
+ max_tokens=10,
198
+ temperature=0.0,
199
+ messages=[{"role": "user", "content": text_content}],
200
+ )
201
+ embeds_resp = await client.chat.completions.create(
202
+ model=MODEL_NAME,
203
+ max_tokens=10,
204
+ temperature=0.0,
205
+ messages=[{"role": "user", "content": embeds_content}],
206
+ )
207
+
208
+ text_out = text_resp.choices[0].message.content
209
+ embeds_out = embeds_resp.choices[0].message.content
210
+ assert text_out is not None and len(text_out) > 0
211
+ assert embeds_out is not None and len(embeds_out) > 0
212
+ assert text_out == embeds_out
tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py
new
+293
@@ -0,0 +1,293 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""E2E tests for `prompt_embeds` content parts in the Chat Completions API."""
5
+
6
+import asyncio
7
+import io
8
+
9
+import openai
10
+import pybase64 as base64
11
+import pytest
12
+import pytest_asyncio
13
+import torch
14
+from openai import BadRequestError
15
+
16
+from tests.utils import VLLM_PATH, RemoteOpenAIServer
17
+
18
+MODEL_NAME = "facebook/opt-125m"
19
+CHAT_TEMPLATE = VLLM_PATH / "examples/template_chatml.jinja"
20
+# Matches `--dtype` in `server_args` to avoid an implicit cast in
21
+# `safe_load_prompt_embeds` (mismatched floating-point dtypes are cast to the
22
+# model's dtype automatically, we match here just to skip the conversion).
23
+SERVER_DTYPE: torch.dtype = torch.bfloat16
24
+
25
+
26
+@pytest.fixture(scope="module")
27
+def server_args() -> list[str]:
28
+ return [
29
+ "--dtype",
30
+ "bfloat16",
31
+ "--max-model-len",
32
+ "2048",
33
+ "--max-num-seqs",
34
+ "128",
35
+ "--enforce-eager",
36
+ "--chat-template",
37
+ str(CHAT_TEMPLATE),
38
+ # Prompt Embeds server args
39
+ "--enable-prompt-embeds",
40
+ ]
41
+
42
+
43
+@pytest.fixture(scope="module")
44
+def server(server_args):
45
+ with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
46
+ yield remote_server
47
+
48
+
49
+@pytest_asyncio.fixture
50
+async def client(server):
51
+ async with server.get_async_client() as async_client:
52
+ yield async_client
53
+
54
+
55
+def _encode_embeds(embeds: torch.Tensor) -> str:
56
+ buf = io.BytesIO()
57
+ torch.save(embeds, buf)
58
+ return base64.b64encode(buf.getvalue()).decode("utf-8")
59
+
60
+
61
+@pytest.fixture(scope="module")
62
+def prompt_embeds_b64(hf_runner) -> list[str]:
63
+ """Pre-compute embeddings for two short prompts and return as base64."""
64
+ prompts = ["Hello, my name is", "What is an LLM?"]
65
+ with hf_runner(MODEL_NAME) as hf_model:
66
+ embeddings = hf_model.get_prompt_embeddings(prompts)
67
+ # Cast to the server's dtype so `safe_load_prompt_embeds` doesn't need to
68
+ # convert on its own, the function accepts any floating-point dtype and
69
+ # will cast to the model's dtype, but matching up front skips the work.
70
+ return [_encode_embeds(e.to(SERVER_DTYPE)) for e in embeddings]
71
+
72
+
73
+@pytest.mark.asyncio
74
+async def test_single_prompt_embeds_part(
75
+ client: openai.AsyncOpenAI,
76
+ prompt_embeds_b64: list[str],
77
+):
78
+ """A user message with one prompt_embeds part + text."""
79
+ b64 = prompt_embeds_b64[0]
80
+ chat = await client.chat.completions.create(
81
+ model=MODEL_NAME,
82
+ max_tokens=5,
83
+ temperature=0.0,
84
+ messages=[
85
+ {
86
+ "role": "user",
87
+ "content": [
88
+ {"type": "prompt_embeds", "data": b64},
89
+ {"type": "text", "text": "Continue:"},
90
+ ],
91
+ }
92
+ ],
93
+ )
94
+ assert chat.choices[0].message.content is not None
95
+ assert len(chat.choices[0].message.content) > 0
96
+
97
+
98
+@pytest.mark.asyncio
99
+async def test_multiple_prompt_embeds_parts(
100
+ client: openai.AsyncOpenAI,
101
+ prompt_embeds_b64: list[str],
102
+):
103
+ """Multiple prompt_embeds parts in a single message."""
104
+ b64_a, b64_b = prompt_embeds_b64
105
+ chat = await client.chat.completions.create(
106
+ model=MODEL_NAME,
107
+ max_tokens=5,
108
+ temperature=0.0,
109
+ messages=[
110
+ {
111
+ "role": "user",
112
+ "content": [
113
+ {"type": "prompt_embeds", "data": b64_a},
114
+ {"type": "text", "text": " and "},
115
+ {"type": "prompt_embeds", "data": b64_b},
116
+ ],
117
+ }
118
+ ],
119
+ )
120
+ assert chat.choices[0].message.content is not None
121
+ assert len(chat.choices[0].message.content) > 0
122
+
123
+
124
+@pytest.mark.asyncio
125
+async def test_multi_message_conversation(
126
+ client: openai.AsyncOpenAI,
127
+ prompt_embeds_b64: list[str],
128
+):
129
+ """prompt_embeds in both system and user messages."""
130
+ b64_sys, b64_usr = prompt_embeds_b64
131
+ chat = await client.chat.completions.create(
132
+ model=MODEL_NAME,
133
+ max_tokens=5,
134
+ temperature=0.0,
135
+ messages=[
136
+ {
137
+ "role": "system",
138
+ "content": [
139
+ {"type": "text", "text": "You are helpful."},
140
+ {"type": "prompt_embeds", "data": b64_sys},
141
+ ],
142
+ },
143
+ {
144
+ "role": "user",
145
+ "content": [
146
+ {"type": "prompt_embeds", "data": b64_usr},
147
+ {"type": "text", "text": "Summarize."},
148
+ ],
149
+ },
150
+ ],
151
+ )
152
+ assert chat.choices[0].message.content is not None
153
+ assert len(chat.choices[0].message.content) > 0
154
+
155
+
156
+@pytest.mark.asyncio
157
+async def test_streaming(
158
+ client: openai.AsyncOpenAI,
159
+ prompt_embeds_b64: list[str],
160
+):
161
+ """Streaming chat completion with prompt_embeds."""
162
+ b64 = prompt_embeds_b64[0]
163
+
164
+ # Non-streaming baseline.
165
+ baseline = await client.chat.completions.create(
166
+ model=MODEL_NAME,
167
+ max_tokens=5,
168
+ temperature=0.0,
169
+ messages=[
170
+ {
171
+ "role": "user",
172
+ "content": [
173
+ {"type": "prompt_embeds", "data": b64},
174
+ {"type": "text", "text": "Continue:"},
175
+ ],
176
+ }
177
+ ],
178
+ )
179
+ expected = baseline.choices[0].message.content
180
+
181
+ # Streaming.
182
+ stream = await client.chat.completions.create(
183
+ model=MODEL_NAME,
184
+ max_tokens=5,
185
+ temperature=0.0,
186
+ stream=True,
187
+ messages=[
188
+ {
189
+ "role": "user",
190
+ "content": [
191
+ {"type": "prompt_embeds", "data": b64},
192
+ {"type": "text", "text": "Continue:"},
193
+ ],
194
+ }
195
+ ],
196
+ )
197
+ chunks: list[str] = []
198
+ async for chunk in stream:
199
+ delta = chunk.choices[0].delta.content
200
+ if delta:
201
+ chunks.append(delta)
202
+ assert "".join(chunks) == expected
203
+
204
+
205
+@pytest.fixture(scope="module")
206
+def aligned_content_and_embeds_b64(hf_runner) -> tuple[str, str]:
207
+ """Return `(content, base64_embeds)` where the embeddings are the model's
208
+ embedding of `content` tokenized WITHOUT special tokens.
209
+ """
210
+ content = "Hello, my name is"
211
+ with hf_runner(MODEL_NAME) as hf_model:
212
+ ids = hf_model.tokenizer(
213
+ content, add_special_tokens=False, return_tensors="pt"
214
+ ).input_ids
215
+ ids = hf_model.wrap_device({"input_ids": ids})["input_ids"]
216
+ embed_layer = hf_model.model.get_input_embeddings()
217
+ embeds = embed_layer(ids).squeeze(0).to(SERVER_DTYPE).cpu()
218
+ return content, _encode_embeds(embeds)
219
+
220
+
221
+@pytest.mark.asyncio
222
+async def test_text_content_and_prompt_embeds_match(
223
+ client: openai.AsyncOpenAI,
224
+ aligned_content_and_embeds_b64: tuple[str, str],
225
+):
226
+ """Equal content in text and `prompt_embeds` should yield identical
227
+ Chat Completions output under greedy decoding.
228
+ """
229
+ content, encoded_embeds = aligned_content_and_embeds_b64
230
+
231
+ text_resp, embeds_resp = await asyncio.gather(
232
+ client.chat.completions.create(
233
+ model=MODEL_NAME,
234
+ max_tokens=10,
235
+ temperature=0.0,
236
+ messages=[{"role": "user", "content": content}],
237
+ ),
238
+ client.chat.completions.create(
239
+ model=MODEL_NAME,
240
+ max_tokens=10,
241
+ temperature=0.0,
242
+ messages=[
243
+ {
244
+ "role": "user",
245
+ "content": [{"type": "prompt_embeds", "data": encoded_embeds}],
246
+ }
247
+ ],
248
+ ),
249
+ )
250
+
251
+ text_out = text_resp.choices[0].message.content
252
+ embeds_out = embeds_resp.choices[0].message.content
253
+ assert text_out is not None and len(text_out) > 0
254
+ assert embeds_out is not None and len(embeds_out) > 0
255
+ assert text_out == embeds_out
256
+
257
+
258
+@pytest.mark.asyncio
259
+async def test_missing_data_field(
260
+ client: openai.AsyncOpenAI,
261
+):
262
+ """A prompt_embeds part without `data` should return a clear error."""
263
+ with pytest.raises(BadRequestError):
264
+ await client.chat.completions.create(
265
+ model=MODEL_NAME,
266
+ max_tokens=5,
267
+ messages=[
268
+ {
269
+ "role": "user",
270
+ "content": [{"type": "prompt_embeds"}],
271
+ }
272
+ ],
273
+ )
274
+
275
+
276
+@pytest.mark.asyncio
277
+async def test_invalid_base64(
278
+ client: openai.AsyncOpenAI,
279
+):
280
+ """Invalid base64 in the `data` field should return a clear error."""
281
+ with pytest.raises(BadRequestError):
282
+ await client.chat.completions.create(
283
+ model=MODEL_NAME,
284
+ max_tokens=5,
285
+ messages=[
286
+ {
287
+ "role": "user",
288
+ "content": [
289
+ {"type": "prompt_embeds", "data": "not_valid_base64!!"},
290
+ ],
291
+ }
292
+ ],
293
+ )
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+1
@@ -538,6 +538,7 @@ class MockModelConfig:
538
is_encoder_decoder: bool = False
539
is_multimodal_model: bool = False
540
renderer_num_workers: int = 1
541
+ enable_prompt_embeds: bool = False
542
543
def get_diff_sampling_param(self):
544
return self.diff_sampling_param or {}
tests/entrypoints/openai/completion/test_prompt_validation.py
+2
@@ -62,6 +62,8 @@ def test_load_prompt_embeds(
62
):
63
model_config = Mock(spec=ModelConfig)
64
model_config.enable_prompt_embeds = True
65
+ model_config.get_hidden_size.return_value = hidden_size
66
+ model_config.dtype = dtype
67
68
# construct arbitrary tensors of various dtypes, layouts, and sizes.
69
# We need to check against different layouts to make sure that if a user
tests/renderers/test_chat_utils_prompt_embeds.py
new
+576
@@ -0,0 +1,576 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Offline unit tests for `prompt_embeds` chat-completion content parts."""
5
+
6
+from __future__ import annotations
7
+
8
+import inspect
9
+import io
10
+from typing import Final
11
+from unittest import mock
12
+
13
+import pybase64 as base64
14
+import pytest
15
+import regex as re
16
+import torch
17
+from transformers import AutoTokenizer
18
+
19
+from vllm.entrypoints.chat_utils import (
20
+ _ENABLE_PROMPT_EMBEDS_ERROR,
21
+ _PROMPT_EMBEDS_MISSING_DATA_ERROR,
22
+ _RESERVED_PLACEHOLDER_IN_TEXT_ERROR,
23
+ MM_PARSER_MAP,
24
+ MODALITY_PLACEHOLDERS_MAP,
25
+ PROMPT_EMBEDS_PLACEHOLDER_TOKEN,
26
+ parse_chat_messages,
27
+ parse_chat_messages_async,
28
+)
29
+from vllm.renderers.hf import (
30
+ _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR,
31
+ _build_mixed_prompt_embeds,
32
+ _build_prompt_embeds_positions,
33
+ _build_prompt_embeds_updates,
34
+ _ensure_prompt_embeds_placeholder_token,
35
+ _expand_prompt_embeds_placeholders,
36
+)
37
+
38
+# Cover distinct tokenizer families:
39
+# GPT2TokenizerFast (BPE, OpenAI-style)
40
+# Qwen2TokenizerFast (SentencePiece BPE variant)
41
+# BertTokenizerFast (WordPiece)
42
+TOKENIZER_IDS: Final[list[str]] = [
43
+ "gpt2",
44
+ "Qwen/Qwen2.5-1.5B-Instruct",
45
+ "bert-base-uncased",
46
+]
47
+
48
+
49
+@pytest.fixture(params=TOKENIZER_IDS, ids=TOKENIZER_IDS)
50
+def tokenizer(request):
51
+ """A fresh tokenizer instance per tokenizer family."""
52
+ return AutoTokenizer.from_pretrained(request.param)
53
+
54
+
55
+# Minimal chat template that works with any tokenizer. Iterates
56
+# `message.content` as either a string or a list of dicts (openai format).
57
+_SIMPLE_CHAT_TEMPLATE: Final[str] = (
58
+ "{% for m in messages %}"
59
+ "{% if m['content'] is string %}{{m['content']}}"
60
+ "{% else %}{% for p in m['content'] %}{{p['text']}}{% endfor %}"
61
+ "{% endif %}\n{% endfor %}"
62
+)
63
+
64
+
65
+async def _maybe_await(fn, *args, **kwargs):
66
+ """Call *fn* and `await` the result if it's a coroutine."""
67
+ result = fn(*args, **kwargs)
68
+ if inspect.iscoroutine(result):
69
+ result = await result
70
+ return result
71
+
72
+
73
+# Parametrize over sync / async parse paths so every end-to-end test
74
+# exercises both.
75
+_PARSE_FUNCTIONS = [parse_chat_messages, parse_chat_messages_async]
76
+
77
+
78
+@pytest.fixture(params=_PARSE_FUNCTIONS, ids=["sync", "async"])
79
+def parse_fn(request):
80
+ """Either the sync or async `parse_chat_messages` callable."""
81
+ return request.param
82
+
83
+
84
+def _encode_tensor(t: torch.Tensor) -> str:
85
+ buf = io.BytesIO()
86
+ torch.save(t, buf)
87
+ return base64.b64encode(buf.getvalue()).decode("utf-8")
88
+
89
+
90
+_MOCK_HIDDEN_SIZE: Final[int] = 8
91
+_MOCK_DTYPE: Final[torch.dtype] = torch.float32
92
+
93
+
94
+def _make_mock_model_config(*, enable_prompt_embeds: bool = True) -> mock.MagicMock:
95
+ mc = mock.MagicMock()
96
+ mc.enable_prompt_embeds = enable_prompt_embeds
97
+ mc.multimodal_config = None
98
+ mc.allowed_local_media_path = None
99
+ mc.allowed_media_domains = None
100
+ # Test text-only code path in `MultiModalItemTracker.resolve_items`.
101
+ mc.is_multimodal_model = False
102
+ # `safe_load_prompt_embeds` pins each tensor to the model's hidden_size
103
+ # and dtype, so the mock must return concrete values.
104
+ mc.get_hidden_size.return_value = _MOCK_HIDDEN_SIZE
105
+ mc.dtype = _MOCK_DTYPE
106
+ return mc
107
+
108
+
109
+def test_prompt_embeds_keys_registered():
110
+ assert "prompt_embeds" in MODALITY_PLACEHOLDERS_MAP
111
+ assert MODALITY_PLACEHOLDERS_MAP["prompt_embeds"] == "<##PROMPT_EMBEDS##>"
112
+ assert "prompt_embeds" in MM_PARSER_MAP
113
+
114
+
115
+def test_ensure_placeholder_token_is_single_token_and_idempotent(tokenizer):
116
+ """Ensure the placeholder token is a single token and that multiple calls to
117
+ "ensure" are idempotent, across all tokenizer families."""
118
+ tid1 = _ensure_prompt_embeds_placeholder_token(tokenizer)
119
+ tid2 = _ensure_prompt_embeds_placeholder_token(tokenizer)
120
+ assert tid1 == tid2
121
+
122
+ ids = tokenizer.encode(PROMPT_EMBEDS_PLACEHOLDER_TOKEN, add_special_tokens=False)
123
+ assert ids == [tid1]
124
+
125
+ # Repeating it in a string N times must produce exactly that many tokens.
126
+ N = 5
127
+ ids_rep = tokenizer.encode(
128
+ PROMPT_EMBEDS_PLACEHOLDER_TOKEN * N, add_special_tokens=False
129
+ )
130
+ assert ids_rep == [tid1] * N
131
+
132
+
133
+def test_parse_chat_messages_openai_format():
134
+ NUM_TOKENS = 3
135
+ t = torch.randn(NUM_TOKENS, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE)
136
+ b64 = _encode_tensor(t)
137
+ mc = _make_mock_model_config()
138
+
139
+ messages = [
140
+ {
141
+ "role": "user",
142
+ "content": [
143
+ {"type": "text", "text": "Hello "},
144
+ {"type": "prompt_embeds", "data": b64},
145
+ {"type": "text", "text": " world"},
146
+ ],
147
+ }
148
+ ]
149
+ conv, mm_data, _ = parse_chat_messages(
150
+ messages,
151
+ mc,
152
+ content_format="openai",
153
+ )
154
+ # The middle content part is rewritten to a single placeholder-token
155
+ # sentinel.
156
+ texts = [p["text"] for p in conv[0]["content"]]
157
+ assert texts == [
158
+ "Hello ",
159
+ PROMPT_EMBEDS_PLACEHOLDER_TOKEN,
160
+ " world",
161
+ ]
162
+ assert mm_data is not None and "prompt_embeds" in mm_data
163
+ assert torch.equal(mm_data["prompt_embeds"][0], t)
164
+
165
+
166
+# Each layout entry is one content part:
167
+# ("text", "A") -> {"type": "text", "text": "A"}
168
+# ("embed", N) -> {"type": "prompt_embeds", "data": <base64 of (N, H) tensor>}
169
+@pytest.mark.parametrize(
170
+ "layout",
171
+ [
172
+ # Case: Single embed only.
173
+ [("embed", 2)],
174
+ # Case: Embed at the start of the message.
175
+ [("embed", 3), ("text", "B")],
176
+ # Case: Embed at the end of the message.
177
+ [("text", "A"), ("embed", 1)],
178
+ # Case: Embed sandwiched between text spans.
179
+ [("text", "A"), ("embed", 2), ("text", "B")],
180
+ # Case: Multiple embeds with text in between.
181
+ [("text", "A"), ("embed", 2), ("text", "B"), ("embed", 3)],
182
+ # Case: Adjacent embeds with no separating text.
183
+ [("embed", 1), ("embed", 2)],
184
+ # Case: Multiple text spans before a trailing embed.
185
+ [("text", "A"), ("text", "B"), ("embed", 1)],
186
+ # Case: Long-ish run mixing both kinds.
187
+ [
188
+ ("text", "head"),
189
+ ("embed", 4),
190
+ ("text", "mid"),
191
+ ("embed", 1),
192
+ ("embed", 2),
193
+ ("text", "tail"),
194
+ ],
195
+ ],
196
+ ids=[
197
+ "single-embed",
198
+ "embed-then-text",
199
+ "text-then-embed",
200
+ "text-embed-text",
201
+ "text-embed-text-embed",
202
+ "adjacent-embeds",
203
+ "text-text-embed",
204
+ "long-mixed-run",
205
+ ],
206
+)
207
+@pytest.mark.parametrize(
208
+ "interleave_mm_strings",
209
+ # `None`: text-only path where `multimodal_config` is absent.
210
+ # `False`: non-interleave multimodal path (the common default).
211
+ # `True`: sentinel-substitution interleave path.
212
+ # All three must preserve the request ordering of prompt_embeds
213
+ # relative to surrounding text because prompt_embeds are spliced at the
214
+ # token offset during rendering.
215
+ [None, False, True],
216
+ ids=["text-only", "interleave-off", "interleave-on"],
217
+)
218
+def test_parse_chat_messages_string_format_preserves_position(
219
+ layout, interleave_mm_strings
220
+):
221
+ mc = _make_mock_model_config()
222
+ if interleave_mm_strings is not None:
223
+ mm_cfg = mock.MagicMock()
224
+ mm_cfg.interleave_mm_strings = interleave_mm_strings
225
+ mc.multimodal_config = mm_cfg
226
+
227
+ content: list[dict] = []
228
+ expected_parts: list[str] = []
229
+ expected_embeds: list[torch.Tensor] = []
230
+ for kind, value in layout:
231
+ if kind == "text":
232
+ content.append({"type": "text", "text": value})
233
+ expected_parts.append(value)
234
+ else: # prompt embeds
235
+ num_tokens = value
236
+ t = torch.randn(num_tokens, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE)
237
+ expected_embeds.append(t)
238
+ content.append({"type": "prompt_embeds", "data": _encode_tensor(t)})
239
+ # Parser emits ONE sentinel per part.
240
+ expected_parts.append(PROMPT_EMBEDS_PLACEHOLDER_TOKEN)
241
+
242
+ messages = [{"role": "user", "content": content}]
243
+ conv, mm_data, _ = parse_chat_messages(
244
+ messages,
245
+ mc,
246
+ content_format="string",
247
+ )
248
+
249
+ assert conv[0]["content"] == "\n".join(expected_parts)
250
+ assert mm_data is not None and "prompt_embeds" in mm_data
251
+ assert len(mm_data["prompt_embeds"]) == len(expected_embeds)
252
+ for got, want in zip(mm_data["prompt_embeds"], expected_embeds, strict=True):
253
+ assert torch.equal(got, want)
254
+
255
+
256
+def test_parse_chat_messages_requires_flag():
257
+ t = torch.randn(2, 4)
258
+ b64 = _encode_tensor(t)
259
+ mc = _make_mock_model_config(enable_prompt_embeds=False)
260
+
261
+ messages = [
262
+ {
263
+ "role": "user",
264
+ "content": [{"type": "prompt_embeds", "data": b64}],
265
+ }
266
+ ]
267
+ with pytest.raises(ValueError, match=_ENABLE_PROMPT_EMBEDS_ERROR):
268
+ parse_chat_messages(
269
+ messages,
270
+ mc,
271
+ content_format="openai",
272
+ )
273
+
274
+
275
+def test_parse_chat_messages_rejects_missing_data():
276
+ # `data` is marked `Required` on `ChatCompletionContentPartPromptEmbedsParam`;
277
+ # malformed requests without `data` must surface a clear validation error
278
+ # rather than being silently dropped.
279
+ mc = _make_mock_model_config()
280
+ messages = [
281
+ {
282
+ "role": "user",
283
+ "content": [{"type": "prompt_embeds"}], # no `data`
284
+ }
285
+ ]
286
+ with pytest.raises(ValueError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR):
287
+ parse_chat_messages(
288
+ messages,
289
+ mc,
290
+ content_format="openai",
291
+ )
292
+
293
+
294
+# Reserved placeholder guard: when `enable_prompt_embeds=True` the tokenizer is
295
+# mutated to make `<prompt_embeds>` a single unsplittable token. Any user text
296
+# containing that literal sequence would tokenize to the same sentinel ID and
297
+# be mistaken for a splice point, so we reject it at parse time.
298
+_PLACEHOLDER_ERROR_PATTERN: Final[str] = re.sub(
299
+ r"\\{[^}]*\\}", ".*", re.escape(_RESERVED_PLACEHOLDER_IN_TEXT_ERROR)
300
+)
301
+
302
+
303
+@pytest.mark.parametrize(
304
+ "content",
305
+ [
306
+ # Case: Top-level string content (wrapped as a single text part).
307
+ f"hello {PROMPT_EMBEDS_PLACEHOLDER_TOKEN} world",
308
+ # Case: List with a typed text part containing the placeholder.
309
+ [{"type": "text", "text": f"leading {PROMPT_EMBEDS_PLACEHOLDER_TOKEN}"}],
310
+ # Case: List with a plain-string part (no wrapping dict).
311
+ [f"raw string {PROMPT_EMBEDS_PLACEHOLDER_TOKEN}"],
312
+ ],
313
+ ids=["top-level-string", "typed-text-part", "plain-string-part"],
314
+)
315
+def test_parse_chat_messages_rejects_placeholder_in_user_text(content):
316
+ mc = _make_mock_model_config() # enable_prompt_embeds=True by default
317
+ messages = [{"role": "user", "content": content}]
318
+ with pytest.raises(ValueError, match=_PLACEHOLDER_ERROR_PATTERN):
319
+ parse_chat_messages(messages, mc, content_format="openai")
320
+
321
+
322
+def test_parse_chat_messages_allows_placeholder_in_text_when_feature_disabled():
323
+ # When `enable_prompt_embeds=False` the tokenizer is never mutated, so the
324
+ # literal `<prompt_embeds>` is just ordinary text and must pass through.
325
+ mc = _make_mock_model_config(enable_prompt_embeds=False)
326
+ messages = [
327
+ {
328
+ "role": "user",
329
+ "content": f"benign mention of {PROMPT_EMBEDS_PLACEHOLDER_TOKEN} here",
330
+ }
331
+ ]
332
+ conv, mm_data, _ = parse_chat_messages(messages, mc, content_format="openai")
333
+ assert mm_data is None or "prompt_embeds" not in mm_data
334
+ # Text reaches the rendered conversation unchanged.
335
+ texts = [p["text"] for p in conv[0]["content"]]
336
+ assert PROMPT_EMBEDS_PLACEHOLDER_TOKEN in "".join(texts)
337
+
338
+
339
+# Token-stream spec: ints are regular token IDs, tuples `(N,)` expand to
340
+# a placeholder span of length N (creates corresponding `(N, H)` tensor).
341
+# `expected` lists the `(start_idx, length)` pairs that
342
+# `_build_prompt_embeds_positions` should return.
343
+@pytest.mark.parametrize(
344
+ "stream, expected",
345
+ [
346
+ # Case: Single run in the middle.
347
+ ([10, 20, (3,), 30], [(2, 3)]),
348
+ # Case: Single run at the start.
349
+ ([(2,), 10, 20], [(0, 2)]),
350
+ # Case: Single run at the end.
351
+ ([10, 20, (4,)], [(2, 4)]),
352
+ # Case: Two runs with tokens between.
353
+ ([1, (2,), 2, 3, (3,), 4], [(1, 2), (5, 3)]),
354
+ # Case: Adjacent runs (no separating tokens).
355
+ ([(1,), (2,)], [(0, 1), (1, 2)]),
356
+ # Case: Three runs.
357
+ ([5, (2,), 6, (1,), 7, (3,), 8], [(1, 2), (4, 1), (6, 3)]),
358
+ ],
359
+ ids=[
360
+ "single-middle",
361
+ "single-start",
362
+ "single-end",
363
+ "two-runs-separated",
364
+ "two-runs-adjacent",
365
+ "three-runs",
366
+ ],
367
+)
368
+def test_build_positions(tokenizer, stream, expected):
369
+ H = 4
370
+ tid = _ensure_prompt_embeds_placeholder_token(tokenizer)
371
+ tensors: list[torch.Tensor] = []
372
+ token_ids: list[int] = []
373
+ for item in stream:
374
+ if isinstance(item, tuple):
375
+ length = item[0]
376
+ tensors.append(torch.randn(length, H))
377
+ token_ids.extend([tid] * length)
378
+ else:
379
+ token_ids.append(item)
380
+ mm_updates = _build_prompt_embeds_updates(tensors, tid)
381
+ positions = _build_prompt_embeds_positions(token_ids, len(tensors), mm_updates)
382
+ assert positions == expected
383
+
384
+
385
+def test_build_positions_length_mismatch(tokenizer):
386
+ N1, H1 = 2, 4
387
+ N2, H2 = 3, 4
388
+ tid = _ensure_prompt_embeds_placeholder_token(tokenizer)
389
+ # 2 tensors expected but only a single placeholder run in the token
390
+ # stream (simulating dropping the second one).
391
+ tensors = [torch.randn(N1, H1), torch.randn(N2, H2)]
392
+ token_ids = [1, tid, tid, 2, 3]
393
+ mm_updates = _build_prompt_embeds_updates(tensors, tid)
394
+ # The error constant is a `str.format` template, escape it and turn
395
+ # the `{field}` placeholders into `.*` so it matches any substitution.
396
+ pattern = re.sub(
397
+ r"\\{[^}]*\\}", ".*", re.escape(_PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR)
398
+ )
399
+ with pytest.raises(ValueError, match=pattern):
400
+ _build_prompt_embeds_positions(token_ids, len(tensors), mm_updates)
401
+
402
+
403
+# ints = regular token IDs (any value)
404
+# (N,) = embed span of length N
405
+@pytest.mark.parametrize(
406
+ "stream",
407
+ [
408
+ [10, 20, (3,), 30],
409
+ [(2,), 10, 20],
410
+ [10, 20, (4,)],
411
+ [1, (2,), 2, 3, (3,), 4],
412
+ [(1,), (2,)],
413
+ [5, (2,), 6, (1,), 7, (3,), 8],
414
+ ],
415
+ ids=[
416
+ "single-middle",
417
+ "single-start",
418
+ "single-end",
419
+ "two-spans-separated",
420
+ "two-spans-adjacent",
421
+ "three-spans",
422
+ ],
423
+)
424
+def test_build_mixed_prompt_embeds(stream):
425
+ H = 8
426
+ _PLACEHOLDER = 0 # sentinel for embed positions in token_ids
427
+
428
+ tensors: list[torch.Tensor] = []
429
+ token_ids: list[int] = []
430
+ positions: list[tuple[int, int]] = []
431
+ cursor = 0
432
+ for item in stream:
433
+ if isinstance(item, tuple):
434
+ length = item[0]
435
+ tensors.append(torch.randn(length, H))
436
+ positions.append((cursor, length))
437
+ token_ids.extend([_PLACEHOLDER] * length)
438
+ cursor += length
439
+ else:
440
+ token_ids.append(item)
441
+ cursor += 1
442
+
443
+ embeds, mask = _build_mixed_prompt_embeds(token_ids, tensors, positions)
444
+
445
+ assert embeds.shape == (len(token_ids), H)
446
+ assert len(mask) == len(token_ids)
447
+
448
+ # Mask: False exactly at embed positions, True everywhere else.
449
+ expected_mask = torch.ones(len(token_ids), dtype=torch.bool)
450
+ for start, length in positions:
451
+ expected_mask[start : start + length] = False
452
+ assert mask == expected_mask.tolist()
453
+
454
+ # Embed rows match input tensors at the right positions.
455
+ for tensor, (start, length) in zip(tensors, positions):
456
+ assert torch.equal(embeds[start : start + length], tensor)
457
+
458
+ # Non-embed positions remain zero-filled.
459
+ assert torch.all(embeds[expected_mask] == 0)
460
+
461
+
462
+# End-to-end tests: each runs both sync and async parse paths via the
463
+# `parse_fn` fixture.
464
+
465
+
466
+@pytest.mark.asyncio
467
+@pytest.mark.parametrize("role", ["user", "system"])
468
+async def test_end_to_end_expand_and_build(tokenizer, parse_fn, role):
469
+ """Full renderer pipeline: parse -> chat template -> expand -> locate
470
+ -> build mixed prompt, across tokenizers, roles, and sync/async."""
471
+ tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE
472
+ tid = _ensure_prompt_embeds_placeholder_token(tokenizer)
473
+
474
+ LEN_A, LEN_B = 3, 2
475
+ t_a = torch.randn(LEN_A, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE)
476
+ t_b = torch.randn(LEN_B, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE)
477
+ NUM_TENSORS = 2
478
+
479
+ mc = _make_mock_model_config()
480
+
481
+ messages = [
482
+ {
483
+ "role": role,
484
+ "content": [
485
+ {"type": "text", "text": "Hello "},
486
+ {"type": "prompt_embeds", "data": _encode_tensor(t_a)},
487
+ {"type": "text", "text": " world "},
488
+ {"type": "prompt_embeds", "data": _encode_tensor(t_b)},
489
+ {"type": "text", "text": "!"},
490
+ ],
491
+ }
492
+ ]
493
+
494
+ conv, mm_data, _ = await _maybe_await(
495
+ parse_fn, messages, mc, content_format="openai"
496
+ )
497
+ tensors = list(mm_data["prompt_embeds"])
498
+ assert len(tensors) == NUM_TENSORS
499
+
500
+ # Tokenize: each prompt_embeds part becomes 1 placeholder token.
501
+ # `return_dict=False` to get a flat `list[int]` on transformers v5
502
+ # (where the default flipped to True and yields a `BatchEncoding` dict).
503
+ token_ids = tokenizer.apply_chat_template(conv, tokenize=True, return_dict=False)
504
+ assert sum(t == tid for t in token_ids) == NUM_TENSORS
505
+
506
+ # Expand, locate, and build.
507
+ mm_updates = _build_prompt_embeds_updates(tensors, tid)
508
+ expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates)
509
+ assert len(expanded) == len(token_ids) + LEN_A + LEN_B - NUM_TENSORS
510
+
511
+ positions = _build_prompt_embeds_positions(expanded, len(tensors), mm_updates)
512
+ assert positions[0][1] == LEN_A
513
+ assert positions[1][1] == LEN_B
514
+
515
+ embeds, mask = _build_mixed_prompt_embeds(expanded, tensors, positions)
516
+ assert embeds.shape == (len(expanded), _MOCK_HIDDEN_SIZE)
517
+ assert mask.count(False) == LEN_A + LEN_B
518
+ assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_A], t_a)
519
+ assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_B], t_b)
520
+
521
+
522
+@pytest.mark.asyncio
523
+async def test_end_to_end_multi_message_conversation(tokenizer, parse_fn):
524
+ """Full pipeline with prompt_embeds spread across system + user messages,
525
+ verifying ordering and positioning in the final token stream."""
526
+ tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE
527
+ tid = _ensure_prompt_embeds_placeholder_token(tokenizer)
528
+
529
+ LEN_SYS, LEN_USR = 4, 3
530
+ t_sys = torch.randn(LEN_SYS, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE)
531
+ t_usr = torch.randn(LEN_USR, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE)
532
+ NUM_TENSORS = 2 # t_sys and t_usr.
533
+
534
+ mc = _make_mock_model_config()
535
+
536
+ messages = [
537
+ {
538
+ "role": "system",
539
+ "content": [
540
+ {"type": "text", "text": "You are helpful."},
541
+ {"type": "prompt_embeds", "data": _encode_tensor(t_sys)},
542
+ ],
543
+ },
544
+ {
545
+ "role": "user",
546
+ "content": [
547
+ {"type": "prompt_embeds", "data": _encode_tensor(t_usr)},
548
+ {"type": "text", "text": "Summarize."},
549
+ ],
550
+ },
551
+ ]
552
+
553
+ conv, mm_data, _ = await _maybe_await(
554
+ parse_fn, messages, mc, content_format="openai"
555
+ )
556
+ tensors = list(mm_data["prompt_embeds"])
557
+ assert len(tensors) == NUM_TENSORS
558
+
559
+ # Tokenize, expand, locate, and build.
560
+ # `return_dict=False` to get a flat `list[int]` on transformers v5
561
+ # (where the default flipped to True and yields a `BatchEncoding` dict).
562
+ token_ids = tokenizer.apply_chat_template(conv, tokenize=True, return_dict=False)
563
+ mm_updates = _build_prompt_embeds_updates(tensors, tid)
564
+ expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates)
565
+ positions = _build_prompt_embeds_positions(expanded, len(tensors), mm_updates)
566
+
567
+ assert positions[0][1] == LEN_SYS
568
+ assert positions[1][1] == LEN_USR
569
+ # System embed must appear before user embed in the token stream.
570
+ assert positions[0][0] < positions[1][0]
571
+
572
+ embeds, mask = _build_mixed_prompt_embeds(expanded, tensors, positions)
573
+ assert embeds.shape == (len(expanded), _MOCK_HIDDEN_SIZE)
574
+ assert mask.count(False) == LEN_SYS + LEN_USR
575
+ assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_SYS], t_sys)
576
+ assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_USR], t_usr)
tests/renderers/test_completions.py
+17
-8
@@ -39,6 +39,11 @@ class MockModelConfig:
39
is_encoder_decoder: bool = False
40
is_multimodal_model: bool = False
41
renderer_num_workers: int = 1
42
+ hidden_size: int = 768
43
+ dtype: torch.dtype = torch.float32
44
+
45
+ def get_hidden_size(self) -> int:
46
+ return self.hidden_size
47
48
49
@dataclass
@@ -384,12 +389,13 @@ class TestRenderEmbedPrompt:
389
assert torch.equal(results[0]["prompt_embeds"], tensor_input)
390
391
def test_multiple_prompt_embeds(self):
387
- renderer = _build_renderer(MockModelConfig())
392
+ hidden_size = 512
393
+ renderer = _build_renderer(MockModelConfig(hidden_size=hidden_size))
394
395
# Create multiple test tensors
396
tensor_inputs = [
391
- torch.randn(8, 512, dtype=torch.float32),
392
- torch.randn(12, 512, dtype=torch.float32),
397
+ torch.randn(8, hidden_size, dtype=torch.float32),
398
+ torch.randn(12, hidden_size, dtype=torch.float32),
399
]
400
401
prompts = renderer.render_prompts(
@@ -432,13 +438,15 @@ class TestRenderEmbedPrompt:
438
assert torch.equal(results[0]["prompt_embeds"], expected)
439
440
def test_prompt_embed_different_dtypes(self):
435
- renderer = _build_renderer(MockModelConfig())
436
-
441
+ hidden_size = 256
442
# Test different supported dtypes
443
dtypes = [torch.float32, torch.float16, torch.bfloat16]
444
445
for dtype in dtypes:
441
- tensor_input = torch.randn(5, 256, dtype=dtype)
446
+ renderer = _build_renderer(
447
+ MockModelConfig(hidden_size=hidden_size, dtype=dtype)
448
+ )
449
+ tensor_input = torch.randn(5, hidden_size, dtype=dtype)
450
451
prompts = renderer.render_prompts(
452
_preprocess_prompt(
@@ -474,10 +482,11 @@ class TestRenderEmbedPrompt:
482
assert results[0]["prompt_embeds"].shape == (10, 768)
483
484
def test_both_prompts_and_embeds(self):
477
- renderer = _build_renderer(MockModelConfig())
485
+ hidden_size = 256
486
+ renderer = _build_renderer(MockModelConfig(hidden_size=hidden_size))
487
488
text_input = "Hello world"
480
- tensor_input = torch.randn(5, 256, dtype=torch.float32)
489
+ tensor_input = torch.randn(5, hidden_size, dtype=torch.float32)
490
491
prompts = renderer.render_prompts(
492
_preprocess_prompt(
tests/renderers/test_sparse_tensor_validation.py
+67
-4
@@ -12,6 +12,7 @@ import pybase64 as base64
12
import pytest
13
import torch
14
15
+from vllm.exceptions import VLLMValidationError
16
from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO
17
from vllm.renderers.embed_utils import safe_load_prompt_embeds
18
@@ -53,8 +54,14 @@ def _create_malicious_sparse_tensor() -> torch.Tensor:
54
values = torch.tensor([1.0])
55
shape = (3, 3)
56
56
- # Create sparse tensor (this will be invalid)
57
- sparse_tensor = torch.sparse_coo_tensor(indices, values, shape, dtype=torch.float32)
57
+ # Create sparse tensor (this will be invalid). Pass `check_invariants=False`
58
+ # explicitly so this fixture is robust to process-wide invariant-check state
59
+ # left enabled by other tests (the global flag isn't thread-local, and
60
+ # concurrent users of the `check_sparse_tensor_invariants` context manager
61
+ # can leak the "enabled" state across tests).
62
+ sparse_tensor = torch.sparse_coo_tensor(
63
+ indices, values, shape, dtype=torch.float32, check_invariants=False
64
+ )
65
return sparse_tensor
66
67
@@ -117,7 +124,7 @@ class TestPromptEmbedsValidation:
124
shape = (10, 10)
125
126
malicious_tensor = torch.sparse_coo_tensor(
120
- indices, values, shape, dtype=torch.float32
127
+ indices, values, shape, dtype=torch.float32, check_invariants=False
128
)
129
encoded = _encode_tensor(malicious_tensor)
130
@@ -132,13 +139,69 @@ class TestPromptEmbedsValidation:
139
shape = (10, 10)
140
141
malicious_tensor = torch.sparse_coo_tensor(
135
- indices, values, shape, dtype=torch.float32
142
+ indices, values, shape, dtype=torch.float32, check_invariants=False
143
)
144
encoded = _encode_tensor(malicious_tensor)
145
146
with pytest.raises((RuntimeError, ValueError)):
147
safe_load_prompt_embeds(model_config, encoded)
148
149
+ def test_hidden_size_mismatch_rejected(self, model_config):
150
+ """Tensors whose trailing dim doesn't match the model's hidden_size
151
+ must be rejected at parse time."""
152
+ # opt-125m has hidden_size=768, passing 512 triggers the check.
153
+ wrong_hidden = torch.randn(10, 512, dtype=torch.float32)
154
+ encoded = _encode_tensor(wrong_hidden)
155
+
156
+ with pytest.raises(VLLMValidationError, match="hidden_size"):
157
+ safe_load_prompt_embeds(model_config, encoded)
158
+
159
+ def test_float_dtype_mismatch_cast_to_model_dtype(self, model_config):
160
+ """Tensors whose dtype doesn't match the model's dtype but are still
161
+ floating-point are cast, since API clients generally can't know the
162
+ server's `--dtype` setting ahead of time."""
163
+ # Fixture pins model dtype to float32, upload a bfloat16 tensor.
164
+ mismatched_float = torch.randn(10, 768, dtype=torch.bfloat16)
165
+ encoded = _encode_tensor(mismatched_float)
166
+
167
+ result = safe_load_prompt_embeds(model_config, encoded)
168
+
169
+ assert result.dtype == torch.float32
170
+ assert result.shape == mismatched_float.shape
171
+
172
+ def test_non_float_dtype_rejected(self, model_config):
173
+ """Non-floating-point dtypes cannot be safely cast for embeddings
174
+ (e.g. integer tensors almost certainly indicate caller confusion),
175
+ so they are rejected at parse time."""
176
+ non_float = torch.randint(0, 100, (10, 768), dtype=torch.int32)
177
+ encoded = _encode_tensor(non_float)
178
+
179
+ with pytest.raises(VLLMValidationError, match="floating-point"):
180
+ safe_load_prompt_embeds(model_config, encoded)
181
+
182
+ def test_non_2d_tensor_rejected(self, model_config):
183
+ """Tensors that aren't 2D (even after squeezing a leading dim)
184
+ must be rejected with a clear error."""
185
+ # A 1D tensor cannot be interpreted as (num_tokens, hidden_size).
186
+ bad = torch.randn(768, dtype=torch.float32)
187
+ encoded = _encode_tensor(bad)
188
+
189
+ with pytest.raises(VLLMValidationError, match="2D tensor"):
190
+ safe_load_prompt_embeds(model_config, encoded)
191
+
192
+ def test_non_tensor_payload_rejected(self, model_config):
193
+ """Deserializing to a non-Tensor object must raise a clear error
194
+ instead of propagating an AssertionError."""
195
+ # `torch.save` will serialize a plain dict; `weights_only=True` allows
196
+ # loading built-in containers, so this exercises the isinstance check.
197
+ buffer = io.BytesIO()
198
+ torch.save({"not": "a tensor"}, buffer)
199
+ buffer.seek(0)
200
+ encoded = base64.b64encode(buffer.read())
201
+
202
+ with pytest.raises(VLLMValidationError, match="torch.Tensor"):
203
+ safe_load_prompt_embeds(model_config, encoded)
204
+
205
206
class TestImageEmbedsValidation:
207
"""Test sparse tensor validation in image embeddings (Chat API)."""
tests/tokenizers_/test_deepseek_v4.py
+1
@@ -40,6 +40,7 @@ def _model_config():
40
multimodal_config=None,
41
allowed_local_media_path="",
42
allowed_media_domains=None,
43
+ enable_prompt_embeds=False,
44
)
45
46
tests/v1/engine/test_async_llm.py
+7
-4
@@ -256,8 +256,10 @@ async def test_multi_abort(output_kind: RequestOutputKind):
256
)
257
)
258
259
- # Let requests start
260
- await asyncio.sleep(0.5)
259
+ # Let requests start generating, use a longer sleep to ensure all
260
+ # requests have exited prefill and produced at least one
261
+ # decode token before we abort.
262
+ await asyncio.sleep(1.0)
263
264
# Use multi-abort to abort multiple requests at once
265
abort_request_ids = [request_ids[i] for i in REQUEST_IDS_TO_ABORT]
@@ -369,9 +371,10 @@ async def test_mid_stream_cancellation(
371
# Wait for all tasks to complete
372
results = await asyncio.gather(*tasks)
373
372
- # Verify all tasks were cancelled at the expected point
374
+ # Verify all tasks were cancelled at the expected point.
375
+ # Uses >= because the cancel check is `count >= cancel_after`.
376
for num_generated_tokens, request_id in results:
374
- assert num_generated_tokens == NUM_EXPECTED_TOKENS, (
377
+ assert num_generated_tokens >= NUM_EXPECTED_TOKENS, (
378
f"{request_id} generated {num_generated_tokens} tokens but "
379
f"expected to cancel after {NUM_EXPECTED_TOKENS}"
380
)
vllm/entrypoints/chat_utils.py
+201
-9
@@ -11,7 +11,7 @@ from dataclasses import dataclass
11
from functools import cached_property, lru_cache, partial
12
from itertools import accumulate
13
from pathlib import Path
14
-from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar, cast
14
+from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeAlias, TypeVar, cast
15
16
from openai.types.chat import (
17
ChatCompletionAssistantMessageParam,
@@ -36,7 +36,7 @@ from PIL import Image
36
from pydantic import BaseModel, ConfigDict, TypeAdapter
37
38
# pydantic needs the TypedDict from typing_extensions
39
-from typing_extensions import Required, TypedDict
39
+from typing_extensions import Required, TypedDict, override
40
41
from vllm import envs
42
from vllm.config import ModelConfig
@@ -55,6 +55,10 @@ from vllm.multimodal.inputs import (
55
)
56
from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector
57
from vllm.multimodal.processing import BaseMultiModalProcessor
58
+from vllm.renderers.embed_utils import (
59
+ safe_load_prompt_embeds,
60
+ safe_load_prompt_embeds_async,
61
+)
62
from vllm.utils import random_uuid
63
from vllm.utils.collection_utils import is_list_of
64
from vllm.utils.import_utils import LazyLoader
@@ -98,9 +102,40 @@ MODALITY_PLACEHOLDERS_MAP = {
102
"image": "<##IMAGE##>",
103
"audio": "<##AUDIO##>",
104
"video": "<##VIDEO##>",
105
+ "prompt_embeds": "<##PROMPT_EMBEDS##>",
106
}
107
108
109
+PROMPT_EMBEDS_PLACEHOLDER_TOKEN: Final[str] = "<prompt_embeds>"
110
+"""The special token used as a placeholder for each embedding
111
+position during chat template rendering.
112
+
113
+Registered as an additional special token when `--enable-prompt-embeds` is set.
114
+See `_ensure_prompt_embeds_placeholder_token` in `vllm/renderers/hf.py`.
115
+"""
116
+
117
+
118
+_REQUIRE_MM_PROCESSOR_ERROR: Final[str] = (
119
+ "Resolving modality {modality!r} requires a multimodal processor "
120
+ "but none is available."
121
+)
122
+
123
+_ENABLE_PROMPT_EMBEDS_ERROR: Final[str] = (
124
+ "You must set `--enable-prompt-embeds` to input `prompt_embeds`"
125
+)
126
+
127
+_PROMPT_EMBEDS_MISSING_DATA_ERROR: Final[str] = (
128
+ "prompt_embeds content part requires a non-empty `data` field "
129
+ "with base64-encoded tensor bytes."
130
+)
131
+
132
+_RESERVED_PLACEHOLDER_IN_TEXT_ERROR: Final[str] = (
133
+ "Text content may not contain the reserved placeholder {token!r}. "
134
+ "This placeholder is used internally to mark `prompt_embeds` splice "
135
+ "positions in the tokenized prompt."
136
+)
137
+
138
+
139
class AudioURL(TypedDict, total=False):
140
url: Required[str]
141
"""
@@ -147,6 +182,17 @@ class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False):
182
"""
183
184
185
+class ChatCompletionContentPartPromptEmbedsParam(TypedDict, total=False):
186
+ data: Required[str]
187
+ """
188
+ Base64-encoded bytes of a serialized `torch.Tensor` of shape
189
+ `(num_tokens, hidden_size)`. The tensor's `dtype` and `hidden_size` must
190
+ match the model's input embedding layer.
191
+ """
192
+ type: Required[Literal["prompt_embeds"]]
193
+ """The type of the content part."""
194
+
195
+
196
class VideoURL(TypedDict, total=False):
197
url: Required[str]
198
"""
@@ -282,6 +328,7 @@ ChatCompletionContentPartParam: TypeAlias = (
328
| CustomChatCompletionContentSimpleImageParam
329
| ChatCompletionContentPartImageEmbedsParam
330
| ChatCompletionContentPartAudioEmbedsParam
331
+ | ChatCompletionContentPartPromptEmbedsParam
332
| CustomChatCompletionContentSimpleAudioParam
333
| CustomChatCompletionContentSimpleVideoParam
334
| CustomChatCompletionContentToolReferenceParam
@@ -367,7 +414,13 @@ ChatTemplateContentFormat = Literal["string", "openai"]
414
415
416
ModalityStr = Literal[
370
- "image", "audio", "video", "image_embeds", "audio_embeds", "vision_chunk"
417
+ "image",
418
+ "audio",
419
+ "video",
420
+ "image_embeds",
421
+ "audio_embeds",
422
+ "vision_chunk",
423
+ "prompt_embeds",
424
]
425
_T = TypeVar("_T")
426
@@ -549,7 +602,17 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
602
603
An optional uuid can be added which serves as a unique identifier of the
604
media.
605
+
606
+ Note:
607
+ `prompt_embeds` bypass MM-processor validation because they are
608
+ pre-computed embeddings that do not go through any HF processor, encoder,
609
+ or model-specific placeholder logic. The corresponding placeholder string is
610
+ managed by the parser via `_add_placeholder`, so we return None here.
611
"""
612
+ if modality == "prompt_embeds":
613
+ self._items_by_modality["prompt_embeds"].append(item)
614
+ return None
615
+
616
input_modality = modality.replace("_embeds", "")
617
original_modality = modality
618
use_vision_chunk = (
@@ -660,17 +723,32 @@ def _resolve_vision_chunk_items(
723
724
def _resolve_items(
725
items_by_modality: dict[str, list[tuple[object, str | None]]],
663
- mm_processor: BaseMultiModalProcessor,
726
+ mm_processor: BaseMultiModalProcessor | None,
727
modality_order: dict[str, list[str]],
728
) -> tuple[MultiModalDataDict, MultiModalUUIDDict]:
729
+ """
730
+ Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`.
731
+
732
+ Note:
733
+ `mm_processor` is `None` for text-only models (no registered HF
734
+ processor) whose only modality is `prompt_embeds`. Every other
735
+ modality requires a processor, enforced by the guard below.
736
+ """
737
if "image" in items_by_modality and "image_embeds" in items_by_modality:
738
raise ValueError("Mixing raw image and embedding inputs is not allowed")
739
if "audio" in items_by_modality and "audio_embeds" in items_by_modality:
740
raise ValueError("Mixing raw audio and embedding inputs is not allowed")
741
+ # `prompt_embeds` bypasses HF MM processors. Every other modality requires one.
742
+ processor_modalities = items_by_modality.keys() - {"prompt_embeds"}
743
+ if processor_modalities and mm_processor is None:
744
+ raise RuntimeError(
745
+ _REQUIRE_MM_PROCESSOR_ERROR.format(modality=processor_modalities)
746
+ )
747
748
mm_data = {}
749
mm_uuids = {}
750
if "image_embeds" in items_by_modality:
751
+ assert mm_processor is not None
752
mm_data["image"] = _get_embeds_data(
753
"image",
754
[data for data, uuid in items_by_modality["image_embeds"]],
@@ -681,6 +759,7 @@ def _resolve_items(
759
mm_data["image"] = [data for data, uuid in items_by_modality["image"]]
760
mm_uuids["image"] = [uuid for data, uuid in items_by_modality["image"]]
761
if "audio_embeds" in items_by_modality:
762
+ assert mm_processor is not None
763
mm_data["audio"] = _get_embeds_data(
764
"audio",
765
[data for data, uuid in items_by_modality["audio_embeds"]],
@@ -694,6 +773,7 @@ def _resolve_items(
773
mm_data["video"] = [data for data, uuid in items_by_modality["video"]]
774
mm_uuids["video"] = [uuid for data, uuid in items_by_modality["video"]]
775
if "vision_chunk" in items_by_modality:
776
+ assert mm_processor is not None
777
# Process vision_chunk items - extract from (data, modality) tuples
778
# and convert to VisionChunk types with proper UUID handling
779
processed_chunks, vision_chunk_uuids = _resolve_vision_chunk_items(
@@ -703,6 +783,10 @@ def _resolve_items(
783
)
784
mm_data["vision_chunk"] = processed_chunks
785
mm_uuids["vision_chunk"] = vision_chunk_uuids
786
+ if "prompt_embeds" in items_by_modality:
787
+ mm_data["prompt_embeds"] = [
788
+ data for data, _uuid in items_by_modality["prompt_embeds"]
789
+ ]
790
791
return mm_data, mm_uuids
792
@@ -714,8 +798,16 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[tuple[object, str | None]]
798
if not self._items_by_modality:
799
return None, None
800
801
+ # Text-only models (`is_multimodal_model=False`) with inputs of
802
+ # modality `prompt_embeds` have no MM processor since `prompt_embeds` are
803
+ # pre-computed and require no processing, so we pass `None`.
804
+ mm_processor = (
805
+ self.mm_processor if self._model_config.is_multimodal_model else None
806
+ )
807
return _resolve_items(
718
- dict(self._items_by_modality), self.mm_processor, self._modality_order
808
+ dict(self._items_by_modality),
809
+ mm_processor,
810
+ self._modality_order,
811
)
812
813
def create_parser(
@@ -738,8 +830,13 @@ class AsyncMultiModalItemTracker(
830
for modality, coros in self._items_by_modality.items()
831
}
832
833
+ mm_processor = (
834
+ self.mm_processor if self._model_config.is_multimodal_model else None
835
+ )
836
return _resolve_items(
742
- resolved_items_by_modality, self.mm_processor, self._modality_order
837
+ resolved_items_by_modality,
838
+ mm_processor,
839
+ self._modality_order,
840
)
841
842
def create_parser(
@@ -758,10 +855,16 @@ class BaseMultiModalContentParser(ABC):
855
# general MM placeholder:
856
# {
857
# "<##IMAGE##>": ["<image>", "<image>", "<image>"],
761
- # "<##AUDIO##>": ["<audio>", "<audio>"]
858
+ # "<##AUDIO##>": ["<audio>", "<audio>"],
859
+ # "<##PROMPT_EMBEDS##>": ["<prompt_embeds>", "<prompt_embeds>"]
860
# }
861
self._placeholder_storage: dict[str, list] = defaultdict(list)
862
863
+ @property
864
+ @abstractmethod
865
+ def model_config(self) -> ModelConfig:
866
+ raise NotImplementedError
867
+
868
def _add_placeholder(self, modality: ModalityStr, placeholder: str | None):
869
mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]
870
if placeholder:
@@ -806,6 +909,10 @@ class BaseMultiModalContentParser(ABC):
909
) -> None:
910
raise NotImplementedError
911
912
+ @abstractmethod
913
+ def parse_prompt_embeds(self, data: str) -> None:
914
+ raise NotImplementedError
915
+
916
@abstractmethod
917
def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
918
raise NotImplementedError
@@ -834,6 +941,21 @@ class MultiModalContentParser(BaseMultiModalContentParser):
941
def model_config(self) -> ModelConfig:
942
return self._tracker.model_config
943
944
+ @override
945
+ def parse_prompt_embeds(self, data: str) -> None:
946
+ """Decode a base64 prompt embeds tensor and store it in the tracker.
947
+
948
+ Emits a single `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` sentinel per
949
+ content part. The renderer later expands each sentinel to a span of
950
+ `tensor.shape[0]` placeholder tokens after tokenization.
951
+ """
952
+ if not self.model_config.enable_prompt_embeds:
953
+ raise ValueError(_ENABLE_PROMPT_EMBEDS_ERROR)
954
+
955
+ tensor = safe_load_prompt_embeds(self.model_config, data.encode())
956
+ self._tracker.add("prompt_embeds", (tensor, None))
957
+ self._add_placeholder("prompt_embeds", PROMPT_EMBEDS_PLACEHOLDER_TOKEN)
958
+
959
def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
960
image = self._connector.fetch_image(image_url) if image_url else None
961
@@ -958,6 +1080,29 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):
1080
def model_config(self) -> ModelConfig:
1081
return self._tracker.model_config
1082
1083
+ @override
1084
+ def parse_prompt_embeds(self, data: str) -> None:
1085
+ """Schedule async prompt embeds decode and store the coroutine in the tracker.
1086
+
1087
+ Like the sync variant, emits a single sentinel `PROMPT_EMBEDS_PLACEHOLDER_TOKEN`
1088
+ per content part. Unlike the sync variant, the tensor decode is deferred to a
1089
+ thread-pool executor via `safe_load_prompt_embeds_async`.
1090
+ """
1091
+ if not self.model_config.enable_prompt_embeds:
1092
+ raise ValueError(_ENABLE_PROMPT_EMBEDS_ERROR)
1093
+
1094
+ coro = self._load_prompt_embeds_async(data.encode())
1095
+ self._tracker.add("prompt_embeds", coro)
1096
+ self._add_placeholder("prompt_embeds", PROMPT_EMBEDS_PLACEHOLDER_TOKEN)
1097
+
1098
+ async def _load_prompt_embeds_async(
1099
+ self, data_bytes: bytes
1100
+ ) -> tuple[torch.Tensor, None]:
1101
+ # Second tuple slot fills the tracker's generic `(item, uuid | None)`
1102
+ # contract. prompt_embeds has no UUID concept, so it's always `None`.
1103
+ tensor = await safe_load_prompt_embeds_async(self.model_config, data_bytes)
1104
+ return tensor, None
1105
+
1106
async def _image_with_uuid_async(self, image_url: str | None, uuid: str | None):
1107
image = (
1108
await self._connector.fetch_image_async(image_url) if image_url else None
@@ -1269,6 +1414,7 @@ def _get_full_multimodal_text_prompt(
1414
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
1415
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
1416
_AudioEmbedsParser = partial(cast, ChatCompletionContentPartAudioEmbedsParam)
1417
+_PromptEmbedsParser = partial(cast, ChatCompletionContentPartPromptEmbedsParam)
1418
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
1419
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
1420
_PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
@@ -1294,6 +1440,7 @@ MM_PARSER_MAP: dict[
1440
"image_url": lambda part: _ImageParser(part).get("image_url", {}).get("url", None),
1441
"image_embeds": lambda part: _ImageEmbedsParser(part).get("image_embeds", None),
1442
"audio_embeds": lambda part: _AudioEmbedsParser(part).get("audio_embeds", None),
1443
+ "prompt_embeds": lambda part: _PromptEmbedsParser(part).get("data", None),
1444
"image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
1445
"audio_url": lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
1446
"input_audio": lambda part: _InputAudioParser(part).get("input_audio", None),
@@ -1372,6 +1519,11 @@ def _parse_chat_message_content_mm_part(
1519
)
1520
audio_embeds = audio_params.get("audio_embeds", None)
1521
return "audio_embeds", audio_embeds
1522
+ if "prompt_embeds" in part:
1523
+ prompt_embeds_params = cast( # type: ignore[assignment]
1524
+ ChatCompletionContentPartPromptEmbedsParam, part
1525
+ )
1526
+ return "prompt_embeds", prompt_embeds_params.get("data", None)
1527
if "audio_url" in part:
1528
audio_params = cast( # type: ignore[assignment]
1529
CustomChatCompletionContentSimpleAudioParam, part
@@ -1455,6 +1607,24 @@ def _parse_chat_message_content_parts(
1607
return [ConversationMessage(role=role, content=text_prompt)]
1608
1609
1610
+def _reject_reserved_placeholder_in_text(text: str, model_config: ModelConfig) -> None:
1611
+ """Reject user-supplied text parts that contains the reserved `prompt_embeds`
1612
+ placeholder sentinel.
1613
+
1614
+ When the server accepts `prompt_embeds`, the placeholder token is
1615
+ registered as a single unsplittable special token on the tokenizer. Any
1616
+ user text that happens to contain the literal sequence would tokenize to
1617
+ the same ID and be mistaken for a splice point by the renderer, letting a
1618
+ caller move or inject splice positions via plain text content.
1619
+ """
1620
+ if model_config.enable_prompt_embeds and PROMPT_EMBEDS_PLACEHOLDER_TOKEN in text:
1621
+ raise ValueError(
1622
+ _RESERVED_PLACEHOLDER_IN_TEXT_ERROR.format(
1623
+ token=PROMPT_EMBEDS_PLACEHOLDER_TOKEN
1624
+ )
1625
+ )
1626
+
1627
+
1628
def _parse_chat_message_content_part(
1629
part: ChatCompletionContentPartParam,
1630
mm_parser: BaseMultiModalContentParser,
@@ -1470,6 +1640,7 @@ def _parse_chat_message_content_part(
1640
with multimodal placeholders.
1641
"""
1642
if isinstance(part, str): # Handle plain text parts
1643
+ _reject_reserved_placeholder_in_text(part, mm_parser.model_config)
1644
if wrap_dicts:
1645
return {"type": "text", "text": part}
1646
return part
@@ -1488,6 +1659,7 @@ def _parse_chat_message_content_part(
1659
1660
if part_type in ("text", "input_text", "output_text", "refusal", "thinking"):
1661
str_content = cast(str, content)
1662
+ _reject_reserved_placeholder_in_text(str_content, mm_parser.model_config)
1663
if wrap_dicts:
1664
return {"type": "text", "text": str_content}
1665
else:
@@ -1516,6 +1688,11 @@ def _parse_chat_message_content_part(
1688
content = cast(str | dict[str, str], content) if content is not None else None
1689
mm_parser.parse_audio_embeds(content, uuid)
1690
modality = "audio"
1691
+ elif part_type == "prompt_embeds":
1692
+ if not content:
1693
+ raise ValueError(_PROMPT_EMBEDS_MISSING_DATA_ERROR)
1694
+ mm_parser.parse_prompt_embeds(cast(str, content))
1695
+ modality = "prompt_embeds"
1696
elif part_type == "audio_url":
1697
str_content = cast(str, content)
1698
mm_parser.parse_audio(str_content, uuid)
@@ -1544,7 +1721,18 @@ def _parse_chat_message_content_part(
1721
)
1722
1723
if wrap_dicts:
1724
+ if modality == "prompt_embeds":
1725
+ # Chat templates don't know about the "prompt_embeds" modality,
1726
+ # emit the single sentinel token as text so the template renders
1727
+ # it inline. The renderer later expands it to N tokens post-tokenize.
1728
+ return {"type": "text", "text": PROMPT_EMBEDS_PLACEHOLDER_TOKEN}
1729
return {"type": modality}
1730
+ if modality == "prompt_embeds":
1731
+ # Emit the renderer token inline regardless of `interleave_strings`,
1732
+ # prompt_embeds are spliced at the token offset so position matters.
1733
+ # Falling back to front-padding via `missing_placeholders` would
1734
+ # reorder them relative to surrounding text.
1735
+ return PROMPT_EMBEDS_PLACEHOLDER_TOKEN
1736
return MODALITY_PLACEHOLDERS_MAP[modality] if interleave_strings else None
1737
1738
@@ -1668,7 +1856,10 @@ def parse_chat_messages(
1856
MultiModalUUIDDict | None,
1857
]:
1858
conversation: list[ConversationMessage] = []
1671
- mm_tracker = MultiModalItemTracker(model_config, media_io_kwargs=media_io_kwargs)
1859
+ mm_tracker = MultiModalItemTracker(
1860
+ model_config,
1861
+ media_io_kwargs=media_io_kwargs,
1862
+ )
1863
1864
for msg in messages:
1865
sub_messages = _parse_chat_message_content(
@@ -1705,7 +1896,8 @@ async def parse_chat_messages_async(
1896
]:
1897
conversation: list[ConversationMessage] = []
1898
mm_tracker = AsyncMultiModalItemTracker(
1708
- model_config, media_io_kwargs=media_io_kwargs
1899
+ model_config,
1900
+ media_io_kwargs=media_io_kwargs,
1901
)
1902
1903
for msg in messages:
vllm/entrypoints/llm.py
+4
-1
@@ -926,7 +926,10 @@ class LLM:
926
add_generation_prompt=add_generation_prompt,
927
continue_final_message=continue_final_message,
928
tools=tools,
929
- tokenize=is_mistral_tokenizer(renderer.tokenizer),
929
+ tokenize=(
930
+ is_mistral_tokenizer(renderer.tokenizer)
931
+ or self.model_config.enable_prompt_embeds
932
+ ),
933
),
934
),
935
mm_processor_kwargs=mm_processor_kwargs,
vllm/entrypoints/serve/render/serving.py
+4
-1
@@ -541,7 +541,10 @@ class OpenAIServingRender:
541
default_template_kwargs,
542
dict(
543
tools=tool_dicts,
544
- tokenize=is_mistral_tokenizer(renderer.tokenizer),
544
+ tokenize=(
545
+ is_mistral_tokenizer(renderer.tokenizer)
546
+ or self.model_config.enable_prompt_embeds
547
+ ),
548
),
549
)
550
vllm/inputs/engine.py
+19
@@ -71,12 +71,27 @@ class EmbedsInput(_InputOptions):
71
prompt: NotRequired[str]
72
"""The prompt text corresponding to the token IDs, if available."""
73
74
+ prompt_token_ids: NotRequired[list[int]]
75
+ """Token IDs of the rendered prompt. Only set for mixed-mode inputs
76
+ (chat completion with `prompt_embeds` content parts). When present,
77
+ `is_token_ids` MUST also be present and have the same length.
78
+ For pure-embeds inputs this field is absent."""
79
+
80
+ is_token_ids: NotRequired[list[bool]]
81
+ """Per-position mask for mixed-mode inputs. `True` means the position
82
+ is a real token ID (use the model's embedding layer); `False` means
83
+ the position uses a pre-computed embedding row from `prompt_embeds`.
84
+ Length MUST equal `len(prompt_token_ids)`.
85
+ For pure-embeds inputs this field is absent."""
86
+
87
88
def embeds_input(
89
prompt_embeds: "torch.Tensor",
90
*,
91
prompt: str | None = None,
92
cache_salt: str | None = None,
93
+ prompt_token_ids: list[int] | None = None,
94
+ is_token_ids: list[bool] | None = None,
95
) -> EmbedsInput:
96
"""
97
Construct [`EmbedsInput`][vllm.inputs.engine.EmbedsInput]
@@ -88,6 +103,10 @@ def embeds_input(
103
inputs["prompt"] = prompt
104
if cache_salt is not None:
105
inputs["cache_salt"] = cache_salt
106
+ if prompt_token_ids is not None:
107
+ inputs["prompt_token_ids"] = prompt_token_ids
108
+ if is_token_ids is not None:
109
+ inputs["is_token_ids"] = is_token_ids
110
111
return inputs
112
vllm/inputs/llm.py
+11
@@ -125,6 +125,17 @@ class EmbedsPrompt(_PromptOptions):
125
prompt: NotRequired[str]
126
"""The prompt text corresponding to the token embeddings, if available."""
127
128
+ prompt_token_ids: NotRequired[list[int]]
129
+ """Token IDs for mixed-mode inputs (chat completion with
130
+ `prompt_embeds` content parts). The tokens at positions where
131
+ `prompt_is_token_ids` is `False` are placeholder tokens that
132
+ get replaced by entries from `prompt_embeds` in the forward pass."""
133
+
134
+ prompt_is_token_ids: NotRequired[list[bool]]
135
+ """Per-position mask, `True` uses the real token ID, `False` uses
136
+ the corresponding entry from `prompt_embeds`.
137
+ Must be the same length as `prompt_token_ids` when both are set."""
138
+
139
140
DecoderOnlyPrompt: TypeAlias = (
141
str | TextPrompt | list[int] | TokensPrompt | EmbedsPrompt
vllm/renderers/base.py
+2
@@ -769,6 +769,8 @@ class BaseRenderer(ABC, Generic[_T]):
769
return embeds_input(
770
prompt_embeds=prompt_embeds,
771
cache_salt=prompt.get("cache_salt"),
772
+ prompt_token_ids=prompt.get("prompt_token_ids"),
773
+ is_token_ids=prompt.get("prompt_is_token_ids"),
774
)
775
776
async def _process_tokens_async(
vllm/renderers/embed_utils.py
+45
-6
@@ -7,6 +7,7 @@ import pybase64
7
import torch
8
9
from vllm.exceptions import VLLMValidationError
10
+from vllm.utils.async_utils import make_async
11
12
if TYPE_CHECKING:
13
from vllm.config import ModelConfig
@@ -30,15 +31,53 @@ def safe_load_prompt_embeds(
31
weights_only=True,
32
map_location=torch.device("cpu"),
33
)
33
- assert isinstance(tensor, torch.Tensor) and tensor.dtype in (
34
- torch.float32,
35
- torch.bfloat16,
36
- torch.float16,
37
- )
34
+ if not isinstance(tensor, torch.Tensor):
35
+ raise VLLMValidationError(
36
+ "`prompt_embeds` payload did not deserialize to a torch.Tensor.",
37
+ parameter="prompt_embeds",
38
+ )
39
tensor = tensor.to_dense()
40
41
if tensor.dim() > 2:
42
tensor = tensor.squeeze(0)
42
- assert tensor.dim() == 2
43
+ if tensor.dim() != 2:
44
+ raise VLLMValidationError(
45
+ "`prompt_embeds` must be a 2D tensor of shape "
46
+ f"(num_tokens, hidden_size); got shape {tuple(tensor.shape)}.",
47
+ parameter="prompt_embeds",
48
+ )
49
+
50
+ # Pin each tensor to the model's hidden_size. Validating here
51
+ # also transitively guarantees cross-tensor consistency for requests that
52
+ # include multiple `prompt_embeds` parts, which is required by downstream
53
+ # concatenation in `_build_mixed_prompt_embeds`.
54
+ expected_hidden_size = model_config.get_hidden_size()
55
+ if tensor.shape[1] != expected_hidden_size:
56
+ raise VLLMValidationError(
57
+ f"`prompt_embeds` hidden_size {tensor.shape[1]} does not match "
58
+ f"the model's hidden_size {expected_hidden_size}.",
59
+ parameter="prompt_embeds",
60
+ )
61
+
62
+ # Cast to the model's dtype so API clients don't need to know the server's
63
+ # `--dtype` setting ahead of time. Only floating-point source dtypes are
64
+ # allowed. integer / bool / complex inputs almost certainly indicate caller
65
+ # error (e.g. quantized payloads, wrong tensor), and a silent `.to()`
66
+ # could hide a real mistake.
67
+ expected_dtype = model_config.dtype
68
+ if tensor.dtype != expected_dtype:
69
+ if not tensor.is_floating_point():
70
+ raise VLLMValidationError(
71
+ f"`prompt_embeds` dtype {tensor.dtype} is not a floating-point "
72
+ f"type, cannot safely cast to the model's dtype {expected_dtype}.",
73
+ parameter="prompt_embeds",
74
+ )
75
+ tensor = tensor.to(expected_dtype)
76
77
return tensor
78
+
79
+
80
+safe_load_prompt_embeds_async = make_async(safe_load_prompt_embeds)
81
+"""Async variant of `safe_load_prompt_embeds` that defers the decode to a
82
+thread-pool executor, so the asyncio event loop is not blocked by the base64
83
+decode + `torch.load` work."""
vllm/renderers/hf.py
+440
-18
@@ -1,11 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from __future__ import annotations
4
+
5
import inspect
6
import itertools
7
+import weakref
8
from collections import defaultdict, deque
6
-from collections.abc import Set
9
+from collections.abc import Sequence
10
from functools import lru_cache
8
-from typing import Any, Literal, cast, overload
11
+from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
12
13
import jinja2
14
import jinja2.ext
@@ -13,20 +16,32 @@ import jinja2.meta
16
import jinja2.nodes
17
import jinja2.parser
18
import jinja2.sandbox
19
+import torch
20
+from typing_extensions import override
21
17
-from vllm.config import ModelConfig, VllmConfig
22
from vllm.entrypoints.chat_utils import (
19
- ChatCompletionMessageParam,
20
- ChatTemplateContentFormat,
21
- ChatTemplateContentFormatOption,
23
+ PROMPT_EMBEDS_PLACEHOLDER_TOKEN,
24
ChatTemplateResolutionError,
23
- ConversationMessage,
25
load_chat_template,
26
parse_chat_messages,
27
parse_chat_messages_async,
28
)
28
-from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict
29
+from vllm.inputs import EmbedsPrompt
30
+from vllm.inputs.engine import MultiModalInput
31
from vllm.logger import init_logger
32
+from vllm.multimodal.hasher import MultiModalHasher
33
+from vllm.multimodal.inputs import (
34
+ MultiModalFieldElem,
35
+ MultiModalKwargsItem,
36
+ MultiModalKwargsItems,
37
+ MultiModalSharedField,
38
+ PlaceholderRange,
39
+)
40
+from vllm.multimodal.processing.processor import (
41
+ PromptReplacement,
42
+ apply_token_matches,
43
+ find_mm_placeholders,
44
+)
45
from vllm.tokenizers.hf import HfTokenizer
46
from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
47
from vllm.transformers_utils.processor import cached_get_processor
@@ -34,13 +49,166 @@ from vllm.utils.async_utils import make_async
49
from vllm.utils.func_utils import supports_kw
50
51
from .base import BaseRenderer
37
-from .inputs import DictPrompt
52
from .inputs.preprocess import parse_dec_only_prompt
39
-from .params import ChatParams
53
+
54
+if TYPE_CHECKING:
55
+ from collections.abc import Set
56
+
57
+ from vllm.config import ModelConfig, VllmConfig
58
+ from vllm.entrypoints.chat_utils import (
59
+ ChatCompletionMessageParam,
60
+ ChatTemplateContentFormat,
61
+ ChatTemplateContentFormatOption,
62
+ ConversationMessage,
63
+ )
64
+ from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict, TokensPrompt
65
+ from vllm.inputs.engine import TokensInput
66
+ from vllm.multimodal.processing.processor import (
67
+ MultiModalPromptUpdates,
68
+ ResolvedPromptUpdate,
69
+ )
70
+
71
+ from .inputs import DictPrompt
72
+ from .params import ChatParams
73
74
logger = init_logger(__name__)
75
76
77
+# Cache of `tokenizer -> prompt_embeds placeholder token ID`. Keyed by the
78
+# tokenizer object (not `id(tokenizer)`) so a fresh tokenizer landing at a
79
+# recycled memory address can't pick up a stale tid. Entries evict atomically
80
+# with the tokenizer's garbage-collection.
81
+_PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE: Final[
82
+ weakref.WeakKeyDictionary[HfTokenizer, int]
83
+] = weakref.WeakKeyDictionary()
84
+_PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_ERROR: Final[str] = (
85
+ "Expected {token!r} to tokenize to exactly 1 token, got {num_ids} ({ids!r})."
86
+)
87
+_PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR: Final[str] = (
88
+ "Expected {expected} prompt_embeds placeholder spans in the "
89
+ "tokenized prompt, found {actual}."
90
+)
91
+_MISSING_PROMPT_TOKEN_IDS_ERROR: Final[str] = (
92
+ "Expected prompt_token_ids in rendered prompt when prompt_embeds "
93
+ "are present. This indicates the chat template was invoked with "
94
+ "tokenize=False."
95
+)
96
+_TOKENIZE_OVERRIDE_WARNING: Final[str] = (
97
+ "Overriding `tokenize=False` to `True` because `prompt_embeds` "
98
+ "post-processing requires tokenized IDs."
99
+)
100
+
101
+
102
+def _ensure_prompt_embeds_placeholder_token(tokenizer: HfTokenizer) -> int:
103
+ """Register `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` as a special token and return
104
+ its token ID."""
105
+ cached = _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE.get(tokenizer)
106
+ if cached is not None:
107
+ return cached
108
+
109
+ tokenizer.add_special_tokens(
110
+ {"additional_special_tokens": [PROMPT_EMBEDS_PLACEHOLDER_TOKEN]}
111
+ )
112
+
113
+ ids = tokenizer.encode(PROMPT_EMBEDS_PLACEHOLDER_TOKEN, add_special_tokens=False)
114
+ if len(ids) != 1:
115
+ raise RuntimeError(
116
+ _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_ERROR.format(
117
+ token=PROMPT_EMBEDS_PLACEHOLDER_TOKEN,
118
+ num_ids=len(ids),
119
+ ids=ids,
120
+ )
121
+ )
122
+
123
+ token_id = ids[0]
124
+ _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE[tokenizer] = token_id
125
+ return token_id
126
+
127
+
128
+def _build_prompt_embeds_updates(
129
+ prompt_embeds_tensors: Sequence[torch.Tensor],
130
+ placeholder_token_id: int,
131
+) -> MultiModalPromptUpdates:
132
+ """Build `MultiModalPromptUpdates` for `prompt_embeds` expansion.
133
+
134
+ Each tensor produces a `PromptReplacement` that maps
135
+ `[placeholder_token_id]` -> `[placeholder_token_id] x N`
136
+ (where `N = tensor.shape[0]`).
137
+ """
138
+ updates: list[Sequence[ResolvedPromptUpdate]] = []
139
+ for i, tensor in enumerate(prompt_embeds_tensors):
140
+ update = PromptReplacement(
141
+ modality="prompt_embeds",
142
+ target=[placeholder_token_id],
143
+ replacement=[placeholder_token_id] * tensor.shape[0],
144
+ )
145
+ updates.append([update.resolve(item_idx=i)])
146
+ return {"prompt_embeds": updates}
147
+
148
+
149
+def _expand_prompt_embeds_placeholders(
150
+ token_ids: list[int],
151
+ mm_prompt_updates: MultiModalPromptUpdates,
152
+) -> list[int]:
153
+ """Expand each 1-token `prompt_embeds` sentinel into an N-token span.
154
+
155
+ Uses `apply_token_matches`. Each single placeholder token in
156
+ `token_ids` is replaced with a consecutive span of
157
+ `tensor.shape[0]` copies, following tensors in order.
158
+ """
159
+ expanded, _ = apply_token_matches(token_ids, mm_prompt_updates, tokenizer=None)
160
+ return expanded
161
+
162
+
163
+def _build_prompt_embeds_positions(
164
+ token_ids: list[int],
165
+ num_tensors: int,
166
+ mm_prompt_updates: MultiModalPromptUpdates,
167
+) -> list[tuple[int, int]]:
168
+ """Locate each prompt_embeds placeholder span in `token_ids`.
169
+
170
+ Expects `token_ids` to already contain expanded N-token spans.
171
+ Returns `[(start_idx, length), ...]` aligned with the tensors.
172
+ """
173
+ placeholders = find_mm_placeholders(
174
+ prompt=token_ids,
175
+ mm_prompt_updates=mm_prompt_updates,
176
+ tokenizer=None,
177
+ )
178
+ features = placeholders.get("prompt_embeds", [])
179
+
180
+ if len(features) != num_tensors:
181
+ raise ValueError(
182
+ _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR.format(
183
+ expected=num_tensors,
184
+ actual=len(features),
185
+ )
186
+ )
187
+
188
+ return [(f.start_idx, f.length) for f in features]
189
+
190
+
191
+def _build_mixed_prompt_embeds(
192
+ token_ids: list[int],
193
+ prompt_embeds_tensors: Sequence[torch.Tensor],
194
+ positions: list[tuple[int, int]],
195
+) -> tuple[torch.Tensor, list[bool]]:
196
+ """Build the full-length `prompt_embeds` tensor and the `is_token_ids`
197
+ mask aligned to `token_ids`."""
198
+ total_len = len(token_ids)
199
+ hidden_size = prompt_embeds_tensors[0].shape[1]
200
+ dtype = prompt_embeds_tensors[0].dtype
201
+
202
+ full_embeds = torch.zeros(total_len, hidden_size, dtype=dtype)
203
+ is_token_ids = torch.ones(total_len, dtype=torch.bool)
204
+
205
+ for (start, length), tensor in zip(positions, prompt_embeds_tensors, strict=True):
206
+ full_embeds[start : start + length] = tensor
207
+ is_token_ids[start : start + length] = False
208
+
209
+ return full_embeds, is_token_ids.tolist()
210
+
211
+
212
_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()
213
"""
214
Used in `_try_get_processor_chat_template` to avoid calling
@@ -98,7 +266,7 @@ def resolve_chat_template(
266
chat_template: str | None,
267
tools: list[dict[str, Any]] | None,
268
*,
101
- model_config: "ModelConfig",
269
+ model_config: ModelConfig,
270
) -> str | None:
271
# 1st priority: The given chat template
272
if chat_template is not None:
@@ -281,7 +449,7 @@ def _resolve_chat_template_content_format(
449
tools: list[dict[str, Any]] | None,
450
tokenizer: HfTokenizer,
451
*,
284
- model_config: "ModelConfig",
452
+ model_config: ModelConfig,
453
) -> ChatTemplateContentFormat:
454
resolved_chat_template = resolve_chat_template(
455
tokenizer,
@@ -335,7 +503,7 @@ def resolve_chat_template_content_format(
503
given_format: ChatTemplateContentFormatOption,
504
tokenizer: HfTokenizer,
505
*,
338
- model_config: "ModelConfig",
506
+ model_config: ModelConfig,
507
) -> ChatTemplateContentFormat:
508
if given_format != "auto":
509
return given_format
@@ -437,7 +605,7 @@ def resolve_chat_template_kwargs(
605
606
@overload
607
def safe_apply_chat_template(
440
- model_config: "ModelConfig",
608
+ model_config: ModelConfig,
609
tokenizer: HfTokenizer,
610
conversation: list[ConversationMessage],
611
*,
@@ -448,7 +616,7 @@ def safe_apply_chat_template(
616
) -> list[int]: ...
617
@overload
618
def safe_apply_chat_template(
451
- model_config: "ModelConfig",
619
+ model_config: ModelConfig,
620
tokenizer: HfTokenizer,
621
conversation: list[ConversationMessage],
622
*,
@@ -458,7 +626,7 @@ def safe_apply_chat_template(
626
**kwargs,
627
) -> str: ...
628
def safe_apply_chat_template(
461
- model_config: "ModelConfig",
629
+ model_config: ModelConfig,
630
tokenizer: HfTokenizer,
631
conversation: list[ConversationMessage],
632
*,
@@ -486,6 +654,14 @@ def safe_apply_chat_template(
654
chat_template_kwargs=kwargs,
655
)
656
657
+ # transformers v5 changed the default of `return_dict` to True, which
658
+ # makes `apply_chat_template(tokenize=True)` return a `BatchEncoding`
659
+ # instead of `list[int]`. Force `return_dict=False` so downstream code
660
+ # that expects a flat token list (e.g. `parse_dec_only_prompt`) works
661
+ # consistently across v4 and v5.
662
+ if tokenize and "return_dict" not in resolved_kwargs:
663
+ resolved_kwargs["return_dict"] = False
664
+
665
try:
666
return tokenizer.apply_chat_template(
667
conversation=conversation, # type: ignore[arg-type]
@@ -627,6 +803,12 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
803
model_config = self.model_config
804
tokenizer = self.get_tokenizer()
805
806
+ prompt_embeds_placeholder_token_id: int | None = None
807
+ if model_config.enable_prompt_embeds:
808
+ prompt_embeds_placeholder_token_id = (
809
+ _ensure_prompt_embeds_placeholder_token(tokenizer)
810
+ )
811
+
812
conversation, mm_data, mm_uuids = parse_chat_messages(
813
messages,
814
model_config,
@@ -641,11 +823,30 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
823
mm_processor_kwargs=params.mm_processor_kwargs,
824
)
825
826
+ # prompt_embeds tensors are carried by the tracker through mm_data,
827
+ # but they must NOT be fed to the MM processor (which would reject
828
+ # the unknown key). Extract them here.
829
+ prompt_embeds_tensors: list[torch.Tensor] | None = None
830
+ if mm_data is not None and "prompt_embeds" in mm_data:
831
+ prompt_embeds_tensors = list(
832
+ cast(Sequence[torch.Tensor], mm_data["prompt_embeds"])
833
+ )
834
+ mm_data = {k: v for k, v in mm_data.items() if k != "prompt_embeds"}
835
+ if not mm_data:
836
+ mm_data = None
837
+
838
+ chat_template_kwargs = params.get_apply_chat_template_kwargs()
839
+ if prompt_embeds_tensors:
840
+ # prompt_embeds post-processing requires prompt_token_ids.
841
+ if chat_template_kwargs.get("tokenize") is False:
842
+ logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
843
+ chat_template_kwargs["tokenize"] = True
844
+
845
prompt_raw = safe_apply_chat_template(
846
model_config,
847
tokenizer,
848
conversation,
648
- **params.get_apply_chat_template_kwargs(),
849
+ **chat_template_kwargs,
850
)
851
852
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
@@ -671,6 +872,29 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
872
)
873
874
prompt = parse_dec_only_prompt(prompt_raw)
875
+
876
+ # When `prompt_embeds` is mixed with other modality data,
877
+ # `_process_tokens` runs `_process_multimodal` first (expanding
878
+ # `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then
879
+ # `_apply_prompt_embeds_to_engine_input` augments the result.
880
+ # Stash the tensors and placeholder ID for that override to consume.
881
+ if prompt_embeds_tensors and mm_data:
882
+ assert prompt_embeds_placeholder_token_id is not None
883
+ cast(dict, prompt)["_prompt_embeds"] = (
884
+ prompt_embeds_tensors,
885
+ prompt_embeds_placeholder_token_id,
886
+ )
887
+ if params.mm_processor_kwargs:
888
+ cast(dict, prompt)["mm_processor_kwargs"] = params.mm_processor_kwargs
889
+ elif prompt_embeds_tensors:
890
+ # Pure mode: no other MM data, mutate prompt to EmbedsPrompt shape.
891
+ assert prompt_embeds_placeholder_token_id is not None
892
+ self._apply_prompt_embeds_to_prompt(
893
+ prompt,
894
+ prompt_embeds_tensors,
895
+ prompt_embeds_placeholder_token_id,
896
+ )
897
+
898
if mm_data is not None:
899
prompt["multi_modal_data"] = mm_data
900
if mm_uuids is not None:
@@ -686,6 +910,12 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
910
model_config = self.model_config
911
tokenizer = self.get_tokenizer()
912
913
+ prompt_embeds_placeholder_token_id: int | None = None
914
+ if model_config.enable_prompt_embeds:
915
+ prompt_embeds_placeholder_token_id = (
916
+ _ensure_prompt_embeds_placeholder_token(tokenizer)
917
+ )
918
+
919
conversation, mm_data, mm_uuids = await parse_chat_messages_async(
920
messages,
921
model_config,
@@ -700,11 +930,27 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
930
mm_processor_kwargs=params.mm_processor_kwargs,
931
)
932
933
+ prompt_embeds_tensors: list[torch.Tensor] | None = None
934
+ if mm_data is not None and "prompt_embeds" in mm_data:
935
+ prompt_embeds_tensors = list(
936
+ cast(Sequence[torch.Tensor], mm_data["prompt_embeds"])
937
+ )
938
+ mm_data = {k: v for k, v in mm_data.items() if k != "prompt_embeds"}
939
+ if not mm_data:
940
+ mm_data = None
941
+
942
+ chat_template_kwargs = params.get_apply_chat_template_kwargs()
943
+ if prompt_embeds_tensors:
944
+ # prompt_embeds post-processing requires prompt_token_ids.
945
+ if chat_template_kwargs.get("tokenize") is False:
946
+ logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
947
+ chat_template_kwargs["tokenize"] = True
948
+
949
prompt_raw = await self._apply_chat_template_async(
950
model_config,
951
tokenizer,
952
conversation,
707
- **params.get_apply_chat_template_kwargs(),
953
+ **chat_template_kwargs,
954
)
955
956
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
@@ -728,9 +974,185 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
974
)
975
976
prompt = parse_dec_only_prompt(prompt_raw)
977
+
978
+ # See `render_messages` for the rationale.
979
+ if prompt_embeds_tensors and mm_data:
980
+ assert prompt_embeds_placeholder_token_id is not None
981
+ cast(dict, prompt)["_prompt_embeds"] = (
982
+ prompt_embeds_tensors,
983
+ prompt_embeds_placeholder_token_id,
984
+ )
985
+ if params.mm_processor_kwargs:
986
+ cast(dict, prompt)["mm_processor_kwargs"] = params.mm_processor_kwargs
987
+ elif prompt_embeds_tensors:
988
+ assert prompt_embeds_placeholder_token_id is not None
989
+ self._apply_prompt_embeds_to_prompt(
990
+ prompt,
991
+ prompt_embeds_tensors,
992
+ prompt_embeds_placeholder_token_id,
993
+ )
994
+
995
if mm_data is not None:
996
prompt["multi_modal_data"] = mm_data
997
if mm_uuids is not None:
998
prompt["multi_modal_uuids"] = mm_uuids
999
1000
return conversation, prompt
1001
+
1002
+ @override
1003
+ def _process_tokens(
1004
+ self,
1005
+ prompt: TokensPrompt,
1006
+ *,
1007
+ skip_mm_cache: bool = False,
1008
+ ) -> TokensInput | MultiModalInput:
1009
+ """Pre-expand `prompt_embeds` sentinels before delegating to the MM
1010
+ processor, then attach `prompt_embeds` modality data to the result.
1011
+
1012
+ Mixed mode only: the `_prompt_embeds` stash is set by
1013
+ `render_messages` when `prompt_embeds` co-exist with other MM data
1014
+ (images, audio, …). We expand each 1-token sentinel to an N-token
1015
+ span *before* calling `super()._process_tokens()` so the MM
1016
+ processor records all placeholder offsets in the final (post-expansion)
1017
+ coordinate space, no offset shifting needed afterwards.
1018
+ """
1019
+ prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
1020
+ if prompt_embeds_info is not None:
1021
+ tensors, placeholder_token_id = prompt_embeds_info
1022
+ mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
1023
+ cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
1024
+ list(prompt["prompt_token_ids"]), mm_updates
1025
+ )
1026
+ engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)
1027
+ if prompt_embeds_info is not None:
1028
+ tensors, _ = prompt_embeds_info
1029
+ self._apply_prompt_embeds_to_engine_input(
1030
+ cast(MultiModalInput, engine_input),
1031
+ tensors,
1032
+ mm_updates,
1033
+ )
1034
+ return engine_input
1035
+
1036
+ @override
1037
+ async def _process_tokens_async(
1038
+ self,
1039
+ prompt: TokensPrompt,
1040
+ *,
1041
+ skip_mm_cache: bool = False,
1042
+ ) -> TokensInput | MultiModalInput:
1043
+ """Async equivalent of `_process_tokens`."""
1044
+ prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
1045
+ if prompt_embeds_info is not None:
1046
+ tensors, placeholder_token_id = prompt_embeds_info
1047
+ mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
1048
+ cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
1049
+ list(prompt["prompt_token_ids"]), mm_updates
1050
+ )
1051
+ engine_input = await super()._process_tokens_async(
1052
+ prompt, skip_mm_cache=skip_mm_cache
1053
+ )
1054
+ if prompt_embeds_info is not None:
1055
+ tensors, _ = prompt_embeds_info
1056
+ self._apply_prompt_embeds_to_engine_input(
1057
+ cast(MultiModalInput, engine_input),
1058
+ tensors,
1059
+ mm_updates,
1060
+ )
1061
+ return engine_input
1062
+
1063
+ @staticmethod
1064
+ def _apply_prompt_embeds_to_prompt(
1065
+ prompt: DictPrompt,
1066
+ prompt_embeds_tensors: list[torch.Tensor],
1067
+ placeholder_token_id: int,
1068
+ ) -> None:
1069
+ """Mutate `prompt` from `TokensPrompt` to `EmbedsPrompt` shape.
1070
+
1071
+ Pure `prompt_embeds` path only (no other MM modalities). Expands
1072
+ each `<prompt_embeds>` sentinel token into an N-token span and builds
1073
+ the full-length `prompt_embeds` tensor + `prompt_is_token_ids` mask
1074
+ that the engine's `enable_prompt_embeds` worker branch consumes.
1075
+ """
1076
+ token_ids = cast(list[int] | None, prompt.get("prompt_token_ids"))
1077
+ if token_ids is None:
1078
+ raise RuntimeError(_MISSING_PROMPT_TOKEN_IDS_ERROR)
1079
+
1080
+ embeds_orig_positions: list[int] = [
1081
+ i for i, tok in enumerate(token_ids) if tok == placeholder_token_id
1082
+ ]
1083
+ if len(embeds_orig_positions) != len(prompt_embeds_tensors):
1084
+ raise ValueError(
1085
+ f"Expected {len(prompt_embeds_tensors)} prompt_embeds "
1086
+ f"placeholder tokens in the rendered prompt, found "
1087
+ f"{len(embeds_orig_positions)}."
1088
+ )
1089
+
1090
+ mm_updates = _build_prompt_embeds_updates(
1091
+ prompt_embeds_tensors, placeholder_token_id
1092
+ )
1093
+ expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates)
1094
+ positions = _build_prompt_embeds_positions(
1095
+ expanded, len(prompt_embeds_tensors), mm_updates
1096
+ )
1097
+
1098
+ embeds_prompt = cast(EmbedsPrompt, prompt)
1099
+ embeds_prompt["prompt_token_ids"] = expanded
1100
+ full_embeds, is_token_ids_mask = _build_mixed_prompt_embeds(
1101
+ expanded, prompt_embeds_tensors, positions
1102
+ )
1103
+ embeds_prompt["prompt_embeds"] = full_embeds
1104
+ embeds_prompt["prompt_is_token_ids"] = is_token_ids_mask
1105
+
1106
+ @staticmethod
1107
+ def _apply_prompt_embeds_to_engine_input(
1108
+ engine_input: MultiModalInput,
1109
+ prompt_embeds_tensors: list[torch.Tensor],
1110
+ mm_updates: MultiModalPromptUpdates,
1111
+ ) -> None:
1112
+ """Augment `engine_input` in-place with a `prompt_embeds` modality.
1113
+
1114
+ Mixed mode: called after `_process_multimodal` has already run on the
1115
+ pre-expanded token IDs (expansion was done in `_process_tokens` before
1116
+ calling `super()`). Locates the already-expanded `prompt_embeds` spans
1117
+ and adds `prompt_embeds` entries to `mm_kwargs`, `mm_hashes`, and
1118
+ `mm_placeholders`.
1119
+ """
1120
+ # token_ids already contain the pre-expanded N-token spans.
1121
+ token_ids = list(engine_input["prompt_token_ids"])
1122
+
1123
+ positions = _build_prompt_embeds_positions(
1124
+ token_ids, len(prompt_embeds_tensors), mm_updates
1125
+ )
1126
+
1127
+ pe_kwargs_items: list[MultiModalKwargsItem] = []
1128
+ pe_hashes: list[str] = []
1129
+ pe_placeholders: list[PlaceholderRange] = []
1130
+ for tensor, (start, length) in zip(
1131
+ prompt_embeds_tensors, positions, strict=True
1132
+ ):
1133
+ pe_kwargs_items.append(
1134
+ MultiModalKwargsItem(
1135
+ {
1136
+ "embedding": MultiModalFieldElem(
1137
+ data=tensor,
1138
+ field=MultiModalSharedField(batch_size=1),
1139
+ )
1140
+ }
1141
+ )
1142
+ )
1143
+ pe_hashes.append(MultiModalHasher.hash_kwargs(prompt_embeds=tensor))
1144
+ # `is_embed=None` matches the existing image_embeds-style
1145
+ # "no encoder, just splice the tensor directly" semantics.
1146
+ pe_placeholders.append(
1147
+ PlaceholderRange(offset=start, length=length, is_embed=None)
1148
+ )
1149
+
1150
+ cast(
1151
+ MultiModalKwargsItems[MultiModalKwargsItem | None],
1152
+ engine_input["mm_kwargs"],
1153
+ )["prompt_embeds"] = pe_kwargs_items
1154
+ engine_input["mm_hashes"] = {
1155
+ **engine_input["mm_hashes"],
1156
+ "prompt_embeds": pe_hashes,
1157
+ }
1158
+ cast(dict, engine_input["mm_placeholders"])["prompt_embeds"] = pe_placeholders
vllm/v1/core/sched/output.py
+2
@@ -38,6 +38,7 @@ class NewRequestData:
38
num_computed_tokens: int
39
lora_request: LoRARequest | None
40
prompt_embeds: "torch.Tensor | None" = None
41
+ prompt_is_token_ids: list[bool] | None = None
42
43
# Only used for v2 model runner.
44
prefill_token_ids: list[int] | None = None
@@ -59,6 +60,7 @@ class NewRequestData:
60
num_computed_tokens=request.num_computed_tokens,
61
lora_request=request.lora_request,
62
prompt_embeds=request.prompt_embeds,
63
+ prompt_is_token_ids=request.prompt_is_token_ids,
64
prefill_token_ids=prefill_token_ids,
65
)
66
vllm/v1/engine/__init__.py
+6
@@ -94,6 +94,12 @@ class EngineCoreRequest(
94
data_parallel_rank: int | None
95
prompt_embeds: torch.Tensor | None = None
96
97
+ # Per-position mask for mixed-mode inputs (e.g chat completion with
98
+ # prompt_embeds content parts). `True` means the position is a real
99
+ # token ID; `False` means the position uses a pre-computed entry from
100
+ # `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
101
+ prompt_is_token_ids: list[bool] | None = None
102
+
103
# Index of the client, used to ensure outputs are sent back to the same
104
# client for this request when scaling out the front-end.
105
client_index: int = 0
vllm/v1/engine/input_processor.py
+4
-1
@@ -292,11 +292,13 @@ class InputProcessor:
292
293
# Mypy can be conservative for TypedDict unions; normalize access.
294
if decoder_inputs["type"] == "embeds":
295
- prompt_token_ids = None
295
prompt_embeds = decoder_inputs["prompt_embeds"]
296
+ prompt_token_ids = decoder_inputs.get("prompt_token_ids")
297
+ prompt_is_token_ids = decoder_inputs.get("is_token_ids")
298
else:
299
prompt_token_ids = decoder_inputs["prompt_token_ids"]
300
prompt_embeds = None
301
+ prompt_is_token_ids = None
302
303
sampling_params = None
304
pooling_params = None
@@ -361,6 +363,7 @@ class InputProcessor:
363
request_id=request_id,
364
prompt_token_ids=prompt_token_ids,
365
prompt_embeds=prompt_embeds,
366
+ prompt_is_token_ids=prompt_is_token_ids,
367
mm_features=mm_features,
368
sampling_params=sampling_params,
369
pooling_params=pooling_params,
vllm/v1/request.py
+6
@@ -66,6 +66,7 @@ class Request:
66
client_index: int = 0,
67
arrival_time: float | None = None,
68
prompt_embeds: torch.Tensor | None = None,
69
+ prompt_is_token_ids: list[bool] | None = None,
70
mm_features: list[MultiModalFeatureSpec] | None = None,
71
lora_request: "LoRARequest | None" = None,
72
cache_salt: str | None = None,
@@ -114,6 +115,10 @@ class Request:
115
116
self.prompt_token_ids = prompt_token_ids
117
self.prompt_embeds = prompt_embeds
118
+ # Per-position mask used in mixed-mode (chat completion with
119
+ # prompt_embeds). `None` except when both `prompt_token_ids` and
120
+ # `prompt_embeds` are set and their positions are interleaved.
121
+ self.prompt_is_token_ids = prompt_is_token_ids
122
# Cache per-block prompt-embed hashes to avoid rehashing the same
123
# tensor slices when generating extra keys.
124
self._prompt_embeds_per_block_hashes: dict[tuple[int, int], bytes] = {}
@@ -184,6 +189,7 @@ class Request:
189
client_index=request.client_index,
190
prompt_token_ids=request.prompt_token_ids,
191
prompt_embeds=request.prompt_embeds,
192
+ prompt_is_token_ids=request.prompt_is_token_ids,
193
mm_features=request.mm_features,
194
sampling_params=request.sampling_params,
195
pooling_params=request.pooling_params,
vllm/v1/worker/gpu_input_batch.py
+10
-1
@@ -50,6 +50,10 @@ class CachedRequestState:
50
lora_request: LoRARequest | None = None
51
prompt_embeds: torch.Tensor | None = None
52
53
+ # Per-position mask for mixed-mode inputs (e.g chat completion with
54
+ # prompt_embeds content parts). See `Request.prompt_is_token_ids`.
55
+ prompt_is_token_ids: list[bool] | None = None
56
+
57
# Used when both async_scheduling and spec_decode are enabled.
58
prev_num_draft_len: int = 0
59
@@ -356,7 +360,12 @@ class InputBatch:
360
end_idx = start_idx + len(request.output_token_ids)
361
if request.prompt_token_ids is not None:
362
self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids
359
- self.is_token_ids[req_index, :num_prompt_tokens] = True
363
+ if request.prompt_is_token_ids is not None:
364
+ self.is_token_ids[req_index, :num_prompt_tokens] = (
365
+ request.prompt_is_token_ids
366
+ )
367
+ else:
368
+ self.is_token_ids[req_index, :num_prompt_tokens] = True
369
else:
370
self.is_token_ids[req_index, :num_prompt_tokens] = False
371
if request.prompt_embeds is not None:
vllm/v1/worker/gpu_model_runner.py
+35
-1
@@ -1163,6 +1163,7 @@ class GPUModelRunner(
1163
req_id=req_id,
1164
prompt_token_ids=new_req_data.prompt_token_ids,
1165
prompt_embeds=new_req_data.prompt_embeds,
1166
+ prompt_is_token_ids=new_req_data.prompt_is_token_ids,
1167
mm_features=new_req_data.mm_features,
1168
sampling_params=sampling_params,
1169
pooling_params=pooling_params,
@@ -1505,10 +1506,16 @@ class GPUModelRunner(
1506
)
1507
mrope_model = cast(SupportsMRoPE, model)
1508
1509
+ # `prompt_embeds` is a passthrough modality (no grid_thw), models'
1510
+ # M-RoPE code assumes per-feature grid info, so filter it out. The
1511
+ # prompt_embeds positions are treated as text positions for M-RoPE.
1512
+ mrope_features = [
1513
+ f for f in req_state.mm_features if f.modality != "prompt_embeds"
1514
+ ]
1515
req_state.mrope_positions, req_state.mrope_position_delta = (
1516
mrope_model.get_mrope_input_positions(
1517
req_state.prompt_token_ids,
1511
- req_state.mm_features,
1518
+ mrope_features,
1519
)
1520
)
1521
@@ -2744,6 +2751,33 @@ class GPUModelRunner(
2751
if not mm_kwargs:
2752
return []
2753
2754
+ # `prompt_embeds` is a passthrough modality, the tensor is already in
2755
+ # the model embedding space, so no encoder runs. Inject each
2756
+ # `prompt_embeds` tensor directly into the encoder cache here so that
2757
+ # `_gather_mm_embeddings` can splice it via the standard `is_mm_embed`
2758
+ # path.
2759
+ pe_indices = [
2760
+ i
2761
+ for i, (modality, _) in enumerate(mm_kwargs)
2762
+ if modality == "prompt_embeds"
2763
+ ]
2764
+ if pe_indices:
2765
+ for i in pe_indices:
2766
+ pe_tensor = mm_kwargs[i][1]["embedding"].data
2767
+ assert isinstance(pe_tensor, torch.Tensor)
2768
+
2769
+ self.encoder_cache[mm_hashes[i]] = pe_tensor.to(self.device)
2770
+ self.maybe_save_ec_to_connector(self.encoder_cache, mm_hashes[i])
2771
+ # Filter out `prompt_embeds` items from mm_kwargs/mm_hashes/mm_lora_refs
2772
+ # since they don't require further encoder processing.
2773
+ mm_hashes = [h for i, h in enumerate(mm_hashes) if i not in pe_indices]
2774
+ mm_kwargs = [k for i, k in enumerate(mm_kwargs) if i not in pe_indices]
2775
+ mm_lora_refs = [
2776
+ r for i, r in enumerate(mm_lora_refs) if i not in pe_indices
2777
+ ]
2778
+ if not mm_kwargs:
2779
+ return [] # nothing left to encode after filtering out `prompt_embeds`
2780
+
2781
should_time = bool(
2782
self.observability_config
2783
and self.observability_config.enable_mm_processor_stats