diff --git a/docs/features/README.md b/docs/features/README.md index e62d9cddee7..28362f40147 100644 --- a/docs/features/README.md +++ b/docs/features/README.md @@ -52,10 +52,10 @@ th:not(:first-child) { | [mm](multimodal_inputs.md) | ✅ | ✅ | [🟠](https://github.com/vllm-project/vllm/pull/4194)^ | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | | | | best-of | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/7968) | ✅ | ✅ | | | | beam-search | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/7968) | ❔ | ✅ | ✅ | | -| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ❌ | ❔ | ❔ | ✅ | +| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ✅ | ❔ | ❔ | ✅ | \* Chunked prefill and prefix caching are only applicable to last-token or all pooling with causal attention. -^ LoRA is only applicable to the language backbone of multimodal models. +^ LoRA is only applicable to the language backbone of multimodal models. ### Feature x Hardware diff --git a/docs/features/prompt_embeds.md b/docs/features/prompt_embeds.md index 3d68b07a3ac..dd0b4d62c42 100644 --- a/docs/features/prompt_embeds.md +++ b/docs/features/prompt_embeds.md @@ -20,12 +20,47 @@ You can pass prompt embeddings from Hugging Face Transformers models to the `'p ## Online Serving -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`. +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`. + +### Completions API + +Prompt embeddings inputs are added via a `'prompt_embeds'` key in the JSON request body. When a mixture of `'prompt_embeds'` and `'prompt'` inputs are provided in a single request, the prompt embeds are always returned first. Prompt embeddings are passed in as base64 encoded torch tensors. +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. + +### Chat Completions API + +Prompt embeddings can be included as content parts in chat messages, interleaved with text: + +```json +{ + "messages": [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "prompt_embeds", "data": ""} + ] + }, + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": ""}, + {"type": "text", "text": "Summarize the above."} + ] + } + ] +} +``` + +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. + +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. + !!! warning The vLLM engine may crash if incorrect shape of embeddings is passed. Only enable this flag for trusted users! diff --git a/examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py b/examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py index 40eae0c062d..f3204645d0a 100644 --- a/examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py +++ b/examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py @@ -1,12 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -vLLM OpenAI-Compatible Client with Prompt Embeddings +"""vLLM OpenAI-Compatible Client with Prompt Embeddings. This script demonstrates how to: -1. Generate prompt embeddings using Hugging Face Transformers -2. Encode them in base64 format -3. Send them to a vLLM server via the OpenAI-compatible Completions API +1. Generate prompt embeddings using Hugging Face Transformers. +2. Encode them in base64 format. +3. Send them to a vLLM server for inference via both: + - OpenAI-compatible Chat Completions API + - OpenAI-compatible Completions API + +Important distinction between the two APIs: + +- Chat Completions API: `prompt_embeds` content parts should encode ONLY + the user-provided content, not a templated conversation. The server + renders the surrounding 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 likely produce undesirable results. + +- Completions API: the server does NOT apply a chat template to + `prompt_embeds`. The caller is responsible for producing embeddings for + the full, already-templated prompt (i.e. apply the chat template first, + 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. Run the vLLM server first: vllm serve meta-llama/Llama-3.2-1B-Instruct \ @@ -34,34 +51,68 @@ from openai import OpenAI from vllm.utils.serial_utils import tensor2base64 -def main(): - client = OpenAI( - api_key="EMPTY", - base_url="http://localhost:8000/v1", +def run_chat_completion_prompt_embeds( + client: OpenAI, + model_name: str, + tokenizer: transformers.PreTrainedTokenizerBase, + embedding_layer, + messages: list[dict], +) -> None: + """Run a Chat Completions API request using prompt_embeds content parts. + + This example embeds ONLY the user-provided content of the final user turn, the + vLLM server applies the chat template around it at request time. + """ + user_content = messages[-1]["content"] + content_token_ids = tokenizer( + user_content, return_tensors="pt", add_special_tokens=False + ).input_ids + content_prompt_embeds = embedding_layer(content_token_ids).squeeze(0) + encoded_embeds = tensor2base64(content_prompt_embeds) + + api_messages = [ + *messages[:-1], + { + "role": messages[-1]["role"], + "content": [{"type": "prompt_embeds", "data": encoded_embeds}], + }, + ] + + chat_completion = client.chat.completions.create( + model=model_name, + max_tokens=6, + temperature=0.0, + messages=api_messages, ) - model_name = "meta-llama/Llama-3.2-1B-Instruct" + print("-" * 30) + print("Chat Completions API") + print(chat_completion.choices[0].message.content) + print("-" * 30) - # Transformers - tokenizer = transformers.AutoTokenizer.from_pretrained(model_name) - transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name) - # Refer to the HuggingFace repo for the correct format to use - chat = [{"role": "user", "content": "Please tell me about the capital of France."}] - token_ids = tokenizer.apply_chat_template( - chat, add_generation_prompt=True, return_tensors="pt", return_dict=True +def run_completion_prompt_embeds( + client: OpenAI, + model_name: str, + tokenizer: transformers.PreTrainedTokenizerBase, + embedding_layer, + messages: list[dict], +) -> None: + """Run a Completions API request using prompt embeddings. + + The Completions endpoint does not apply a chat template, + so the caller must apply it and embed the full templated prompt. + """ + templated_token_ids = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt", return_dict=True ).input_ids - - embedding_layer = transformers_model.get_input_embeddings() - prompt_embeds = embedding_layer(token_ids).squeeze(0) - - # Prompt embeddings - encoded_embeds = tensor2base64(prompt_embeds) + templated_prompt_embeds = embedding_layer(templated_token_ids).squeeze(0) + encoded_embeds = tensor2base64(templated_prompt_embeds) completion = client.completions.create( model=model_name, prompt=None, - max_tokens=5, + max_tokens=6, temperature=0.0, # NOTE: The OpenAI client allows passing in extra JSON body via the # `extra_body` argument. @@ -69,9 +120,39 @@ def main(): ) print("-" * 30) + print("Completions API") print(completion.choices[0].text) print("-" * 30) +def main() -> None: + client = OpenAI( + api_key="EMPTY", + base_url="http://localhost:8000/v1", + ) + + model_name = "meta-llama/Llama-3.2-1B-Instruct" + + tokenizer = transformers.AutoTokenizer.from_pretrained(model_name) + transformers_model = transformers.AutoModelForCausalLM.from_pretrained(model_name) + embedding_layer = transformers_model.get_input_embeddings() + + messages = [ + {"role": "user", "content": "Please tell me about the capital of France."} + ] + + # Chat Completions API: embed ONLY the user content. The server wraps + # the embedding in the chat template when it renders the messages. + run_chat_completion_prompt_embeds( + client, model_name, tokenizer, embedding_layer, messages + ) + + # Completions API: embed the FULL templated prompt. The caller must + # apply the chat template up-front. + run_completion_prompt_embeds( + client, model_name, tokenizer, embedding_layer, messages + ) + + if __name__ == "__main__": main() diff --git a/tests/entrypoints/llm/test_mm_processor_kwargs.py b/tests/entrypoints/llm/test_mm_processor_kwargs.py index 19cf91230ca..1b0092df011 100644 --- a/tests/entrypoints/llm/test_mm_processor_kwargs.py +++ b/tests/entrypoints/llm/test_mm_processor_kwargs.py @@ -11,7 +11,9 @@ from vllm import LLM, SamplingParams def _make_mock_llm() -> LLM: llm = object.__new__(LLM) - llm.model_config = SimpleNamespace(runner_type="generate") + llm.model_config = SimpleNamespace( + runner_type="generate", enable_prompt_embeds=False + ) return llm diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py new file mode 100644 index 00000000000..d005edc950c --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_audio_embeds.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E test for mixing `prompt_embeds` with `audio_embeds` in a single +Chat Completions request.""" + +import json + +import openai +import pytest +import pytest_asyncio +import safetensors +import torch +import torch.nn as nn +from huggingface_hub import hf_hub_download +from transformers import AutoConfig, AutoTokenizer + +from tests.utils import RemoteOpenAIServer +from vllm.utils.serial_utils import tensor2base64 + +QWEN2AUDIO_MODEL = "Qwen/Qwen2-Audio-7B-Instruct" + +# Use the model's native dtype to avoid an implicit cast inside +# `safe_load_prompt_embeds` (mismatched floating-point dtypes are cast to the +# model's dtype automatically, matching here just skips the conversion). +QWEN2AUDIO_DTYPE = torch.bfloat16 + + +@pytest.fixture(scope="module") +def qwen2audio_server_args() -> list[str]: + return [ + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "4", + "--enforce-eager", + "--trust-remote-code", + "--gpu-memory-utilization", + "0.85", + "--limit-mm-per-prompt", + json.dumps({"audio": 1}), + "--enable-prompt-embeds", + "--enable-mm-embeds", + ] + + +@pytest.fixture(scope="module") +def qwen2audio_server(qwen2audio_server_args): + with RemoteOpenAIServer( + QWEN2AUDIO_MODEL, + qwen2audio_server_args, + max_wait_seconds=600, + ) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def qwen2audio_client(qwen2audio_server): + async with qwen2audio_server.get_async_client() as async_client: + yield async_client + + +@pytest.fixture(scope="module") +def qwen2audio_hidden_size() -> int: + config = AutoConfig.from_pretrained(QWEN2AUDIO_MODEL, trust_remote_code=True) + return config.text_config.hidden_size + + +@pytest.fixture(scope="module") +def qwen2audio_prompt_embeds_b64(qwen2audio_hidden_size: int) -> str: + tensor = torch.randn(4, qwen2audio_hidden_size, dtype=QWEN2AUDIO_DTYPE) + return tensor2base64(tensor) + + +@pytest.fixture(scope="module") +def qwen2audio_audio_embeds_b64(qwen2audio_hidden_size: int) -> str: + # Shape matches the `audio_embeds` unit-test fixture. + torch.manual_seed(0) + tensor = torch.randn(1, 128, qwen2audio_hidden_size, dtype=QWEN2AUDIO_DTYPE) + return tensor2base64(tensor) + + +@pytest.mark.asyncio +async def test_prompt_embeds_plus_audio_embeds( + qwen2audio_client: openai.AsyncOpenAI, + qwen2audio_prompt_embeds_b64: str, + qwen2audio_audio_embeds_b64: str, +): + """Single user message carrying both prompt_embeds and audio_embeds parts.""" + chat = await qwen2audio_client.chat.completions.create( + model=QWEN2AUDIO_MODEL, + max_tokens=5, + temperature=0.0, + messages=[ + { + "role": "user", + "content": [ + { + "type": "prompt_embeds", + "data": qwen2audio_prompt_embeds_b64, + }, + { + "type": "audio_embeds", + "audio_embeds": qwen2audio_audio_embeds_b64, + }, + {"type": "text", "text": "Continue."}, + ], + } + ], + ) + assert chat.choices[0].message.content is not None + assert len(chat.choices[0].message.content) > 0 + + +@pytest.fixture(scope="module") +def qwen2audio_aligned_content_and_embeds_b64() -> tuple[str, str]: + """Return `(content, base64_embeds)` where the embeddings are the model's + embedding of `content` tokenized WITHOUT special tokens. + + Loads only the `embed_tokens` shard from disk on CPU (~1.1 GB of host + RAM) instead of the full 7B model on GPU. + """ + content = "Describe this audio." + tokenizer = AutoTokenizer.from_pretrained(QWEN2AUDIO_MODEL, trust_remote_code=True) + + index_path = hf_hub_download(QWEN2AUDIO_MODEL, "model.safetensors.index.json") + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight")) + shard_path = hf_hub_download(QWEN2AUDIO_MODEL, weight_map[embed_key]) + with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f: + embed_weight = f.get_tensor(embed_key) + embed_layer = nn.Embedding.from_pretrained(embed_weight.to(QWEN2AUDIO_DTYPE)) + + ids = tokenizer(content, add_special_tokens=False, return_tensors="pt").input_ids + embeds = embed_layer(ids).squeeze(0) + return content, tensor2base64(embeds) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "audio_first", + [True, False], + ids=["audio_embeds-then-text", "text-then-audio_embeds"], +) +async def test_text_content_and_prompt_embeds_match_with_audio_embeds( + qwen2audio_client: openai.AsyncOpenAI, + qwen2audio_audio_embeds_b64: str, + qwen2audio_aligned_content_and_embeds_b64: tuple[str, str], + audio_first: bool, +): + """Same content as text vs `prompt_embeds` should yield identical Chat + Completions output when mixed with `audio_embeds` in the same message. + """ + content, encoded_text_embeds = qwen2audio_aligned_content_and_embeds_b64 + + audio_part = { + "type": "audio_embeds", + "audio_embeds": qwen2audio_audio_embeds_b64, + } + text_part = {"type": "text", "text": content} + embeds_part = {"type": "prompt_embeds", "data": encoded_text_embeds} + + if audio_first: + text_content = [audio_part, text_part] + embeds_content = [audio_part, embeds_part] + else: + text_content = [text_part, audio_part] + embeds_content = [embeds_part, audio_part] + + text_resp = await qwen2audio_client.chat.completions.create( + model=QWEN2AUDIO_MODEL, + max_tokens=10, + temperature=0.0, + messages=[{"role": "user", "content": text_content}], + ) + embeds_resp = await qwen2audio_client.chat.completions.create( + model=QWEN2AUDIO_MODEL, + max_tokens=10, + temperature=0.0, + messages=[{"role": "user", "content": embeds_content}], + ) + + text_out = text_resp.choices[0].message.content + embeds_out = embeds_resp.choices[0].message.content + assert text_out is not None and len(text_out) > 0 + assert embeds_out is not None and len(embeds_out) > 0 + assert text_out == embeds_out diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py new file mode 100644 index 00000000000..dbbed3c4712 --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_mixed_image_embeds.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E tests for mixing `prompt_embeds` with image content parts in a single +Chat Completions request. +""" + +import json + +import openai +import pytest +import pytest_asyncio +import safetensors +import torch +import torch.nn as nn +from huggingface_hub import hf_hub_download +from transformers import AutoTokenizer + +from tests.utils import RemoteOpenAIServer +from vllm.assets.image import ImageAsset +from vllm.multimodal.utils import encode_image_url +from vllm.utils.serial_utils import tensor2base64 + +MODEL_NAME = "Qwen/Qwen2-VL-2B-Instruct" + +# Use the model's native dtype to skip the implicit cast inside +# `safe_load_prompt_embeds` (mismatched floating-point dtypes are cast to the +# model's dtype automatically). +MODEL_DTYPE = torch.bfloat16 + + +@pytest.fixture(scope="module") +def server_args() -> list[str]: + return [ + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "4", + "--enforce-eager", + "--gpu-memory-utilization", + "0.4", + "--limit-mm-per-prompt", + json.dumps({"image": 1}), + "--enable-prompt-embeds", + "--enable-mm-embeds", + ] + + +@pytest.fixture(scope="module") +def server(server_args): + with RemoteOpenAIServer( + MODEL_NAME, + server_args, + max_wait_seconds=600, + ) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server): + async with server.get_async_client() as async_client: + yield async_client + + +@pytest.fixture(scope="module") +def image_url() -> str: + """Stable real image as a data URL, kept identical across both the + text and prompt_embeds requests so any output difference must come from + how the text content is delivered.""" + return encode_image_url(ImageAsset("stop_sign").pil_image) + + +@pytest.fixture(scope="module") +def aligned_content_and_embeds_b64() -> tuple[str, str]: + """`(content, base64_embeds)` where the embeddings are the model's + embedding of `content` tokenized WITHOUT special tokens. + + Loads only the `embed_tokens` shard from disk on CPU instead of the full + model on GPU, so the fixture has zero VRAM footprint and won't contend + with the running vLLM server. + """ + content = "Describe this image." + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + + index_path = hf_hub_download(MODEL_NAME, "model.safetensors.index.json") + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + embed_key = next(k for k in weight_map if k.endswith("embed_tokens.weight")) + shard_path = hf_hub_download(MODEL_NAME, weight_map[embed_key]) + with safetensors.safe_open(shard_path, framework="pt", device="cpu") as f: + embed_weight = f.get_tensor(embed_key) + embed_layer = nn.Embedding.from_pretrained(embed_weight.to(MODEL_DTYPE)) + + ids = tokenizer(content, add_special_tokens=False, return_tensors="pt").input_ids + embeds = embed_layer(ids).squeeze(0) + return content, tensor2base64(embeds) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "image_first", + [True, False], + ids=["image_url-then-text", "text-then-image_url"], +) +async def test_text_content_and_prompt_embeds_match_with_image_url( + client: openai.AsyncOpenAI, + image_url: str, + aligned_content_and_embeds_b64: tuple[str, str], + image_first: bool, +): + """Same content as text vs `prompt_embeds` should yield identical Chat + Completions output when mixed with an `image_url` part in the same + message under greedy decoding. + """ + content, encoded_text_embeds = aligned_content_and_embeds_b64 + + image_part = {"type": "image_url", "image_url": {"url": image_url}} + text_part = {"type": "text", "text": content} + embeds_part = {"type": "prompt_embeds", "data": encoded_text_embeds} + + if image_first: + text_content = [image_part, text_part] + embeds_content = [image_part, embeds_part] + else: + text_content = [text_part, image_part] + embeds_content = [embeds_part, image_part] + + text_resp = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=10, + temperature=0.0, + messages=[{"role": "user", "content": text_content}], + ) + embeds_resp = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=10, + temperature=0.0, + messages=[{"role": "user", "content": embeds_content}], + ) + + text_out = text_resp.choices[0].message.content + embeds_out = embeds_resp.choices[0].message.content + assert text_out is not None and len(text_out) > 0 + assert embeds_out is not None and len(embeds_out) > 0 + assert text_out == embeds_out + + +@pytest.fixture(scope="module") +def image_embeds_b64() -> dict[str, str]: + """Synthetic but stable `image_embeds` for Qwen2-VL.""" + grid = (1, 4, 4) + spatial_merge_size = 2 + num_patches = (grid[1] // spatial_merge_size) * (grid[2] // spatial_merge_size) + text_hidden_size = 1536 # Qwen2-VL-2B + torch.manual_seed(0) + return { + "image_embeds": tensor2base64( + torch.randn(num_patches, text_hidden_size, dtype=MODEL_DTYPE) + ), + "image_grid_thw": tensor2base64(torch.tensor(grid)), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "image_first", + [True, False], + ids=["image_embeds-then-text", "text-then-image_embeds"], +) +async def test_text_content_and_prompt_embeds_match_with_image_embeds( + client: openai.AsyncOpenAI, + image_embeds_b64: dict[str, str], + aligned_content_and_embeds_b64: tuple[str, str], + image_first: bool, +): + """Same content as text vs `prompt_embeds` should yield identical Chat + Completions output when mixed with a precomputed `image_embeds` part in + the same message under greedy decoding. + """ + content, encoded_text_embeds = aligned_content_and_embeds_b64 + + image_part = {"type": "image_embeds", "image_embeds": image_embeds_b64} + text_part = {"type": "text", "text": content} + embeds_part = {"type": "prompt_embeds", "data": encoded_text_embeds} + + if image_first: + text_content = [image_part, text_part] + embeds_content = [image_part, embeds_part] + else: + text_content = [text_part, image_part] + embeds_content = [embeds_part, image_part] + + text_resp = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=10, + temperature=0.0, + messages=[{"role": "user", "content": text_content}], + ) + embeds_resp = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=10, + temperature=0.0, + messages=[{"role": "user", "content": embeds_content}], + ) + + text_out = text_resp.choices[0].message.content + embeds_out = embeds_resp.choices[0].message.content + assert text_out is not None and len(text_out) > 0 + assert embeds_out is not None and len(embeds_out) > 0 + assert text_out == embeds_out diff --git a/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py new file mode 100644 index 00000000000..1813d74798d --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_chat_completion_with_prompt_embeds.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E tests for `prompt_embeds` content parts in the Chat Completions API.""" + +import asyncio +import io + +import openai +import pybase64 as base64 +import pytest +import pytest_asyncio +import torch +from openai import BadRequestError + +from tests.utils import VLLM_PATH, RemoteOpenAIServer + +MODEL_NAME = "facebook/opt-125m" +CHAT_TEMPLATE = VLLM_PATH / "examples/template_chatml.jinja" +# Matches `--dtype` in `server_args` to avoid an implicit cast in +# `safe_load_prompt_embeds` (mismatched floating-point dtypes are cast to the +# model's dtype automatically, we match here just to skip the conversion). +SERVER_DTYPE: torch.dtype = torch.bfloat16 + + +@pytest.fixture(scope="module") +def server_args() -> list[str]: + return [ + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "128", + "--enforce-eager", + "--chat-template", + str(CHAT_TEMPLATE), + # Prompt Embeds server args + "--enable-prompt-embeds", + ] + + +@pytest.fixture(scope="module") +def server(server_args): + with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server): + async with server.get_async_client() as async_client: + yield async_client + + +def _encode_embeds(embeds: torch.Tensor) -> str: + buf = io.BytesIO() + torch.save(embeds, buf) + return base64.b64encode(buf.getvalue()).decode("utf-8") + + +@pytest.fixture(scope="module") +def prompt_embeds_b64(hf_runner) -> list[str]: + """Pre-compute embeddings for two short prompts and return as base64.""" + prompts = ["Hello, my name is", "What is an LLM?"] + with hf_runner(MODEL_NAME) as hf_model: + embeddings = hf_model.get_prompt_embeddings(prompts) + # Cast to the server's dtype so `safe_load_prompt_embeds` doesn't need to + # convert on its own, the function accepts any floating-point dtype and + # will cast to the model's dtype, but matching up front skips the work. + return [_encode_embeds(e.to(SERVER_DTYPE)) for e in embeddings] + + +@pytest.mark.asyncio +async def test_single_prompt_embeds_part( + client: openai.AsyncOpenAI, + prompt_embeds_b64: list[str], +): + """A user message with one prompt_embeds part + text.""" + b64 = prompt_embeds_b64[0] + chat = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=5, + temperature=0.0, + messages=[ + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": b64}, + {"type": "text", "text": "Continue:"}, + ], + } + ], + ) + assert chat.choices[0].message.content is not None + assert len(chat.choices[0].message.content) > 0 + + +@pytest.mark.asyncio +async def test_multiple_prompt_embeds_parts( + client: openai.AsyncOpenAI, + prompt_embeds_b64: list[str], +): + """Multiple prompt_embeds parts in a single message.""" + b64_a, b64_b = prompt_embeds_b64 + chat = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=5, + temperature=0.0, + messages=[ + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": b64_a}, + {"type": "text", "text": " and "}, + {"type": "prompt_embeds", "data": b64_b}, + ], + } + ], + ) + assert chat.choices[0].message.content is not None + assert len(chat.choices[0].message.content) > 0 + + +@pytest.mark.asyncio +async def test_multi_message_conversation( + client: openai.AsyncOpenAI, + prompt_embeds_b64: list[str], +): + """prompt_embeds in both system and user messages.""" + b64_sys, b64_usr = prompt_embeds_b64 + chat = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=5, + temperature=0.0, + messages=[ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are helpful."}, + {"type": "prompt_embeds", "data": b64_sys}, + ], + }, + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": b64_usr}, + {"type": "text", "text": "Summarize."}, + ], + }, + ], + ) + assert chat.choices[0].message.content is not None + assert len(chat.choices[0].message.content) > 0 + + +@pytest.mark.asyncio +async def test_streaming( + client: openai.AsyncOpenAI, + prompt_embeds_b64: list[str], +): + """Streaming chat completion with prompt_embeds.""" + b64 = prompt_embeds_b64[0] + + # Non-streaming baseline. + baseline = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=5, + temperature=0.0, + messages=[ + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": b64}, + {"type": "text", "text": "Continue:"}, + ], + } + ], + ) + expected = baseline.choices[0].message.content + + # Streaming. + stream = await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=5, + temperature=0.0, + stream=True, + messages=[ + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": b64}, + {"type": "text", "text": "Continue:"}, + ], + } + ], + ) + chunks: list[str] = [] + async for chunk in stream: + delta = chunk.choices[0].delta.content + if delta: + chunks.append(delta) + assert "".join(chunks) == expected + + +@pytest.fixture(scope="module") +def aligned_content_and_embeds_b64(hf_runner) -> tuple[str, str]: + """Return `(content, base64_embeds)` where the embeddings are the model's + embedding of `content` tokenized WITHOUT special tokens. + """ + content = "Hello, my name is" + with hf_runner(MODEL_NAME) as hf_model: + ids = hf_model.tokenizer( + content, add_special_tokens=False, return_tensors="pt" + ).input_ids + ids = hf_model.wrap_device({"input_ids": ids})["input_ids"] + embed_layer = hf_model.model.get_input_embeddings() + embeds = embed_layer(ids).squeeze(0).to(SERVER_DTYPE).cpu() + return content, _encode_embeds(embeds) + + +@pytest.mark.asyncio +async def test_text_content_and_prompt_embeds_match( + client: openai.AsyncOpenAI, + aligned_content_and_embeds_b64: tuple[str, str], +): + """Equal content in text and `prompt_embeds` should yield identical + Chat Completions output under greedy decoding. + """ + content, encoded_embeds = aligned_content_and_embeds_b64 + + text_resp, embeds_resp = await asyncio.gather( + client.chat.completions.create( + model=MODEL_NAME, + max_tokens=10, + temperature=0.0, + messages=[{"role": "user", "content": content}], + ), + client.chat.completions.create( + model=MODEL_NAME, + max_tokens=10, + temperature=0.0, + messages=[ + { + "role": "user", + "content": [{"type": "prompt_embeds", "data": encoded_embeds}], + } + ], + ), + ) + + text_out = text_resp.choices[0].message.content + embeds_out = embeds_resp.choices[0].message.content + assert text_out is not None and len(text_out) > 0 + assert embeds_out is not None and len(embeds_out) > 0 + assert text_out == embeds_out + + +@pytest.mark.asyncio +async def test_missing_data_field( + client: openai.AsyncOpenAI, +): + """A prompt_embeds part without `data` should return a clear error.""" + with pytest.raises(BadRequestError): + await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=5, + messages=[ + { + "role": "user", + "content": [{"type": "prompt_embeds"}], + } + ], + ) + + +@pytest.mark.asyncio +async def test_invalid_base64( + client: openai.AsyncOpenAI, +): + """Invalid base64 in the `data` field should return a clear error.""" + with pytest.raises(BadRequestError): + await client.chat.completions.create( + model=MODEL_NAME, + max_tokens=5, + messages=[ + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": "not_valid_base64!!"}, + ], + } + ], + ) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 39d59d28f85..df4d5ad47ca 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -538,6 +538,7 @@ class MockModelConfig: is_encoder_decoder: bool = False is_multimodal_model: bool = False renderer_num_workers: int = 1 + enable_prompt_embeds: bool = False def get_diff_sampling_param(self): return self.diff_sampling_param or {} diff --git a/tests/entrypoints/openai/completion/test_prompt_validation.py b/tests/entrypoints/openai/completion/test_prompt_validation.py index f44d13c555c..81204b27bc0 100644 --- a/tests/entrypoints/openai/completion/test_prompt_validation.py +++ b/tests/entrypoints/openai/completion/test_prompt_validation.py @@ -62,6 +62,8 @@ def test_load_prompt_embeds( ): model_config = Mock(spec=ModelConfig) model_config.enable_prompt_embeds = True + model_config.get_hidden_size.return_value = hidden_size + model_config.dtype = dtype # construct arbitrary tensors of various dtypes, layouts, and sizes. # We need to check against different layouts to make sure that if a user diff --git a/tests/renderers/test_chat_utils_prompt_embeds.py b/tests/renderers/test_chat_utils_prompt_embeds.py new file mode 100644 index 00000000000..e33cc304710 --- /dev/null +++ b/tests/renderers/test_chat_utils_prompt_embeds.py @@ -0,0 +1,576 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Offline unit tests for `prompt_embeds` chat-completion content parts.""" + +from __future__ import annotations + +import inspect +import io +from typing import Final +from unittest import mock + +import pybase64 as base64 +import pytest +import regex as re +import torch +from transformers import AutoTokenizer + +from vllm.entrypoints.chat_utils import ( + _ENABLE_PROMPT_EMBEDS_ERROR, + _PROMPT_EMBEDS_MISSING_DATA_ERROR, + _RESERVED_PLACEHOLDER_IN_TEXT_ERROR, + MM_PARSER_MAP, + MODALITY_PLACEHOLDERS_MAP, + PROMPT_EMBEDS_PLACEHOLDER_TOKEN, + parse_chat_messages, + parse_chat_messages_async, +) +from vllm.renderers.hf import ( + _PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR, + _build_mixed_prompt_embeds, + _build_prompt_embeds_positions, + _build_prompt_embeds_updates, + _ensure_prompt_embeds_placeholder_token, + _expand_prompt_embeds_placeholders, +) + +# Cover distinct tokenizer families: +# GPT2TokenizerFast (BPE, OpenAI-style) +# Qwen2TokenizerFast (SentencePiece BPE variant) +# BertTokenizerFast (WordPiece) +TOKENIZER_IDS: Final[list[str]] = [ + "gpt2", + "Qwen/Qwen2.5-1.5B-Instruct", + "bert-base-uncased", +] + + +@pytest.fixture(params=TOKENIZER_IDS, ids=TOKENIZER_IDS) +def tokenizer(request): + """A fresh tokenizer instance per tokenizer family.""" + return AutoTokenizer.from_pretrained(request.param) + + +# Minimal chat template that works with any tokenizer. Iterates +# `message.content` as either a string or a list of dicts (openai format). +_SIMPLE_CHAT_TEMPLATE: Final[str] = ( + "{% for m in messages %}" + "{% if m['content'] is string %}{{m['content']}}" + "{% else %}{% for p in m['content'] %}{{p['text']}}{% endfor %}" + "{% endif %}\n{% endfor %}" +) + + +async def _maybe_await(fn, *args, **kwargs): + """Call *fn* and `await` the result if it's a coroutine.""" + result = fn(*args, **kwargs) + if inspect.iscoroutine(result): + result = await result + return result + + +# Parametrize over sync / async parse paths so every end-to-end test +# exercises both. +_PARSE_FUNCTIONS = [parse_chat_messages, parse_chat_messages_async] + + +@pytest.fixture(params=_PARSE_FUNCTIONS, ids=["sync", "async"]) +def parse_fn(request): + """Either the sync or async `parse_chat_messages` callable.""" + return request.param + + +def _encode_tensor(t: torch.Tensor) -> str: + buf = io.BytesIO() + torch.save(t, buf) + return base64.b64encode(buf.getvalue()).decode("utf-8") + + +_MOCK_HIDDEN_SIZE: Final[int] = 8 +_MOCK_DTYPE: Final[torch.dtype] = torch.float32 + + +def _make_mock_model_config(*, enable_prompt_embeds: bool = True) -> mock.MagicMock: + mc = mock.MagicMock() + mc.enable_prompt_embeds = enable_prompt_embeds + mc.multimodal_config = None + mc.allowed_local_media_path = None + mc.allowed_media_domains = None + # Test text-only code path in `MultiModalItemTracker.resolve_items`. + mc.is_multimodal_model = False + # `safe_load_prompt_embeds` pins each tensor to the model's hidden_size + # and dtype, so the mock must return concrete values. + mc.get_hidden_size.return_value = _MOCK_HIDDEN_SIZE + mc.dtype = _MOCK_DTYPE + return mc + + +def test_prompt_embeds_keys_registered(): + assert "prompt_embeds" in MODALITY_PLACEHOLDERS_MAP + assert MODALITY_PLACEHOLDERS_MAP["prompt_embeds"] == "<##PROMPT_EMBEDS##>" + assert "prompt_embeds" in MM_PARSER_MAP + + +def test_ensure_placeholder_token_is_single_token_and_idempotent(tokenizer): + """Ensure the placeholder token is a single token and that multiple calls to + "ensure" are idempotent, across all tokenizer families.""" + tid1 = _ensure_prompt_embeds_placeholder_token(tokenizer) + tid2 = _ensure_prompt_embeds_placeholder_token(tokenizer) + assert tid1 == tid2 + + ids = tokenizer.encode(PROMPT_EMBEDS_PLACEHOLDER_TOKEN, add_special_tokens=False) + assert ids == [tid1] + + # Repeating it in a string N times must produce exactly that many tokens. + N = 5 + ids_rep = tokenizer.encode( + PROMPT_EMBEDS_PLACEHOLDER_TOKEN * N, add_special_tokens=False + ) + assert ids_rep == [tid1] * N + + +def test_parse_chat_messages_openai_format(): + NUM_TOKENS = 3 + t = torch.randn(NUM_TOKENS, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) + b64 = _encode_tensor(t) + mc = _make_mock_model_config() + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Hello "}, + {"type": "prompt_embeds", "data": b64}, + {"type": "text", "text": " world"}, + ], + } + ] + conv, mm_data, _ = parse_chat_messages( + messages, + mc, + content_format="openai", + ) + # The middle content part is rewritten to a single placeholder-token + # sentinel. + texts = [p["text"] for p in conv[0]["content"]] + assert texts == [ + "Hello ", + PROMPT_EMBEDS_PLACEHOLDER_TOKEN, + " world", + ] + assert mm_data is not None and "prompt_embeds" in mm_data + assert torch.equal(mm_data["prompt_embeds"][0], t) + + +# Each layout entry is one content part: +# ("text", "A") -> {"type": "text", "text": "A"} +# ("embed", N) -> {"type": "prompt_embeds", "data": } +@pytest.mark.parametrize( + "layout", + [ + # Case: Single embed only. + [("embed", 2)], + # Case: Embed at the start of the message. + [("embed", 3), ("text", "B")], + # Case: Embed at the end of the message. + [("text", "A"), ("embed", 1)], + # Case: Embed sandwiched between text spans. + [("text", "A"), ("embed", 2), ("text", "B")], + # Case: Multiple embeds with text in between. + [("text", "A"), ("embed", 2), ("text", "B"), ("embed", 3)], + # Case: Adjacent embeds with no separating text. + [("embed", 1), ("embed", 2)], + # Case: Multiple text spans before a trailing embed. + [("text", "A"), ("text", "B"), ("embed", 1)], + # Case: Long-ish run mixing both kinds. + [ + ("text", "head"), + ("embed", 4), + ("text", "mid"), + ("embed", 1), + ("embed", 2), + ("text", "tail"), + ], + ], + ids=[ + "single-embed", + "embed-then-text", + "text-then-embed", + "text-embed-text", + "text-embed-text-embed", + "adjacent-embeds", + "text-text-embed", + "long-mixed-run", + ], +) +@pytest.mark.parametrize( + "interleave_mm_strings", + # `None`: text-only path where `multimodal_config` is absent. + # `False`: non-interleave multimodal path (the common default). + # `True`: sentinel-substitution interleave path. + # All three must preserve the request ordering of prompt_embeds + # relative to surrounding text because prompt_embeds are spliced at the + # token offset during rendering. + [None, False, True], + ids=["text-only", "interleave-off", "interleave-on"], +) +def test_parse_chat_messages_string_format_preserves_position( + layout, interleave_mm_strings +): + mc = _make_mock_model_config() + if interleave_mm_strings is not None: + mm_cfg = mock.MagicMock() + mm_cfg.interleave_mm_strings = interleave_mm_strings + mc.multimodal_config = mm_cfg + + content: list[dict] = [] + expected_parts: list[str] = [] + expected_embeds: list[torch.Tensor] = [] + for kind, value in layout: + if kind == "text": + content.append({"type": "text", "text": value}) + expected_parts.append(value) + else: # prompt embeds + num_tokens = value + t = torch.randn(num_tokens, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) + expected_embeds.append(t) + content.append({"type": "prompt_embeds", "data": _encode_tensor(t)}) + # Parser emits ONE sentinel per part. + expected_parts.append(PROMPT_EMBEDS_PLACEHOLDER_TOKEN) + + messages = [{"role": "user", "content": content}] + conv, mm_data, _ = parse_chat_messages( + messages, + mc, + content_format="string", + ) + + assert conv[0]["content"] == "\n".join(expected_parts) + assert mm_data is not None and "prompt_embeds" in mm_data + assert len(mm_data["prompt_embeds"]) == len(expected_embeds) + for got, want in zip(mm_data["prompt_embeds"], expected_embeds, strict=True): + assert torch.equal(got, want) + + +def test_parse_chat_messages_requires_flag(): + t = torch.randn(2, 4) + b64 = _encode_tensor(t) + mc = _make_mock_model_config(enable_prompt_embeds=False) + + messages = [ + { + "role": "user", + "content": [{"type": "prompt_embeds", "data": b64}], + } + ] + with pytest.raises(ValueError, match=_ENABLE_PROMPT_EMBEDS_ERROR): + parse_chat_messages( + messages, + mc, + content_format="openai", + ) + + +def test_parse_chat_messages_rejects_missing_data(): + # `data` is marked `Required` on `ChatCompletionContentPartPromptEmbedsParam`; + # malformed requests without `data` must surface a clear validation error + # rather than being silently dropped. + mc = _make_mock_model_config() + messages = [ + { + "role": "user", + "content": [{"type": "prompt_embeds"}], # no `data` + } + ] + with pytest.raises(ValueError, match=_PROMPT_EMBEDS_MISSING_DATA_ERROR): + parse_chat_messages( + messages, + mc, + content_format="openai", + ) + + +# Reserved placeholder guard: when `enable_prompt_embeds=True` the tokenizer is +# mutated to make `` a single unsplittable token. Any user text +# containing that literal sequence would tokenize to the same sentinel ID and +# be mistaken for a splice point, so we reject it at parse time. +_PLACEHOLDER_ERROR_PATTERN: Final[str] = re.sub( + r"\\{[^}]*\\}", ".*", re.escape(_RESERVED_PLACEHOLDER_IN_TEXT_ERROR) +) + + +@pytest.mark.parametrize( + "content", + [ + # Case: Top-level string content (wrapped as a single text part). + f"hello {PROMPT_EMBEDS_PLACEHOLDER_TOKEN} world", + # Case: List with a typed text part containing the placeholder. + [{"type": "text", "text": f"leading {PROMPT_EMBEDS_PLACEHOLDER_TOKEN}"}], + # Case: List with a plain-string part (no wrapping dict). + [f"raw string {PROMPT_EMBEDS_PLACEHOLDER_TOKEN}"], + ], + ids=["top-level-string", "typed-text-part", "plain-string-part"], +) +def test_parse_chat_messages_rejects_placeholder_in_user_text(content): + mc = _make_mock_model_config() # enable_prompt_embeds=True by default + messages = [{"role": "user", "content": content}] + with pytest.raises(ValueError, match=_PLACEHOLDER_ERROR_PATTERN): + parse_chat_messages(messages, mc, content_format="openai") + + +def test_parse_chat_messages_allows_placeholder_in_text_when_feature_disabled(): + # When `enable_prompt_embeds=False` the tokenizer is never mutated, so the + # literal `` is just ordinary text and must pass through. + mc = _make_mock_model_config(enable_prompt_embeds=False) + messages = [ + { + "role": "user", + "content": f"benign mention of {PROMPT_EMBEDS_PLACEHOLDER_TOKEN} here", + } + ] + conv, mm_data, _ = parse_chat_messages(messages, mc, content_format="openai") + assert mm_data is None or "prompt_embeds" not in mm_data + # Text reaches the rendered conversation unchanged. + texts = [p["text"] for p in conv[0]["content"]] + assert PROMPT_EMBEDS_PLACEHOLDER_TOKEN in "".join(texts) + + +# Token-stream spec: ints are regular token IDs, tuples `(N,)` expand to +# a placeholder span of length N (creates corresponding `(N, H)` tensor). +# `expected` lists the `(start_idx, length)` pairs that +# `_build_prompt_embeds_positions` should return. +@pytest.mark.parametrize( + "stream, expected", + [ + # Case: Single run in the middle. + ([10, 20, (3,), 30], [(2, 3)]), + # Case: Single run at the start. + ([(2,), 10, 20], [(0, 2)]), + # Case: Single run at the end. + ([10, 20, (4,)], [(2, 4)]), + # Case: Two runs with tokens between. + ([1, (2,), 2, 3, (3,), 4], [(1, 2), (5, 3)]), + # Case: Adjacent runs (no separating tokens). + ([(1,), (2,)], [(0, 1), (1, 2)]), + # Case: Three runs. + ([5, (2,), 6, (1,), 7, (3,), 8], [(1, 2), (4, 1), (6, 3)]), + ], + ids=[ + "single-middle", + "single-start", + "single-end", + "two-runs-separated", + "two-runs-adjacent", + "three-runs", + ], +) +def test_build_positions(tokenizer, stream, expected): + H = 4 + tid = _ensure_prompt_embeds_placeholder_token(tokenizer) + tensors: list[torch.Tensor] = [] + token_ids: list[int] = [] + for item in stream: + if isinstance(item, tuple): + length = item[0] + tensors.append(torch.randn(length, H)) + token_ids.extend([tid] * length) + else: + token_ids.append(item) + mm_updates = _build_prompt_embeds_updates(tensors, tid) + positions = _build_prompt_embeds_positions(token_ids, len(tensors), mm_updates) + assert positions == expected + + +def test_build_positions_length_mismatch(tokenizer): + N1, H1 = 2, 4 + N2, H2 = 3, 4 + tid = _ensure_prompt_embeds_placeholder_token(tokenizer) + # 2 tensors expected but only a single placeholder run in the token + # stream (simulating dropping the second one). + tensors = [torch.randn(N1, H1), torch.randn(N2, H2)] + token_ids = [1, tid, tid, 2, 3] + mm_updates = _build_prompt_embeds_updates(tensors, tid) + # The error constant is a `str.format` template, escape it and turn + # the `{field}` placeholders into `.*` so it matches any substitution. + pattern = re.sub( + r"\\{[^}]*\\}", ".*", re.escape(_PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR) + ) + with pytest.raises(ValueError, match=pattern): + _build_prompt_embeds_positions(token_ids, len(tensors), mm_updates) + + +# ints = regular token IDs (any value) +# (N,) = embed span of length N +@pytest.mark.parametrize( + "stream", + [ + [10, 20, (3,), 30], + [(2,), 10, 20], + [10, 20, (4,)], + [1, (2,), 2, 3, (3,), 4], + [(1,), (2,)], + [5, (2,), 6, (1,), 7, (3,), 8], + ], + ids=[ + "single-middle", + "single-start", + "single-end", + "two-spans-separated", + "two-spans-adjacent", + "three-spans", + ], +) +def test_build_mixed_prompt_embeds(stream): + H = 8 + _PLACEHOLDER = 0 # sentinel for embed positions in token_ids + + tensors: list[torch.Tensor] = [] + token_ids: list[int] = [] + positions: list[tuple[int, int]] = [] + cursor = 0 + for item in stream: + if isinstance(item, tuple): + length = item[0] + tensors.append(torch.randn(length, H)) + positions.append((cursor, length)) + token_ids.extend([_PLACEHOLDER] * length) + cursor += length + else: + token_ids.append(item) + cursor += 1 + + embeds, mask = _build_mixed_prompt_embeds(token_ids, tensors, positions) + + assert embeds.shape == (len(token_ids), H) + assert len(mask) == len(token_ids) + + # Mask: False exactly at embed positions, True everywhere else. + expected_mask = torch.ones(len(token_ids), dtype=torch.bool) + for start, length in positions: + expected_mask[start : start + length] = False + assert mask == expected_mask.tolist() + + # Embed rows match input tensors at the right positions. + for tensor, (start, length) in zip(tensors, positions): + assert torch.equal(embeds[start : start + length], tensor) + + # Non-embed positions remain zero-filled. + assert torch.all(embeds[expected_mask] == 0) + + +# End-to-end tests: each runs both sync and async parse paths via the +# `parse_fn` fixture. + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role", ["user", "system"]) +async def test_end_to_end_expand_and_build(tokenizer, parse_fn, role): + """Full renderer pipeline: parse -> chat template -> expand -> locate + -> build mixed prompt, across tokenizers, roles, and sync/async.""" + tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE + tid = _ensure_prompt_embeds_placeholder_token(tokenizer) + + LEN_A, LEN_B = 3, 2 + t_a = torch.randn(LEN_A, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) + t_b = torch.randn(LEN_B, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) + NUM_TENSORS = 2 + + mc = _make_mock_model_config() + + messages = [ + { + "role": role, + "content": [ + {"type": "text", "text": "Hello "}, + {"type": "prompt_embeds", "data": _encode_tensor(t_a)}, + {"type": "text", "text": " world "}, + {"type": "prompt_embeds", "data": _encode_tensor(t_b)}, + {"type": "text", "text": "!"}, + ], + } + ] + + conv, mm_data, _ = await _maybe_await( + parse_fn, messages, mc, content_format="openai" + ) + tensors = list(mm_data["prompt_embeds"]) + assert len(tensors) == NUM_TENSORS + + # Tokenize: each prompt_embeds part becomes 1 placeholder token. + # `return_dict=False` to get a flat `list[int]` on transformers v5 + # (where the default flipped to True and yields a `BatchEncoding` dict). + token_ids = tokenizer.apply_chat_template(conv, tokenize=True, return_dict=False) + assert sum(t == tid for t in token_ids) == NUM_TENSORS + + # Expand, locate, and build. + mm_updates = _build_prompt_embeds_updates(tensors, tid) + expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates) + assert len(expanded) == len(token_ids) + LEN_A + LEN_B - NUM_TENSORS + + positions = _build_prompt_embeds_positions(expanded, len(tensors), mm_updates) + assert positions[0][1] == LEN_A + assert positions[1][1] == LEN_B + + embeds, mask = _build_mixed_prompt_embeds(expanded, tensors, positions) + assert embeds.shape == (len(expanded), _MOCK_HIDDEN_SIZE) + assert mask.count(False) == LEN_A + LEN_B + assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_A], t_a) + assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_B], t_b) + + +@pytest.mark.asyncio +async def test_end_to_end_multi_message_conversation(tokenizer, parse_fn): + """Full pipeline with prompt_embeds spread across system + user messages, + verifying ordering and positioning in the final token stream.""" + tokenizer.chat_template = _SIMPLE_CHAT_TEMPLATE + tid = _ensure_prompt_embeds_placeholder_token(tokenizer) + + LEN_SYS, LEN_USR = 4, 3 + t_sys = torch.randn(LEN_SYS, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) + t_usr = torch.randn(LEN_USR, _MOCK_HIDDEN_SIZE, dtype=_MOCK_DTYPE) + NUM_TENSORS = 2 # t_sys and t_usr. + + mc = _make_mock_model_config() + + messages = [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are helpful."}, + {"type": "prompt_embeds", "data": _encode_tensor(t_sys)}, + ], + }, + { + "role": "user", + "content": [ + {"type": "prompt_embeds", "data": _encode_tensor(t_usr)}, + {"type": "text", "text": "Summarize."}, + ], + }, + ] + + conv, mm_data, _ = await _maybe_await( + parse_fn, messages, mc, content_format="openai" + ) + tensors = list(mm_data["prompt_embeds"]) + assert len(tensors) == NUM_TENSORS + + # Tokenize, expand, locate, and build. + # `return_dict=False` to get a flat `list[int]` on transformers v5 + # (where the default flipped to True and yields a `BatchEncoding` dict). + token_ids = tokenizer.apply_chat_template(conv, tokenize=True, return_dict=False) + mm_updates = _build_prompt_embeds_updates(tensors, tid) + expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates) + positions = _build_prompt_embeds_positions(expanded, len(tensors), mm_updates) + + assert positions[0][1] == LEN_SYS + assert positions[1][1] == LEN_USR + # System embed must appear before user embed in the token stream. + assert positions[0][0] < positions[1][0] + + embeds, mask = _build_mixed_prompt_embeds(expanded, tensors, positions) + assert embeds.shape == (len(expanded), _MOCK_HIDDEN_SIZE) + assert mask.count(False) == LEN_SYS + LEN_USR + assert torch.equal(embeds[positions[0][0] : positions[0][0] + LEN_SYS], t_sys) + assert torch.equal(embeds[positions[1][0] : positions[1][0] + LEN_USR], t_usr) diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index ccc806ba137..00d604afdcf 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -39,6 +39,11 @@ class MockModelConfig: is_encoder_decoder: bool = False is_multimodal_model: bool = False renderer_num_workers: int = 1 + hidden_size: int = 768 + dtype: torch.dtype = torch.float32 + + def get_hidden_size(self) -> int: + return self.hidden_size @dataclass @@ -384,12 +389,13 @@ class TestRenderEmbedPrompt: assert torch.equal(results[0]["prompt_embeds"], tensor_input) def test_multiple_prompt_embeds(self): - renderer = _build_renderer(MockModelConfig()) + hidden_size = 512 + renderer = _build_renderer(MockModelConfig(hidden_size=hidden_size)) # Create multiple test tensors tensor_inputs = [ - torch.randn(8, 512, dtype=torch.float32), - torch.randn(12, 512, dtype=torch.float32), + torch.randn(8, hidden_size, dtype=torch.float32), + torch.randn(12, hidden_size, dtype=torch.float32), ] prompts = renderer.render_prompts( @@ -432,13 +438,15 @@ class TestRenderEmbedPrompt: assert torch.equal(results[0]["prompt_embeds"], expected) def test_prompt_embed_different_dtypes(self): - renderer = _build_renderer(MockModelConfig()) - + hidden_size = 256 # Test different supported dtypes dtypes = [torch.float32, torch.float16, torch.bfloat16] for dtype in dtypes: - tensor_input = torch.randn(5, 256, dtype=dtype) + renderer = _build_renderer( + MockModelConfig(hidden_size=hidden_size, dtype=dtype) + ) + tensor_input = torch.randn(5, hidden_size, dtype=dtype) prompts = renderer.render_prompts( _preprocess_prompt( @@ -474,10 +482,11 @@ class TestRenderEmbedPrompt: assert results[0]["prompt_embeds"].shape == (10, 768) def test_both_prompts_and_embeds(self): - renderer = _build_renderer(MockModelConfig()) + hidden_size = 256 + renderer = _build_renderer(MockModelConfig(hidden_size=hidden_size)) text_input = "Hello world" - tensor_input = torch.randn(5, 256, dtype=torch.float32) + tensor_input = torch.randn(5, hidden_size, dtype=torch.float32) prompts = renderer.render_prompts( _preprocess_prompt( diff --git a/tests/renderers/test_sparse_tensor_validation.py b/tests/renderers/test_sparse_tensor_validation.py index 5c51cd30a33..642867086fc 100644 --- a/tests/renderers/test_sparse_tensor_validation.py +++ b/tests/renderers/test_sparse_tensor_validation.py @@ -12,6 +12,7 @@ import pybase64 as base64 import pytest import torch +from vllm.exceptions import VLLMValidationError from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO from vllm.renderers.embed_utils import safe_load_prompt_embeds @@ -53,8 +54,14 @@ def _create_malicious_sparse_tensor() -> torch.Tensor: values = torch.tensor([1.0]) shape = (3, 3) - # Create sparse tensor (this will be invalid) - sparse_tensor = torch.sparse_coo_tensor(indices, values, shape, dtype=torch.float32) + # Create sparse tensor (this will be invalid). Pass `check_invariants=False` + # explicitly so this fixture is robust to process-wide invariant-check state + # left enabled by other tests (the global flag isn't thread-local, and + # concurrent users of the `check_sparse_tensor_invariants` context manager + # can leak the "enabled" state across tests). + sparse_tensor = torch.sparse_coo_tensor( + indices, values, shape, dtype=torch.float32, check_invariants=False + ) return sparse_tensor @@ -117,7 +124,7 @@ class TestPromptEmbedsValidation: shape = (10, 10) malicious_tensor = torch.sparse_coo_tensor( - indices, values, shape, dtype=torch.float32 + indices, values, shape, dtype=torch.float32, check_invariants=False ) encoded = _encode_tensor(malicious_tensor) @@ -132,13 +139,69 @@ class TestPromptEmbedsValidation: shape = (10, 10) malicious_tensor = torch.sparse_coo_tensor( - indices, values, shape, dtype=torch.float32 + indices, values, shape, dtype=torch.float32, check_invariants=False ) encoded = _encode_tensor(malicious_tensor) with pytest.raises((RuntimeError, ValueError)): safe_load_prompt_embeds(model_config, encoded) + def test_hidden_size_mismatch_rejected(self, model_config): + """Tensors whose trailing dim doesn't match the model's hidden_size + must be rejected at parse time.""" + # opt-125m has hidden_size=768, passing 512 triggers the check. + wrong_hidden = torch.randn(10, 512, dtype=torch.float32) + encoded = _encode_tensor(wrong_hidden) + + with pytest.raises(VLLMValidationError, match="hidden_size"): + safe_load_prompt_embeds(model_config, encoded) + + def test_float_dtype_mismatch_cast_to_model_dtype(self, model_config): + """Tensors whose dtype doesn't match the model's dtype but are still + floating-point are cast, since API clients generally can't know the + server's `--dtype` setting ahead of time.""" + # Fixture pins model dtype to float32, upload a bfloat16 tensor. + mismatched_float = torch.randn(10, 768, dtype=torch.bfloat16) + encoded = _encode_tensor(mismatched_float) + + result = safe_load_prompt_embeds(model_config, encoded) + + assert result.dtype == torch.float32 + assert result.shape == mismatched_float.shape + + def test_non_float_dtype_rejected(self, model_config): + """Non-floating-point dtypes cannot be safely cast for embeddings + (e.g. integer tensors almost certainly indicate caller confusion), + so they are rejected at parse time.""" + non_float = torch.randint(0, 100, (10, 768), dtype=torch.int32) + encoded = _encode_tensor(non_float) + + with pytest.raises(VLLMValidationError, match="floating-point"): + safe_load_prompt_embeds(model_config, encoded) + + def test_non_2d_tensor_rejected(self, model_config): + """Tensors that aren't 2D (even after squeezing a leading dim) + must be rejected with a clear error.""" + # A 1D tensor cannot be interpreted as (num_tokens, hidden_size). + bad = torch.randn(768, dtype=torch.float32) + encoded = _encode_tensor(bad) + + with pytest.raises(VLLMValidationError, match="2D tensor"): + safe_load_prompt_embeds(model_config, encoded) + + def test_non_tensor_payload_rejected(self, model_config): + """Deserializing to a non-Tensor object must raise a clear error + instead of propagating an AssertionError.""" + # `torch.save` will serialize a plain dict; `weights_only=True` allows + # loading built-in containers, so this exercises the isinstance check. + buffer = io.BytesIO() + torch.save({"not": "a tensor"}, buffer) + buffer.seek(0) + encoded = base64.b64encode(buffer.read()) + + with pytest.raises(VLLMValidationError, match="torch.Tensor"): + safe_load_prompt_embeds(model_config, encoded) + class TestImageEmbedsValidation: """Test sparse tensor validation in image embeddings (Chat API).""" diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index f8099a54634..358732eabf4 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -40,6 +40,7 @@ def _model_config(): multimodal_config=None, allowed_local_media_path="", allowed_media_domains=None, + enable_prompt_embeds=False, ) diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 21a651c62ab..92de5a7e981 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -256,8 +256,10 @@ async def test_multi_abort(output_kind: RequestOutputKind): ) ) - # Let requests start - await asyncio.sleep(0.5) + # Let requests start generating, use a longer sleep to ensure all + # requests have exited prefill and produced at least one + # decode token before we abort. + await asyncio.sleep(1.0) # Use multi-abort to abort multiple requests at once abort_request_ids = [request_ids[i] for i in REQUEST_IDS_TO_ABORT] @@ -369,9 +371,10 @@ async def test_mid_stream_cancellation( # Wait for all tasks to complete results = await asyncio.gather(*tasks) - # Verify all tasks were cancelled at the expected point + # Verify all tasks were cancelled at the expected point. + # Uses >= because the cancel check is `count >= cancel_after`. for num_generated_tokens, request_id in results: - assert num_generated_tokens == NUM_EXPECTED_TOKENS, ( + assert num_generated_tokens >= NUM_EXPECTED_TOKENS, ( f"{request_id} generated {num_generated_tokens} tokens but " f"expected to cancel after {NUM_EXPECTED_TOKENS}" ) diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index c6bf6f0e6e7..cfe0857b679 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -11,7 +11,7 @@ from dataclasses import dataclass from functools import cached_property, lru_cache, partial from itertools import accumulate from pathlib import Path -from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeAlias, TypeVar, cast from openai.types.chat import ( ChatCompletionAssistantMessageParam, @@ -36,7 +36,7 @@ from PIL import Image from pydantic import BaseModel, ConfigDict, TypeAdapter # pydantic needs the TypedDict from typing_extensions -from typing_extensions import Required, TypedDict +from typing_extensions import Required, TypedDict, override from vllm import envs from vllm.config import ModelConfig @@ -55,6 +55,10 @@ from vllm.multimodal.inputs import ( ) from vllm.multimodal.media import MEDIA_CONNECTOR_REGISTRY, MediaConnector from vllm.multimodal.processing import BaseMultiModalProcessor +from vllm.renderers.embed_utils import ( + safe_load_prompt_embeds, + safe_load_prompt_embeds_async, +) from vllm.utils import random_uuid from vllm.utils.collection_utils import is_list_of from vllm.utils.import_utils import LazyLoader @@ -98,9 +102,40 @@ MODALITY_PLACEHOLDERS_MAP = { "image": "<##IMAGE##>", "audio": "<##AUDIO##>", "video": "<##VIDEO##>", + "prompt_embeds": "<##PROMPT_EMBEDS##>", } +PROMPT_EMBEDS_PLACEHOLDER_TOKEN: Final[str] = "" +"""The special token used as a placeholder for each embedding +position during chat template rendering. + +Registered as an additional special token when `--enable-prompt-embeds` is set. +See `_ensure_prompt_embeds_placeholder_token` in `vllm/renderers/hf.py`. +""" + + +_REQUIRE_MM_PROCESSOR_ERROR: Final[str] = ( + "Resolving modality {modality!r} requires a multimodal processor " + "but none is available." +) + +_ENABLE_PROMPT_EMBEDS_ERROR: Final[str] = ( + "You must set `--enable-prompt-embeds` to input `prompt_embeds`" +) + +_PROMPT_EMBEDS_MISSING_DATA_ERROR: Final[str] = ( + "prompt_embeds content part requires a non-empty `data` field " + "with base64-encoded tensor bytes." +) + +_RESERVED_PLACEHOLDER_IN_TEXT_ERROR: Final[str] = ( + "Text content may not contain the reserved placeholder {token!r}. " + "This placeholder is used internally to mark `prompt_embeds` splice " + "positions in the tokenized prompt." +) + + class AudioURL(TypedDict, total=False): url: Required[str] """ @@ -147,6 +182,17 @@ class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False): """ +class ChatCompletionContentPartPromptEmbedsParam(TypedDict, total=False): + data: Required[str] + """ + Base64-encoded bytes of a serialized `torch.Tensor` of shape + `(num_tokens, hidden_size)`. The tensor's `dtype` and `hidden_size` must + match the model's input embedding layer. + """ + type: Required[Literal["prompt_embeds"]] + """The type of the content part.""" + + class VideoURL(TypedDict, total=False): url: Required[str] """ @@ -282,6 +328,7 @@ ChatCompletionContentPartParam: TypeAlias = ( | CustomChatCompletionContentSimpleImageParam | ChatCompletionContentPartImageEmbedsParam | ChatCompletionContentPartAudioEmbedsParam + | ChatCompletionContentPartPromptEmbedsParam | CustomChatCompletionContentSimpleAudioParam | CustomChatCompletionContentSimpleVideoParam | CustomChatCompletionContentToolReferenceParam @@ -367,7 +414,13 @@ ChatTemplateContentFormat = Literal["string", "openai"] ModalityStr = Literal[ - "image", "audio", "video", "image_embeds", "audio_embeds", "vision_chunk" + "image", + "audio", + "video", + "image_embeds", + "audio_embeds", + "vision_chunk", + "prompt_embeds", ] _T = TypeVar("_T") @@ -549,7 +602,17 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]): An optional uuid can be added which serves as a unique identifier of the media. + + Note: + `prompt_embeds` bypass MM-processor validation because they are + pre-computed embeddings that do not go through any HF processor, encoder, + or model-specific placeholder logic. The corresponding placeholder string is + managed by the parser via `_add_placeholder`, so we return None here. """ + if modality == "prompt_embeds": + self._items_by_modality["prompt_embeds"].append(item) + return None + input_modality = modality.replace("_embeds", "") original_modality = modality use_vision_chunk = ( @@ -660,17 +723,32 @@ def _resolve_vision_chunk_items( def _resolve_items( items_by_modality: dict[str, list[tuple[object, str | None]]], - mm_processor: BaseMultiModalProcessor, + mm_processor: BaseMultiModalProcessor | None, modality_order: dict[str, list[str]], ) -> tuple[MultiModalDataDict, MultiModalUUIDDict]: + """ + Materialize the tracker's per-modality items into `mm_data` / `mm_uuids`. + + Note: + `mm_processor` is `None` for text-only models (no registered HF + processor) whose only modality is `prompt_embeds`. Every other + modality requires a processor, enforced by the guard below. + """ if "image" in items_by_modality and "image_embeds" in items_by_modality: raise ValueError("Mixing raw image and embedding inputs is not allowed") if "audio" in items_by_modality and "audio_embeds" in items_by_modality: raise ValueError("Mixing raw audio and embedding inputs is not allowed") + # `prompt_embeds` bypasses HF MM processors. Every other modality requires one. + processor_modalities = items_by_modality.keys() - {"prompt_embeds"} + if processor_modalities and mm_processor is None: + raise RuntimeError( + _REQUIRE_MM_PROCESSOR_ERROR.format(modality=processor_modalities) + ) mm_data = {} mm_uuids = {} if "image_embeds" in items_by_modality: + assert mm_processor is not None mm_data["image"] = _get_embeds_data( "image", [data for data, uuid in items_by_modality["image_embeds"]], @@ -681,6 +759,7 @@ def _resolve_items( mm_data["image"] = [data for data, uuid in items_by_modality["image"]] mm_uuids["image"] = [uuid for data, uuid in items_by_modality["image"]] if "audio_embeds" in items_by_modality: + assert mm_processor is not None mm_data["audio"] = _get_embeds_data( "audio", [data for data, uuid in items_by_modality["audio_embeds"]], @@ -694,6 +773,7 @@ def _resolve_items( mm_data["video"] = [data for data, uuid in items_by_modality["video"]] mm_uuids["video"] = [uuid for data, uuid in items_by_modality["video"]] if "vision_chunk" in items_by_modality: + assert mm_processor is not None # Process vision_chunk items - extract from (data, modality) tuples # and convert to VisionChunk types with proper UUID handling processed_chunks, vision_chunk_uuids = _resolve_vision_chunk_items( @@ -703,6 +783,10 @@ def _resolve_items( ) mm_data["vision_chunk"] = processed_chunks mm_uuids["vision_chunk"] = vision_chunk_uuids + if "prompt_embeds" in items_by_modality: + mm_data["prompt_embeds"] = [ + data for data, _uuid in items_by_modality["prompt_embeds"] + ] return mm_data, mm_uuids @@ -714,8 +798,16 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[tuple[object, str | None]] if not self._items_by_modality: return None, None + # Text-only models (`is_multimodal_model=False`) with inputs of + # modality `prompt_embeds` have no MM processor since `prompt_embeds` are + # pre-computed and require no processing, so we pass `None`. + mm_processor = ( + self.mm_processor if self._model_config.is_multimodal_model else None + ) return _resolve_items( - dict(self._items_by_modality), self.mm_processor, self._modality_order + dict(self._items_by_modality), + mm_processor, + self._modality_order, ) def create_parser( @@ -738,8 +830,13 @@ class AsyncMultiModalItemTracker( for modality, coros in self._items_by_modality.items() } + mm_processor = ( + self.mm_processor if self._model_config.is_multimodal_model else None + ) return _resolve_items( - resolved_items_by_modality, self.mm_processor, self._modality_order + resolved_items_by_modality, + mm_processor, + self._modality_order, ) def create_parser( @@ -758,10 +855,16 @@ class BaseMultiModalContentParser(ABC): # general MM placeholder: # { # "<##IMAGE##>": ["", "", ""], - # "<##AUDIO##>": ["