@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2c764c089a
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Added real /v1/embeddings support for messages + chat_template_kw (#45173)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

Taneem Ibrahim committed Jun 14, 2026 at 20:08 UTC 2c764c089ae7ea2132d0c530b22a8f85fbb90af6
5 files changed +400 -10
tests/entrypoints/pooling/embed/test_io_processor.py
+201
@@ -3,6 +3,7 @@
3 """Unit tests for EmbedIOProcessor."""
4
5 import pytest
6 +from pydantic import TypeAdapter
7
8 from vllm import PoolingParams
9 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
@@ -10,10 +11,100 @@ from vllm.entrypoints.pooling.embed.protocol import (
11 CohereEmbedContent,
12 CohereEmbedInput,
13 CohereEmbedRequest,
14 + EmbeddingBatchChatInputRequest,
15 + EmbeddingBatchChatRequest,
16 + EmbeddingChatInputRequest,
17 + EmbeddingChatRequest,
18 + EmbeddingCompletionRequest,
19 + EmbeddingRequest,
20 )
21 from vllm.entrypoints.pooling.typing import PoolingServeContext
22
23
24 +class TestEmbeddingRequestParsing:
25 + """Unit tests for OpenAI embedding request parsing."""
26 +
27 + def test_input_messages_parses_as_chat_request(self):
28 + request = TypeAdapter(EmbeddingRequest).validate_python(
29 + {
30 + "model": "test",
31 + "input": [{"role": "user", "content": "hello"}],
32 + "chat_template_kwargs": {"instruction": "Represent the query: "},
33 + }
34 + )
35 +
36 + assert isinstance(request, EmbeddingChatInputRequest)
37 + assert request.input == [{"role": "user", "content": "hello"}]
38 + assert request.messages == [{"role": "user", "content": "hello"}]
39 + assert request.chat_template_kwargs == {"instruction": "Represent the query: "}
40 +
41 + def test_batched_input_messages_parses_as_batch_chat_input_request(self):
42 + request = TypeAdapter(EmbeddingRequest).validate_python(
43 + {
44 + "model": "test",
45 + "input": [
46 + [{"role": "user", "content": "hello"}],
47 + [{"role": "user", "content": "goodbye"}],
48 + ],
49 + "chat_template_kwargs": {"instruction": "Represent the query: "},
50 + }
51 + )
52 +
53 + assert isinstance(request, EmbeddingBatchChatInputRequest)
54 + assert request.input == [
55 + [{"role": "user", "content": "hello"}],
56 + [{"role": "user", "content": "goodbye"}],
57 + ]
58 + assert request.messages == [
59 + [{"role": "user", "content": "hello"}],
60 + [{"role": "user", "content": "goodbye"}],
61 + ]
62 + assert request.chat_template_kwargs == {"instruction": "Represent the query: "}
63 +
64 + def test_token_ids_still_parse_as_completion_request(self):
65 + request = TypeAdapter(EmbeddingRequest).validate_python(
66 + {
67 + "model": "test",
68 + "input": [[1, 2, 3], [4, 5]],
69 + }
70 + )
71 +
72 + assert isinstance(request, EmbeddingCompletionRequest)
73 + assert request.input == [[1, 2, 3], [4, 5]]
74 +
75 + def test_messages_still_parses_as_chat_request(self):
76 + request = TypeAdapter(EmbeddingRequest).validate_python(
77 + {
78 + "model": "test",
79 + "messages": [{"role": "user", "content": "hello"}],
80 + "chat_template_kwargs": {"instruction": "Represent the query: "},
81 + }
82 + )
83 +
84 + assert isinstance(request, EmbeddingChatRequest)
85 + assert request.messages == [{"role": "user", "content": "hello"}]
86 + assert request.chat_template_kwargs == {"instruction": "Represent the query: "}
87 +
88 + def test_batched_messages_parses_as_batch_chat_request(self):
89 + request = TypeAdapter(EmbeddingRequest).validate_python(
90 + {
91 + "model": "test",
92 + "messages": [
93 + [{"role": "user", "content": "hello"}],
94 + [{"role": "user", "content": "goodbye"}],
95 + ],
96 + "chat_template_kwargs": {"instruction": "Represent the query: "},
97 + }
98 + )
99 +
100 + assert isinstance(request, EmbeddingBatchChatRequest)
101 + assert request.messages == [
102 + [{"role": "user", "content": "hello"}],
103 + [{"role": "user", "content": "goodbye"}],
104 + ]
105 + assert request.chat_template_kwargs == {"instruction": "Represent the query: "}
106 +
107 +
108 class TestResolveTruncation:
109 """Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
110
@@ -324,3 +415,113 @@ class TestPreProcessCohereOnline:
415 },
416 )
417 ]
418 +
419 +
420 +class TestPreProcessOpenAIEmbeddingChatOnline:
421 + """Unit tests for OpenAI embedding chat preprocessing."""
422 +
423 + class _FakeModelConfig:
424 + max_model_len = 128
425 + encoder_config: dict[str, object] = {}
426 + pooler_config = None
427 + multimodal_config = None
428 + is_encoder_decoder = False
429 +
430 + class _FakeRenderer:
431 + tokenizer = object()
432 +
433 + def __init__(self):
434 + self.calls = []
435 +
436 + def render_chat(
437 + self,
438 + all_messages,
439 + chat_params,
440 + tok_params,
441 + prompt_extras=None,
442 + ):
443 + self.calls.append(
444 + {
445 + "all_messages": all_messages,
446 + "chat_params": chat_params,
447 + "tok_params": tok_params,
448 + "prompt_extras": prompt_extras,
449 + }
450 + )
451 + return all_messages, [
452 + {"prompt_token_ids": [index]} for index, _ in enumerate(all_messages)
453 + ]
454 +
455 + @classmethod
456 + def _make_handler(cls, renderer):
457 + handler = object.__new__(EmbedIOProcessor)
458 + handler.renderer = renderer
459 + handler.model_config = cls._FakeModelConfig()
460 + handler.chat_template = "template"
461 + handler.chat_template_content_format = "auto"
462 + handler.trust_request_chat_template = False
463 + handler.enable_chunked_processing = False
464 + return handler
465 +
466 + @staticmethod
467 + def _make_context(
468 + request: (
469 + EmbeddingChatRequest
470 + | EmbeddingBatchChatRequest
471 + | EmbeddingChatInputRequest
472 + | EmbeddingBatchChatInputRequest
473 + ),
474 + ) -> PoolingServeContext[
475 + EmbeddingChatRequest
476 + | EmbeddingBatchChatRequest
477 + | EmbeddingChatInputRequest
478 + | EmbeddingBatchChatInputRequest
479 + ]:
480 + return PoolingServeContext(
481 + request=request,
482 + pooling_params=PoolingParams(),
483 + model_name="test",
484 + request_id="embd-test",
485 + )
486 +
487 + def test_chat_template_kwargs_forwarded_for_batched_input_messages(self):
488 + request = TypeAdapter(EmbeddingRequest).validate_python(
489 + {
490 + "model": "test",
491 + "input": [
492 + [{"role": "user", "content": "hello"}],
493 + [{"role": "user", "content": "goodbye"}],
494 + ],
495 + "add_generation_prompt": True,
496 + "chat_template_kwargs": {"instruction": "Represent the query: "},
497 + "mm_processor_kwargs": {"max_pixels": 1},
498 + "cache_salt": "salt",
499 + }
500 + )
501 + assert isinstance(request, EmbeddingBatchChatInputRequest)
502 +
503 + renderer = self._FakeRenderer()
504 + handler = self._make_handler(renderer)
505 + ctx = self._make_context(request)
506 +
507 + handler.pre_process_online(ctx)
508 +
509 + assert ctx.engine_inputs == [
510 + {"prompt_token_ids": [0]},
511 + {"prompt_token_ids": [1]},
512 + ]
513 + assert len(renderer.calls) == 1
514 +
515 + call = renderer.calls[0]
516 + assert call["all_messages"] == request.messages
517 + assert call["prompt_extras"] == {
518 + "mm_processor_kwargs": {"max_pixels": 1},
519 + "cache_salt": "salt",
520 + }
521 +
522 + chat_template_kwargs = call["chat_params"].chat_template_kwargs
523 + assert chat_template_kwargs["instruction"] == "Represent the query: "
524 + assert chat_template_kwargs["add_generation_prompt"] is True
525 + assert chat_template_kwargs["continue_final_message"] is False
526 + assert "tools" not in chat_template_kwargs
527 + assert chat_template_kwargs["tokenize"] is False
vllm/entrypoints/pooling/base/protocol.py
+7 -5
@@ -168,11 +168,7 @@ class CompletionRequestMixin(OpenAIBaseModel):
168 # --8<-- [end:completion-extra-params]
169
170
171 -class ChatRequestMixin(OpenAIBaseModel):
172 - # --8<-- [start:chat-params]
173 - messages: list[ChatCompletionMessageParam]
174 - # --8<-- [end:chat-params]
175 -
171 +class ChatRequestOptionsMixin(OpenAIBaseModel):
172 # --8<-- [start:chat-extra-params]
173 add_generation_prompt: bool = Field(
174 default=False,
@@ -256,6 +252,12 @@ class ChatRequestMixin(OpenAIBaseModel):
252 )
253
254
255 +class ChatRequestMixin(ChatRequestOptionsMixin):
256 + # --8<-- [start:chat-params]
257 + messages: list[ChatCompletionMessageParam]
258 + # --8<-- [end:chat-params]
259 +
260 +
261 class EncodingRequestMixin(OpenAIBaseModel):
262 # --8<-- [start:encoding-params]
263 encoding_format: EncodingFormat = "float"
vllm/entrypoints/pooling/embed/io_processor.py
+77
@@ -36,6 +36,9 @@ from .protocol import (
36 CohereEmbedContent,
37 CohereEmbedInput,
38 CohereEmbedRequest,
39 + EmbeddingBatchChatInputRequest,
40 + EmbeddingBatchChatRequest,
41 + EmbeddingChatInputRequest,
42 EmbeddingChatRequest,
43 EmbeddingCompletionRequest,
44 )
@@ -66,6 +69,16 @@ class EmbedIOProcessor(PoolingIOProcessor):
69 def pre_process_online(self, ctx: PoolingServeContext):
70 if isinstance(ctx.request, CohereEmbedRequest):
71 self._pre_process_cohere_online(ctx)
72 + elif isinstance(
73 + ctx.request,
74 + (
75 + EmbeddingChatRequest,
76 + EmbeddingBatchChatRequest,
77 + EmbeddingChatInputRequest,
78 + EmbeddingBatchChatInputRequest,
79 + ),
80 + ):
81 + self._pre_process_openai_chat_online(ctx)
82 else:
83 super().pre_process_online(ctx)
84
@@ -367,6 +380,70 @@ class EmbedIOProcessor(PoolingIOProcessor):
380 )
381 return super().create_pooling_params(request)
382
383 + def _pre_process_openai_chat_online(
384 + self,
385 + ctx: PoolingServeContext[
386 + EmbeddingChatRequest
387 + | EmbeddingBatchChatRequest
388 + | EmbeddingChatInputRequest
389 + | EmbeddingBatchChatInputRequest
390 + ],
391 + ) -> None:
392 + request = ctx.request
393 + self._validate_chat_template(
394 + request_chat_template=request.chat_template,
395 + chat_template_kwargs=request.chat_template_kwargs,
396 + trust_request_chat_template=self.trust_request_chat_template,
397 + )
398 +
399 + if isinstance(
400 + request, (EmbeddingBatchChatRequest, EmbeddingBatchChatInputRequest)
401 + ):
402 + all_messages = request.messages
403 + else:
404 + all_messages = [request.messages]
405 + ctx.engine_inputs = self._batch_render_openai_chat(request, all_messages)
406 +
407 + def _batch_render_openai_chat(
408 + self,
409 + request: (
410 + EmbeddingChatRequest
411 + | EmbeddingBatchChatRequest
412 + | EmbeddingChatInputRequest
413 + | EmbeddingBatchChatInputRequest
414 + ),
415 + all_messages: Sequence[list[ChatCompletionMessageParam]],
416 + ) -> list[EngineInput]:
417 + renderer = self.renderer
418 + mm_config = self.model_config.multimodal_config
419 +
420 + tok_params = request.build_tok_params(self.model_config)
421 + chat_params = request.build_chat_params(
422 + self.chat_template,
423 + self.chat_template_content_format,
424 + ).with_defaults(
425 + merge_kwargs(
426 + None,
427 + dict(
428 + tools=None,
429 + tokenize=is_mistral_tokenizer(renderer.tokenizer),
430 + ),
431 + ),
432 + default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
433 + )
434 +
435 + _, engine_inputs = renderer.render_chat(
436 + all_messages,
437 + chat_params,
438 + tok_params,
439 + prompt_extras={
440 + k: v
441 + for k in ("mm_processor_kwargs", "cache_salt")
442 + if (v := getattr(request, k, None)) is not None
443 + },
444 + )
445 + return engine_inputs
446 +
447 def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None:
448 """Convert a ``CohereEmbedRequest`` into engine prompts.
449
vllm/entrypoints/pooling/embed/protocol.py
+107 -3
@@ -10,17 +10,19 @@ import builtins
10 import struct
11 import time
12 from collections.abc import Sequence
13 -from typing import Literal, TypeAlias
13 +from typing import Annotated, Any, Literal, TypeAlias
14
15 import pybase64 as base64
16 -from pydantic import BaseModel, Field
16 +from pydantic import BaseModel, Field, model_validator
17
18 from vllm import PoolingParams
19 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
20 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
21 from vllm.utils import random_uuid
22
23 from ..base.protocol import (
24 ChatRequestMixin,
25 + ChatRequestOptionsMixin,
26 CompletionRequestMixin,
27 EmbeddingTokenizeParamsMixin,
28 EmbedRequestMixin,
@@ -42,12 +44,34 @@ class EmbeddingCompletionRequest(
44 )
45
46
47 +def _is_chat_message(value: Any) -> bool:
48 + return isinstance(value, dict) and isinstance(value.get("role"), str)
49 +
50 +
51 +def _is_chat_messages(value: Any) -> bool:
52 + return (
53 + isinstance(value, list)
54 + and bool(value)
55 + and all(_is_chat_message(item) for item in value)
56 + )
57 +
58 +
59 +def _is_batched_chat_messages(value: Any) -> bool:
60 + return (
61 + isinstance(value, list)
62 + and bool(value)
63 + and all(_is_chat_messages(item) for item in value)
64 + )
65 +
66 +
67 class EmbeddingChatRequest(
68 PoolingBasicRequestMixin,
69 ChatRequestMixin,
70 EmbedRequestMixin,
71 EmbeddingTokenizeParamsMixin,
72 ):
73 + """OpenAI embeddings request with one top-level chat conversation."""
74 +
75 def to_pooling_params(self):
76 return PoolingParams(
77 task="embed",
@@ -56,7 +80,87 @@ class EmbeddingChatRequest(
80 )
81
82
59 -EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest
83 +class EmbeddingBatchChatRequest(
84 + PoolingBasicRequestMixin,
85 + ChatRequestOptionsMixin,
86 + EmbedRequestMixin,
87 + EmbeddingTokenizeParamsMixin,
88 +):
89 + """OpenAI embeddings request with batched top-level chat conversations.
90 +
91 + Mirrors ``BatchChatCompletionRequest`` by keeping batched conversations in
92 + ``messages`` instead of introducing a separate batch-specific field.
93 + """
94 +
95 + messages: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = (
96 + Field(..., min_length=1)
97 + )
98 +
99 + def to_pooling_params(self):
100 + return PoolingParams(
101 + task="embed",
102 + dimensions=self.dimensions,
103 + use_activation=self.use_activation,
104 + )
105 +
106 +
107 +class EmbeddingChatInputRequest(
108 + EmbeddingChatRequest,
109 +):
110 + """OpenAI embeddings request with one chat conversation in ``input``."""
111 +
112 + input: list[ChatCompletionMessageParam]
113 +
114 + @model_validator(mode="before")
115 + @classmethod
116 + def normalize_input_messages(cls, data):
117 + if not isinstance(data, dict):
118 + return data
119 +
120 + if "messages" in data or "input" not in data:
121 + return data
122 +
123 + input_data = data["input"]
124 + if not _is_chat_messages(input_data):
125 + return data
126 +
127 + normalized = dict(data)
128 + normalized["messages"] = input_data
129 + return normalized
130 +
131 +
132 +class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest):
133 + """OpenAI embeddings request with batched chat conversations in ``input``."""
134 +
135 + input: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = (
136 + Field(..., min_length=1)
137 + )
138 +
139 + @model_validator(mode="before")
140 + @classmethod
141 + def normalize_input_messages(cls, data):
142 + if not isinstance(data, dict):
143 + return data
144 +
145 + if "messages" in data or "input" not in data:
146 + return data
147 +
148 + input_data = data["input"]
149 + if not _is_batched_chat_messages(input_data):
150 + return data
151 +
152 + normalized = dict(data)
153 + normalized["messages"] = input_data
154 + return normalized
155 +
156 +
157 +EmbeddingRequest: TypeAlias = (
158 + EmbeddingCompletionRequest
159 + | EmbeddingChatRequest
160 + | EmbeddingBatchChatRequest
161 + | EmbeddingChatInputRequest
162 + | EmbeddingBatchChatInputRequest
163 +)
164
165
166 # ---------------------------------------------------------------------------
vllm/entrypoints/pooling/typing.py
+8 -2
@@ -20,8 +20,10 @@ from .classify.protocol import (
20 from .embed.protocol import (
21 CohereEmbedRequest,
22 EmbeddingBytesResponse,
23 + EmbeddingChatInputRequest,
24 EmbeddingChatRequest,
25 EmbeddingCompletionRequest,
26 + EmbeddingRequest,
27 EmbeddingResponse,
28 )
29 from .pooling.protocol import (
@@ -41,11 +43,15 @@ PoolingCompletionLikeRequest: TypeAlias = (
43 )
44
45 PoolingChatLikeRequest: TypeAlias = (
44 - EmbeddingChatRequest | ClassificationChatRequest | PoolingChatRequest
46 + EmbeddingChatRequest
47 + | EmbeddingChatInputRequest
48 + | ClassificationChatRequest
49 + | PoolingChatRequest
50 )
51
52 AnyPoolingRequest: TypeAlias = (
48 - PoolingCompletionLikeRequest
53 + EmbeddingRequest
54 + | PoolingCompletionLikeRequest
55 | PoolingChatLikeRequest
56 | IOProcessorRequest
57 | ScoringRequest