mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-20 20:50:15 +00:00
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:
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
+190
@@ -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
|
||||
+212
@@ -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)
|
||||
@@ -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)."""
|
||||
|
||||
@@ -40,6 +40,7 @@ def _model_config():
|
||||
multimodal_config=None,
|
||||
allowed_local_media_path="",
|
||||
allowed_media_domains=None,
|
||||
enable_prompt_embeds=False,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user