@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4e8c3f1c19
Fork 0

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 5, 02:15 UTC

[Frontend][last/5] Improve pooling entrypoints | clean up. (#39675)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

wang.yuqi committed Apr 16, 2026 at 22:53 UTC 4e8c3f1c197952454f244e5585a5de5990865939
29 files changed +465 -427
docs/models/pooling_models/README.md
+10
@@ -59,6 +59,16 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
59 Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
60 are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
61
62 +### Pooling Types
63 +
64 +| Pooling Tasks | Granularity | Description |
65 +|----------------|---------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
66 +| `CLS` pooling | Sequence-wise | For BERT‑like (bidirectional self‑attention) models, CLS pooling is used by default. This means the last_hidden_states corresponding to the first token (the [CLS] token) is taken as the output. |
67 +| `LAST` pooling | Sequence-wise | For GPT‑like (causal self‑attention) models, LAST pooling is used by default. This means the last_hidden_states corresponding to the last token is taken as the output. |
68 +| `MEAN` pooling | Sequence-wise | Many studies have shown that averaging the last_hidden_states over all input tokens performs better on certain downstream tasks. Therefore, more and more models are using MEAN pooling. |
69 +| `ALL` pooling | Token-wise | Outputs the last_hidden_states for all input tokens. |
70 +| `STEP` pooling | Token-wise | Filters and outputs the last_hidden_states corresponding to the token IDs returned by returned_token_ids. |
71 +
72 ### Score Types
73
74 The scoring models is designed to compute similarity scores between two input prompts. It supports three model types
docs/models/pooling_models/scoring.md
+4
@@ -160,6 +160,8 @@ The following Score API parameters are supported:
160 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
161 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
162 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
163 +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
164 +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:score-request-params"
165 ```
166
167 #### Examples
@@ -370,6 +372,8 @@ The following rerank api parameters are supported:
372 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
373 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
374 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
375 +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:scoring-common-params"
376 +--8<-- "vllm/entrypoints/pooling/scoring/protocol.py:rerank-request-params"
377 ```
378
379 #### Examples
docs/models/pooling_models/token_classify.md
+1 -1
@@ -68,7 +68,7 @@ If your model is not in the above list, we will try to automatically convert the
68 Forced alignment usage requires `--hf-overrides '{"architectures": ["Qwen3ASRForcedAlignerForTokenClassification"]}'`.
69 Please refer to [examples/pooling/token_classify/forced_alignment_offline.py](../../../examples/pooling/token_classify/forced_alignment_offline.py).
70
71 -### As Reward Models
71 +### Reward Models
72
73 Using token classification models as reward models. For details on reward models, see [Reward Models](reward.md).
74
docs/serving/openai_compatible_server.md
+1 -18
@@ -467,28 +467,11 @@ It consists of two endpoints:
467 - `/tokenize` corresponds to calling `tokenizer.encode()`.
468 - `/detokenize` corresponds to calling `tokenizer.decode()`.
469
470 -### Score API
471 -
472 -#### Score Template
473 -
474 -Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](#chat-template)).
475 -
476 -Score templates are supported for **cross-encoder** models only. If you are using an **embedding** model for scoring, vLLM does not apply a score template.
477 -
478 -Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter:
479 -
480 -- **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}`
481 -- **Document**: `{{ (messages | selectattr("role", "eq", "document") | first).content }}`
482 -
483 -This approach is more robust than index-based access (`messages[0]`, `messages[1]`) because it selects messages by their semantic role. It also avoids assumptions about message ordering if additional message types are added to `messages` in the future.
484 -
485 -Example template file: [examples/pooling/score/template/nemotron-rerank.jinja](../../examples/pooling/score/template/nemotron-rerank.jinja)
486 -
470 ### Generative Scoring API
471
472 The `/generative_scoring` endpoint uses a CausalLM model (e.g., Llama, Qwen, Mistral) to compute the probability of specified token IDs appearing as the next token. Each item (document) is concatenated with the query to form a prompt, and the model predicts how likely each label token is as the next token after that prompt. This lets you score items against a query — for example, asking "Is this the capital of France?" and scoring each city by how likely the model is to answer "Yes".
473
491 -This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
474 +This endpoint is automatically available when the server is started with a generative model (task `"generate"`). It is separate from the pooling-based [Score API](../models/pooling_models/scoring.md#score-api), which uses cross-encoder, bi-encoder, or late-interaction models.
475
476 **Requirements:**
477
vllm/entrypoints/llm.py
+1 -1
@@ -48,7 +48,7 @@ from vllm.entrypoints.chat_utils import (
48 ChatTemplateContentFormatOption,
49 load_chat_template,
50 )
51 -from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
51 +from vllm.entrypoints.pooling.factories import init_pooling_io_processors
52 from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessor
53 from vllm.entrypoints.pooling.scoring.typing import ScoreInput
54 from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext
vllm/entrypoints/openai/api_server.py
+14 -16
@@ -220,6 +220,12 @@ def build_app(
220
221 elastic_ep_attach_router(app)
222
223 + from vllm.entrypoints.openai.generative_scoring.api_router import (
224 + register_generative_scoring_api_router,
225 + )
226 +
227 + register_generative_scoring_api_router(app)
228 +
229 if "generate" in supported_tasks or "render" in supported_tasks:
230 from vllm.entrypoints.serve.render.api_router import (
231 attach_router as attach_render_router,
@@ -242,17 +248,10 @@ def build_app(
248 register_realtime_api_router(app)
249
250 if any(task in POOLING_TASKS for task in supported_tasks):
245 - from vllm.entrypoints.pooling import register_pooling_api_routers
251 + from vllm.entrypoints.pooling.factories import register_pooling_api_routers
252
253 register_pooling_api_routers(app, supported_tasks, model_config)
254
249 - if "generate" in supported_tasks:
250 - from vllm.entrypoints.openai.generative_scoring.api_router import (
251 - register_generative_scoring_api_router,
252 - )
253 -
254 - register_generative_scoring_api_router(app)
255 -
255 app.root_path = args.root_path
256 app.add_middleware(
257 CORSMiddleware,
@@ -401,6 +400,12 @@ async def init_app_state(
400 engine_client, state, args, request_logger, supported_tasks
401 )
402
403 + from vllm.entrypoints.openai.generative_scoring.api_router import (
404 + init_generative_scoring_state,
405 + )
406 +
407 + await init_generative_scoring_state(engine_client, state, args, request_logger)
408 +
409 if "transcription" in supported_tasks:
410 from vllm.entrypoints.openai.speech_to_text.api_router import (
411 init_transcription_state,
@@ -416,17 +421,10 @@ async def init_app_state(
421 init_realtime_state(engine_client, state, args, request_logger, supported_tasks)
422
423 if any(task in POOLING_TASKS for task in supported_tasks):
419 - from vllm.entrypoints.pooling import init_pooling_state
424 + from vllm.entrypoints.pooling.factories import init_pooling_state
425
426 init_pooling_state(engine_client, state, args, request_logger, supported_tasks)
427
423 - if "generate" in supported_tasks:
424 - from vllm.entrypoints.openai.generative_scoring.api_router import (
425 - init_generative_scoring_state,
426 - )
427 -
428 - await init_generative_scoring_state(engine_client, state, args, request_logger)
429 -
428 state.enable_server_load_tracking = args.enable_server_load_tracking
429 state.server_load_metrics = 0
430
vllm/entrypoints/openai/generate/factories.py new
+42
@@ -0,0 +1,42 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import TYPE_CHECKING
4 +
5 +from vllm.config import ModelConfig
6 +from vllm.tasks import SupportedTask
7 +
8 +if TYPE_CHECKING:
9 + from vllm.entrypoints.sagemaker.api_router import (
10 + EndpointFn,
11 + GetHandlerFn,
12 + RequestType,
13 + )
14 +
15 +
16 +def get_generate_invocation_types(
17 + supported_tasks: tuple["SupportedTask", ...],
18 + model_config: ModelConfig | None = None,
19 +):
20 + # NOTE: Items defined earlier take higher priority
21 + invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
22 +
23 + if "generate" in supported_tasks:
24 + from vllm.entrypoints.openai.chat_completion.api_router import (
25 + chat,
26 + create_chat_completion,
27 + )
28 + from vllm.entrypoints.openai.chat_completion.protocol import (
29 + ChatCompletionRequest,
30 + )
31 + from vllm.entrypoints.openai.completion.api_router import (
32 + completion,
33 + create_completion,
34 + )
35 + from vllm.entrypoints.openai.completion.protocol import CompletionRequest
36 +
37 + invocation_types += [
38 + (ChatCompletionRequest, (chat, create_chat_completion)),
39 + (CompletionRequest, (completion, create_completion)),
40 + ]
41 +
42 + return invocation_types
vllm/entrypoints/pooling/__init__.py
-130
@@ -1,130 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from typing import TYPE_CHECKING
5 -
6 -from fastapi import FastAPI
7 -
8 -from vllm.config import ModelConfig
9 -from vllm.entrypoints.pooling.utils import enable_scoring_api
10 -from vllm.logger import init_logger
11 -
12 -if TYPE_CHECKING:
13 - from argparse import Namespace
14 -
15 - from starlette.datastructures import State
16 -
17 - from vllm.engine.protocol import EngineClient
18 - from vllm.entrypoints.logger import RequestLogger
19 - from vllm.tasks import SupportedTask
20 -else:
21 - RequestLogger = object
22 - SupportedTask = object
23 -
24 -logger = init_logger(__name__)
25 -
26 -
27 -def register_pooling_api_routers(
28 - app: FastAPI,
29 - supported_tasks: tuple["SupportedTask", ...],
30 - model_config: ModelConfig | None = None,
31 -):
32 - if model_config is None:
33 - return
34 -
35 - pooling_task = model_config.get_pooling_task(supported_tasks)
36 -
37 - if pooling_task is not None:
38 - from vllm.entrypoints.pooling.pooling.api_router import router as pooling_router
39 -
40 - app.include_router(pooling_router)
41 -
42 - if "classify" in supported_tasks:
43 - from vllm.entrypoints.pooling.classify.api_router import (
44 - router as classify_router,
45 - )
46 -
47 - app.include_router(classify_router)
48 -
49 - if "embed" in supported_tasks:
50 - from vllm.entrypoints.pooling.embed.api_router import router as embed_router
51 -
52 - app.include_router(embed_router)
53 -
54 - if enable_scoring_api(supported_tasks, model_config):
55 - from vllm.entrypoints.pooling.scoring.api_router import router as score_router
56 -
57 - app.include_router(score_router)
58 -
59 -
60 -def init_pooling_state(
61 - engine_client: "EngineClient",
62 - state: "State",
63 - args: "Namespace",
64 - request_logger: RequestLogger | None,
65 - supported_tasks: tuple["SupportedTask", ...],
66 -):
67 - from vllm.entrypoints.chat_utils import load_chat_template
68 - from vllm.entrypoints.pooling.classify.serving import ServingClassification
69 - from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
70 - from vllm.entrypoints.pooling.pooling.serving import ServingPooling
71 - from vllm.entrypoints.pooling.scoring.serving import ServingScores
72 - from vllm.tasks import POOLING_TASKS
73 -
74 - model_config = engine_client.model_config
75 - resolved_chat_template = load_chat_template(args.chat_template)
76 -
77 - state.serving_pooling = (
78 - (
79 - ServingPooling(
80 - engine_client,
81 - state.openai_serving_models,
82 - supported_tasks=supported_tasks,
83 - request_logger=request_logger,
84 - chat_template=resolved_chat_template,
85 - chat_template_content_format=args.chat_template_content_format,
86 - trust_request_chat_template=args.trust_request_chat_template,
87 - )
88 - )
89 - if any(t in supported_tasks for t in POOLING_TASKS)
90 - else None
91 - )
92 - state.serving_embedding = (
93 - ServingEmbedding(
94 - engine_client,
95 - state.openai_serving_models,
96 - request_logger=request_logger,
97 - chat_template=resolved_chat_template,
98 - chat_template_content_format=args.chat_template_content_format,
99 - trust_request_chat_template=args.trust_request_chat_template,
100 - )
101 - if "embed" in supported_tasks
102 - else None
103 - )
104 - state.serving_classification = (
105 - ServingClassification(
106 - engine_client,
107 - state.openai_serving_models,
108 - request_logger=request_logger,
109 - chat_template=resolved_chat_template,
110 - chat_template_content_format=args.chat_template_content_format,
111 - trust_request_chat_template=args.trust_request_chat_template,
112 - )
113 - if "classify" in supported_tasks
114 - else None
115 - )
116 - state.serving_scores = (
117 - ServingScores(
118 - engine_client,
119 - state.openai_serving_models,
120 - request_logger=request_logger,
121 - chat_template=resolved_chat_template,
122 - chat_template_content_format=args.chat_template_content_format,
123 - trust_request_chat_template=args.trust_request_chat_template,
124 - enable_flash_late_interaction=getattr(
125 - args, "enable_flash_late_interaction", True
126 - ),
127 - )
128 - if enable_scoring_api(supported_tasks, model_config)
129 - else None
130 - )
vllm/entrypoints/pooling/base/io_processor.py
+15 -8
@@ -13,22 +13,29 @@ from vllm.entrypoints.chat_utils import (
13 ConversationMessage,
14 )
15 from vllm.entrypoints.openai.engine.serving import RendererChatRequest, RendererRequest
16 -from vllm.entrypoints.pooling.scoring.typing import ScoringData
17 -from vllm.entrypoints.pooling.typing import (
16 +from vllm.inputs import EngineInput, SingletonPrompt
17 +from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs
18 +from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
19 +from vllm.tool_parsers import ToolParser
20 +from vllm.utils.mistral import is_mistral_tokenizer
21 +
22 +from ..scoring.typing import ScoringData
23 +from ..typing import (
24 OfflineInputsContext,
25 OfflineOutputsContext,
26 PoolingChatLikeRequest,
27 PoolingCompletionLikeRequest,
28 PoolingServeContext,
29 )
24 -from vllm.inputs import EngineInput, SingletonPrompt
25 -from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs
26 -from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
27 -from vllm.tool_parsers import ToolParser
28 -from vllm.utils.mistral import is_mistral_tokenizer
30
31
32 class PoolingIOProcessor:
33 + """Processor for handling preprocessing & postprocessing ops for pooling requests.
34 +
35 + This class manages both online (serving) and offline (batch) processing of pooling
36 + requests, handling chat and completion formats.
37 + """
38 +
39 name: str
40
41 def __init__(
@@ -58,7 +65,7 @@ class PoolingIOProcessor:
65 def pre_process_online(self, ctx: PoolingServeContext):
66 request = ctx.request
67
61 - if isinstance(ctx.request, PoolingChatLikeRequest):
68 + if isinstance(request, PoolingChatLikeRequest):
69 self._validate_chat_template(
70 request_chat_template=request.chat_template,
71 chat_template_kwargs=request.chat_template_kwargs,
vllm/entrypoints/pooling/base/serving.py
+1 -1
@@ -22,7 +22,6 @@ from vllm.entrypoints.chat_utils import (
22 from vllm.entrypoints.logger import RequestLogger
23 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
24 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
25 -from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext
25 from vllm.exceptions import VLLMNotFoundError
26 from vllm.inputs import EngineInput
27 from vllm.lora.request import LoRARequest
@@ -36,6 +35,7 @@ from vllm.tracing import (
35 from vllm.utils import random_uuid
36 from vllm.utils.async_utils import make_async, merge_async_iterators
37
38 +from ..typing import AnyPoolingRequest, PoolingServeContext
39 from .io_processor import PoolingIOProcessor
40
41
vllm/entrypoints/pooling/classify/api_router.py
+3 -2
@@ -5,13 +5,14 @@ from fastapi import APIRouter, Depends, Request
5 from fastapi.responses import Response
6
7 from vllm.entrypoints.openai.utils import validate_json_request
8 -from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest
9 -from vllm.entrypoints.pooling.classify.serving import ServingClassification
8 from vllm.entrypoints.utils import (
9 load_aware_call,
10 with_cancellation,
11 )
12
13 +from .protocol import ClassificationRequest
14 +from .serving import ServingClassification
15 +
16 router = APIRouter()
17
18
vllm/entrypoints/pooling/classify/io_processor.py
+1 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
4 +from ..base.io_processor import PoolingIOProcessor
5
6
7 class ClassifyIOProcessor(PoolingIOProcessor):
vllm/entrypoints/pooling/classify/protocol.py
+5 -4
@@ -9,15 +9,16 @@ from pydantic import Field
9 from vllm import PoolingParams
10 from vllm.config import ModelConfig
11 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
12 -from vllm.entrypoints.pooling.base.protocol import (
12 +from vllm.logger import init_logger
13 +from vllm.renderers import TokenizeParams
14 +from vllm.utils import random_uuid
15 +
16 +from ..base.protocol import (
17 ChatRequestMixin,
18 ClassifyRequestMixin,
19 CompletionRequestMixin,
20 PoolingBasicRequestMixin,
21 )
18 -from vllm.logger import init_logger
19 -from vllm.renderers import TokenizeParams
20 -from vllm.utils import random_uuid
22
23 logger = init_logger(__name__)
24
vllm/entrypoints/pooling/classify/serving.py
+2 -2
@@ -7,11 +7,11 @@ import numpy as np
7 from fastapi.responses import JSONResponse
8
9 from vllm.entrypoints.openai.engine.protocol import UsageInfo
10 -from vllm.entrypoints.pooling.base.serving import PoolingServing
11 -from vllm.entrypoints.pooling.typing import PoolingServeContext
10 from vllm.logger import init_logger
11 from vllm.outputs import ClassificationOutput
12
13 +from ..base.serving import PoolingServing
14 +from ..typing import PoolingServeContext
15 from .io_processor import ClassifyIOProcessor
16 from .protocol import (
17 ClassificationData,
vllm/entrypoints/pooling/embed/api_router.py
+3 -5
@@ -7,13 +7,11 @@ from fastapi import APIRouter, Depends, Request
7
8 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
9 from vllm.entrypoints.openai.utils import validate_json_request
10 -from vllm.entrypoints.pooling.embed.protocol import (
11 - CohereEmbedRequest,
12 - EmbeddingRequest,
13 -)
14 -from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
10 from vllm.entrypoints.utils import load_aware_call, with_cancellation
11
12 +from .protocol import CohereEmbedRequest, EmbeddingRequest
13 +from .serving import ServingEmbedding
14 +
15 router = APIRouter()
16
17
vllm/entrypoints/pooling/embed/io_processor.py
+20 -19
@@ -16,21 +16,6 @@ from vllm.entrypoints.chat_utils import (
16 ChatCompletionMessageParam,
17 CustomChatCompletionMessageParam,
18 )
19 -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
20 -from vllm.entrypoints.pooling.embed.protocol import (
21 - CohereEmbedContent,
22 - CohereEmbedInput,
23 - CohereEmbedRequest,
24 - EmbeddingChatRequest,
25 - EmbeddingCompletionRequest,
26 -)
27 -from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessorMixin
28 -from vllm.entrypoints.pooling.typing import (
29 - OfflineInputsContext,
30 - PoolingChatLikeRequest,
31 - PoolingCompletionLikeRequest,
32 - PoolingServeContext,
33 -)
19 from vllm.inputs import EngineInput, tokens_input
20 from vllm.logger import init_logger
21 from vllm.outputs import PoolingOutput, PoolingRequestOutput
@@ -39,6 +24,22 @@ from vllm.renderers.hf import resolve_chat_template
24 from vllm.utils.collection_utils import chunk_list
25 from vllm.utils.mistral import is_mistral_tokenizer
26
27 +from ..base.io_processor import PoolingIOProcessor
28 +from ..scoring.io_processor import JinaRankingIOProcessorMixin
29 +from ..typing import (
30 + OfflineInputsContext,
31 + PoolingChatLikeRequest,
32 + PoolingCompletionLikeRequest,
33 + PoolingServeContext,
34 +)
35 +from .protocol import (
36 + CohereEmbedContent,
37 + CohereEmbedInput,
38 + CohereEmbedRequest,
39 + EmbeddingChatRequest,
40 + EmbeddingCompletionRequest,
41 +)
42 +
43 logger = init_logger(__name__)
44
45
@@ -94,7 +95,7 @@ class EmbedIOProcessor(PoolingIOProcessor):
95 if ctx.engine_inputs is None:
96 raise ValueError("Engine prompts not available")
97
97 - ctx.intermediates = ctx.engine_inputs
98 + ctx.original_engine_inputs = ctx.engine_inputs
99 request_id = ctx.request_id
100 max_model_len = self.model_config.max_model_len
101 chunked_engine_inputs: list[EngineInput] = []
@@ -189,10 +190,10 @@ class EmbedIOProcessor(PoolingIOProcessor):
190 aggregator["total_weight"] += weight
191 aggregator["chunk_count"] += 1
192
192 - if ctx.intermediates is None:
193 - raise ValueError("Original prompts inputs not available")
193 + if ctx.original_engine_inputs is None:
194 + raise ValueError("Original engine inputs not available")
195
195 - original_engine_inputs = cast(list[EngineInput], ctx.intermediates)
196 + original_engine_inputs = ctx.original_engine_inputs
197 num_prompts = len(original_engine_inputs)
198
199 # Finalize aggregated results
vllm/entrypoints/pooling/embed/protocol.py
+4 -3
@@ -18,14 +18,15 @@ from pydantic import BaseModel, Field
18 from vllm import PoolingParams
19 from vllm.config import ModelConfig
20 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
21 -from vllm.entrypoints.pooling.base.protocol import (
21 +from vllm.renderers import TokenizeParams
22 +from vllm.utils import random_uuid
23 +
24 +from ..base.protocol import (
25 ChatRequestMixin,
26 CompletionRequestMixin,
27 EmbedRequestMixin,
28 PoolingBasicRequestMixin,
29 )
27 -from vllm.renderers import TokenizeParams
28 -from vllm.utils import random_uuid
30
31 # ---------------------------------------------------------------------------
32 # OpenAI /v1/embeddings — request models
vllm/entrypoints/pooling/embed/serving.py
+14 -13
@@ -9,9 +9,20 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse
9 from typing_extensions import assert_never
10
11 from vllm.entrypoints.openai.engine.protocol import UsageInfo
12 -from vllm.entrypoints.pooling.base.serving import PoolingServing
13 -from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
14 -from vllm.entrypoints.pooling.embed.protocol import (
12 +from vllm.logger import init_logger
13 +from vllm.outputs import PoolingRequestOutput
14 +from vllm.utils.serial_utils import EmbedDType, Endianness
15 +
16 +from ..base.serving import PoolingServing
17 +from ..typing import PoolingServeContext
18 +from ..utils import (
19 + encode_pooling_bytes,
20 + encode_pooling_output_base64,
21 + encode_pooling_output_float,
22 + get_json_response_cls,
23 +)
24 +from .io_processor import EmbedIOProcessor
25 +from .protocol import (
26 CohereBilledUnits,
27 CohereEmbedRequest,
28 CohereEmbedResponse,
@@ -22,16 +33,6 @@ from vllm.entrypoints.pooling.embed.protocol import (
33 EmbeddingResponseData,
34 build_typed_embeddings,
35 )
25 -from vllm.entrypoints.pooling.typing import PoolingServeContext
26 -from vllm.entrypoints.pooling.utils import (
27 - encode_pooling_bytes,
28 - encode_pooling_output_base64,
29 - encode_pooling_output_float,
30 - get_json_response_cls,
31 -)
32 -from vllm.logger import init_logger
33 -from vllm.outputs import PoolingRequestOutput
34 -from vllm.utils.serial_utils import EmbedDType, Endianness
36
37 logger = init_logger(__name__)
38
vllm/entrypoints/pooling/factories.py new
+256
@@ -0,0 +1,256 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import TYPE_CHECKING
5 +
6 +from fastapi import FastAPI
7 +
8 +from vllm.config import ModelConfig, VllmConfig
9 +from vllm.entrypoints.chat_utils import ChatTemplateConfig
10 +from vllm.logger import init_logger
11 +from vllm.plugins.io_processors import has_io_processor
12 +from vllm.renderers import BaseRenderer
13 +from vllm.tasks import POOLING_TASKS, SupportedTask
14 +
15 +from .base.io_processor import PoolingIOProcessor
16 +from .utils import enable_scoring_api
17 +
18 +if TYPE_CHECKING:
19 + from argparse import Namespace
20 +
21 + from starlette.datastructures import State
22 +
23 + from vllm.engine.protocol import EngineClient
24 + from vllm.entrypoints.logger import RequestLogger
25 + from vllm.entrypoints.sagemaker.api_router import (
26 + EndpointFn,
27 + GetHandlerFn,
28 + RequestType,
29 + )
30 +
31 +else:
32 + RequestLogger = object
33 +
34 +
35 +logger = init_logger(__name__)
36 +
37 +
38 +def init_pooling_io_processors(
39 + supported_tasks: tuple[SupportedTask, ...],
40 + vllm_config: VllmConfig,
41 + renderer: BaseRenderer,
42 + chat_template_config: ChatTemplateConfig,
43 +) -> dict[str, PoolingIOProcessor]:
44 + model_config = vllm_config.model_config
45 + processors: dict[str, type[PoolingIOProcessor]] = {}
46 +
47 + if "classify" in supported_tasks:
48 + from .classify.io_processor import ClassifyIOProcessor
49 +
50 + processors["classify"] = ClassifyIOProcessor
51 +
52 + if "token_classify" in supported_tasks:
53 + from .classify.io_processor import TokenClassifyIOProcessor
54 +
55 + processors["token_classify"] = TokenClassifyIOProcessor
56 +
57 + if "embed" in supported_tasks:
58 + from .embed.io_processor import EmbedIOProcessor
59 +
60 + processors["embed"] = EmbedIOProcessor
61 +
62 + if "token_embed" in supported_tasks:
63 + from .embed.io_processor import TokenEmbedIOProcessor
64 +
65 + processors["token_embed"] = TokenEmbedIOProcessor
66 +
67 + if has_io_processor(
68 + vllm_config,
69 + model_config.io_processor_plugin,
70 + ):
71 + from .pooling.io_processor import PluginWithIOProcessorPlugins
72 +
73 + processors["plugin"] = PluginWithIOProcessorPlugins
74 + elif "plugin" in supported_tasks:
75 + from .pooling.io_processor import PluginWithoutIOProcessorPlugins
76 +
77 + processors["plugin"] = PluginWithoutIOProcessorPlugins
78 +
79 + if enable_scoring_api(supported_tasks, model_config):
80 + score_type = model_config.score_type
81 + from .scoring.io_processor import ScoringIOProcessors
82 +
83 + if score_type is not None and score_type in ScoringIOProcessors:
84 + processors[score_type] = ScoringIOProcessors[score_type]
85 +
86 + if model_config.architecture == "JinaForRanking":
87 + from .embed.io_processor import JinaRankingTokenEmbedIOProcessor
88 + from .scoring.io_processor import ScoringIOProcessors
89 +
90 + processors["token_embed"] = JinaRankingTokenEmbedIOProcessor
91 + processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"]
92 +
93 + return {
94 + task: processor_cls(
95 + vllm_config=vllm_config,
96 + renderer=renderer,
97 + chat_template_config=chat_template_config,
98 + )
99 + for task, processor_cls in processors.items()
100 + }
101 +
102 +
103 +def register_pooling_api_routers(
104 + app: FastAPI,
105 + supported_tasks: tuple["SupportedTask", ...],
106 + model_config: ModelConfig | None = None,
107 +):
108 + if model_config is None:
109 + return
110 +
111 + pooling_task = model_config.get_pooling_task(supported_tasks)
112 +
113 + if pooling_task is not None:
114 + from .pooling.api_router import router as pooling_router
115 +
116 + app.include_router(pooling_router)
117 +
118 + if "classify" in supported_tasks:
119 + from .classify.api_router import (
120 + router as classify_router,
121 + )
122 +
123 + app.include_router(classify_router)
124 +
125 + if "embed" in supported_tasks:
126 + from .embed.api_router import router as embed_router
127 +
128 + app.include_router(embed_router)
129 +
130 + if enable_scoring_api(supported_tasks, model_config):
131 + from .scoring.api_router import router as score_router
132 +
133 + app.include_router(score_router)
134 +
135 +
136 +def init_pooling_state(
137 + engine_client: "EngineClient",
138 + state: "State",
139 + args: "Namespace",
140 + request_logger: RequestLogger | None,
141 + supported_tasks: tuple["SupportedTask", ...],
142 +):
143 + from vllm.entrypoints.chat_utils import load_chat_template
144 + from vllm.tasks import POOLING_TASKS
145 +
146 + from .classify.serving import ServingClassification
147 + from .embed.serving import ServingEmbedding
148 + from .pooling.serving import ServingPooling
149 + from .scoring.serving import ServingScores
150 +
151 + model_config = engine_client.model_config
152 + resolved_chat_template = load_chat_template(args.chat_template)
153 +
154 + state.serving_pooling = (
155 + (
156 + ServingPooling(
157 + engine_client,
158 + state.openai_serving_models,
159 + supported_tasks=supported_tasks,
160 + request_logger=request_logger,
161 + chat_template=resolved_chat_template,
162 + chat_template_content_format=args.chat_template_content_format,
163 + trust_request_chat_template=args.trust_request_chat_template,
164 + )
165 + )
166 + if any(t in supported_tasks for t in POOLING_TASKS)
167 + else None
168 + )
169 + state.serving_embedding = (
170 + ServingEmbedding(
171 + engine_client,
172 + state.openai_serving_models,
173 + request_logger=request_logger,
174 + chat_template=resolved_chat_template,
175 + chat_template_content_format=args.chat_template_content_format,
176 + trust_request_chat_template=args.trust_request_chat_template,
177 + )
178 + if "embed" in supported_tasks
179 + else None
180 + )
181 + state.serving_classification = (
182 + ServingClassification(
183 + engine_client,
184 + state.openai_serving_models,
185 + request_logger=request_logger,
186 + chat_template=resolved_chat_template,
187 + chat_template_content_format=args.chat_template_content_format,
188 + trust_request_chat_template=args.trust_request_chat_template,
189 + )
190 + if "classify" in supported_tasks
191 + else None
192 + )
193 + state.serving_scores = (
194 + ServingScores(
195 + engine_client,
196 + state.openai_serving_models,
197 + request_logger=request_logger,
198 + chat_template=resolved_chat_template,
199 + chat_template_content_format=args.chat_template_content_format,
200 + trust_request_chat_template=args.trust_request_chat_template,
201 + enable_flash_late_interaction=getattr(
202 + args, "enable_flash_late_interaction", True
203 + ),
204 + )
205 + if enable_scoring_api(supported_tasks, model_config)
206 + else None
207 + )
208 +
209 +
210 +def get_pooling_invocation_types(
211 + supported_tasks: tuple["SupportedTask", ...],
212 + model_config: ModelConfig | None = None,
213 +):
214 + # NOTE: Items defined earlier take higher priority
215 + invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
216 +
217 + if "embed" in supported_tasks:
218 + from .embed.api_router import create_embedding, embedding
219 + from .embed.protocol import EmbeddingRequest
220 +
221 + invocation_types += [
222 + (EmbeddingRequest, (embedding, create_embedding)),
223 + ]
224 +
225 + if "classify" in supported_tasks:
226 + from .classify.api_router import classify, create_classify
227 + from .classify.protocol import ClassificationRequest
228 +
229 + invocation_types += [
230 + (ClassificationRequest, (classify, create_classify)),
231 + ]
232 +
233 + if enable_scoring_api(supported_tasks, model_config):
234 + from .scoring.api_router import do_rerank, rerank
235 + from .scoring.protocol import RerankRequest
236 +
237 + invocation_types += [
238 + (RerankRequest, (rerank, do_rerank)),
239 + ]
240 +
241 + from .scoring.api_router import create_score, score
242 + from .scoring.protocol import ScoreRequest
243 +
244 + invocation_types += [
245 + (ScoreRequest, (score, create_score)),
246 + ]
247 +
248 + if any(task in POOLING_TASKS for task in supported_tasks):
249 + from .pooling.api_router import create_pooling, pooling
250 + from .pooling.protocol import PoolingRequest
251 +
252 + invocation_types += [
253 + (PoolingRequest, (pooling, create_pooling)),
254 + ]
255 +
256 + return invocation_types
vllm/entrypoints/pooling/io_processor_factories.py deleted
-76
@@ -1,76 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -from vllm.config import VllmConfig
5 -from vllm.entrypoints.chat_utils import ChatTemplateConfig
6 -from vllm.plugins.io_processors import has_io_processor
7 -from vllm.renderers import BaseRenderer
8 -from vllm.tasks import SupportedTask
9 -
10 -from .base.io_processor import PoolingIOProcessor
11 -from .utils import enable_scoring_api
12 -
13 -
14 -def init_pooling_io_processors(
15 - supported_tasks: tuple[SupportedTask, ...],
16 - vllm_config: VllmConfig,
17 - renderer: BaseRenderer,
18 - chat_template_config: ChatTemplateConfig,
19 -) -> dict[str, PoolingIOProcessor]:
20 - model_config = vllm_config.model_config
21 - processors: dict[str, type[PoolingIOProcessor]] = {}
22 -
23 - if "classify" in supported_tasks:
24 - from .classify.io_processor import ClassifyIOProcessor
25 -
26 - processors["classify"] = ClassifyIOProcessor
27 -
28 - if "token_classify" in supported_tasks:
29 - from .classify.io_processor import TokenClassifyIOProcessor
30 -
31 - processors["token_classify"] = TokenClassifyIOProcessor
32 -
33 - if "embed" in supported_tasks:
34 - from .embed.io_processor import EmbedIOProcessor
35 -
36 - processors["embed"] = EmbedIOProcessor
37 -
38 - if "token_embed" in supported_tasks:
39 - from .embed.io_processor import TokenEmbedIOProcessor
40 -
41 - processors["token_embed"] = TokenEmbedIOProcessor
42 -
43 - if has_io_processor(
44 - vllm_config,
45 - model_config.io_processor_plugin,
46 - ):
47 - from .pooling.io_processor import PluginWithIOProcessorPlugins
48 -
49 - processors["plugin"] = PluginWithIOProcessorPlugins
50 - elif "plugin" in supported_tasks:
51 - from .pooling.io_processor import PluginWithoutIOProcessorPlugins
52 -
53 - processors["plugin"] = PluginWithoutIOProcessorPlugins
54 -
55 - if enable_scoring_api(supported_tasks, model_config):
56 - score_type = model_config.score_type
57 - from .scoring.io_processor import ScoringIOProcessors
58 -
59 - if score_type is not None and score_type in ScoringIOProcessors:
60 - processors[score_type] = ScoringIOProcessors[score_type]
61 -
62 - if model_config.architecture == "JinaForRanking":
63 - from .embed.io_processor import JinaRankingTokenEmbedIOProcessor
64 - from .scoring.io_processor import ScoringIOProcessors
65 -
66 - processors["token_embed"] = JinaRankingTokenEmbedIOProcessor
67 - processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"]
68 -
69 - return {
70 - task: processor_cls(
71 - vllm_config=vllm_config,
72 - renderer=renderer,
73 - chat_template_config=chat_template_config,
74 - )
75 - for task, processor_cls in processors.items()
76 - }
vllm/entrypoints/pooling/pooling/api_router.py
+3 -2
@@ -6,10 +6,11 @@ from fastapi import APIRouter, Depends, Request
6
7 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
8 from vllm.entrypoints.openai.utils import validate_json_request
9 -from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest
10 -from vllm.entrypoints.pooling.pooling.serving import ServingPooling
9 from vllm.entrypoints.utils import load_aware_call, with_cancellation
10
11 +from .protocol import PoolingRequest
12 +from .serving import ServingPooling
13 +
14 router = APIRouter()
15
16
vllm/entrypoints/pooling/pooling/io_processor.py
+3 -1
@@ -4,12 +4,12 @@ from collections.abc import Sequence
4 from typing import Any
5
6 from vllm import PoolingParams, PoolingRequestOutput
7 -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
7 from vllm.inputs import EngineInput
8 from vllm.logger import init_logger
9 from vllm.plugins.io_processors import get_io_processor
10 from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
11
12 +from ..base.io_processor import PoolingIOProcessor
13 from ..typing import OfflineInputsContext, OfflineOutputsContext, PoolingServeContext
14 from .protocol import IOProcessorRequest, IOProcessorResponse
15
@@ -17,6 +17,8 @@ logger = init_logger(__name__)
17
18
19 class PluginWithoutIOProcessorPlugins(PoolingIOProcessor):
20 + # Some models, such as Terratorch (tests/models/test_terratorch.py),
21 + # use plugin tasks in the pooler but do not use IO Processor plugins.
22 name = "plugin"
23
24
vllm/entrypoints/pooling/pooling/protocol.py
+5 -4
@@ -8,7 +8,11 @@ from pydantic import Field
8 from vllm import PoolingParams
9 from vllm.config import ModelConfig
10 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
11 -from vllm.entrypoints.pooling.base.protocol import (
11 +from vllm.renderers import TokenizeParams
12 +from vllm.tasks import PoolingTask
13 +from vllm.utils import random_uuid
14 +
15 +from ..base.protocol import (
16 ChatRequestMixin,
17 ClassifyRequestMixin,
18 CompletionRequestMixin,
@@ -16,9 +20,6 @@ from vllm.entrypoints.pooling.base.protocol import (
20 EncodingRequestMixin,
21 PoolingBasicRequestMixin,
22 )
19 -from vllm.renderers import TokenizeParams
20 -from vllm.tasks import PoolingTask
21 -from vllm.utils import random_uuid
23
24
25 class PoolingCompletionRequest(
vllm/entrypoints/pooling/pooling/serving.py
+16 -15
@@ -11,27 +11,28 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse
11 from typing_extensions import assert_never
12
13 from vllm.entrypoints.openai.engine.protocol import UsageInfo
14 -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
15 -from vllm.entrypoints.pooling.base.serving import PoolingServingBase
16 -from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
17 -from vllm.entrypoints.pooling.pooling.protocol import (
14 +from vllm.logger import init_logger
15 +from vllm.outputs import PoolingRequestOutput
16 +from vllm.tasks import SupportedTask
17 +from vllm.utils.serial_utils import EmbedDType, Endianness
18 +
19 +from ..base.io_processor import PoolingIOProcessor
20 +from ..base.serving import PoolingServingBase
21 +from ..factories import init_pooling_io_processors
22 +from ..typing import AnyPoolingRequest, PoolingServeContext
23 +from ..utils import (
24 + encode_pooling_bytes,
25 + encode_pooling_output_base64,
26 + encode_pooling_output_float,
27 + get_json_response_cls,
28 +)
29 +from .protocol import (
30 IOProcessorRequest,
31 PoolingBytesResponse,
32 PoolingRequest,
33 PoolingResponse,
34 PoolingResponseData,
35 )
24 -from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext
25 -from vllm.entrypoints.pooling.utils import (
26 - encode_pooling_bytes,
27 - encode_pooling_output_base64,
28 - encode_pooling_output_float,
29 - get_json_response_cls,
30 -)
31 -from vllm.logger import init_logger
32 -from vllm.outputs import PoolingRequestOutput
33 -from vllm.tasks import SupportedTask
34 -from vllm.utils.serial_utils import EmbedDType, Endianness
36
37 logger = init_logger(__name__)
38
vllm/entrypoints/pooling/scoring/io_processor.py
+6 -6
@@ -7,12 +7,6 @@ from typing import Any, TypeAlias
7 import torch.nn.functional as F
8
9 from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt
10 -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
11 -from vllm.entrypoints.pooling.typing import (
12 - OfflineInputsContext,
13 - OfflineOutputsContext,
14 - PoolingServeContext,
15 -)
10 from vllm.inputs import EngineInput
11 from vllm.renderers import TokenizeParams
12 from vllm.renderers.hf import safe_apply_chat_template
@@ -20,6 +14,12 @@ from vllm.tasks import PoolingTask
14 from vllm.utils.mistral import is_mistral_tokenizer
15
16 from ...chat_utils import ChatTemplateResolutionError
17 +from ..base.io_processor import PoolingIOProcessor
18 +from ..typing import (
19 + OfflineInputsContext,
20 + OfflineOutputsContext,
21 + PoolingServeContext,
22 +)
23 from .protocol import RerankRequest, ScoreRequest, ScoringRequest
24 from .typing import ScoreData, ScoreInput, ScoringData
25 from .utils import (
vllm/entrypoints/pooling/scoring/protocol.py
+13 -10
@@ -8,18 +8,16 @@ from pydantic import BaseModel, Field
8 from vllm import PoolingParams
9 from vllm.config import ModelConfig
10 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
11 -from vllm.entrypoints.pooling.base.protocol import (
12 - ClassifyRequestMixin,
13 - PoolingBasicRequestMixin,
14 -)
11 from vllm.renderers import TokenizeParams
12 from vllm.tasks import PoolingTask
13 from vllm.utils import random_uuid
14
15 +from ..base.protocol import ClassifyRequestMixin, PoolingBasicRequestMixin
16 from .typing import ScoreContentPartParam, ScoreInput
17
18
22 -class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
19 +class ScoringRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
20 + # --8<-- [start:scoring-common-params]
21 max_tokens_per_query: int = Field(
22 default=0,
23 description=(
@@ -37,6 +35,7 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
35 "applies to the combined query+document)."
36 ),
37 )
38 + # --8<-- [end:scoring-common-params]
39
40 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
41 encoder_config = model_config.encoder_config or {}
@@ -57,14 +56,16 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
56 )
57
58
60 -class ScoreDataRequest(ScoreRequestMixin):
59 +class ScoreDataRequest(ScoringRequestMixin):
60 data_1: ScoreInput | list[ScoreInput]
61 data_2: ScoreInput | list[ScoreInput]
62
63
65 -class ScoreQueriesDocumentsRequest(ScoreRequestMixin):
64 +class ScoreQueriesDocumentsRequest(ScoringRequestMixin):
65 + # --8<-- [start:score-request-params]
66 queries: ScoreInput | list[ScoreInput]
67 documents: ScoreInput | list[ScoreInput]
68 + # --8<-- [end:score-request-params]
69
70 @property
71 def data_1(self):
@@ -75,7 +76,7 @@ class ScoreQueriesDocumentsRequest(ScoreRequestMixin):
76 return self.documents
77
78
78 -class ScoreQueriesItemsRequest(ScoreRequestMixin):
79 +class ScoreQueriesItemsRequest(ScoringRequestMixin):
80 queries: ScoreInput | list[ScoreInput]
81 items: ScoreInput | list[ScoreInput]
82
@@ -88,7 +89,7 @@ class ScoreQueriesItemsRequest(ScoreRequestMixin):
89 return self.items
90
91
91 -class ScoreTextRequest(ScoreRequestMixin):
92 +class ScoreTextRequest(ScoringRequestMixin):
93 text_1: ScoreInput | list[ScoreInput]
94 text_2: ScoreInput | list[ScoreInput]
95
@@ -109,10 +110,12 @@ ScoreRequest: TypeAlias = (
110 )
111
112
112 -class RerankRequest(ScoreRequestMixin):
113 +class RerankRequest(ScoringRequestMixin):
114 + # --8<-- [start:rerank-request-params]
115 query: ScoreInput
116 documents: ScoreInput | list[ScoreInput]
117 top_n: int = Field(default_factory=lambda: 0)
118 + # --8<-- [end:rerank-request-params]
119
120
121 ScoringRequest: TypeAlias = ScoreRequest | RerankRequest
vllm/entrypoints/pooling/scoring/serving.py
+3 -3
@@ -6,8 +6,6 @@ from fastapi.responses import JSONResponse, Response
6 from vllm import PoolingParams
7 from vllm.engine.protocol import EngineClient
8 from vllm.entrypoints.openai.engine.protocol import UsageInfo
9 -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
10 -from vllm.entrypoints.pooling.base.serving import PoolingServing
9 from vllm.logger import init_logger
10 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
11 from vllm.v1.pool.late_interaction import (
@@ -15,6 +13,8 @@ from vllm.v1.pool.late_interaction import (
13 build_late_interaction_query_params,
14 )
15
16 +from ..base.io_processor import PoolingIOProcessor
17 +from ..base.serving import PoolingServing
18 from .io_processor import ScoringIOProcessors, ScoringServeContext
19 from .protocol import (
20 RerankDocument,
@@ -90,7 +90,7 @@ class ServingScores(PoolingServing):
90 ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch),
91 )
92 else:
93 - raise NotImplementedError("")
93 + raise ValueError(f"Invalid {self.request_id_prefix} request type")
94
95 def _request_output_to_score_response(
96 self,
vllm/entrypoints/pooling/typing.py
+12 -9
@@ -9,29 +9,30 @@ from fastapi import Request
9 from pydantic import ConfigDict
10
11 from vllm import PoolingParams, PoolingRequestOutput, PromptType
12 -from vllm.entrypoints.pooling.classify.protocol import (
12 +from vllm.inputs import DataPrompt, EngineInput
13 +from vllm.lora.request import LoRARequest
14 +
15 +from .classify.protocol import (
16 ClassificationChatRequest,
17 ClassificationCompletionRequest,
18 ClassificationResponse,
19 )
17 -from vllm.entrypoints.pooling.embed.protocol import (
20 +from .embed.protocol import (
21 CohereEmbedRequest,
22 EmbeddingBytesResponse,
23 EmbeddingChatRequest,
24 EmbeddingCompletionRequest,
25 EmbeddingResponse,
26 )
24 -from vllm.entrypoints.pooling.pooling.protocol import (
27 +from .pooling.protocol import (
28 IOProcessorRequest,
29 PoolingBytesResponse,
30 PoolingChatRequest,
31 PoolingCompletionRequest,
32 PoolingResponse,
33 )
31 -from vllm.entrypoints.pooling.scoring.protocol import ScoringRequest, ScoringResponse
32 -from vllm.entrypoints.pooling.scoring.typing import ScoringData
33 -from vllm.inputs import DataPrompt, EngineInput
34 -from vllm.lora.request import LoRARequest
34 +from .scoring.protocol import ScoringRequest, ScoringResponse
35 +from .scoring.typing import ScoringData
36
37 PoolingCompletionLikeRequest: TypeAlias = (
38 EmbeddingCompletionRequest
@@ -65,6 +66,8 @@ PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest)
66
67 @dataclass(kw_only=True)
68 class PoolingServeContext(Generic[PoolingRequestT]):
69 + model_config = ConfigDict(arbitrary_types_allowed=True)
70 +
71 request: PoolingRequestT
72 raw_request: Request | None = None
73 model_name: str
@@ -74,14 +77,14 @@ class PoolingServeContext(Generic[PoolingRequestT]):
77 lora_request: LoRARequest | None = None
78 engine_inputs: Sequence[EngineInput] | None = None
79 prompt_request_ids: list[str] | None = None
77 - intermediates: Any | None = None
80
81 result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = (
82 None
83 )
84 final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)
85
84 - model_config = ConfigDict(arbitrary_types_allowed=True)
86 + ## for Long Text Embedding with Chunked Processing
87 + original_engine_inputs: Sequence[EngineInput] | None = None
88
89 ## for bi-encoder & late-interaction
90 n_queries: int | None = None
vllm/entrypoints/sagemaker/api_router.py
+7 -77
@@ -13,12 +13,13 @@ from fastapi.responses import JSONResponse, Response
13 from vllm.config import ModelConfig
14 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
15 from vllm.entrypoints.openai.engine.serving import OpenAIServing
16 +from vllm.entrypoints.openai.generate.factories import get_generate_invocation_types
17 from vllm.entrypoints.openai.utils import validate_json_request
18 from vllm.entrypoints.pooling.base.serving import PoolingServingBase
18 -from vllm.entrypoints.pooling.utils import enable_scoring_api
19 +from vllm.entrypoints.pooling.factories import get_pooling_invocation_types
20 from vllm.entrypoints.serve.instrumentator.basic import base
21 from vllm.entrypoints.serve.instrumentator.health import health
21 -from vllm.tasks import POOLING_TASKS, SupportedTask
22 +from vllm.tasks import SupportedTask
23
24 # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers
25 # (requires typing_extensions >= 4.13)
@@ -27,80 +28,6 @@ GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServingBase | None]
28 EndpointFn = Callable[[RequestType, Request], Awaitable[Any]]
29
30
30 -def get_invocation_types(
31 - supported_tasks: tuple["SupportedTask", ...],
32 - model_config: ModelConfig | None = None,
33 -):
34 - # NOTE: Items defined earlier take higher priority
35 - INVOCATION_TYPES: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
36 -
37 - if "generate" in supported_tasks:
38 - from vllm.entrypoints.openai.chat_completion.api_router import (
39 - chat,
40 - create_chat_completion,
41 - )
42 - from vllm.entrypoints.openai.chat_completion.protocol import (
43 - ChatCompletionRequest,
44 - )
45 - from vllm.entrypoints.openai.completion.api_router import (
46 - completion,
47 - create_completion,
48 - )
49 - from vllm.entrypoints.openai.completion.protocol import CompletionRequest
50 -
51 - INVOCATION_TYPES += [
52 - (ChatCompletionRequest, (chat, create_chat_completion)),
53 - (CompletionRequest, (completion, create_completion)),
54 - ]
55 -
56 - if "embed" in supported_tasks:
57 - from vllm.entrypoints.pooling.embed.api_router import (
58 - create_embedding,
59 - embedding,
60 - )
61 - from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest
62 -
63 - INVOCATION_TYPES += [
64 - (EmbeddingRequest, (embedding, create_embedding)),
65 - ]
66 -
67 - if "classify" in supported_tasks:
68 - from vllm.entrypoints.pooling.classify.api_router import (
69 - classify,
70 - create_classify,
71 - )
72 - from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest
73 -
74 - INVOCATION_TYPES += [
75 - (ClassificationRequest, (classify, create_classify)),
76 - ]
77 -
78 - if enable_scoring_api(supported_tasks, model_config):
79 - from vllm.entrypoints.pooling.scoring.api_router import do_rerank, rerank
80 - from vllm.entrypoints.pooling.scoring.protocol import RerankRequest
81 -
82 - INVOCATION_TYPES += [
83 - (RerankRequest, (rerank, do_rerank)),
84 - ]
85 -
86 - from vllm.entrypoints.pooling.scoring.api_router import create_score, score
87 - from vllm.entrypoints.pooling.scoring.protocol import ScoreRequest
88 -
89 - INVOCATION_TYPES += [
90 - (ScoreRequest, (score, create_score)),
91 - ]
92 -
93 - if any(task in POOLING_TASKS for task in supported_tasks):
94 - from vllm.entrypoints.pooling.pooling.api_router import create_pooling, pooling
95 - from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest
96 -
97 - INVOCATION_TYPES += [
98 - (PoolingRequest, (pooling, create_pooling)),
99 - ]
100 -
101 - return INVOCATION_TYPES
102 -
103 -
31 def attach_router(
32 app: FastAPI,
33 supported_tasks: tuple["SupportedTask", ...],
@@ -109,7 +36,10 @@ def attach_router(
36 router = APIRouter()
37
38 # NOTE: Construct the TypeAdapters only once
112 - INVOCATION_TYPES = get_invocation_types(supported_tasks, model_config)
39 + INVOCATION_TYPES = get_generate_invocation_types(
40 + supported_tasks, model_config
41 + ) + get_pooling_invocation_types(supported_tasks, model_config)
42 +
43 INVOCATION_VALIDATORS = [
44 (pydantic.TypeAdapter(request_type), (get_handler, endpoint))
45 for request_type, (get_handler, endpoint) in INVOCATION_TYPES