A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 01:30 UTC
[Frontend][2/n] Improve pooling entrypoints | embed. (#36110)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Mar 9, 2026 at 11:42 UTC
fff3711a244dd9e2915323e31c20768d922e90b5
20 files changed
+509
-756
tests/entrypoints/pooling/score/test_online_score_vision.py
+10
-7
@@ -25,7 +25,7 @@ ROCM_ATTN_BACKENDS = [
25
"FLEX_ATTENTION",
26
]
27
28
-ATTN_BACKENDS = ROCM_ATTN_BACKENDS if current_platform.is_rocm() else []
28
+ATTN_BACKENDS = ROCM_ATTN_BACKENDS if current_platform.is_rocm() else ["auto"]
29
30
# Per-backend tolerance with explicit entries; "default" is the fallback
31
BACKEND_TOL: dict[str, float] = {
@@ -105,13 +105,16 @@ def server(request):
105
"8192",
106
"--chat-template",
107
str(VLLM_PATH / "examples/pooling/score/template/qwen3_vl_reranker.jinja"),
108
- "--attention-config",
109
- json.dumps({"backend": backend}),
110
- ] + ROCM_EXTRA_ARGS
108
+ ]
109
112
- env = dict(ROCM_ENV_OVERRIDES)
113
- if backend != "ROCM_AITER_FA":
114
- env["VLLM_ROCM_USE_AITER"] = "0"
110
+ env = dict()
111
+ if backend != "auto":
112
+ args += ["--attention-config", json.dumps({"backend": backend})]
113
+ args += ROCM_EXTRA_ARGS
114
+
115
+ env = dict(ROCM_ENV_OVERRIDES)
116
+ if backend != "ROCM_AITER_FA":
117
+ env["VLLM_ROCM_USE_AITER"] = "0"
118
119
with RemoteOpenAIServer(
120
MODEL_NAME, args, override_hf_configs=HF_OVERRIDES, env_dict=env
vllm/entrypoints/llm.py
+4
-4
@@ -397,7 +397,7 @@ class LLM:
397
self.io_processor = self.llm_engine.io_processor
398
self.input_processor = self.llm_engine.input_processor
399
self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)
400
- self.init_pooling_io_processors = init_pooling_io_processors(
400
+ self.pooling_io_processors = init_pooling_io_processors(
401
supported_tasks=supported_tasks,
402
model_config=self.model_config,
403
renderer=self.renderer,
@@ -1174,8 +1174,8 @@ class LLM:
1174
)
1175
raise ValueError(msg)
1176
1177
- if pooling_task in self.init_pooling_io_processors:
1178
- io_processor = self.init_pooling_io_processors[pooling_task]
1177
+ if pooling_task in self.pooling_io_processors:
1178
+ io_processor = self.pooling_io_processors[pooling_task]
1179
processor_inputs = io_processor.pre_process_offline(
1180
prompts_seq, tokenization_kwargs
1181
)
@@ -1194,7 +1194,7 @@ class LLM:
1194
outputs = self._run_engine(
1195
use_tqdm=use_tqdm, output_type=PoolingRequestOutput
1196
)
1197
- outputs = io_processor.post_process(outputs)
1197
+ outputs = io_processor.post_process_offline(outputs)
1198
else:
1199
outputs = self._run_completion(
1200
prompts=prompts_seq,
vllm/entrypoints/openai/engine/serving.py
+4
-21
@@ -60,12 +60,6 @@ from vllm.entrypoints.openai.speech_to_text.protocol import (
60
TranscriptionResponse,
61
TranslationRequest,
62
)
63
-from vllm.entrypoints.pooling.embed.protocol import (
64
- EmbeddingBytesResponse,
65
- EmbeddingChatRequest,
66
- EmbeddingCompletionRequest,
67
- EmbeddingResponse,
68
-)
63
from vllm.entrypoints.pooling.pooling.protocol import (
64
IOProcessorRequest,
65
PoolingChatRequest,
@@ -144,17 +138,13 @@ CompletionLikeRequest: TypeAlias = (
138
CompletionRequest
139
| TokenizeCompletionRequest
140
| DetokenizeRequest
147
- | EmbeddingCompletionRequest
141
| RerankRequest
142
| ScoreRequest
143
| PoolingCompletionRequest
144
)
145
146
ChatLikeRequest: TypeAlias = (
154
- ChatCompletionRequest
155
- | TokenizeChatRequest
156
- | EmbeddingChatRequest
157
- | PoolingChatRequest
147
+ ChatCompletionRequest | TokenizeChatRequest | PoolingChatRequest
148
)
149
150
SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest
@@ -171,8 +161,6 @@ AnyRequest: TypeAlias = (
161
AnyResponse: TypeAlias = (
162
CompletionResponse
163
| ChatCompletionResponse
174
- | EmbeddingResponse
175
- | EmbeddingBytesResponse
164
| TranscriptionResponse
165
| TokenizeResponse
166
| PoolingResponse
@@ -203,8 +191,7 @@ class ServeContext(Generic[RequestT]):
191
192
class OpenAIServing:
193
request_id_prefix: ClassVar[str] = """
206
- A short string prepended to every request’s ID (e.g. "embd")
207
- so you can easily tell “this ID came from Embedding.”
194
+ A short string prepended to every request’s ID.
195
"""
196
197
def __init__(
@@ -432,8 +419,7 @@ class OpenAIServing:
419
ctx: ServeContext,
420
) -> ErrorResponse | None:
421
"""
435
- Default preprocessing hook. Subclasses may override
436
- to prepare `ctx` (embedding, etc.).
422
+ Default preprocessing hook. Subclasses may override to prepare `ctx`.
423
"""
424
return None
425
@@ -730,13 +716,10 @@ class OpenAIServing:
716
token_num = len(input_ids)
717
max_model_len = self.model_config.max_model_len
718
733
- # Note: EmbeddingRequest,
734
- # and ScoreRequest doesn't have max_tokens
719
+ # Note: ScoreRequest doesn't have max_tokens
720
if isinstance(
721
request,
722
(
738
- EmbeddingChatRequest,
739
- EmbeddingCompletionRequest,
723
ScoreDataRequest,
724
ScoreTextRequest,
725
ScoreQueriesDocumentsRequest,
vllm/entrypoints/openai/run_batch.py
+10
-12
@@ -53,6 +53,7 @@ from vllm.entrypoints.pooling.score.protocol import (
53
ScoreRequest,
54
ScoreResponse,
55
)
56
+from vllm.entrypoints.utils import create_error_response
57
from vllm.logger import init_logger
58
from vllm.reasoning import ReasoningParserManager
59
from vllm.utils import random_uuid
@@ -503,7 +504,10 @@ async def run_request(
504
request: BatchRequestInput,
505
tracker: BatchProgressTracker,
506
) -> BatchRequestOutput:
506
- response = await serving_engine_func(request.body)
507
+ try:
508
+ response = await serving_engine_func(request.body)
509
+ except Exception as e:
510
+ response = create_error_response(e)
511
512
if isinstance(
513
response,
@@ -678,10 +682,10 @@ async def build_endpoint_registry(
682
683
# Get serving objects from state (defaulting to None if not set)
684
openai_serving_chat = getattr(state, "openai_serving_chat", None)
681
- openai_serving_embedding = getattr(state, "openai_serving_embedding", None)
682
- openai_serving_scores = getattr(state, "openai_serving_scores", None)
685
openai_serving_transcription = getattr(state, "openai_serving_transcription", None)
686
openai_serving_translation = getattr(state, "openai_serving_translation", None)
687
+ serving_embedding = getattr(state, "serving_embedding", None)
688
+ serving_scores = getattr(state, "serving_scores", None)
689
690
# Registry of endpoint configurations
691
endpoint_registry: dict[str, dict[str, Any]] = {
@@ -697,27 +701,21 @@ async def build_endpoint_registry(
701
"embeddings": {
702
"url_matcher": lambda url: url == "/v1/embeddings",
703
"handler_getter": lambda: (
700
- openai_serving_embedding.create_embedding
701
- if openai_serving_embedding is not None
702
- else None
704
+ serving_embedding if serving_embedding is not None else None
705
),
706
"wrapper_fn": None,
707
},
708
"score": {
709
"url_matcher": lambda url: url.endswith("/score"),
710
"handler_getter": lambda: (
709
- openai_serving_scores.create_score
710
- if openai_serving_scores is not None
711
- else None
711
+ serving_scores.create_score if serving_scores is not None else None
712
),
713
"wrapper_fn": None,
714
},
715
"rerank": {
716
"url_matcher": lambda url: url.endswith("/rerank"),
717
"handler_getter": lambda: (
718
- openai_serving_scores.do_rerank
719
- if openai_serving_scores is not None
720
- else None
718
+ serving_scores.do_rerank if serving_scores is not None else None
719
),
720
"wrapper_fn": None,
721
},
vllm/entrypoints/pooling/__init__.py
+6
-6
@@ -56,14 +56,14 @@ def init_pooling_state(
56
):
57
from vllm.entrypoints.chat_utils import load_chat_template
58
from vllm.entrypoints.pooling.classify.serving import ServingClassification
59
- from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
59
+ from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
60
from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
61
from vllm.entrypoints.pooling.score.serving import ServingScores
62
from vllm.tasks import POOLING_TASKS
63
64
resolved_chat_template = load_chat_template(args.chat_template)
65
66
- state.openai_serving_pooling = (
66
+ state.serving_pooling = (
67
(
68
OpenAIServingPooling(
69
engine_client,
@@ -77,8 +77,8 @@ def init_pooling_state(
77
if any(t in supported_tasks for t in POOLING_TASKS)
78
else None
79
)
80
- state.openai_serving_embedding = (
81
- OpenAIServingEmbedding(
80
+ state.serving_embedding = (
81
+ ServingEmbedding(
82
engine_client,
83
state.openai_serving_models,
84
request_logger=request_logger,
@@ -89,7 +89,7 @@ def init_pooling_state(
89
if "embed" in supported_tasks
90
else None
91
)
92
- state.openai_serving_classification = (
92
+ state.serving_classification = (
93
ServingClassification(
94
engine_client,
95
state.openai_serving_models,
@@ -105,7 +105,7 @@ def init_pooling_state(
105
# - "score" task (cross-encoder models)
106
# - "embed" task (bi-encoder models)
107
# - "token_embed" task (late interaction models like ColBERT)
108
- state.openai_serving_scores = (
108
+ state.serving_scores = (
109
ServingScores(
110
engine_client,
111
state.openai_serving_models,
vllm/entrypoints/pooling/base/io_processor.py
+76
-20
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable, Sequence
5
-from concurrent.futures import ThreadPoolExecutor
5
from typing import Any, Final
6
7
from vllm import PoolingRequestOutput, PromptType
@@ -14,9 +13,13 @@ from vllm.entrypoints.chat_utils import (
13
ConversationMessage,
14
)
15
from vllm.entrypoints.openai.engine.serving import RendererChatRequest, RendererRequest
17
-from vllm.inputs import ProcessorInputs, SingletonPrompt
16
+from vllm.entrypoints.pooling.typing import (
17
+ PoolingChatLikeRequest,
18
+ PoolingCompletionLikeRequest,
19
+ PoolingServeContext,
20
+)
21
+from vllm.inputs.data import ProcessorInputs, SingletonPrompt
22
from vllm.renderers import BaseRenderer, merge_kwargs
19
-from vllm.renderers.inputs import TokPrompt
23
from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
24
from vllm.tokenizers import TokenizerLike
25
from vllm.tool_parsers import ToolParser
@@ -24,14 +27,14 @@ from vllm.utils.mistral import is_mistral_tokenizer
27
28
29
class PoolingIOProcessor:
30
+ name: str
31
+
32
def __init__(
33
self,
34
model_config: ModelConfig,
35
renderer: BaseRenderer,
36
chat_template_config: ChatTemplateConfig,
37
):
33
- self._tokenizer_executor = ThreadPoolExecutor(max_workers=1)
34
-
38
self.model_config = model_config
39
self.renderer = renderer
40
@@ -43,37 +46,90 @@ class PoolingIOProcessor:
46
chat_template_config.trust_request_chat_template
47
)
48
46
- def pre_process_online(self, *args, **kwargs):
47
- raise NotImplementedError
49
+ def create_pooling_params(self, request):
50
+ return request.to_pooling_params()
51
+
52
+ #######################################
53
+ # online APIs
54
+
55
+ def pre_process_online(self, ctx: PoolingServeContext):
56
+ request = ctx.request
57
+
58
+ if isinstance(ctx.request, PoolingChatLikeRequest):
59
+ self._validate_chat_template(
60
+ request_chat_template=request.chat_template,
61
+ chat_template_kwargs=request.chat_template_kwargs,
62
+ trust_request_chat_template=self.trust_request_chat_template,
63
+ )
64
+ _, engine_prompts = self._preprocess_chat_online(
65
+ request,
66
+ request.messages,
67
+ default_template=self.chat_template,
68
+ default_template_content_format=self.chat_template_content_format,
69
+ default_template_kwargs=None,
70
+ )
71
+ elif isinstance(request, PoolingCompletionLikeRequest):
72
+ engine_prompts = self._preprocess_completion_online(
73
+ request,
74
+ prompt_input=request.input,
75
+ prompt_embeds=None,
76
+ )
77
+ else:
78
+ raise ValueError(f"Invalid {self.name} request type")
79
+
80
+ ctx.engine_prompts = engine_prompts
81
+
82
+ async def pre_process_online_async(self, ctx: PoolingServeContext):
83
+ self.pre_process_online(ctx)
84
+
85
+ def post_process_online(
86
+ self,
87
+ ctx: PoolingServeContext,
88
+ ):
89
+ pass
90
49
- async def pre_process_online_async(self, *args, **kwargs):
50
- return self.pre_process_online(*args, **kwargs)
91
+ async def post_process_online_async(
92
+ self,
93
+ ctx: PoolingServeContext,
94
+ ):
95
+ self.post_process_online(ctx)
96
52
- def pre_process_offline(self, *args, **kwargs):
53
- raise NotImplementedError
97
+ #######################################
98
+ # offline APIs
99
+
100
+ def pre_process_offline(
101
+ self,
102
+ prompts: PromptType | Sequence[PromptType],
103
+ tokenization_kwargs: dict[str, Any] | None = None,
104
+ ) -> Sequence[ProcessorInputs]:
105
+ return self._preprocess_completion_offline(
106
+ prompts=prompts, tokenization_kwargs=tokenization_kwargs
107
+ )
108
109
async def pre_process_offline_async(self, *args, **kwargs):
110
return self.pre_process_offline(*args, **kwargs)
111
58
- def post_process(
59
- self, outputs: list[PoolingRequestOutput]
112
+ def post_process_offline(
113
+ self,
114
+ outputs: list[PoolingRequestOutput],
115
) -> list[PoolingRequestOutput]:
116
return outputs
117
63
- async def post_process_async(
64
- self, outputs: list[PoolingRequestOutput]
118
+ async def post_process_offline_async(
119
+ self,
120
+ outputs: list[PoolingRequestOutput],
121
) -> list[PoolingRequestOutput]:
66
- return self.post_process(outputs)
122
+ return self.post_process_offline(outputs)
123
68
- def create_pooling_params(self, request):
69
- return request.to_pooling_params()
124
+ #######################################
125
+ # helpers
126
127
def _preprocess_completion_online(
128
self,
129
request: RendererRequest,
130
prompt_input: str | list[str] | list[int] | list[list[int]] | None,
131
prompt_embeds: bytes | list[bytes] | None,
76
- ) -> list[TokPrompt]:
132
+ ) -> list[ProcessorInputs]:
133
renderer = self.renderer
134
model_config = self.model_config
135
@@ -112,7 +168,7 @@ class PoolingIOProcessor:
168
default_template_kwargs: dict[str, Any] | None,
169
tool_dicts: list[dict[str, Any]] | None = None,
170
tool_parser: Callable[[TokenizerLike], ToolParser] | None = None,
115
- ) -> tuple[list[ConversationMessage], list[TokPrompt]]:
171
+ ) -> tuple[list[ConversationMessage], list[ProcessorInputs]]:
172
renderer = self.renderer
173
174
default_template_kwargs = merge_kwargs(
vllm/entrypoints/pooling/base/serving.py
+26
-59
@@ -1,23 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import time
3
from collections.abc import AsyncGenerator, Mapping
5
-from dataclasses import dataclass, field
4
from http import HTTPStatus
7
-from typing import ClassVar, Generic, TypeVar
5
+from typing import ClassVar
6
7
from fastapi import Request
10
-from pydantic import ConfigDict
8
+from fastapi.responses import Response
9
from starlette.datastructures import Headers
12
-from starlette.responses import JSONResponse
13
-
14
-from vllm import (
15
- PoolingParams,
16
- PoolingRequestOutput,
17
- PromptType,
18
- SamplingParams,
19
- envs,
20
-)
10
+
11
+from vllm import PoolingParams, PoolingRequestOutput, envs
12
from vllm.config import ModelConfig
13
from vllm.engine.protocol import EngineClient
14
from vllm.entrypoints.chat_utils import (
@@ -27,12 +18,12 @@ from vllm.entrypoints.chat_utils import (
18
from vllm.entrypoints.logger import RequestLogger
19
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
20
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
30
-from vllm.entrypoints.pooling.typing import AnyPoolingRequest, AnyPoolingResponse
31
-from vllm.inputs import ProcessorInputs
21
+from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext
22
+from vllm.exceptions import VLLMNotFoundError
23
+from vllm.inputs.data import ProcessorInputs
24
from vllm.lora.request import LoRARequest
33
-from vllm.renderers import BaseRenderer
25
+from vllm.renderers.base import BaseRenderer
26
from vllm.renderers.inputs.preprocess import extract_prompt_components
35
-from vllm.sampling_params import BeamSearchParams
27
from vllm.tracing import (
28
contains_trace_headers,
29
extract_trace_headers,
@@ -43,26 +34,6 @@ from vllm.utils.async_utils import merge_async_iterators
34
35
from .io_processor import PoolingIOProcessor
36
46
-PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest)
47
-
48
-
49
-@dataclass(kw_only=True)
50
-class PoolingServeContext(Generic[PoolingRequestT]):
51
- request: PoolingRequestT
52
- raw_request: Request | None = None
53
- model_name: str
54
- request_id: str
55
- created_time: int = field(default_factory=lambda: int(time.time()))
56
- lora_request: LoRARequest | None = None
57
- engine_prompts: list[ProcessorInputs] | None = None
58
-
59
- result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = (
60
- None
61
- )
62
- final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)
63
-
64
- model_config = ConfigDict(arbitrary_types_allowed=True)
65
-
37
38
class PoolingServing:
39
request_id_prefix: ClassVar[str]
@@ -109,8 +80,8 @@ class PoolingServing:
80
async def __call__(
81
self,
82
request: AnyPoolingRequest,
112
- raw_request: Request,
113
- ) -> JSONResponse:
83
+ raw_request: Request | None = None,
84
+ ) -> Response:
85
model_name = self.models.model_name()
86
request_id = f"{self.request_id_prefix}-{self._base_request_id(raw_request)}"
87
@@ -125,19 +96,11 @@ class PoolingServing:
96
97
self._validate_request(ctx)
98
self._maybe_get_adapters(ctx)
128
- await self._preprocess(ctx)
99
+ await self.io_processor.pre_process_online_async(ctx)
100
await self._prepare_generators(ctx)
101
await self._collect_batch(ctx)
131
- response = await self._build_response(ctx)
132
- return JSONResponse(content=response.model_dump())
133
-
134
- async def _preprocess(
135
- self,
136
- ctx: PoolingServeContext,
137
- ):
138
- ctx.engine_prompts = await self.io_processor.pre_process_online_async(
139
- ctx.request
140
- )
102
+ await self.io_processor.post_process_online_async(ctx)
103
+ return await self._build_response(ctx)
104
105
async def _prepare_generators(
106
self,
@@ -157,10 +120,14 @@ class PoolingServing:
120
pooling_params = self.io_processor.create_pooling_params(ctx.request)
121
122
for i, engine_prompt in enumerate(ctx.engine_prompts):
160
- request_id_item = f"{ctx.request_id}-{i}"
123
+ prompt_request_id = (
124
+ f"{ctx.request_id}-{i}"
125
+ if ctx.prompt_request_ids is None
126
+ else ctx.prompt_request_ids[i]
127
+ )
128
129
self._log_inputs(
163
- request_id_item,
130
+ prompt_request_id,
131
engine_prompt,
132
params=pooling_params,
133
lora_request=ctx.lora_request,
@@ -169,7 +136,7 @@ class PoolingServing:
136
generator = self.engine_client.encode(
137
engine_prompt,
138
pooling_params,
172
- request_id_item,
139
+ prompt_request_id,
140
lora_request=ctx.lora_request,
141
trace_headers=trace_headers,
142
priority=getattr(ctx.request, "priority", 0),
@@ -189,9 +156,9 @@ class PoolingServing:
156
if ctx.result_generator is None:
157
raise ValueError("Result generator not available")
158
192
- num_prompts = len(ctx.engine_prompts)
159
+ num_inputs = len(ctx.engine_prompts)
160
final_res_batch: list[PoolingRequestOutput | None]
194
- final_res_batch = [None] * num_prompts
161
+ final_res_batch = [None] * num_inputs
162
163
async for i, res in ctx.result_generator:
164
final_res_batch[i] = res
@@ -204,7 +171,7 @@ class PoolingServing:
171
async def _build_response(
172
self,
173
ctx: PoolingServeContext,
207
- ) -> AnyPoolingResponse:
174
+ ) -> Response:
175
raise NotImplementedError
176
177
@staticmethod
@@ -294,7 +261,7 @@ class PoolingServing:
261
return None
262
263
# if _check_model has been called earlier, this will be unreachable
297
- raise ValueError(f"The model `{request.model}` does not exist.")
264
+ raise VLLMNotFoundError(f"The model `{request.model}` does not exist.")
265
266
def _get_active_default_mm_loras(
267
self, request: AnyPoolingRequest
@@ -349,8 +316,8 @@ class PoolingServing:
316
def _log_inputs(
317
self,
318
request_id: str,
352
- inputs: PromptType | ProcessorInputs,
353
- params: SamplingParams | PoolingParams | BeamSearchParams | None,
319
+ inputs: ProcessorInputs,
320
+ params: PoolingParams,
321
lora_request: LoRARequest | None,
322
) -> None:
323
if self.request_logger is None:
vllm/entrypoints/pooling/classify/api_router.py
+4
-6
@@ -2,12 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from fastapi import APIRouter, Depends, Request
5
-from starlette.responses import JSONResponse
5
+from fastapi.responses import JSONResponse, Response
6
7
from vllm.entrypoints.openai.utils import validate_json_request
8
-from vllm.entrypoints.pooling.classify.protocol import (
9
- ClassificationRequest,
10
-)
8
+from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest
9
from vllm.entrypoints.pooling.classify.serving import ServingClassification
10
from vllm.entrypoints.utils import (
11
create_error_response,
@@ -19,7 +17,7 @@ router = APIRouter()
17
18
19
def classify(request: Request) -> ServingClassification | None:
22
- return request.app.state.openai_serving_classification
20
+ return request.app.state.serving_classification
21
22
23
@router.post("/classify", dependencies=[Depends(validate_json_request)])
@@ -27,7 +25,7 @@ def classify(request: Request) -> ServingClassification | None:
25
@load_aware_call
26
async def create_classify(
27
request: ClassificationRequest, raw_request: Request
30
-) -> JSONResponse:
28
+) -> Response:
29
handler = classify(raw_request)
30
if handler is None:
31
error_response = create_error_response(
vllm/entrypoints/pooling/classify/io_processor.py
+1
-43
@@ -1,50 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Sequence
4
-from typing import Any
3
6
-from vllm import PromptType
4
from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
8
-from vllm.entrypoints.pooling.classify.protocol import (
9
- ClassificationChatRequest,
10
- ClassificationCompletionRequest,
11
-)
12
-from vllm.inputs import ProcessorInputs
13
-from vllm.renderers.inputs import TokPrompt
5
6
7
class ClassifyIOProcessor(PoolingIOProcessor):
17
- def pre_process_online(
18
- self, request: ClassificationCompletionRequest | ClassificationChatRequest
19
- ) -> list[TokPrompt] | None:
20
- if isinstance(request, ClassificationChatRequest):
21
- self._validate_chat_template(
22
- request_chat_template=request.chat_template,
23
- chat_template_kwargs=request.chat_template_kwargs,
24
- trust_request_chat_template=self.trust_request_chat_template,
25
- )
26
- _, engine_prompts = self._preprocess_chat_online(
27
- request,
28
- request.messages,
29
- default_template=self.chat_template,
30
- default_template_content_format=self.chat_template_content_format,
31
- default_template_kwargs=None,
32
- )
33
- elif isinstance(request, ClassificationCompletionRequest):
34
- engine_prompts = self._preprocess_completion_online(
35
- request,
36
- prompt_input=request.input,
37
- prompt_embeds=None,
38
- )
39
- else:
40
- raise ValueError("Invalid classification request type")
41
- return engine_prompts
42
-
43
- def pre_process_offline(
44
- self,
45
- prompts: PromptType | Sequence[PromptType],
46
- tokenization_kwargs: dict[str, Any] | None = None,
47
- ) -> Sequence[ProcessorInputs]:
48
- return self._preprocess_completion_offline(
49
- prompts=prompts, tokenization_kwargs=tokenization_kwargs
50
- )
8
+ name = "classification"
vllm/entrypoints/pooling/classify/serving.py
+9
-9
@@ -4,13 +4,15 @@
4
from typing import TypeAlias
5
6
import numpy as np
7
+from fastapi.responses import JSONResponse
8
8
-from vllm import ClassificationOutput
9
from vllm.config import ModelConfig
10
from vllm.entrypoints.chat_utils import ChatTemplateConfig
11
from vllm.entrypoints.openai.engine.protocol import UsageInfo
12
-from vllm.entrypoints.pooling.base.serving import PoolingServeContext, PoolingServing
12
+from vllm.entrypoints.pooling.base.serving import PoolingServing
13
+from vllm.entrypoints.pooling.typing import PoolingServeContext
14
from vllm.logger import init_logger
15
+from vllm.outputs import ClassificationOutput
16
from vllm.renderers import BaseRenderer
17
18
from .io_processor import ClassifyIOProcessor
@@ -44,15 +46,11 @@ class ServingClassification(PoolingServing):
46
async def _build_response(
47
self,
48
ctx: ClassificationServeContext,
47
- ) -> ClassificationResponse:
48
- final_res_batch_checked = await self.io_processor.post_process_async(
49
- ctx.final_res_batch
50
- )
51
-
49
+ ) -> JSONResponse:
50
id2label = getattr(self.model_config.hf_config, "id2label", {})
51
num_prompt_tokens = 0
52
items: list[ClassificationData] = []
55
- for idx, final_res in enumerate(final_res_batch_checked):
53
+ for idx, final_res in enumerate(ctx.final_res_batch):
54
classify_res = ClassificationOutput.from_base(final_res.outputs)
55
56
probs = classify_res.probs
@@ -75,10 +73,12 @@ class ServingClassification(PoolingServing):
73
total_tokens=num_prompt_tokens,
74
)
75
78
- return ClassificationResponse(
76
+ response = ClassificationResponse(
77
id=ctx.request_id,
78
created=ctx.created_time,
79
model=ctx.model_name,
80
data=items,
81
usage=usage,
82
)
83
+
84
+ return JSONResponse(content=response.model_dump())
vllm/entrypoints/pooling/embed/api_router.py
+14
-44
@@ -1,43 +1,26 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import importlib.util
4
-from functools import lru_cache
3
+
4
from http import HTTPStatus
5
6
from fastapi import APIRouter, Depends, Request
8
-from fastapi.responses import JSONResponse, StreamingResponse
9
-from typing_extensions import assert_never
7
+from fastapi.responses import JSONResponse
8
9
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
10
from vllm.entrypoints.openai.utils import validate_json_request
13
-from vllm.entrypoints.pooling.embed.protocol import (
14
- EmbeddingBytesResponse,
15
- EmbeddingRequest,
16
- EmbeddingResponse,
11
+from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest
12
+from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
13
+from vllm.entrypoints.utils import (
14
+ create_error_response,
15
+ load_aware_call,
16
+ with_cancellation,
17
)
18
-from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding
19
-from vllm.entrypoints.utils import load_aware_call, with_cancellation
20
-from vllm.logger import init_logger
18
19
router = APIRouter()
20
24
-logger = init_logger(__name__)
25
-
26
-
27
-@lru_cache(maxsize=1)
28
-def _get_json_response_cls():
29
- if importlib.util.find_spec("orjson") is not None:
30
- from fastapi.responses import ORJSONResponse
31
-
32
- return ORJSONResponse
33
- logger.warning_once(
34
- "To make v1/embeddings API fast, please install orjson by `pip install orjson`"
35
- )
36
- return JSONResponse
37
-
21
39
-def embedding(request: Request) -> OpenAIServingEmbedding | None:
40
- return request.app.state.openai_serving_embedding
22
+def embedding(request: Request) -> ServingEmbedding | None:
23
+ return request.app.state.serving_embedding
24
25
26
@router.post(
@@ -56,24 +39,11 @@ async def create_embedding(
39
):
40
handler = embedding(raw_request)
41
if handler is None:
59
- base_server = raw_request.app.state.openai_serving_tokenization
60
- return base_server.create_error_response(
42
+ error_response = create_error_response(
43
message="The model does not support Embeddings API"
44
)
63
-
64
- generator = await handler.create_embedding(request, raw_request)
65
-
66
- if isinstance(generator, ErrorResponse):
45
return JSONResponse(
68
- content=generator.model_dump(), status_code=generator.error.code
46
+ content=error_response.model_dump(),
47
+ status_code=error_response.error.code,
48
)
70
- elif isinstance(generator, EmbeddingResponse):
71
- return _get_json_response_cls()(content=generator.model_dump())
72
- elif isinstance(generator, EmbeddingBytesResponse):
73
- return StreamingResponse(
74
- content=generator.content,
75
- headers=generator.headers,
76
- media_type=generator.media_type,
77
- )
78
-
79
- assert_never(generator)
49
+ return await handler(request, raw_request)
vllm/entrypoints/pooling/embed/io_processor.py
new
+198
@@ -0,0 +1,198 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Any, cast
4
+
5
+import torch
6
+
7
+from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
8
+from vllm.entrypoints.pooling.typing import PoolingServeContext
9
+from vllm.inputs.data import ProcessorInputs, token_inputs
10
+from vllm.outputs import PoolingOutput, PoolingRequestOutput
11
+from vllm.utils.collection_utils import chunk_list
12
+
13
+
14
+class EmbedIOProcessor(PoolingIOProcessor):
15
+ name = "embedding"
16
+
17
+ def __init__(self, *args, **kwargs):
18
+ super().__init__(*args, **kwargs)
19
+ assert self.model_config.pooler_config is not None
20
+
21
+ self.pooler_config = self.model_config.pooler_config
22
+ self.enable_chunked_processing = self.pooler_config.enable_chunked_processing
23
+
24
+ #################################################################
25
+ # Long Text Embedding with Chunked Processing
26
+ # PTAL: examples/pooling/embed/openai_embedding_long_text
27
+
28
+ def pre_process_online(self, ctx: PoolingServeContext):
29
+ super().pre_process_online(ctx)
30
+
31
+ if not self.enable_chunked_processing:
32
+ return None
33
+
34
+ if ctx.engine_prompts is None:
35
+ raise ValueError("Engine prompts not available")
36
+
37
+ ctx.intermediates = ctx.engine_prompts
38
+ request_id = ctx.request_id
39
+ max_model_len = self.model_config.max_model_len
40
+ chunked_engine_prompts: list[ProcessorInputs] = []
41
+ prompt_request_ids: list[str] = []
42
+ for prompt_idx, engine_prompt in enumerate(ctx.engine_prompts):
43
+ token_ids = engine_prompt.get("prompt_token_ids", None)
44
+ if token_ids is None:
45
+ raise NotImplementedError(
46
+ "Long Text Embedding with Chunked Processing does "
47
+ "not support EmbedsPrompt and EncoderDecoderInputs."
48
+ )
49
+
50
+ prompt_token_ids = cast(list[int], token_ids)
51
+
52
+ for chunk_idx, chunk_tokens in enumerate(
53
+ chunk_list(prompt_token_ids, max_model_len)
54
+ ):
55
+ chunked_engine_prompts.append(
56
+ token_inputs(prompt_token_ids=chunk_tokens)
57
+ )
58
+ prompt_request_ids.append(
59
+ f"{request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}"
60
+ )
61
+
62
+ ctx.engine_prompts = chunked_engine_prompts
63
+ ctx.prompt_request_ids = prompt_request_ids
64
+ return None
65
+
66
+ def post_process_online(
67
+ self,
68
+ ctx: PoolingServeContext,
69
+ ):
70
+ if ctx.final_res_batch is None:
71
+ raise ValueError("Final response batch not available")
72
+
73
+ if not self.enable_chunked_processing:
74
+ return super().post_process_online(ctx)
75
+
76
+ # Online aggregation for chunked requests to
77
+ # minimize memory usage
78
+ # Track aggregation state for each prompt
79
+ prompt_aggregators: dict[int, dict[str, Any]] = {}
80
+ short_prompts_results: dict[int, PoolingRequestOutput] = {}
81
+ for result_idx, result in enumerate(ctx.final_res_batch):
82
+ if "-chunk-" not in result.request_id:
83
+ # Non-chunked result - extract prompt_idx from request_id
84
+ parts = result.request_id.split("-")
85
+ try:
86
+ # Last part should be prompt index
87
+ prompt_idx = int(parts[-1])
88
+ except (ValueError, IndexError):
89
+ prompt_idx = result_idx # Fallback to result_idx
90
+
91
+ short_prompts_results[prompt_idx] = result
92
+ else:
93
+ # Extract prompt_idx from chunked request_id
94
+ parts = result.request_id.split("-")
95
+ try:
96
+ prompt_idx = int(parts[parts.index("prompt") + 1])
97
+ except (ValueError, IndexError):
98
+ # Fallback: extract from result_idx if parsing fails
99
+ prompt_idx = result_idx
100
+
101
+ # Initialize aggregator for this prompt if needed
102
+ if prompt_idx not in prompt_aggregators:
103
+ prompt_aggregators[prompt_idx] = {
104
+ "weighted_sum": None,
105
+ "total_weight": 0,
106
+ "chunk_count": 0,
107
+ "request_id": result.request_id.split("-chunk-")[0],
108
+ }
109
+
110
+ aggregator = prompt_aggregators[prompt_idx]
111
+
112
+ # MEAN pooling with online weighted averaging
113
+ # Ensure result is PoolingRequestOutput
114
+ # for embedding processing
115
+ if not isinstance(result, PoolingRequestOutput):
116
+ raise ValueError(
117
+ f"Expected PoolingRequestOutput for "
118
+ f"chunked embedding, got "
119
+ f"{type(result).__name__}"
120
+ )
121
+ if result.prompt_token_ids is None:
122
+ raise ValueError(
123
+ "prompt_token_ids cannot be None for chunked processing"
124
+ )
125
+
126
+ weight = len(result.prompt_token_ids)
127
+ embedding_data = result.outputs.data
128
+ weighted_embedding = embedding_data.to(dtype=torch.float32) * weight
129
+
130
+ if aggregator["weighted_sum"] is None:
131
+ # First chunk
132
+ aggregator["weighted_sum"] = weighted_embedding
133
+ else:
134
+ # Accumulate
135
+ aggregator["weighted_sum"] += weighted_embedding
136
+
137
+ aggregator["total_weight"] += weight
138
+ aggregator["chunk_count"] += 1
139
+
140
+ if ctx.intermediates is None:
141
+ raise ValueError("Original prompts inputs not available")
142
+
143
+ original_engine_prompts = cast(list[ProcessorInputs], ctx.intermediates)
144
+ num_prompts = len(original_engine_prompts)
145
+
146
+ # Finalize aggregated results
147
+ final_res_batch: list[PoolingRequestOutput] = []
148
+ for prompt_idx in range(num_prompts):
149
+ if prompt_idx in prompt_aggregators:
150
+ # Finalize MEAN aggregation for this chunked prompt
151
+ aggregator = prompt_aggregators[prompt_idx]
152
+
153
+ weighted_sum = aggregator["weighted_sum"]
154
+ total_weight = aggregator["total_weight"]
155
+
156
+ if (
157
+ weighted_sum is not None
158
+ and isinstance(weighted_sum, torch.Tensor)
159
+ and isinstance(total_weight, (int, float))
160
+ and total_weight > 0
161
+ ):
162
+ # Compute final mean embedding
163
+ final_embedding = weighted_sum / total_weight
164
+
165
+ # Create a PoolingRequestOutput
166
+ # for the aggregated result
167
+ pooling_output_data = PoolingOutput(data=final_embedding)
168
+
169
+ # Get original prompt token IDs for this prompt
170
+ original_prompt = original_engine_prompts[prompt_idx]
171
+ token_ids = original_prompt.get("prompt_token_ids", None)
172
+ if token_ids is None:
173
+ raise NotImplementedError(
174
+ "Long Text Embedding with Chunked Processing does "
175
+ "not support EmbedsPrompt and EncoderDecoderInputs."
176
+ )
177
+
178
+ original_token_ids = cast(list[int], token_ids)
179
+ pooling_request_output = PoolingRequestOutput(
180
+ request_id=aggregator["request_id"],
181
+ prompt_token_ids=original_token_ids,
182
+ outputs=pooling_output_data,
183
+ num_cached_tokens=0,
184
+ finished=True,
185
+ )
186
+
187
+ final_res_batch.append(pooling_request_output)
188
+ else:
189
+ raise ValueError(
190
+ f"Failed to aggregate chunks for prompt {prompt_idx}"
191
+ )
192
+ elif prompt_idx in short_prompts_results:
193
+ final_res_batch.append(short_prompts_results[prompt_idx])
194
+ else:
195
+ raise ValueError(f"Result not found for prompt {prompt_idx}")
196
+
197
+ ctx.final_res_batch = final_res_batch
198
+ return None
vllm/entrypoints/pooling/embed/serving.py
+65
-509
@@ -1,108 +1,95 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import json
4
-from collections.abc import AsyncGenerator, Callable, Mapping
4
+from collections.abc import Callable
5
from functools import partial
6
-from typing import Any, Final, Literal, TypeAlias, cast
6
+from typing import Literal, TypeAlias, cast
7
8
-import torch
9
-from fastapi import Request
8
+from fastapi.responses import JSONResponse, StreamingResponse
9
from typing_extensions import assert_never
10
12
-from vllm.engine.protocol import EngineClient
13
-from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
14
-from vllm.entrypoints.logger import RequestLogger
15
-from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
16
-from vllm.entrypoints.openai.engine.serving import OpenAIServing, ServeContext
17
-from vllm.entrypoints.openai.models.serving import OpenAIServingModels
11
+from vllm.config import ModelConfig
12
+from vllm.entrypoints.chat_utils import ChatTemplateConfig
13
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
14
+from vllm.entrypoints.pooling.base.serving import PoolingServing
15
+from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
16
from vllm.entrypoints.pooling.embed.protocol import (
17
EmbeddingBytesResponse,
20
- EmbeddingChatRequest,
21
- EmbeddingCompletionRequest,
18
EmbeddingRequest,
19
EmbeddingResponse,
20
EmbeddingResponseData,
21
)
22
+from vllm.entrypoints.pooling.typing import PoolingServeContext
23
from vllm.entrypoints.pooling.utils import (
24
encode_pooling_bytes,
25
encode_pooling_output_base64,
26
encode_pooling_output_float,
27
+ get_json_response_cls,
28
)
31
-from vllm.inputs.data import ProcessorInputs, TokensPrompt, token_inputs
32
-from vllm.logger import init_logger
33
-from vllm.outputs import PoolingOutput, PoolingRequestOutput
34
-from vllm.pooling_params import PoolingParams
35
-from vllm.utils.async_utils import merge_async_iterators
36
-from vllm.utils.collection_utils import chunk_list
29
+from vllm.outputs import PoolingRequestOutput
30
+from vllm.renderers import BaseRenderer
31
from vllm.utils.serial_utils import EmbedDType, Endianness
32
39
-logger = init_logger(__name__)
33
+JSONResponseCLS = get_json_response_cls()
34
35
+EmbeddingServeContext: TypeAlias = PoolingServeContext[EmbeddingRequest]
36
42
-EmbeddingServeContext: TypeAlias = ServeContext[EmbeddingRequest]
37
38
+class ServingEmbedding(PoolingServing):
39
+ """
40
+ Embedding API similar to OpenAI's API.
41
+
42
+ See https://platform.openai.com/docs/api-reference/embeddings/create
43
+ for the API specification. This API mimics the OpenAI Embedding API.
44
+ """
45
45
-class OpenAIServingEmbedding(OpenAIServing):
46
request_id_prefix = "embd"
47
48
- def __init__(
48
+ def init_io_processor(
49
self,
50
- engine_client: EngineClient,
51
- models: OpenAIServingModels,
52
- *,
53
- request_logger: RequestLogger | None,
54
- chat_template: str | None,
55
- chat_template_content_format: ChatTemplateContentFormatOption,
56
- trust_request_chat_template: bool = False,
57
- ) -> None:
58
- super().__init__(
59
- engine_client=engine_client,
60
- models=models,
61
- request_logger=request_logger,
50
+ model_config: ModelConfig,
51
+ renderer: BaseRenderer,
52
+ chat_template_config: ChatTemplateConfig,
53
+ ) -> EmbedIOProcessor:
54
+ return EmbedIOProcessor(
55
+ model_config=model_config,
56
+ renderer=renderer,
57
+ chat_template_config=chat_template_config,
58
)
59
64
- self.chat_template = chat_template
65
- self.chat_template_content_format: Final = chat_template_content_format
66
- self.trust_request_chat_template = trust_request_chat_template
67
-
68
- pooler_config = self.model_config.pooler_config
69
- assert pooler_config is not None
70
- self.pooler_config = pooler_config
71
-
72
- async def _preprocess(
60
+ async def _build_response(
61
self,
62
ctx: EmbeddingServeContext,
75
- ) -> ErrorResponse | None:
76
- ctx.lora_request = self._maybe_get_adapters(ctx.request)
63
+ ) -> JSONResponse | StreamingResponse:
64
+ encoding_format = ctx.request.encoding_format
65
+ embed_dtype = ctx.request.embed_dtype
66
+ endianness = ctx.request.endianness
67
78
- if isinstance(ctx.request, EmbeddingChatRequest):
79
- error_check_ret = self._validate_chat_template(
80
- request_chat_template=ctx.request.chat_template,
81
- chat_template_kwargs=ctx.request.chat_template_kwargs,
82
- trust_request_chat_template=self.trust_request_chat_template,
68
+ if encoding_format == "float" or encoding_format == "base64":
69
+ return self._request_output_to_embed_json_response(
70
+ ctx.final_res_batch,
71
+ ctx.request_id,
72
+ ctx.created_time,
73
+ ctx.model_name,
74
+ encoding_format,
75
+ embed_dtype,
76
+ endianness,
77
)
84
- if error_check_ret is not None:
85
- return error_check_ret
78
87
- _, ctx.engine_prompts = await self._preprocess_chat(
88
- ctx.request,
89
- ctx.request.messages,
90
- default_template=self.chat_template,
91
- default_template_content_format=self.chat_template_content_format,
92
- default_template_kwargs=None,
93
- )
94
- elif isinstance(ctx.request, EmbeddingCompletionRequest):
95
- ctx.engine_prompts = await self._preprocess_completion(
96
- ctx.request,
97
- prompt_input=ctx.request.input,
98
- prompt_embeds=None,
79
+ if encoding_format == "bytes" or encoding_format == "bytes_only":
80
+ return self._request_output_to_to_embed_bytes_response(
81
+ ctx.final_res_batch,
82
+ ctx.request_id,
83
+ ctx.created_time,
84
+ ctx.model_name,
85
+ encoding_format,
86
+ embed_dtype,
87
+ endianness,
88
)
100
- else:
101
- return self.create_error_response("Invalid classification request type")
89
103
- return None
90
+ assert_never(encoding_format)
91
105
- def request_output_to_embed_json_response(
92
+ def _request_output_to_embed_json_response(
93
self,
94
final_res_batch: list[PoolingRequestOutput],
95
request_id: str,
@@ -111,7 +98,7 @@ class OpenAIServingEmbedding(OpenAIServing):
98
encoding_format: Literal["float", "base64"],
99
embed_dtype: EmbedDType,
100
endianness: Endianness,
114
- ) -> EmbeddingResponse:
101
+ ) -> JSONResponse:
102
encode_fn = cast(
103
Callable[[PoolingRequestOutput], list[float] | str],
104
(
@@ -143,15 +130,16 @@ class OpenAIServingEmbedding(OpenAIServing):
130
total_tokens=num_prompt_tokens,
131
)
132
146
- return EmbeddingResponse(
133
+ response = EmbeddingResponse(
134
id=request_id,
135
created=created_time,
136
model=model_name,
137
data=items,
138
usage=usage,
139
)
140
+ return JSONResponseCLS(content=response.model_dump())
141
154
- def request_output_to_embed_bytes_response(
142
+ def _request_output_to_to_embed_bytes_response(
143
self,
144
final_res_batch: list[PoolingRequestOutput],
145
request_id: str,
@@ -160,7 +148,7 @@ class OpenAIServingEmbedding(OpenAIServing):
148
encoding_format: Literal["bytes", "bytes_only"],
149
embed_dtype: EmbedDType,
150
endianness: Endianness,
163
- ) -> EmbeddingBytesResponse:
151
+ ) -> StreamingResponse:
152
content, items, usage = encode_pooling_bytes(
153
pooling_outputs=final_res_batch,
154
embed_dtype=embed_dtype,
@@ -183,441 +171,9 @@ class OpenAIServingEmbedding(OpenAIServing):
171
}
172
)
173
186
- return EmbeddingBytesResponse(content=content, headers=headers)
187
-
188
- def _build_response(
189
- self,
190
- ctx: EmbeddingServeContext,
191
- ) -> EmbeddingResponse | EmbeddingBytesResponse | ErrorResponse:
192
- encoding_format = ctx.request.encoding_format
193
- embed_dtype = ctx.request.embed_dtype
194
- endianness = ctx.request.endianness
195
-
196
- if encoding_format == "float" or encoding_format == "base64":
197
- return self.request_output_to_embed_json_response(
198
- ctx.final_res_batch,
199
- ctx.request_id,
200
- ctx.created_time,
201
- ctx.model_name,
202
- encoding_format,
203
- embed_dtype,
204
- endianness,
205
- )
206
-
207
- if encoding_format == "bytes" or encoding_format == "bytes_only":
208
- return self.request_output_to_embed_bytes_response(
209
- ctx.final_res_batch,
210
- ctx.request_id,
211
- ctx.created_time,
212
- ctx.model_name,
213
- encoding_format,
214
- embed_dtype,
215
- endianness,
216
- )
217
-
218
- assert_never(encoding_format)
219
-
220
- def _get_max_position_embeddings(self) -> int:
221
- """Get the model's effective maximum sequence length for chunking."""
222
- return self.model_config.max_model_len
223
-
224
- def _should_use_chunked_processing(self, request) -> bool:
225
- """Check if chunked processing should be used for this request."""
226
- return (
227
- isinstance(request, (EmbeddingCompletionRequest, EmbeddingChatRequest))
228
- and self.pooler_config.enable_chunked_processing
229
- )
230
-
231
- async def _process_chunked_request(
232
- self,
233
- ctx: EmbeddingServeContext,
234
- token_ids: list[int],
235
- pooling_params: PoolingParams,
236
- trace_headers: Mapping[str, str] | None,
237
- prompt_idx: int,
238
- ) -> list[AsyncGenerator[PoolingRequestOutput, None]]:
239
- """Process a single prompt using chunked processing."""
240
- generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
241
-
242
- # Split into chunks using max_position_embeddings
243
- max_pos_embeddings = self._get_max_position_embeddings()
244
- # Process all chunks for MEAN aggregation
245
- for chunk_idx, chunk_tokens in enumerate(
246
- chunk_list(token_ids, max_pos_embeddings)
247
- ):
248
- # Create a request ID for this chunk
249
- chunk_request_id = f"{ctx.request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}"
250
-
251
- # Create engine prompt for this chunk
252
- chunk_engine_prompt = token_inputs(chunk_tokens)
253
-
254
- # Log the chunk
255
- self._log_inputs(
256
- chunk_request_id,
257
- chunk_engine_prompt,
258
- params=pooling_params,
259
- lora_request=ctx.lora_request,
260
- )
261
-
262
- # Create generator for this chunk and wrap it to return indices
263
- original_generator = self.engine_client.encode(
264
- chunk_engine_prompt,
265
- pooling_params,
266
- chunk_request_id,
267
- lora_request=ctx.lora_request,
268
- trace_headers=trace_headers,
269
- priority=ctx.request.priority,
270
- )
271
-
272
- generators.append(original_generator)
273
-
274
- return generators
275
-
276
- def _validate_input(
277
- self,
278
- request: object,
279
- input_ids: list[int],
280
- input_text: str,
281
- ) -> TokensPrompt:
282
- """Override to support chunked processing for embedding requests."""
283
- token_num = len(input_ids)
284
-
285
- # Note: EmbeddingRequest doesn't have max_tokens
286
- if isinstance(request, (EmbeddingCompletionRequest, EmbeddingChatRequest)):
287
- # Check if chunked processing is enabled for pooling models
288
- enable_chunked = self._should_use_chunked_processing(request)
289
-
290
- # Use max_position_embeddings for chunked processing decisions
291
- max_pos_embeddings = self._get_max_position_embeddings()
292
-
293
- # Determine the effective max length for validation
294
- if self.pooler_config.max_embed_len:
295
- # Use max_embed_len for validation instead of max_model_len
296
- length_type = "maximum embedding input length"
297
- max_length_value = self.pooler_config.max_embed_len
298
- else:
299
- # Fall back to max_model_len validation (original behavior)
300
- length_type = "maximum context length"
301
- max_length_value = self.model_config.max_model_len
302
-
303
- validation_error_msg = (
304
- "This model's {length_type} is {max_length_value} tokens. "
305
- "However, you requested {token_num} tokens in the input for "
306
- "embedding generation. Please reduce the length of the input."
307
- )
308
-
309
- chunked_processing_error_msg = (
310
- "This model's {length_type} is {max_length_value} tokens. "
311
- "However, you requested {token_num} tokens in the input for "
312
- "embedding generation. Please reduce the length of the input "
313
- "or enable chunked processing."
314
- )
315
-
316
- # Check if input exceeds max length
317
- if token_num > max_length_value:
318
- raise ValueError(
319
- validation_error_msg.format(
320
- length_type=length_type,
321
- max_length_value=max_length_value,
322
- token_num=token_num,
323
- )
324
- )
325
-
326
- # Check for chunked processing
327
- # when exceeding max_position_embeddings
328
- if token_num > max_pos_embeddings:
329
- if enable_chunked:
330
- # Allow long inputs when chunked processing is enabled
331
- logger.info(
332
- "Input length %s exceeds max_position_embeddings "
333
- "%s, will use chunked processing",
334
- token_num,
335
- max_pos_embeddings,
336
- )
337
- else:
338
- raise ValueError(
339
- chunked_processing_error_msg.format(
340
- length_type="maximum position embeddings length",
341
- max_length_value=max_pos_embeddings,
342
- token_num=token_num,
343
- )
344
- )
345
-
346
- return TokensPrompt(prompt=input_text, prompt_token_ids=input_ids)
347
-
348
- # For other request types, use the parent's implementation
349
- return super()._validate_input(request, input_ids, input_text)
350
-
351
- async def _create_single_prompt_generator(
352
- self,
353
- ctx: EmbeddingServeContext,
354
- engine_prompt: ProcessorInputs,
355
- pooling_params: PoolingParams,
356
- trace_headers: Mapping[str, str] | None,
357
- prompt_index: int,
358
- ) -> AsyncGenerator[PoolingRequestOutput, None]:
359
- """Create a generator for a single prompt using standard processing."""
360
- request_id_item = f"{ctx.request_id}-{prompt_index}"
361
-
362
- self._log_inputs(
363
- request_id_item,
364
- engine_prompt,
365
- params=pooling_params,
366
- lora_request=ctx.lora_request,
367
- )
368
-
369
- # Return the original generator without wrapping
370
- return self.engine_client.encode(
371
- engine_prompt,
372
- pooling_params,
373
- request_id_item,
374
- lora_request=ctx.lora_request,
375
- trace_headers=trace_headers,
376
- priority=ctx.request.priority,
377
- )
378
-
379
- async def _prepare_generators(
380
- self,
381
- ctx: EmbeddingServeContext,
382
- ) -> ErrorResponse | None:
383
- """Override to support chunked processing."""
384
- # Check if we should use chunked processing
385
- use_chunked = self._should_use_chunked_processing(ctx.request)
386
-
387
- # If no chunked processing needed, delegate to parent class
388
- if not use_chunked:
389
- return await super()._prepare_generators(ctx)
390
-
391
- # Custom logic for chunked processing
392
- generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
393
-
394
- trace_headers = (
395
- None
396
- if ctx.raw_request is None
397
- else await self._get_trace_headers(ctx.raw_request.headers)
174
+ response = EmbeddingBytesResponse(content=content, headers=headers)
175
+ return StreamingResponse(
176
+ content=response.content,
177
+ headers=response.headers,
178
+ media_type=response.media_type,
179
)
399
-
400
- pooling_params = self._create_pooling_params(ctx)
401
- if isinstance(pooling_params, ErrorResponse):
402
- return pooling_params
403
-
404
- if ctx.engine_prompts is None:
405
- return self.create_error_response("Engine prompts not available")
406
-
407
- max_pos_embeddings = self._get_max_position_embeddings()
408
-
409
- for i, engine_prompt in enumerate(ctx.engine_prompts):
410
- # Check if this specific prompt needs chunked processing
411
- if "prompt_token_ids" in engine_prompt:
412
- prompt_token_ids = engine_prompt["prompt_token_ids"] # type: ignore[typeddict-item]
413
-
414
- if len(prompt_token_ids) > max_pos_embeddings:
415
- # Use chunked processing for this prompt
416
- chunk_generators = await self._process_chunked_request(
417
- ctx,
418
- prompt_token_ids,
419
- pooling_params,
420
- trace_headers,
421
- i,
422
- )
423
- generators.extend(chunk_generators)
424
- continue
425
-
426
- # Normal processing for short prompts or non-token prompts
427
- generator = await self._create_single_prompt_generator(
428
- ctx, engine_prompt, pooling_params, trace_headers, i
429
- )
430
- generators.append(generator)
431
-
432
- ctx.result_generator = merge_async_iterators(*generators)
433
-
434
- return None
435
-
436
- async def _collect_batch(
437
- self,
438
- ctx: EmbeddingServeContext,
439
- ) -> ErrorResponse | None:
440
- """Collect and aggregate batch results
441
- with support for chunked processing.
442
-
443
- For chunked requests, performs online aggregation to
444
- minimize memory usage.
445
- For regular requests, collects results normally.
446
- """
447
- if ctx.engine_prompts is None:
448
- return self.create_error_response("Engine prompts not available")
449
-
450
- # Check if we used chunked processing
451
- use_chunked = self._should_use_chunked_processing(ctx.request)
452
-
453
- if not use_chunked:
454
- return await super()._collect_batch(ctx=ctx)
455
-
456
- if ctx.result_generator is None:
457
- return self.create_error_response("Result generator not available")
458
-
459
- # Online aggregation for chunked requests to
460
- # minimize memory usage
461
- # Track aggregation state for each prompt
462
- prompt_aggregators: dict[int, dict[str, Any]] = {}
463
- short_prompts_results: dict[int, PoolingRequestOutput] = {}
464
-
465
- async for result_idx, result in ctx.result_generator:
466
- if "-chunk-" in result.request_id:
467
- # Extract prompt_idx from chunked request_id
468
- parts = result.request_id.split("-")
469
- try:
470
- prompt_idx = int(parts[parts.index("prompt") + 1])
471
- except (ValueError, IndexError):
472
- # Fallback: extract from result_idx if parsing fails
473
- prompt_idx = result_idx
474
-
475
- # Initialize aggregator for this prompt if needed
476
- if prompt_idx not in prompt_aggregators:
477
- prompt_aggregators[prompt_idx] = {
478
- "weighted_sum": None,
479
- "total_weight": 0,
480
- "chunk_count": 0,
481
- "request_id": result.request_id.split("-chunk-")[0],
482
- }
483
-
484
- aggregator = prompt_aggregators[prompt_idx]
485
-
486
- # MEAN pooling with online weighted averaging
487
- # Ensure result is PoolingRequestOutput
488
- # for embedding processing
489
- if not isinstance(result, PoolingRequestOutput):
490
- return self.create_error_response(
491
- f"Expected PoolingRequestOutput for "
492
- f"chunked embedding, got "
493
- f"{type(result).__name__}"
494
- )
495
-
496
- # Handle both PoolingOutput and
497
- # EmbeddingOutput types
498
- if hasattr(result.outputs, "data"):
499
- # PoolingOutput case
500
- embedding_data = result.outputs.data
501
- elif hasattr(result.outputs, "embedding"):
502
- # EmbeddingOutput case -
503
- # convert embedding list to tensor
504
- embedding_data = result.outputs.embedding
505
- else:
506
- return self.create_error_response(
507
- f"Unsupported output type: {type(result.outputs).__name__}"
508
- )
509
-
510
- if not isinstance(embedding_data, torch.Tensor):
511
- embedding_data = torch.tensor(embedding_data, dtype=torch.float32)
512
-
513
- if result.prompt_token_ids is None:
514
- return self.create_error_response(
515
- "prompt_token_ids cannot be None for chunked processing"
516
- )
517
- weight = len(result.prompt_token_ids)
518
-
519
- weighted_embedding = embedding_data.to(dtype=torch.float32) * weight
520
-
521
- if aggregator["weighted_sum"] is None:
522
- # First chunk
523
- aggregator["weighted_sum"] = weighted_embedding
524
- else:
525
- # Accumulate
526
- aggregator["weighted_sum"] += weighted_embedding
527
-
528
- aggregator["total_weight"] += weight
529
- aggregator["chunk_count"] += 1
530
- else:
531
- # Non-chunked result - extract prompt_idx from request_id
532
- parts = result.request_id.split("-")
533
- try:
534
- # Last part should be prompt index
535
- prompt_idx = int(parts[-1])
536
- except (ValueError, IndexError):
537
- prompt_idx = result_idx # Fallback to result_idx
538
-
539
- short_prompts_results[prompt_idx] = result
540
-
541
- # Finalize aggregated results
542
- final_res_batch: list[PoolingRequestOutput] = []
543
- num_prompts = len(ctx.engine_prompts)
544
-
545
- for prompt_idx in range(num_prompts):
546
- if prompt_idx in prompt_aggregators:
547
- # Finalize MEAN aggregation for this chunked prompt
548
- aggregator = prompt_aggregators[prompt_idx]
549
-
550
- weighted_sum = aggregator["weighted_sum"]
551
- total_weight = aggregator["total_weight"]
552
-
553
- if (
554
- weighted_sum is not None
555
- and isinstance(weighted_sum, torch.Tensor)
556
- and isinstance(total_weight, (int, float))
557
- and total_weight > 0
558
- ):
559
- # Compute final mean embedding
560
- final_embedding = weighted_sum / total_weight
561
-
562
- # Create a PoolingRequestOutput
563
- # for the aggregated result
564
- pooling_output_data = PoolingOutput(data=final_embedding)
565
-
566
- # Get original prompt token IDs for this prompt
567
- original_prompt = ctx.engine_prompts[prompt_idx]
568
- if "prompt_token_ids" not in original_prompt:
569
- return self.create_error_response(
570
- f"Chunked prompt {prompt_idx} does not contain token IDs"
571
- )
572
-
573
- original_token_ids = original_prompt["prompt_token_ids"] # type: ignore[typeddict-item]
574
-
575
- pooling_request_output = PoolingRequestOutput(
576
- request_id=aggregator["request_id"],
577
- prompt_token_ids=original_token_ids,
578
- outputs=pooling_output_data,
579
- num_cached_tokens=0,
580
- finished=True,
581
- )
582
-
583
- final_res_batch.append(pooling_request_output)
584
- else:
585
- return self.create_error_response(
586
- f"Failed to aggregate chunks for prompt {prompt_idx}"
587
- )
588
- elif prompt_idx in short_prompts_results:
589
- final_res_batch.append(short_prompts_results[prompt_idx])
590
- else:
591
- return self.create_error_response(
592
- f"Result not found for prompt {prompt_idx}"
593
- )
594
-
595
- ctx.final_res_batch = final_res_batch
596
-
597
- return None
598
-
599
- async def create_embedding(
600
- self,
601
- request: EmbeddingRequest,
602
- raw_request: Request | None = None,
603
- ) -> EmbeddingResponse | ErrorResponse:
604
- """
605
- Embedding API similar to OpenAI's API.
606
-
607
- See https://platform.openai.com/docs/api-reference/embeddings/create
608
- for the API specification. This API mimics the OpenAI Embedding API.
609
- """
610
- model_name = self.models.model_name()
611
- request_id = (
612
- f"{self.request_id_prefix}-"
613
- f"{self._base_request_id(raw_request, request.request_id)}"
614
- )
615
-
616
- ctx = EmbeddingServeContext(
617
- request=request,
618
- raw_request=raw_request,
619
- model_name=model_name,
620
- request_id=request_id,
621
- )
622
-
623
- return await self.handle(ctx) # type: ignore[return-value]
vllm/entrypoints/pooling/io_processor_factories.py
+12
-8
@@ -15,17 +15,21 @@ def init_pooling_io_processors(
15
renderer: BaseRenderer,
16
chat_template_config: ChatTemplateConfig,
17
) -> dict[str, PoolingIOProcessor]:
18
- pooling_io_processors: dict[str, PoolingIOProcessor] = {}
19
-
18
+ processors: list[tuple[str, type[PoolingIOProcessor]]] = []
19
if "classify" in supported_tasks:
21
- from vllm.entrypoints.pooling.classify.io_processor import (
22
- ClassifyIOProcessor,
23
- )
20
+ from vllm.entrypoints.pooling.classify.io_processor import ClassifyIOProcessor
21
+
22
+ processors.append(("classify", ClassifyIOProcessor))
23
+ if "embed" in supported_tasks:
24
+ from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
25
25
- pooling_io_processors["classify"] = ClassifyIOProcessor(
26
+ processors.append(("classify", EmbedIOProcessor))
27
+
28
+ return {
29
+ task: processor_cls(
30
model_config=model_config,
31
renderer=renderer,
32
chat_template_config=chat_template_config,
33
)
30
-
31
- return pooling_io_processors
34
+ for task, processor_cls in processors
35
+ }
vllm/entrypoints/pooling/pooling/api_router.py
+1
-1
@@ -21,7 +21,7 @@ router = APIRouter()
21
22
23
def pooling(request: Request) -> OpenAIServingPooling | None:
24
- return request.app.state.openai_serving_pooling
24
+ return request.app.state.serving_pooling
25
26
27
@router.post(
vllm/entrypoints/pooling/score/api_router.py
+2
-2
@@ -24,11 +24,11 @@ logger = init_logger(__name__)
24
25
26
def score(request: Request) -> ServingScores | None:
27
- return request.app.state.openai_serving_scores
27
+ return request.app.state.serving_scores
28
29
30
def rerank(request: Request) -> ServingScores | None:
31
- return request.app.state.openai_serving_scores
31
+ return request.app.state.serving_scores
32
33
34
@router.post(
vllm/entrypoints/pooling/typing.py
+37
-4
@@ -1,8 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import time
4
+from collections.abc import AsyncGenerator
5
+from dataclasses import dataclass, field
6
+from typing import Any, Generic, TypeAlias, TypeVar
7
4
-from typing import TypeAlias
8
+from fastapi import Request
9
+from pydantic import ConfigDict
10
11
+from vllm import PoolingRequestOutput
12
from vllm.entrypoints.pooling.classify.protocol import (
13
ClassificationChatRequest,
14
ClassificationCompletionRequest,
@@ -25,12 +31,12 @@ from vllm.entrypoints.pooling.score.protocol import (
31
ScoreRequest,
32
ScoreResponse,
33
)
34
+from vllm.inputs import ProcessorInputs
35
+from vllm.lora.request import LoRARequest
36
37
PoolingCompletionLikeRequest: TypeAlias = (
38
EmbeddingCompletionRequest
39
| ClassificationCompletionRequest
32
- | RerankRequest
33
- | ScoreRequest
40
| PoolingCompletionRequest
41
)
42
@@ -39,7 +45,11 @@ PoolingChatLikeRequest: TypeAlias = (
45
)
46
47
AnyPoolingRequest: TypeAlias = (
42
- PoolingCompletionLikeRequest | PoolingChatLikeRequest | IOProcessorRequest
48
+ PoolingCompletionLikeRequest
49
+ | PoolingChatLikeRequest
50
+ | IOProcessorRequest
51
+ | RerankRequest
52
+ | ScoreRequest
53
)
54
55
AnyPoolingResponse: TypeAlias = (
@@ -49,3 +59,26 @@ AnyPoolingResponse: TypeAlias = (
59
| PoolingResponse
60
| ScoreResponse
61
)
62
+
63
+PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest)
64
+
65
+
66
+@dataclass(kw_only=True)
67
+class PoolingServeContext(Generic[PoolingRequestT]):
68
+ request: PoolingRequestT
69
+ raw_request: Request | None = None
70
+ model_name: str
71
+ request_id: str
72
+ created_time: int = field(default_factory=lambda: int(time.time()))
73
+ lora_request: LoRARequest | None = None
74
+
75
+ engine_prompts: list[ProcessorInputs] | None = None
76
+ prompt_request_ids: list[str] | None = None
77
+ intermediates: Any | None = None
78
+
79
+ result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = (
80
+ None
81
+ )
82
+ final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)
83
+
84
+ model_config = ConfigDict(arbitrary_types_allowed=True)
vllm/entrypoints/pooling/utils.py
+19
@@ -1,12 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import importlib.util
5
import math
6
from dataclasses import dataclass
7
+from functools import lru_cache
8
from typing import Any
9
10
import pybase64
11
import torch
12
+from fastapi.responses import JSONResponse
13
14
+from vllm.logger import init_logger
15
from vllm.outputs import PoolingRequestOutput
16
from vllm.utils.serial_utils import (
17
EMBED_DTYPES,
@@ -16,6 +21,8 @@ from vllm.utils.serial_utils import (
21
tensor2binary,
22
)
23
24
+logger = init_logger(__name__)
25
+
26
27
@dataclass
28
class MetadataItem:
@@ -122,3 +129,15 @@ def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.
129
)
130
for item in sorted(items, key=lambda x: x.index)
131
]
132
+
133
+
134
+@lru_cache(maxsize=1)
135
+def get_json_response_cls() -> type[JSONResponse]:
136
+ if importlib.util.find_spec("orjson") is not None:
137
+ from fastapi.responses import ORJSONResponse
138
+
139
+ return ORJSONResponse
140
+ logger.warning_once(
141
+ "To make v1/embeddings API fast, please install orjson by `pip install orjson`"
142
+ )
143
+ return JSONResponse
vllm/entrypoints/utils.py
+5
-1
@@ -303,12 +303,16 @@ def create_error_response(
303
if isinstance(message, Exception):
304
exc = message
305
306
- from vllm.exceptions import VLLMValidationError
306
+ from vllm.exceptions import VLLMNotFoundError, VLLMValidationError
307
308
if isinstance(exc, VLLMValidationError):
309
err_type = "BadRequestError"
310
status_code = HTTPStatus.BAD_REQUEST
311
param = exc.parameter
312
+ elif isinstance(exc, VLLMNotFoundError):
313
+ err_type = "NotFoundError"
314
+ status_code = HTTPStatus.NOT_FOUND
315
+ param = None
316
elif isinstance(exc, (ValueError, TypeError, OverflowError)):
317
# Common validation errors from user input
318
err_type = "BadRequestError"
vllm/exceptions.py
+6
@@ -34,3 +34,9 @@ class VLLMValidationError(ValueError):
34
if self.value is not None:
35
extras.append(f"value={self.value}")
36
return f"{base} ({', '.join(extras)})" if extras else base
37
+
38
+
39
+class VLLMNotFoundError(ValueError):
40
+ """vLLM-specific NotFoundError"""
41
+
42
+ pass