feat: Enable prompt_embeds Content Part Support in vLLM Chat Completions API (#40720)

Signed-off-by: Luis Robaina <[email protected]>
Signed-off-by: Luis Robaina 🚀 <[email protected]>
Signed-off-by: LuisRobaina <[email protected]>
Co-authored-by: Andrew Sansom <[email protected]>
This commit is contained in:
Luis 🚀
2026-05-01 10:05:55 +08:00
committed by GitHub
co-authored by Andrew Sansom
parent 1adaa5056b
commit 14043dfecd
28 changed files with 2301 additions and 82 deletions
+2 -2
View File
@@ -52,10 +52,10 @@ th:not(:first-child) {
| [mm](multimodal_inputs.md) | ✅ | ✅ | [🟠](https://github.com/vllm-project/vllm/pull/4194)<sup>^</sup> | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | | |
| 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.
<sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
<sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
### Feature x Hardware
+36 -1
View File
@@ -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": "<base64_encoded_tensor>"}
]
},
{
"role": "user",
"content": [
{"type": "prompt_embeds", "data": "<base64_encoded_tensor>"},
{"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!
@@ -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()
@@ -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
@@ -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
@@ -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
@@ -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!!"},
],
}
],
)
@@ -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 {}
@@ -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
@@ -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": <base64 of (N, H) tensor>}
@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 `<prompt_embeds>` 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 `<prompt_embeds>` 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)
+17 -8
View File
@@ -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(
@@ -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)."""
+1
View File
@@ -40,6 +40,7 @@ def _model_config():
multimodal_config=None,
allowed_local_media_path="",
allowed_media_domains=None,
enable_prompt_embeds=False,
)
+7 -4
View File
@@ -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}"
)
+201 -9
View File
@@ -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] = "<prompt_embeds>"
"""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##>": ["<image>", "<image>", "<image>"],
# "<##AUDIO##>": ["<audio>", "<audio>"]
# "<##AUDIO##>": ["<audio>", "<audio>"],
# "<##PROMPT_EMBEDS##>": ["<prompt_embeds>", "<prompt_embeds>"]
# }
self._placeholder_storage: dict[str, list] = defaultdict(list)
@property
@abstractmethod
def model_config(self) -> ModelConfig:
raise NotImplementedError
def _add_placeholder(self, modality: ModalityStr, placeholder: str | None):
mod_placeholder = MODALITY_PLACEHOLDERS_MAP[modality]
if placeholder:
@@ -806,6 +909,10 @@ class BaseMultiModalContentParser(ABC):
) -> None:
raise NotImplementedError
@abstractmethod
def parse_prompt_embeds(self, data: str) -> None:
raise NotImplementedError
@abstractmethod
def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:
raise NotImplementedError
@@ -834,6 +941,21 @@ class MultiModalContentParser(BaseMultiModalContentParser):
def model_config(self) -> ModelConfig:
return self._tracker.model_config
@override
def parse_prompt_embeds(self, data: str) -> None:
"""Decode a base64 prompt embeds tensor and store it in the tracker.
Emits a single `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` sentinel per
content part. The renderer later expands each sentinel to a span of
`tensor.shape[0]` placeholder tokens after tokenization.
"""
if not self.model_config.enable_prompt_embeds:
raise ValueError(_ENABLE_PROMPT_EMBEDS_ERROR)
tensor = safe_load_prompt_embeds(self.model_config, data.encode())
self._tracker.add("prompt_embeds", (tensor, None))
self._add_placeholder("prompt_embeds", PROMPT_EMBEDS_PLACEHOLDER_TOKEN)
def parse_image(self, image_url: str | None, uuid: str | None = None) -> None:
image = self._connector.fetch_image(image_url) if image_url else None
@@ -958,6 +1080,29 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):
def model_config(self) -> ModelConfig:
return self._tracker.model_config
@override
def parse_prompt_embeds(self, data: str) -> None:
"""Schedule async prompt embeds decode and store the coroutine in the tracker.
Like the sync variant, emits a single sentinel `PROMPT_EMBEDS_PLACEHOLDER_TOKEN`
per content part. Unlike the sync variant, the tensor decode is deferred to a
thread-pool executor via `safe_load_prompt_embeds_async`.
"""
if not self.model_config.enable_prompt_embeds:
raise ValueError(_ENABLE_PROMPT_EMBEDS_ERROR)
coro = self._load_prompt_embeds_async(data.encode())
self._tracker.add("prompt_embeds", coro)
self._add_placeholder("prompt_embeds", PROMPT_EMBEDS_PLACEHOLDER_TOKEN)
async def _load_prompt_embeds_async(
self, data_bytes: bytes
) -> tuple[torch.Tensor, None]:
# Second tuple slot fills the tracker's generic `(item, uuid | None)`
# contract. prompt_embeds has no UUID concept, so it's always `None`.
tensor = await safe_load_prompt_embeds_async(self.model_config, data_bytes)
return tensor, None
async def _image_with_uuid_async(self, image_url: str | None, uuid: str | None):
image = (
await self._connector.fetch_image_async(image_url) if image_url else None
@@ -1269,6 +1414,7 @@ def _get_full_multimodal_text_prompt(
_TextParser = partial(cast, ChatCompletionContentPartTextParam)
_ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
_AudioEmbedsParser = partial(cast, ChatCompletionContentPartAudioEmbedsParam)
_PromptEmbedsParser = partial(cast, ChatCompletionContentPartPromptEmbedsParam)
_InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
_RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
_PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
@@ -1294,6 +1440,7 @@ MM_PARSER_MAP: dict[
"image_url": lambda part: _ImageParser(part).get("image_url", {}).get("url", None),
"image_embeds": lambda part: _ImageEmbedsParser(part).get("image_embeds", None),
"audio_embeds": lambda part: _AudioEmbedsParser(part).get("audio_embeds", None),
"prompt_embeds": lambda part: _PromptEmbedsParser(part).get("data", None),
"image_pil": lambda part: _PILImageParser(part).get("image_pil", None),
"audio_url": lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),
"input_audio": lambda part: _InputAudioParser(part).get("input_audio", None),
@@ -1372,6 +1519,11 @@ def _parse_chat_message_content_mm_part(
)
audio_embeds = audio_params.get("audio_embeds", None)
return "audio_embeds", audio_embeds
if "prompt_embeds" in part:
prompt_embeds_params = cast( # type: ignore[assignment]
ChatCompletionContentPartPromptEmbedsParam, part
)
return "prompt_embeds", prompt_embeds_params.get("data", None)
if "audio_url" in part:
audio_params = cast( # type: ignore[assignment]
CustomChatCompletionContentSimpleAudioParam, part
@@ -1455,6 +1607,24 @@ def _parse_chat_message_content_parts(
return [ConversationMessage(role=role, content=text_prompt)]
def _reject_reserved_placeholder_in_text(text: str, model_config: ModelConfig) -> None:
"""Reject user-supplied text parts that contains the reserved `prompt_embeds`
placeholder sentinel.
When the server accepts `prompt_embeds`, the placeholder token is
registered as a single unsplittable special token on the tokenizer. Any
user text that happens to contain the literal sequence would tokenize to
the same ID and be mistaken for a splice point by the renderer, letting a
caller move or inject splice positions via plain text content.
"""
if model_config.enable_prompt_embeds and PROMPT_EMBEDS_PLACEHOLDER_TOKEN in text:
raise ValueError(
_RESERVED_PLACEHOLDER_IN_TEXT_ERROR.format(
token=PROMPT_EMBEDS_PLACEHOLDER_TOKEN
)
)
def _parse_chat_message_content_part(
part: ChatCompletionContentPartParam,
mm_parser: BaseMultiModalContentParser,
@@ -1470,6 +1640,7 @@ def _parse_chat_message_content_part(
with multimodal placeholders.
"""
if isinstance(part, str): # Handle plain text parts
_reject_reserved_placeholder_in_text(part, mm_parser.model_config)
if wrap_dicts:
return {"type": "text", "text": part}
return part
@@ -1488,6 +1659,7 @@ def _parse_chat_message_content_part(
if part_type in ("text", "input_text", "output_text", "refusal", "thinking"):
str_content = cast(str, content)
_reject_reserved_placeholder_in_text(str_content, mm_parser.model_config)
if wrap_dicts:
return {"type": "text", "text": str_content}
else:
@@ -1516,6 +1688,11 @@ def _parse_chat_message_content_part(
content = cast(str | dict[str, str], content) if content is not None else None
mm_parser.parse_audio_embeds(content, uuid)
modality = "audio"
elif part_type == "prompt_embeds":
if not content:
raise ValueError(_PROMPT_EMBEDS_MISSING_DATA_ERROR)
mm_parser.parse_prompt_embeds(cast(str, content))
modality = "prompt_embeds"
elif part_type == "audio_url":
str_content = cast(str, content)
mm_parser.parse_audio(str_content, uuid)
@@ -1544,7 +1721,18 @@ def _parse_chat_message_content_part(
)
if wrap_dicts:
if modality == "prompt_embeds":
# Chat templates don't know about the "prompt_embeds" modality,
# emit the single sentinel token as text so the template renders
# it inline. The renderer later expands it to N tokens post-tokenize.
return {"type": "text", "text": PROMPT_EMBEDS_PLACEHOLDER_TOKEN}
return {"type": modality}
if modality == "prompt_embeds":
# Emit the renderer token inline regardless of `interleave_strings`,
# prompt_embeds are spliced at the token offset so position matters.
# Falling back to front-padding via `missing_placeholders` would
# reorder them relative to surrounding text.
return PROMPT_EMBEDS_PLACEHOLDER_TOKEN
return MODALITY_PLACEHOLDERS_MAP[modality] if interleave_strings else None
@@ -1668,7 +1856,10 @@ def parse_chat_messages(
MultiModalUUIDDict | None,
]:
conversation: list[ConversationMessage] = []
mm_tracker = MultiModalItemTracker(model_config, media_io_kwargs=media_io_kwargs)
mm_tracker = MultiModalItemTracker(
model_config,
media_io_kwargs=media_io_kwargs,
)
for msg in messages:
sub_messages = _parse_chat_message_content(
@@ -1705,7 +1896,8 @@ async def parse_chat_messages_async(
]:
conversation: list[ConversationMessage] = []
mm_tracker = AsyncMultiModalItemTracker(
model_config, media_io_kwargs=media_io_kwargs
model_config,
media_io_kwargs=media_io_kwargs,
)
for msg in messages:
+4 -1
View File
@@ -926,7 +926,10 @@ class LLM:
add_generation_prompt=add_generation_prompt,
continue_final_message=continue_final_message,
tools=tools,
tokenize=is_mistral_tokenizer(renderer.tokenizer),
tokenize=(
is_mistral_tokenizer(renderer.tokenizer)
or self.model_config.enable_prompt_embeds
),
),
),
mm_processor_kwargs=mm_processor_kwargs,
+4 -1
View File
@@ -541,7 +541,10 @@ class OpenAIServingRender:
default_template_kwargs,
dict(
tools=tool_dicts,
tokenize=is_mistral_tokenizer(renderer.tokenizer),
tokenize=(
is_mistral_tokenizer(renderer.tokenizer)
or self.model_config.enable_prompt_embeds
),
),
)
+19
View File
@@ -71,12 +71,27 @@ class EmbedsInput(_InputOptions):
prompt: NotRequired[str]
"""The prompt text corresponding to the token IDs, if available."""
prompt_token_ids: NotRequired[list[int]]
"""Token IDs of the rendered prompt. Only set for mixed-mode inputs
(chat completion with `prompt_embeds` content parts). When present,
`is_token_ids` MUST also be present and have the same length.
For pure-embeds inputs this field is absent."""
is_token_ids: NotRequired[list[bool]]
"""Per-position mask for mixed-mode inputs. `True` means the position
is a real token ID (use the model's embedding layer); `False` means
the position uses a pre-computed embedding row from `prompt_embeds`.
Length MUST equal `len(prompt_token_ids)`.
For pure-embeds inputs this field is absent."""
def embeds_input(
prompt_embeds: "torch.Tensor",
*,
prompt: str | None = None,
cache_salt: str | None = None,
prompt_token_ids: list[int] | None = None,
is_token_ids: list[bool] | None = None,
) -> EmbedsInput:
"""
Construct [`EmbedsInput`][vllm.inputs.engine.EmbedsInput]
@@ -88,6 +103,10 @@ def embeds_input(
inputs["prompt"] = prompt
if cache_salt is not None:
inputs["cache_salt"] = cache_salt
if prompt_token_ids is not None:
inputs["prompt_token_ids"] = prompt_token_ids
if is_token_ids is not None:
inputs["is_token_ids"] = is_token_ids
return inputs
+11
View File
@@ -125,6 +125,17 @@ class EmbedsPrompt(_PromptOptions):
prompt: NotRequired[str]
"""The prompt text corresponding to the token embeddings, if available."""
prompt_token_ids: NotRequired[list[int]]
"""Token IDs for mixed-mode inputs (chat completion with
`prompt_embeds` content parts). The tokens at positions where
`prompt_is_token_ids` is `False` are placeholder tokens that
get replaced by entries from `prompt_embeds` in the forward pass."""
prompt_is_token_ids: NotRequired[list[bool]]
"""Per-position mask, `True` uses the real token ID, `False` uses
the corresponding entry from `prompt_embeds`.
Must be the same length as `prompt_token_ids` when both are set."""
DecoderOnlyPrompt: TypeAlias = (
str | TextPrompt | list[int] | TokensPrompt | EmbedsPrompt
+2
View File
@@ -769,6 +769,8 @@ class BaseRenderer(ABC, Generic[_T]):
return embeds_input(
prompt_embeds=prompt_embeds,
cache_salt=prompt.get("cache_salt"),
prompt_token_ids=prompt.get("prompt_token_ids"),
is_token_ids=prompt.get("prompt_is_token_ids"),
)
async def _process_tokens_async(
+45 -6
View File
@@ -7,6 +7,7 @@ import pybase64
import torch
from vllm.exceptions import VLLMValidationError
from vllm.utils.async_utils import make_async
if TYPE_CHECKING:
from vllm.config import ModelConfig
@@ -30,15 +31,53 @@ def safe_load_prompt_embeds(
weights_only=True,
map_location=torch.device("cpu"),
)
assert isinstance(tensor, torch.Tensor) and tensor.dtype in (
torch.float32,
torch.bfloat16,
torch.float16,
)
if not isinstance(tensor, torch.Tensor):
raise VLLMValidationError(
"`prompt_embeds` payload did not deserialize to a torch.Tensor.",
parameter="prompt_embeds",
)
tensor = tensor.to_dense()
if tensor.dim() > 2:
tensor = tensor.squeeze(0)
assert tensor.dim() == 2
if tensor.dim() != 2:
raise VLLMValidationError(
"`prompt_embeds` must be a 2D tensor of shape "
f"(num_tokens, hidden_size); got shape {tuple(tensor.shape)}.",
parameter="prompt_embeds",
)
# Pin each tensor to the model's hidden_size. Validating here
# also transitively guarantees cross-tensor consistency for requests that
# include multiple `prompt_embeds` parts, which is required by downstream
# concatenation in `_build_mixed_prompt_embeds`.
expected_hidden_size = model_config.get_hidden_size()
if tensor.shape[1] != expected_hidden_size:
raise VLLMValidationError(
f"`prompt_embeds` hidden_size {tensor.shape[1]} does not match "
f"the model's hidden_size {expected_hidden_size}.",
parameter="prompt_embeds",
)
# Cast to the model's dtype so API clients don't need to know the server's
# `--dtype` setting ahead of time. Only floating-point source dtypes are
# allowed. integer / bool / complex inputs almost certainly indicate caller
# error (e.g. quantized payloads, wrong tensor), and a silent `.to()`
# could hide a real mistake.
expected_dtype = model_config.dtype
if tensor.dtype != expected_dtype:
if not tensor.is_floating_point():
raise VLLMValidationError(
f"`prompt_embeds` dtype {tensor.dtype} is not a floating-point "
f"type, cannot safely cast to the model's dtype {expected_dtype}.",
parameter="prompt_embeds",
)
tensor = tensor.to(expected_dtype)
return tensor
safe_load_prompt_embeds_async = make_async(safe_load_prompt_embeds)
"""Async variant of `safe_load_prompt_embeds` that defers the decode to a
thread-pool executor, so the asyncio event loop is not blocked by the base64
decode + `torch.load` work."""
+440 -18
View File
@@ -1,11 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import inspect
import itertools
import weakref
from collections import defaultdict, deque
from collections.abc import Set
from collections.abc import Sequence
from functools import lru_cache
from typing import Any, Literal, cast, overload
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
import jinja2
import jinja2.ext
@@ -13,20 +16,32 @@ import jinja2.meta
import jinja2.nodes
import jinja2.parser
import jinja2.sandbox
import torch
from typing_extensions import override
from vllm.config import ModelConfig, VllmConfig
from vllm.entrypoints.chat_utils import (
ChatCompletionMessageParam,
ChatTemplateContentFormat,
ChatTemplateContentFormatOption,
PROMPT_EMBEDS_PLACEHOLDER_TOKEN,
ChatTemplateResolutionError,
ConversationMessage,
load_chat_template,
parse_chat_messages,
parse_chat_messages_async,
)
from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict
from vllm.inputs import EmbedsPrompt
from vllm.inputs.engine import MultiModalInput
from vllm.logger import init_logger
from vllm.multimodal.hasher import MultiModalHasher
from vllm.multimodal.inputs import (
MultiModalFieldElem,
MultiModalKwargsItem,
MultiModalKwargsItems,
MultiModalSharedField,
PlaceholderRange,
)
from vllm.multimodal.processing.processor import (
PromptReplacement,
apply_token_matches,
find_mm_placeholders,
)
from vllm.tokenizers.hf import HfTokenizer
from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
from vllm.transformers_utils.processor import cached_get_processor
@@ -34,13 +49,166 @@ from vllm.utils.async_utils import make_async
from vllm.utils.func_utils import supports_kw
from .base import BaseRenderer
from .inputs import DictPrompt
from .inputs.preprocess import parse_dec_only_prompt
from .params import ChatParams
if TYPE_CHECKING:
from collections.abc import Set
from vllm.config import ModelConfig, VllmConfig
from vllm.entrypoints.chat_utils import (
ChatCompletionMessageParam,
ChatTemplateContentFormat,
ChatTemplateContentFormatOption,
ConversationMessage,
)
from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict, TokensPrompt
from vllm.inputs.engine import TokensInput
from vllm.multimodal.processing.processor import (
MultiModalPromptUpdates,
ResolvedPromptUpdate,
)
from .inputs import DictPrompt
from .params import ChatParams
logger = init_logger(__name__)
# Cache of `tokenizer -> prompt_embeds placeholder token ID`. Keyed by the
# tokenizer object (not `id(tokenizer)`) so a fresh tokenizer landing at a
# recycled memory address can't pick up a stale tid. Entries evict atomically
# with the tokenizer's garbage-collection.
_PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE: Final[
weakref.WeakKeyDictionary[HfTokenizer, int]
] = weakref.WeakKeyDictionary()
_PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_ERROR: Final[str] = (
"Expected {token!r} to tokenize to exactly 1 token, got {num_ids} ({ids!r})."
)
_PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR: Final[str] = (
"Expected {expected} prompt_embeds placeholder spans in the "
"tokenized prompt, found {actual}."
)
_MISSING_PROMPT_TOKEN_IDS_ERROR: Final[str] = (
"Expected prompt_token_ids in rendered prompt when prompt_embeds "
"are present. This indicates the chat template was invoked with "
"tokenize=False."
)
_TOKENIZE_OVERRIDE_WARNING: Final[str] = (
"Overriding `tokenize=False` to `True` because `prompt_embeds` "
"post-processing requires tokenized IDs."
)
def _ensure_prompt_embeds_placeholder_token(tokenizer: HfTokenizer) -> int:
"""Register `PROMPT_EMBEDS_PLACEHOLDER_TOKEN` as a special token and return
its token ID."""
cached = _PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE.get(tokenizer)
if cached is not None:
return cached
tokenizer.add_special_tokens(
{"additional_special_tokens": [PROMPT_EMBEDS_PLACEHOLDER_TOKEN]}
)
ids = tokenizer.encode(PROMPT_EMBEDS_PLACEHOLDER_TOKEN, add_special_tokens=False)
if len(ids) != 1:
raise RuntimeError(
_PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_ERROR.format(
token=PROMPT_EMBEDS_PLACEHOLDER_TOKEN,
num_ids=len(ids),
ids=ids,
)
)
token_id = ids[0]
_PROMPT_EMBEDS_PLACEHOLDER_TOKEN_ID_CACHE[tokenizer] = token_id
return token_id
def _build_prompt_embeds_updates(
prompt_embeds_tensors: Sequence[torch.Tensor],
placeholder_token_id: int,
) -> MultiModalPromptUpdates:
"""Build `MultiModalPromptUpdates` for `prompt_embeds` expansion.
Each tensor produces a `PromptReplacement` that maps
`[placeholder_token_id]` -> `[placeholder_token_id] x N`
(where `N = tensor.shape[0]`).
"""
updates: list[Sequence[ResolvedPromptUpdate]] = []
for i, tensor in enumerate(prompt_embeds_tensors):
update = PromptReplacement(
modality="prompt_embeds",
target=[placeholder_token_id],
replacement=[placeholder_token_id] * tensor.shape[0],
)
updates.append([update.resolve(item_idx=i)])
return {"prompt_embeds": updates}
def _expand_prompt_embeds_placeholders(
token_ids: list[int],
mm_prompt_updates: MultiModalPromptUpdates,
) -> list[int]:
"""Expand each 1-token `prompt_embeds` sentinel into an N-token span.
Uses `apply_token_matches`. Each single placeholder token in
`token_ids` is replaced with a consecutive span of
`tensor.shape[0]` copies, following tensors in order.
"""
expanded, _ = apply_token_matches(token_ids, mm_prompt_updates, tokenizer=None)
return expanded
def _build_prompt_embeds_positions(
token_ids: list[int],
num_tensors: int,
mm_prompt_updates: MultiModalPromptUpdates,
) -> list[tuple[int, int]]:
"""Locate each prompt_embeds placeholder span in `token_ids`.
Expects `token_ids` to already contain expanded N-token spans.
Returns `[(start_idx, length), ...]` aligned with the tensors.
"""
placeholders = find_mm_placeholders(
prompt=token_ids,
mm_prompt_updates=mm_prompt_updates,
tokenizer=None,
)
features = placeholders.get("prompt_embeds", [])
if len(features) != num_tensors:
raise ValueError(
_PROMPT_EMBEDS_PLACEHOLDER_SPAN_MISMATCH_ERROR.format(
expected=num_tensors,
actual=len(features),
)
)
return [(f.start_idx, f.length) for f in features]
def _build_mixed_prompt_embeds(
token_ids: list[int],
prompt_embeds_tensors: Sequence[torch.Tensor],
positions: list[tuple[int, int]],
) -> tuple[torch.Tensor, list[bool]]:
"""Build the full-length `prompt_embeds` tensor and the `is_token_ids`
mask aligned to `token_ids`."""
total_len = len(token_ids)
hidden_size = prompt_embeds_tensors[0].shape[1]
dtype = prompt_embeds_tensors[0].dtype
full_embeds = torch.zeros(total_len, hidden_size, dtype=dtype)
is_token_ids = torch.ones(total_len, dtype=torch.bool)
for (start, length), tensor in zip(positions, prompt_embeds_tensors, strict=True):
full_embeds[start : start + length] = tensor
is_token_ids[start : start + length] = False
return full_embeds, is_token_ids.tolist()
_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]()
"""
Used in `_try_get_processor_chat_template` to avoid calling
@@ -98,7 +266,7 @@ def resolve_chat_template(
chat_template: str | None,
tools: list[dict[str, Any]] | None,
*,
model_config: "ModelConfig",
model_config: ModelConfig,
) -> str | None:
# 1st priority: The given chat template
if chat_template is not None:
@@ -281,7 +449,7 @@ def _resolve_chat_template_content_format(
tools: list[dict[str, Any]] | None,
tokenizer: HfTokenizer,
*,
model_config: "ModelConfig",
model_config: ModelConfig,
) -> ChatTemplateContentFormat:
resolved_chat_template = resolve_chat_template(
tokenizer,
@@ -335,7 +503,7 @@ def resolve_chat_template_content_format(
given_format: ChatTemplateContentFormatOption,
tokenizer: HfTokenizer,
*,
model_config: "ModelConfig",
model_config: ModelConfig,
) -> ChatTemplateContentFormat:
if given_format != "auto":
return given_format
@@ -437,7 +605,7 @@ def resolve_chat_template_kwargs(
@overload
def safe_apply_chat_template(
model_config: "ModelConfig",
model_config: ModelConfig,
tokenizer: HfTokenizer,
conversation: list[ConversationMessage],
*,
@@ -448,7 +616,7 @@ def safe_apply_chat_template(
) -> list[int]: ...
@overload
def safe_apply_chat_template(
model_config: "ModelConfig",
model_config: ModelConfig,
tokenizer: HfTokenizer,
conversation: list[ConversationMessage],
*,
@@ -458,7 +626,7 @@ def safe_apply_chat_template(
**kwargs,
) -> str: ...
def safe_apply_chat_template(
model_config: "ModelConfig",
model_config: ModelConfig,
tokenizer: HfTokenizer,
conversation: list[ConversationMessage],
*,
@@ -486,6 +654,14 @@ def safe_apply_chat_template(
chat_template_kwargs=kwargs,
)
# transformers v5 changed the default of `return_dict` to True, which
# makes `apply_chat_template(tokenize=True)` return a `BatchEncoding`
# instead of `list[int]`. Force `return_dict=False` so downstream code
# that expects a flat token list (e.g. `parse_dec_only_prompt`) works
# consistently across v4 and v5.
if tokenize and "return_dict" not in resolved_kwargs:
resolved_kwargs["return_dict"] = False
try:
return tokenizer.apply_chat_template(
conversation=conversation, # type: ignore[arg-type]
@@ -627,6 +803,12 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
model_config = self.model_config
tokenizer = self.get_tokenizer()
prompt_embeds_placeholder_token_id: int | None = None
if model_config.enable_prompt_embeds:
prompt_embeds_placeholder_token_id = (
_ensure_prompt_embeds_placeholder_token(tokenizer)
)
conversation, mm_data, mm_uuids = parse_chat_messages(
messages,
model_config,
@@ -641,11 +823,30 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
mm_processor_kwargs=params.mm_processor_kwargs,
)
# prompt_embeds tensors are carried by the tracker through mm_data,
# but they must NOT be fed to the MM processor (which would reject
# the unknown key). Extract them here.
prompt_embeds_tensors: list[torch.Tensor] | None = None
if mm_data is not None and "prompt_embeds" in mm_data:
prompt_embeds_tensors = list(
cast(Sequence[torch.Tensor], mm_data["prompt_embeds"])
)
mm_data = {k: v for k, v in mm_data.items() if k != "prompt_embeds"}
if not mm_data:
mm_data = None
chat_template_kwargs = params.get_apply_chat_template_kwargs()
if prompt_embeds_tensors:
# prompt_embeds post-processing requires prompt_token_ids.
if chat_template_kwargs.get("tokenize") is False:
logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
chat_template_kwargs["tokenize"] = True
prompt_raw = safe_apply_chat_template(
model_config,
tokenizer,
conversation,
**params.get_apply_chat_template_kwargs(),
**chat_template_kwargs,
)
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
@@ -671,6 +872,29 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
)
prompt = parse_dec_only_prompt(prompt_raw)
# When `prompt_embeds` is mixed with other modality data,
# `_process_tokens` runs `_process_multimodal` first (expanding
# `<|AUDIO|>` / `<|IMAGE|>` placeholders) and then
# `_apply_prompt_embeds_to_engine_input` augments the result.
# Stash the tensors and placeholder ID for that override to consume.
if prompt_embeds_tensors and mm_data:
assert prompt_embeds_placeholder_token_id is not None
cast(dict, prompt)["_prompt_embeds"] = (
prompt_embeds_tensors,
prompt_embeds_placeholder_token_id,
)
if params.mm_processor_kwargs:
cast(dict, prompt)["mm_processor_kwargs"] = params.mm_processor_kwargs
elif prompt_embeds_tensors:
# Pure mode: no other MM data, mutate prompt to EmbedsPrompt shape.
assert prompt_embeds_placeholder_token_id is not None
self._apply_prompt_embeds_to_prompt(
prompt,
prompt_embeds_tensors,
prompt_embeds_placeholder_token_id,
)
if mm_data is not None:
prompt["multi_modal_data"] = mm_data
if mm_uuids is not None:
@@ -686,6 +910,12 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
model_config = self.model_config
tokenizer = self.get_tokenizer()
prompt_embeds_placeholder_token_id: int | None = None
if model_config.enable_prompt_embeds:
prompt_embeds_placeholder_token_id = (
_ensure_prompt_embeds_placeholder_token(tokenizer)
)
conversation, mm_data, mm_uuids = await parse_chat_messages_async(
messages,
model_config,
@@ -700,11 +930,27 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
mm_processor_kwargs=params.mm_processor_kwargs,
)
prompt_embeds_tensors: list[torch.Tensor] | None = None
if mm_data is not None and "prompt_embeds" in mm_data:
prompt_embeds_tensors = list(
cast(Sequence[torch.Tensor], mm_data["prompt_embeds"])
)
mm_data = {k: v for k, v in mm_data.items() if k != "prompt_embeds"}
if not mm_data:
mm_data = None
chat_template_kwargs = params.get_apply_chat_template_kwargs()
if prompt_embeds_tensors:
# prompt_embeds post-processing requires prompt_token_ids.
if chat_template_kwargs.get("tokenize") is False:
logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
chat_template_kwargs["tokenize"] = True
prompt_raw = await self._apply_chat_template_async(
model_config,
tokenizer,
conversation,
**params.get_apply_chat_template_kwargs(),
**chat_template_kwargs,
)
# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
@@ -728,9 +974,185 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
)
prompt = parse_dec_only_prompt(prompt_raw)
# See `render_messages` for the rationale.
if prompt_embeds_tensors and mm_data:
assert prompt_embeds_placeholder_token_id is not None
cast(dict, prompt)["_prompt_embeds"] = (
prompt_embeds_tensors,
prompt_embeds_placeholder_token_id,
)
if params.mm_processor_kwargs:
cast(dict, prompt)["mm_processor_kwargs"] = params.mm_processor_kwargs
elif prompt_embeds_tensors:
assert prompt_embeds_placeholder_token_id is not None
self._apply_prompt_embeds_to_prompt(
prompt,
prompt_embeds_tensors,
prompt_embeds_placeholder_token_id,
)
if mm_data is not None:
prompt["multi_modal_data"] = mm_data
if mm_uuids is not None:
prompt["multi_modal_uuids"] = mm_uuids
return conversation, prompt
@override
def _process_tokens(
self,
prompt: TokensPrompt,
*,
skip_mm_cache: bool = False,
) -> TokensInput | MultiModalInput:
"""Pre-expand `prompt_embeds` sentinels before delegating to the MM
processor, then attach `prompt_embeds` modality data to the result.
Mixed mode only: the `_prompt_embeds` stash is set by
`render_messages` when `prompt_embeds` co-exist with other MM data
(images, audio, ). We expand each 1-token sentinel to an N-token
span *before* calling `super()._process_tokens()` so the MM
processor records all placeholder offsets in the final (post-expansion)
coordinate space, no offset shifting needed afterwards.
"""
prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
if prompt_embeds_info is not None:
tensors, placeholder_token_id = prompt_embeds_info
mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
list(prompt["prompt_token_ids"]), mm_updates
)
engine_input = super()._process_tokens(prompt, skip_mm_cache=skip_mm_cache)
if prompt_embeds_info is not None:
tensors, _ = prompt_embeds_info
self._apply_prompt_embeds_to_engine_input(
cast(MultiModalInput, engine_input),
tensors,
mm_updates,
)
return engine_input
@override
async def _process_tokens_async(
self,
prompt: TokensPrompt,
*,
skip_mm_cache: bool = False,
) -> TokensInput | MultiModalInput:
"""Async equivalent of `_process_tokens`."""
prompt_embeds_info = cast(dict, prompt).pop("_prompt_embeds", None)
if prompt_embeds_info is not None:
tensors, placeholder_token_id = prompt_embeds_info
mm_updates = _build_prompt_embeds_updates(tensors, placeholder_token_id)
cast(dict, prompt)["prompt_token_ids"] = _expand_prompt_embeds_placeholders(
list(prompt["prompt_token_ids"]), mm_updates
)
engine_input = await super()._process_tokens_async(
prompt, skip_mm_cache=skip_mm_cache
)
if prompt_embeds_info is not None:
tensors, _ = prompt_embeds_info
self._apply_prompt_embeds_to_engine_input(
cast(MultiModalInput, engine_input),
tensors,
mm_updates,
)
return engine_input
@staticmethod
def _apply_prompt_embeds_to_prompt(
prompt: DictPrompt,
prompt_embeds_tensors: list[torch.Tensor],
placeholder_token_id: int,
) -> None:
"""Mutate `prompt` from `TokensPrompt` to `EmbedsPrompt` shape.
Pure `prompt_embeds` path only (no other MM modalities). Expands
each `<prompt_embeds>` sentinel token into an N-token span and builds
the full-length `prompt_embeds` tensor + `prompt_is_token_ids` mask
that the engine's `enable_prompt_embeds` worker branch consumes.
"""
token_ids = cast(list[int] | None, prompt.get("prompt_token_ids"))
if token_ids is None:
raise RuntimeError(_MISSING_PROMPT_TOKEN_IDS_ERROR)
embeds_orig_positions: list[int] = [
i for i, tok in enumerate(token_ids) if tok == placeholder_token_id
]
if len(embeds_orig_positions) != len(prompt_embeds_tensors):
raise ValueError(
f"Expected {len(prompt_embeds_tensors)} prompt_embeds "
f"placeholder tokens in the rendered prompt, found "
f"{len(embeds_orig_positions)}."
)
mm_updates = _build_prompt_embeds_updates(
prompt_embeds_tensors, placeholder_token_id
)
expanded = _expand_prompt_embeds_placeholders(token_ids, mm_updates)
positions = _build_prompt_embeds_positions(
expanded, len(prompt_embeds_tensors), mm_updates
)
embeds_prompt = cast(EmbedsPrompt, prompt)
embeds_prompt["prompt_token_ids"] = expanded
full_embeds, is_token_ids_mask = _build_mixed_prompt_embeds(
expanded, prompt_embeds_tensors, positions
)
embeds_prompt["prompt_embeds"] = full_embeds
embeds_prompt["prompt_is_token_ids"] = is_token_ids_mask
@staticmethod
def _apply_prompt_embeds_to_engine_input(
engine_input: MultiModalInput,
prompt_embeds_tensors: list[torch.Tensor],
mm_updates: MultiModalPromptUpdates,
) -> None:
"""Augment `engine_input` in-place with a `prompt_embeds` modality.
Mixed mode: called after `_process_multimodal` has already run on the
pre-expanded token IDs (expansion was done in `_process_tokens` before
calling `super()`). Locates the already-expanded `prompt_embeds` spans
and adds `prompt_embeds` entries to `mm_kwargs`, `mm_hashes`, and
`mm_placeholders`.
"""
# token_ids already contain the pre-expanded N-token spans.
token_ids = list(engine_input["prompt_token_ids"])
positions = _build_prompt_embeds_positions(
token_ids, len(prompt_embeds_tensors), mm_updates
)
pe_kwargs_items: list[MultiModalKwargsItem] = []
pe_hashes: list[str] = []
pe_placeholders: list[PlaceholderRange] = []
for tensor, (start, length) in zip(
prompt_embeds_tensors, positions, strict=True
):
pe_kwargs_items.append(
MultiModalKwargsItem(
{
"embedding": MultiModalFieldElem(
data=tensor,
field=MultiModalSharedField(batch_size=1),
)
}
)
)
pe_hashes.append(MultiModalHasher.hash_kwargs(prompt_embeds=tensor))
# `is_embed=None` matches the existing image_embeds-style
# "no encoder, just splice the tensor directly" semantics.
pe_placeholders.append(
PlaceholderRange(offset=start, length=length, is_embed=None)
)
cast(
MultiModalKwargsItems[MultiModalKwargsItem | None],
engine_input["mm_kwargs"],
)["prompt_embeds"] = pe_kwargs_items
engine_input["mm_hashes"] = {
**engine_input["mm_hashes"],
"prompt_embeds": pe_hashes,
}
cast(dict, engine_input["mm_placeholders"])["prompt_embeds"] = pe_placeholders
+2
View File
@@ -38,6 +38,7 @@ class NewRequestData:
num_computed_tokens: int
lora_request: LoRARequest | None
prompt_embeds: "torch.Tensor | None" = None
prompt_is_token_ids: list[bool] | None = None
# Only used for v2 model runner.
prefill_token_ids: list[int] | None = None
@@ -59,6 +60,7 @@ class NewRequestData:
num_computed_tokens=request.num_computed_tokens,
lora_request=request.lora_request,
prompt_embeds=request.prompt_embeds,
prompt_is_token_ids=request.prompt_is_token_ids,
prefill_token_ids=prefill_token_ids,
)
+6
View File
@@ -94,6 +94,12 @@ class EngineCoreRequest(
data_parallel_rank: int | None
prompt_embeds: torch.Tensor | None = None
# Per-position mask for mixed-mode inputs (e.g chat completion with
# prompt_embeds content parts). `True` means the position is a real
# token ID; `False` means the position uses a pre-computed entry from
# `prompt_embeds`. `None` for pure-tokens and pure-embeds requests.
prompt_is_token_ids: list[bool] | None = None
# Index of the client, used to ensure outputs are sent back to the same
# client for this request when scaling out the front-end.
client_index: int = 0
+4 -1
View File
@@ -292,11 +292,13 @@ class InputProcessor:
# Mypy can be conservative for TypedDict unions; normalize access.
if decoder_inputs["type"] == "embeds":
prompt_token_ids = None
prompt_embeds = decoder_inputs["prompt_embeds"]
prompt_token_ids = decoder_inputs.get("prompt_token_ids")
prompt_is_token_ids = decoder_inputs.get("is_token_ids")
else:
prompt_token_ids = decoder_inputs["prompt_token_ids"]
prompt_embeds = None
prompt_is_token_ids = None
sampling_params = None
pooling_params = None
@@ -361,6 +363,7 @@ class InputProcessor:
request_id=request_id,
prompt_token_ids=prompt_token_ids,
prompt_embeds=prompt_embeds,
prompt_is_token_ids=prompt_is_token_ids,
mm_features=mm_features,
sampling_params=sampling_params,
pooling_params=pooling_params,
+6
View File
@@ -66,6 +66,7 @@ class Request:
client_index: int = 0,
arrival_time: float | None = None,
prompt_embeds: torch.Tensor | None = None,
prompt_is_token_ids: list[bool] | None = None,
mm_features: list[MultiModalFeatureSpec] | None = None,
lora_request: "LoRARequest | None" = None,
cache_salt: str | None = None,
@@ -114,6 +115,10 @@ class Request:
self.prompt_token_ids = prompt_token_ids
self.prompt_embeds = prompt_embeds
# Per-position mask used in mixed-mode (chat completion with
# prompt_embeds). `None` except when both `prompt_token_ids` and
# `prompt_embeds` are set and their positions are interleaved.
self.prompt_is_token_ids = prompt_is_token_ids
# Cache per-block prompt-embed hashes to avoid rehashing the same
# tensor slices when generating extra keys.
self._prompt_embeds_per_block_hashes: dict[tuple[int, int], bytes] = {}
@@ -184,6 +189,7 @@ class Request:
client_index=request.client_index,
prompt_token_ids=request.prompt_token_ids,
prompt_embeds=request.prompt_embeds,
prompt_is_token_ids=request.prompt_is_token_ids,
mm_features=request.mm_features,
sampling_params=request.sampling_params,
pooling_params=request.pooling_params,
+10 -1
View File
@@ -50,6 +50,10 @@ class CachedRequestState:
lora_request: LoRARequest | None = None
prompt_embeds: torch.Tensor | None = None
# Per-position mask for mixed-mode inputs (e.g chat completion with
# prompt_embeds content parts). See `Request.prompt_is_token_ids`.
prompt_is_token_ids: list[bool] | None = None
# Used when both async_scheduling and spec_decode are enabled.
prev_num_draft_len: int = 0
@@ -356,7 +360,12 @@ class InputBatch:
end_idx = start_idx + len(request.output_token_ids)
if request.prompt_token_ids is not None:
self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids
self.is_token_ids[req_index, :num_prompt_tokens] = True
if request.prompt_is_token_ids is not None:
self.is_token_ids[req_index, :num_prompt_tokens] = (
request.prompt_is_token_ids
)
else:
self.is_token_ids[req_index, :num_prompt_tokens] = True
else:
self.is_token_ids[req_index, :num_prompt_tokens] = False
if request.prompt_embeds is not None:
+35 -1
View File
@@ -1163,6 +1163,7 @@ class GPUModelRunner(
req_id=req_id,
prompt_token_ids=new_req_data.prompt_token_ids,
prompt_embeds=new_req_data.prompt_embeds,
prompt_is_token_ids=new_req_data.prompt_is_token_ids,
mm_features=new_req_data.mm_features,
sampling_params=sampling_params,
pooling_params=pooling_params,
@@ -1505,10 +1506,16 @@ class GPUModelRunner(
)
mrope_model = cast(SupportsMRoPE, model)
# `prompt_embeds` is a passthrough modality (no grid_thw), models'
# M-RoPE code assumes per-feature grid info, so filter it out. The
# prompt_embeds positions are treated as text positions for M-RoPE.
mrope_features = [
f for f in req_state.mm_features if f.modality != "prompt_embeds"
]
req_state.mrope_positions, req_state.mrope_position_delta = (
mrope_model.get_mrope_input_positions(
req_state.prompt_token_ids,
req_state.mm_features,
mrope_features,
)
)
@@ -2744,6 +2751,33 @@ class GPUModelRunner(
if not mm_kwargs:
return []
# `prompt_embeds` is a passthrough modality, the tensor is already in
# the model embedding space, so no encoder runs. Inject each
# `prompt_embeds` tensor directly into the encoder cache here so that
# `_gather_mm_embeddings` can splice it via the standard `is_mm_embed`
# path.
pe_indices = [
i
for i, (modality, _) in enumerate(mm_kwargs)
if modality == "prompt_embeds"
]
if pe_indices:
for i in pe_indices:
pe_tensor = mm_kwargs[i][1]["embedding"].data
assert isinstance(pe_tensor, torch.Tensor)
self.encoder_cache[mm_hashes[i]] = pe_tensor.to(self.device)
self.maybe_save_ec_to_connector(self.encoder_cache, mm_hashes[i])
# Filter out `prompt_embeds` items from mm_kwargs/mm_hashes/mm_lora_refs
# since they don't require further encoder processing.
mm_hashes = [h for i, h in enumerate(mm_hashes) if i not in pe_indices]
mm_kwargs = [k for i, k in enumerate(mm_kwargs) if i not in pe_indices]
mm_lora_refs = [
r for i, r in enumerate(mm_lora_refs) if i not in pe_indices
]
if not mm_kwargs:
return [] # nothing left to encode after filtering out `prompt_embeds`
should_time = bool(
self.observability_config
and self.observability_config.enable_mm_processor_stats