From effacddab42223932954eede6a77f70b96a48c7b Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 11:53:56 +0800 Subject: [PATCH 01/20] offline part Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 133 +++++------------- vllm/entrypoints/pooling/base/io_processor.py | 25 +++- vllm/entrypoints/pooling/base/serving.py | 6 +- .../pooling/classify/io_processor.py | 6 +- vllm/entrypoints/pooling/classify/serving.py | 16 +-- .../entrypoints/pooling/embed/io_processor.py | 6 +- vllm/entrypoints/pooling/embed/serving.py | 16 +-- .../pooling/io_processor_factories.py | 47 +++++-- .../pooling/pooling/io_processor.py | 73 ++++++++++ vllm/entrypoints/pooling/scoring/serving.py | 17 +-- vllm/plugins/io_processors/__init__.py | 19 +++ vllm/v1/engine/llm_engine.py | 5 - 12 files changed, 208 insertions(+), 161 deletions(-) create mode 100644 vllm/entrypoints/pooling/pooling/io_processor.py diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 1be2cdd5c741..c650805cf69b 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -398,12 +398,11 @@ def _make_config(value: Any, cls: type[_R]) -> _R: self.runner_type = self.model_config.runner_type self.renderer = self.llm_engine.renderer self.chat_template = load_chat_template(chat_template) - self.io_processor = self.llm_engine.io_processor self.input_processor = self.llm_engine.input_processor self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template) self.pooling_io_processors = init_pooling_io_processors( supported_tasks=supported_tasks, - model_config=self.model_config, + vllm_config=self.llm_engine.vllm_config, renderer=self.renderer, chat_template_config=self.chat_template_config, ) @@ -1084,7 +1083,10 @@ def encode( self._verify_pooling_task(pooling_task) if isinstance(prompts, dict) and "data" in prompts: - if self.io_processor is None: + if pooling_task != "plugin": + raise ValueError() + + if "plugin" not in self.pooling_io_processors: raise ValueError( "No IOProcessor plugin installed. Please refer " "to the documentation and to the " @@ -1092,107 +1094,42 @@ def encode( "offline inference example for more details." ) - # Validate the request data is valid for the loaded plugin - prompt_data = prompts.get("data") - if prompt_data is None: - raise ValueError( - "The 'data' field of the prompt is expected to contain " - "the prompt data and it cannot be None. " - "Refer to the documentation of the IOProcessor " - "in use for more details." - ) - validated_prompt = self.io_processor.parse_data(prompt_data) + if pooling_params is None: + # Use default pooling params. + pooling_params = PoolingParams() - # obtain the actual model prompts from the pre-processor - prompts = self.io_processor.pre_process(prompt=validated_prompt) - prompts_seq = prompt_to_seq(prompts) - - params_seq: Sequence[PoolingParams] = [ - self.io_processor.merge_pooling_params(param) - for param in self._params_to_seq( - pooling_params, - len(prompts_seq), - ) - ] - for p in params_seq: - if p.task is None: - p.task = "plugin" - - outputs = self._run_completion( - prompts=prompts_seq, - params=params_seq, - output_type=PoolingRequestOutput, - use_tqdm=use_tqdm, - lora_request=lora_request, - tokenization_kwargs=tokenization_kwargs, - ) + prompts_seq = prompt_to_seq(prompts) + params_seq = self._params_to_seq(pooling_params, len(prompts_seq)) - # get the post-processed model outputs - assert self.io_processor is not None - processed_outputs = self.io_processor.post_process(outputs) + for param in params_seq: + if param.task is None: + param.task = pooling_task + elif param.task != pooling_task: + msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" + raise ValueError(msg) - return [ - PoolingRequestOutput[Any]( - request_id="", - outputs=processed_outputs, - num_cached_tokens=getattr( - processed_outputs, "num_cached_tokens", 0 - ), - prompt_token_ids=[], - finished=True, - ) - ] - else: - if pooling_params is None: - # Use default pooling params. - pooling_params = PoolingParams() - - prompts_seq = prompt_to_seq(prompts) - params_seq = self._params_to_seq(pooling_params, len(prompts_seq)) - - for param in params_seq: - if param.task is None: - param.task = pooling_task - elif param.task != pooling_task: - msg = ( - f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" - ) - raise ValueError(msg) + assert pooling_task is not None and pooling_task in self.pooling_io_processors - if pooling_task in self.pooling_io_processors: - io_processor = self.pooling_io_processors[pooling_task] - processor_inputs = io_processor.pre_process_offline( - ctx=OfflineInputsContext( - prompts=prompts_seq, tokenization_kwargs=tokenization_kwargs - ) - ) - seq_lora_requests = self._lora_request_to_seq( - lora_request, len(prompts_seq) - ) - seq_priority = self._priority_to_seq(None, len(prompts)) + io_processor = self.pooling_io_processors[pooling_task] + processor_inputs = io_processor.pre_process_offline( + ctx=OfflineInputsContext( + prompts=prompts_seq, tokenization_kwargs=tokenization_kwargs + ) + ) + seq_lora_requests = self._lora_request_to_seq(lora_request, len(prompts_seq)) + seq_priority = self._priority_to_seq(None, len(prompts)) - self._render_and_add_requests( - prompts=processor_inputs, - params=params_seq, - lora_requests=seq_lora_requests, - priorities=seq_priority, - ) + self._render_and_add_requests( + prompts=processor_inputs, + params=params_seq, + lora_requests=seq_lora_requests, + priorities=seq_priority, + ) - outputs = self._run_engine( - use_tqdm=use_tqdm, output_type=PoolingRequestOutput - ) - outputs = io_processor.post_process_offline( - ctx=OfflineOutputsContext(outputs=outputs) - ) - else: - outputs = self._run_completion( - prompts=prompts_seq, - params=params_seq, - output_type=PoolingRequestOutput, - use_tqdm=use_tqdm, - lora_request=lora_request, - tokenization_kwargs=tokenization_kwargs, - ) + outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput) + outputs = io_processor.post_process_offline( + ctx=OfflineOutputsContext(outputs=outputs) + ) return outputs def _verify_pooling_task(self, pooling_task: PoolingTask | None): diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index fd4c076cdda0..473a66da5b89 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -4,8 +4,8 @@ from collections.abc import Sequence from typing import Any, Final -from vllm import PoolingRequestOutput, PromptType -from vllm.config import ModelConfig +from vllm import PoolingParams, PoolingRequestOutput, PromptType +from vllm.config import VllmConfig from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateConfig, @@ -33,11 +33,12 @@ class PoolingIOProcessor: def __init__( self, - model_config: ModelConfig, + vllm_config: VllmConfig, renderer: BaseRenderer, chat_template_config: ChatTemplateConfig, ): - self.model_config = model_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config self.renderer = renderer self.chat_template = chat_template_config.chat_template @@ -243,3 +244,19 @@ def _validate_chat_template( "Refused request with untrusted chat template." ) return None + + def _params_to_seq( + self, + params: PoolingParams | Sequence[PoolingParams], + num_requests: int, + ) -> Sequence[PoolingParams]: + if isinstance(params, Sequence): + if len(params) != num_requests: + raise ValueError( + f"The lengths of prompts ({params}) " + f"and params ({len(params)}) must be the same." + ) + + return params + + return [params] * num_requests diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 90554aa634b4..4fa77e79b709 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -9,7 +9,7 @@ from starlette.datastructures import Headers from vllm import PoolingParams, PoolingRequestOutput, envs -from vllm.config import ModelConfig +from vllm.config import VllmConfig from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ( ChatTemplateConfig, @@ -64,14 +64,14 @@ def __init__( trust_request_chat_template=trust_request_chat_template, ) self.io_processor = self.init_io_processor( - model_config=models.model_config, + vllm_config=engine_client.vllm_config, renderer=models.renderer, chat_template_config=self.chat_template_config, ) def init_io_processor( self, - model_config: ModelConfig, + vllm_config: VllmConfig, renderer: BaseRenderer, chat_template_config: ChatTemplateConfig, ) -> PoolingIOProcessor: diff --git a/vllm/entrypoints/pooling/classify/io_processor.py b/vllm/entrypoints/pooling/classify/io_processor.py index ee73207dff5f..9bb3774ab0b2 100644 --- a/vllm/entrypoints/pooling/classify/io_processor.py +++ b/vllm/entrypoints/pooling/classify/io_processor.py @@ -5,4 +5,8 @@ class ClassifyIOProcessor(PoolingIOProcessor): - name = "classification" + name = "classify" + + +class TokenClassifyIOProcessor(PoolingIOProcessor): + name = "token_classify" diff --git a/vllm/entrypoints/pooling/classify/serving.py b/vllm/entrypoints/pooling/classify/serving.py index 24d4f9aacffc..0a729075bce0 100644 --- a/vllm/entrypoints/pooling/classify/serving.py +++ b/vllm/entrypoints/pooling/classify/serving.py @@ -6,14 +6,11 @@ import numpy as np from fastapi.responses import JSONResponse -from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.entrypoints.pooling.typing import PoolingServeContext from vllm.logger import init_logger from vllm.outputs import ClassificationOutput -from vllm.renderers import BaseRenderer from .io_processor import ClassifyIOProcessor from .protocol import ( @@ -31,17 +28,8 @@ class ServingClassification(PoolingServing): request_id_prefix = "classify" - def init_io_processor( - self, - model_config: ModelConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, - ) -> ClassifyIOProcessor: - return ClassifyIOProcessor( - model_config=model_config, - renderer=renderer, - chat_template_config=chat_template_config, - ) + def init_io_processor(self, *args, **kwargs) -> ClassifyIOProcessor: + return ClassifyIOProcessor(*args, **kwargs) async def _build_response( self, diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 614f8e0d9d02..623fee4fd3b1 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -37,7 +37,7 @@ class EmbedIOProcessor(PoolingIOProcessor): - name = "embedding" + name = "embed" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -549,3 +549,7 @@ def _enforce_cohere_max_tokens(self, ctx: PoolingServeContext) -> None: request = ctx.request if request.truncate == "NONE" and request.max_tokens is not None: self._check_cohere_max_tokens(ctx.final_res_batch, request.max_tokens) + + +class TokenEmbedIOProcessor(PoolingIOProcessor): + name = "token_embed" diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index f0c331645910..8a477825e70f 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -8,8 +8,6 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never -from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor @@ -33,7 +31,6 @@ ) from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput -from vllm.renderers import BaseRenderer from vllm.utils.serial_utils import EmbedDType, Endianness logger = init_logger(__name__) @@ -49,17 +46,8 @@ class ServingEmbedding(PoolingServing): request_id_prefix = "embd" io_processor: EmbedIOProcessor - def init_io_processor( - self, - model_config: ModelConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, - ) -> EmbedIOProcessor: - return EmbedIOProcessor( - model_config=model_config, - renderer=renderer, - chat_template_config=chat_template_config, - ) + def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor: + return EmbedIOProcessor(*args, **kwargs) async def _build_response( self, diff --git a/vllm/entrypoints/pooling/io_processor_factories.py b/vllm/entrypoints/pooling/io_processor_factories.py index 71033bd2398f..d23d9c7562db 100644 --- a/vllm/entrypoints/pooling/io_processor_factories.py +++ b/vllm/entrypoints/pooling/io_processor_factories.py @@ -1,42 +1,71 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm.config import ModelConfig +from vllm.config import VllmConfig from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessors from vllm.entrypoints.pooling.utils import enable_scoring_api +from vllm.plugins.io_processors import has_io_processor from vllm.renderers import BaseRenderer from vllm.tasks import SupportedTask def init_pooling_io_processors( supported_tasks: tuple[SupportedTask, ...], - model_config: ModelConfig, + vllm_config: VllmConfig, renderer: BaseRenderer, chat_template_config: ChatTemplateConfig, ) -> dict[str, PoolingIOProcessor]: - processors: list[tuple[str, type[PoolingIOProcessor]]] = [] + model_config = vllm_config.model_config + processors: dict[str, type[PoolingIOProcessor]] = {} + if "classify" in supported_tasks: from vllm.entrypoints.pooling.classify.io_processor import ClassifyIOProcessor - processors.append(("classify", ClassifyIOProcessor)) + processors["classify"] = ClassifyIOProcessor + + if "token_classify" in supported_tasks: + from vllm.entrypoints.pooling.classify.io_processor import ( + TokenClassifyIOProcessor, + ) + + processors["token_classify"] = TokenClassifyIOProcessor + if "embed" in supported_tasks: from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor - processors.append(("embed", EmbedIOProcessor)) + processors["embed"] = EmbedIOProcessor + + if "token_embed" in supported_tasks: + from vllm.entrypoints.pooling.embed.io_processor import TokenEmbedIOProcessor + + processors["token_embed"] = TokenEmbedIOProcessor + + if "token_embed" in supported_tasks: + from vllm.entrypoints.pooling.embed.io_processor import TokenEmbedIOProcessor + + processors["token_embed"] = TokenEmbedIOProcessor + + if has_io_processor( + vllm_config, + model_config.io_processor_plugin, + ): + from vllm.entrypoints.pooling.pooling.io_processor import PluginIOProcessor + + processors["plugin"] = PluginIOProcessor if enable_scoring_api(supported_tasks, model_config): score_type = model_config.score_type + if score_type is not None and score_type in ScoringIOProcessors: - processors.append((score_type, ScoringIOProcessors[score_type])) + processors[score_type] = ScoringIOProcessors[score_type] return { task: processor_cls( - model_config=model_config, + vllm_config=vllm_config, renderer=renderer, chat_template_config=chat_template_config, ) - for task, processor_cls in processors + for task, processor_cls in processors.items() } diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py new file mode 100644 index 000000000000..350f268f7af9 --- /dev/null +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence +from typing import Any + +from vllm import PoolingParams, PoolingRequestOutput +from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext +from vllm.inputs import EngineInput +from vllm.plugins.io_processors import get_io_processor +from vllm.renderers.inputs.preprocess import prompt_to_seq + + +class PluginIOProcessor(PoolingIOProcessor): + name = "plugin" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.io_processor = get_io_processor( + self.vllm_config, + self.renderer, + self.model_config.io_processor_plugin, + ) + + ####################################### + # offline APIs + + def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: + # Validate the request data is valid for the loaded plugin + prompt_data = ctx.prompts.get("data") + if prompt_data is None: + raise ValueError( + "The 'data' field of the prompt is expected to contain " + "the prompt data and it cannot be None. " + "Refer to the documentation of the IOProcessor " + "in use for more details." + ) + validated_prompt = self.io_processor.parse_data(prompt_data) + + # obtain the actual model prompts from the pre-processor + prompts = self.io_processor.pre_process(prompt=validated_prompt) + prompts_seq = prompt_to_seq(prompts) + + params_seq: list[PoolingParams] = [ + self.io_processor.merge_pooling_params(param) + for param in self._params_to_seq( + ctx.pooling_params, + len(prompts_seq), + ) + ] + for p in params_seq: + if p.task is None: + p.task = "plugin" + ctx.pooling_params = params_seq + ctx.prompts = prompts_seq + return super().pre_process_offline(ctx) + + def post_process_offline( + self, + ctx: OfflineOutputsContext, + ) -> list[PoolingRequestOutput]: + processed_outputs = self.io_processor.post_process(ctx.outputs) + + return [ + PoolingRequestOutput[Any]( + request_id="", + outputs=processed_outputs, + num_cached_tokens=getattr(processed_outputs, "num_cached_tokens", 0), + prompt_token_ids=[], + finished=True, + ) + ] diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index de5b5797ce49..0cde1daf6b85 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -4,15 +4,13 @@ from fastapi.responses import JSONResponse, Response from vllm import PoolingParams -from vllm.config import ModelConfig +from vllm.config import VllmConfig from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput -from vllm.renderers import BaseRenderer from vllm.v1.pool.late_interaction import ( build_late_interaction_doc_params, build_late_interaction_query_params, @@ -52,22 +50,17 @@ def __init__( super().__init__(engine_client, *args, **kwargs) def init_io_processor( - self, - model_config: ModelConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, + self, vllm_config: VllmConfig, *args, **kwargs ) -> PoolingIOProcessor: + model_config = vllm_config.model_config + score_type: str = model_config.score_type if self.enable_flash_late_interaction: score_type = "flash-late-interaction" assert score_type in ScoringIOProcessors processor_cls = ScoringIOProcessors[score_type] - return processor_cls( - model_config=model_config, - renderer=renderer, - chat_template_config=chat_template_config, - ) + return processor_cls(vllm_config=vllm_config, *args, **kwargs) async def __call__(self, *args, **kwargs) -> Response: if not self.enable_flash_late_interaction: diff --git a/vllm/plugins/io_processors/__init__.py b/vllm/plugins/io_processors/__init__.py index c8cb4f185278..4671aaca267e 100644 --- a/vllm/plugins/io_processors/__init__.py +++ b/vllm/plugins/io_processors/__init__.py @@ -12,6 +12,25 @@ logger = logging.getLogger(__name__) +def has_io_processor( + vllm_config: VllmConfig, + + plugin_from_init: str | None = None, + ): + if plugin_from_init: + model_plugin = plugin_from_init + else: + # A plugin can be specified via the model config + # Retrieve the model specific plugin if available + # This is using a custom field in the hf_config for the model + hf_config = vllm_config.model_config.hf_config.to_dict() + config_plugin = hf_config.get("io_processor_plugin") + model_plugin = config_plugin + + return model_plugin is not None + + + def get_io_processor( vllm_config: VllmConfig, renderer: BaseRenderer, diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 4b6a7ba44e10..55ccdde91339 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -90,11 +90,6 @@ def __init__( self.should_execute_dummy_batch = False self.renderer = renderer = renderer_from_config(self.vllm_config) - self.io_processor = get_io_processor( - self.vllm_config, - self.renderer, - self.model_config.io_processor_plugin, - ) # Convert EngineInput --> EngineCoreRequest. self.input_processor = InputProcessor(self.vllm_config, renderer) From a5666413ff5fb1de7a3527698ae302ccadd81753 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 13:54:51 +0800 Subject: [PATCH 02/20] fix plugin offline Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 51 ++++++++++--------- vllm/entrypoints/pooling/base/io_processor.py | 4 +- .../pooling/pooling/io_processor.py | 1 + 3 files changed, 32 insertions(+), 24 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index c650805cf69b..531f84200b29 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1094,33 +1094,39 @@ def encode( "offline inference example for more details." ) - if pooling_params is None: - # Use default pooling params. - pooling_params = PoolingParams() + assert pooling_task is not None and pooling_task in self.pooling_io_processors + io_processor = self.pooling_io_processors[pooling_task] + + ctx = OfflineInputsContext( + prompts=prompts, + pooling_params=pooling_params, + tokenization_kwargs=tokenization_kwargs, + ) + + engine_inputs = io_processor.pre_process_offline(ctx) + n_inputs = len(engine_inputs) - prompts_seq = prompt_to_seq(prompts) - params_seq = self._params_to_seq(pooling_params, len(prompts_seq)) + if ctx.pooling_params is None: + ctx.pooling_params = PoolingParams() + + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) for param in params_seq: if param.task is None: param.task = pooling_task + elif pooling_task == "plugin": + # `plugin` task uses io_processor.parse_request to verify inputs. + # We actually allow plugin to overwrite pooling_task. + pass elif param.task != pooling_task: msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" raise ValueError(msg) - assert pooling_task is not None and pooling_task in self.pooling_io_processors - - io_processor = self.pooling_io_processors[pooling_task] - processor_inputs = io_processor.pre_process_offline( - ctx=OfflineInputsContext( - prompts=prompts_seq, tokenization_kwargs=tokenization_kwargs - ) - ) - seq_lora_requests = self._lora_request_to_seq(lora_request, len(prompts_seq)) - seq_priority = self._priority_to_seq(None, len(prompts)) + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs) + seq_priority = self._priority_to_seq(None, n_inputs) self._render_and_add_requests( - prompts=processor_inputs, + prompts=engine_inputs, params=params_seq, lora_requests=seq_lora_requests, priorities=seq_priority, @@ -1403,15 +1409,14 @@ def score( n_queries=n_queries, ) - processor_inputs = io_processor.pre_process_offline(ctx) + engine_inputs = io_processor.pre_process_offline(ctx) + n_inputs = len(engine_inputs) - seq_lora_requests = self._lora_request_to_seq( - lora_request, len(processor_inputs) - ) + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs) if ctx.pooling_params is None: ctx.pooling_params = PoolingParams() - params_seq = self._params_to_seq(ctx.pooling_params, len(processor_inputs)) + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) for param in params_seq: if param.task is None: @@ -1420,10 +1425,10 @@ def score( msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" raise ValueError(msg) - seq_priority = self._priority_to_seq(None, len(processor_inputs)) + seq_priority = self._priority_to_seq(None, n_inputs) self._render_and_add_requests( - prompts=processor_inputs, + prompts=engine_inputs, params=params_seq, lora_requests=seq_lora_requests, priorities=seq_priority, diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index 473a66da5b89..9e7f3abae750 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -102,11 +102,13 @@ async def post_process_online_async( def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: assert not isinstance(ctx.prompts, ScoringData) + + prompts_seq = prompt_to_seq(ctx.prompts) tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) ) return self._preprocess_completion_offline( - prompts=ctx.prompts, tok_params=tok_params + prompts=prompts_seq, tok_params=tok_params ) async def pre_process_offline_async(self, ctx: OfflineInputsContext): diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index 350f268f7af9..a365aca539f8 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -52,6 +52,7 @@ def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput for p in params_seq: if p.task is None: p.task = "plugin" + ctx.pooling_params = params_seq ctx.prompts = prompts_seq return super().pre_process_offline(ctx) From 48446095a821235c4883b9ede5fe0b2f286dbc42 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 14:00:58 +0800 Subject: [PATCH 03/20] fix plugin offline Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 6 ++++-- vllm/entrypoints/pooling/base/io_processor.py | 2 +- vllm/entrypoints/pooling/io_processor_factories.py | 5 ----- vllm/plugins/io_processors/__init__.py | 6 ++---- 4 files changed, 7 insertions(+), 12 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 531f84200b29..3fb8648ceda0 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1084,7 +1084,9 @@ def encode( if isinstance(prompts, dict) and "data" in prompts: if pooling_task != "plugin": - raise ValueError() + raise ValueError( + "The 'data' field is only supported for the 'plugin' pooling task." + ) if "plugin" not in self.pooling_io_processors: raise ValueError( @@ -1521,7 +1523,7 @@ def _params_to_seq( if isinstance(params, Sequence): if len(params) != num_requests: raise ValueError( - f"The lengths of prompts ({params}) " + f"The lengths of prompts ({num_requests}) " f"and params ({len(params)}) must be the same." ) diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index 9e7f3abae750..f5de5c295170 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -255,7 +255,7 @@ def _params_to_seq( if isinstance(params, Sequence): if len(params) != num_requests: raise ValueError( - f"The lengths of prompts ({params}) " + f"The lengths of prompts ({num_requests}) " f"and params ({len(params)}) must be the same." ) diff --git a/vllm/entrypoints/pooling/io_processor_factories.py b/vllm/entrypoints/pooling/io_processor_factories.py index d23d9c7562db..49182a9f0c8c 100644 --- a/vllm/entrypoints/pooling/io_processor_factories.py +++ b/vllm/entrypoints/pooling/io_processor_factories.py @@ -42,11 +42,6 @@ def init_pooling_io_processors( processors["token_embed"] = TokenEmbedIOProcessor - if "token_embed" in supported_tasks: - from vllm.entrypoints.pooling.embed.io_processor import TokenEmbedIOProcessor - - processors["token_embed"] = TokenEmbedIOProcessor - if has_io_processor( vllm_config, model_config.io_processor_plugin, diff --git a/vllm/plugins/io_processors/__init__.py b/vllm/plugins/io_processors/__init__.py index 4671aaca267e..c502f4f744d7 100644 --- a/vllm/plugins/io_processors/__init__.py +++ b/vllm/plugins/io_processors/__init__.py @@ -13,10 +13,9 @@ def has_io_processor( - vllm_config: VllmConfig, - + vllm_config: VllmConfig, plugin_from_init: str | None = None, - ): +): if plugin_from_init: model_plugin = plugin_from_init else: @@ -30,7 +29,6 @@ def has_io_processor( return model_plugin is not None - def get_io_processor( vllm_config: VllmConfig, renderer: BaseRenderer, From 41e4e45e47b600820c21ecbacd5a34b98ea22803 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 15:24:41 +0800 Subject: [PATCH 04/20] online part 1 Signed-off-by: wang.yuqi --- vllm/engine/protocol.py | 1 - vllm/entrypoints/openai/engine/serving.py | 185 +---------- vllm/entrypoints/pooling/__init__.py | 29 +- vllm/entrypoints/pooling/base/serving.py | 77 +++-- .../pooling/pooling/io_processor.py | 8 +- vllm/entrypoints/pooling/pooling/serving.py | 301 ++++-------------- vllm/v1/engine/async_llm.py | 5 - 7 files changed, 127 insertions(+), 479 deletions(-) diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 3d466e3fc2af..1cedb86bd17d 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -44,7 +44,6 @@ class EngineClient(ABC): vllm_config: VllmConfig model_config: ModelConfig renderer: BaseRenderer - io_processor: IOProcessor | None input_processor: InputProcessor @property diff --git a/vllm/entrypoints/openai/engine/serving.py b/vllm/entrypoints/openai/engine/serving.py index f5f011a96f27..5bd415b4fbd2 100644 --- a/vllm/entrypoints/openai/engine/serving.py +++ b/vllm/entrypoints/openai/engine/serving.py @@ -44,12 +44,6 @@ TranscriptionResponse, TranslationRequest, ) -from vllm.entrypoints.pooling.pooling.protocol import ( - IOProcessorRequest, - PoolingChatRequest, - PoolingCompletionRequest, - PoolingResponse, -) from vllm.entrypoints.serve.disagg.protocol import GenerateRequest, GenerateResponse from vllm.entrypoints.serve.tokenize.protocol import ( DetokenizeRequest, @@ -62,8 +56,7 @@ from vllm.logger import init_logger from vllm.logprobs import Logprob, PromptLogprobs from vllm.lora.request import LoRARequest -from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput -from vllm.pooling_params import PoolingParams +from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers import ChatParams, TokenizeParams from vllm.renderers.inputs.preprocess import ( extract_prompt_components, @@ -78,10 +71,7 @@ log_tracing_disabled_warning, ) from vllm.utils import random_uuid -from vllm.utils.async_utils import ( - collect_from_async_generator, - merge_async_iterators, -) +from vllm.utils.async_utils import collect_from_async_generator logger = init_logger(__name__) @@ -101,17 +91,11 @@ def build_chat_params( CompletionLikeRequest: TypeAlias = ( - CompletionRequest - | TokenizeCompletionRequest - | DetokenizeRequest - | PoolingCompletionRequest + CompletionRequest | TokenizeCompletionRequest | DetokenizeRequest ) ChatLikeRequest: TypeAlias = ( - ChatCompletionRequest - | BatchChatCompletionRequest - | TokenizeChatRequest - | PoolingChatRequest + ChatCompletionRequest | BatchChatCompletionRequest | TokenizeChatRequest ) SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest @@ -121,7 +105,6 @@ def build_chat_params( | ChatLikeRequest | SpeechToTextRequest | ResponsesRequest - | IOProcessorRequest | GenerateRequest ) @@ -130,7 +113,6 @@ def build_chat_params( | ChatCompletionResponse | TranscriptionResponse | TokenizeResponse - | PoolingResponse | GenerateResponse ) @@ -146,12 +128,6 @@ class ServeContext(Generic[RequestT]): created_time: int = field(default_factory=lambda: int(time.time())) lora_request: LoRARequest | None = None engine_inputs: list[EngineInput] | None = None - - result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = ( - None - ) - final_res_batch: list[PoolingRequestOutput] = field(default_factory=list) - model_config = ConfigDict(arbitrary_types_allowed=True) @@ -171,7 +147,6 @@ def __init__( super().__init__() self.engine_client = engine_client - self.models = models self.request_logger = request_logger @@ -179,7 +154,6 @@ def __init__( self.model_config = engine_client.model_config self.renderer = engine_client.renderer - self.io_processor = engine_client.io_processor self.input_processor = engine_client.input_processor async def beam_search( @@ -381,155 +355,6 @@ async def beam_search( prompt_logprobs=None, ) - async def _preprocess( - self, - ctx: ServeContext, - ) -> ErrorResponse | None: - """ - Default preprocessing hook. Subclasses may override to prepare `ctx`. - """ - return None - - def _build_response( - self, - ctx: ServeContext, - ) -> AnyResponse | ErrorResponse: - """ - Default response builder. Subclass may override this method - to return the appropriate response object. - """ - return self.create_error_response("unimplemented endpoint") - - async def handle( - self, - ctx: ServeContext, - ) -> AnyResponse | ErrorResponse: - async for response in self._pipeline(ctx): - return response - - return self.create_error_response("No response yielded from pipeline") - - async def _pipeline( - self, - ctx: ServeContext, - ) -> AsyncGenerator[AnyResponse | ErrorResponse, None]: - """Execute the request processing pipeline yielding responses.""" - if error := await self._check_model(ctx.request): - yield error - if error := self._validate_request(ctx): - yield error - - preprocess_ret = await self._preprocess(ctx) - if isinstance(preprocess_ret, ErrorResponse): - yield preprocess_ret - - generators_ret = await self._prepare_generators(ctx) - if isinstance(generators_ret, ErrorResponse): - yield generators_ret - - collect_ret = await self._collect_batch(ctx) - if isinstance(collect_ret, ErrorResponse): - yield collect_ret - - yield self._build_response(ctx) - - def _validate_request(self, ctx: ServeContext) -> ErrorResponse | None: - truncate_prompt_tokens = getattr(ctx.request, "truncate_prompt_tokens", None) - - if ( - truncate_prompt_tokens is not None - and truncate_prompt_tokens > self.model_config.max_model_len - ): - return self.create_error_response( - "truncate_prompt_tokens value is " - "greater than max_model_len." - " Please request a smaller truncation size." - ) - return None - - def _create_pooling_params( - self, - ctx: ServeContext, - ) -> PoolingParams | ErrorResponse: - if not hasattr(ctx.request, "to_pooling_params"): - return self.create_error_response( - "Request type does not support pooling parameters" - ) - - return ctx.request.to_pooling_params() - - async def _prepare_generators( - self, - ctx: ServeContext, - ) -> ErrorResponse | None: - """Schedule the request and get the result generator.""" - generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] - - trace_headers = ( - None - if ctx.raw_request is None - else await self._get_trace_headers(ctx.raw_request.headers) - ) - - pooling_params = self._create_pooling_params(ctx) - if isinstance(pooling_params, ErrorResponse): - return pooling_params - - if ctx.engine_inputs is None: - return self.create_error_response("Engine prompts not available") - - for i, engine_input in enumerate(ctx.engine_inputs): - request_id_item = f"{ctx.request_id}-{i}" - - self._log_inputs( - request_id_item, - engine_input, - params=pooling_params, - lora_request=ctx.lora_request, - ) - - generator = self.engine_client.encode( - engine_input, - pooling_params, - request_id_item, - lora_request=ctx.lora_request, - trace_headers=trace_headers, - priority=getattr(ctx.request, "priority", 0), - ) - - generators.append(generator) - - ctx.result_generator = merge_async_iterators(*generators) - - return None - - async def _collect_batch( - self, - ctx: ServeContext, - ) -> ErrorResponse | None: - """Collect batch results from the result generator.""" - if ctx.engine_inputs is None: - return self.create_error_response("Engine prompts not available") - - num_prompts = len(ctx.engine_inputs) - final_res_batch: list[PoolingRequestOutput | None] - final_res_batch = [None] * num_prompts - - if ctx.result_generator is None: - return self.create_error_response("Result generator not available") - - async for i, res in ctx.result_generator: - final_res_batch[i] = res - - if None in final_res_batch: - return self.create_error_response( - "Failed to generate results for all prompts" - ) - - ctx.final_res_batch = [res for res in final_res_batch if res is not None] - - return None - @staticmethod def create_error_response( message: str | Exception, @@ -719,7 +544,7 @@ def _log_inputs( self, request_id: str, inputs: PromptType | EngineInput, - params: SamplingParams | PoolingParams | BeamSearchParams | None, + params: SamplingParams | BeamSearchParams | None, lora_request: LoRARequest | None, ) -> None: if self.request_logger is None: diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index fb0c10e6f4f4..cc8368813612 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -6,6 +6,7 @@ from fastapi import FastAPI from vllm.config import ModelConfig +from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.pooling.utils import enable_scoring_api from vllm.logger import init_logger @@ -67,25 +68,25 @@ def init_pooling_state( from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.pooling.classify.serving import ServingClassification from vllm.entrypoints.pooling.embed.serving import ServingEmbedding - from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling + from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.pooling.scoring.serving import ServingScores from vllm.tasks import POOLING_TASKS model_config = engine_client.model_config - - resolved_chat_template = load_chat_template(args.chat_template) + chat_template_config = ChatTemplateConfig( + chat_template=load_chat_template(args.chat_template), + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + ) state.serving_pooling = ( ( - OpenAIServingPooling( + ServingPooling( engine_client, state.openai_serving_models, - state.openai_serving_render, supported_tasks=supported_tasks, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) ) if any(t in supported_tasks for t in POOLING_TASKS) @@ -96,9 +97,7 @@ def init_pooling_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if "embed" in supported_tasks else None @@ -108,9 +107,7 @@ def init_pooling_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if "classify" in supported_tasks else None @@ -120,9 +117,7 @@ def init_pooling_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, enable_flash_late_interaction=getattr( args, "enable_flash_late_interaction", True ), diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 4fa77e79b709..ab9dc573b763 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -13,7 +13,6 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ( ChatTemplateConfig, - ChatTemplateContentFormatOption, ) from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.engine.protocol import ErrorResponse @@ -35,7 +34,7 @@ from .io_processor import PoolingIOProcessor -class PoolingServing: +class PoolingServingBase: request_id_prefix: ClassVar[str] def __init__( @@ -43,51 +42,28 @@ def __init__( engine_client: EngineClient, models: OpenAIServingModels, *, + chat_template_config: ChatTemplateConfig, request_logger: RequestLogger | None, - chat_template: str | None = None, - chat_template_content_format: ChatTemplateContentFormatOption = "auto", - trust_request_chat_template: bool = False, return_tokens_as_token_ids: bool = False, log_error_stack: bool = False, ): - super().__init__() self.engine_client = engine_client self.models = models self.model_config = models.model_config + self.renderer = models.renderer + self.vllm_config = engine_client.vllm_config self.max_model_len = self.model_config.max_model_len self.request_logger = request_logger self.return_tokens_as_token_ids = return_tokens_as_token_ids self.log_error_stack = log_error_stack - self.chat_template_config = ChatTemplateConfig( - chat_template=chat_template, - chat_template_content_format=chat_template_content_format, - trust_request_chat_template=trust_request_chat_template, - ) - self.io_processor = self.init_io_processor( - vllm_config=engine_client.vllm_config, - renderer=models.renderer, - chat_template_config=self.chat_template_config, - ) - - def init_io_processor( - self, - vllm_config: VllmConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, - ) -> PoolingIOProcessor: - raise NotImplementedError + self.chat_template_config = chat_template_config async def __call__( self, request: AnyPoolingRequest, raw_request: Request | None = None, ) -> Response: - ctx = await self._init_ctx(request, raw_request) - await self.io_processor.pre_process_online_async(ctx) - await self._prepare_generators(ctx) - await self._collect_batch(ctx) - await self.io_processor.post_process_online_async(ctx) - return await self._build_response(ctx) + raise NotImplementedError async def _init_ctx( self, @@ -124,10 +100,8 @@ async def _prepare_generators( else await self._get_trace_headers(ctx.raw_request.headers) ) - if ctx.pooling_params is None: - pooling_params = self.io_processor.create_pooling_params(ctx.request) - else: - pooling_params = ctx.pooling_params + assert ctx.pooling_params is not None + pooling_params = ctx.pooling_params if isinstance(pooling_params, list): for params in pooling_params: @@ -355,3 +329,38 @@ def _log_inputs( params=params, lora_request=lora_request, ) + + +class PoolingServing(PoolingServingBase): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.io_processor = self.init_io_processor( + vllm_config=self.vllm_config, + renderer=self.renderer, + chat_template_config=self.chat_template_config, + ) + + def init_io_processor( + self, + vllm_config: VllmConfig, + renderer: BaseRenderer, + chat_template_config: ChatTemplateConfig, + ) -> PoolingIOProcessor: + raise NotImplementedError + + async def __call__( + self, + request: AnyPoolingRequest, + raw_request: Request | None = None, + ) -> Response: + ctx = await self._init_ctx(request, raw_request) + await self.io_processor.pre_process_online_async(ctx) + + if ctx.pooling_params is None: + ctx.pooling_params = self.io_processor.create_pooling_params(request) + + await self._prepare_generators(ctx) + await self._collect_batch(ctx) + await self.io_processor.post_process_online_async(ctx) + return await self._build_response(ctx) diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index a365aca539f8..f6ef1e7b8135 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -5,11 +5,17 @@ from vllm import PoolingParams, PoolingRequestOutput from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext +from vllm.entrypoints.pooling.typing import ( + OfflineInputsContext, + OfflineOutputsContext, +) from vllm.inputs import EngineInput +from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.renderers.inputs.preprocess import prompt_to_seq +logger = init_logger(__name__) + class PluginIOProcessor(PoolingIOProcessor): name = "plugin" diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 4706684f3637..a552d7dd08e4 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -1,252 +1,105 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import asyncio + import json -import time -from collections.abc import AsyncGenerator, Callable, Sequence +from collections.abc import Callable from functools import partial -from typing import Final, Literal, cast +from typing import Literal, cast from fastapi import Request +from fastapi.responses import Response from typing_extensions import assert_never -from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption -from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo -from vllm.entrypoints.openai.engine.serving import OpenAIServing -from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.engine.protocol import UsageInfo +from vllm.entrypoints.pooling.base.serving import PoolingServingBase +from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors from vllm.entrypoints.pooling.pooling.protocol import ( - IOProcessorRequest, - IOProcessorResponse, PoolingBytesResponse, - PoolingChatRequest, - PoolingCompletionRequest, - PoolingRequest, PoolingResponse, PoolingResponseData, ) +from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext from vllm.entrypoints.pooling.utils import ( encode_pooling_bytes, encode_pooling_output_base64, encode_pooling_output_float, ) -from vllm.entrypoints.serve.render.serving import OpenAIServingRender -from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput -from vllm.renderers.inputs.preprocess import prompt_to_seq from vllm.tasks import SupportedTask -from vllm.utils.async_utils import merge_async_iterators -from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness +from vllm.utils.serial_utils import EmbedDType, Endianness logger = init_logger(__name__) -class OpenAIServingPooling(OpenAIServing): +class ServingPooling(PoolingServingBase): + request_id_prefix = "pooling" + def __init__( self, - engine_client: EngineClient, - models: OpenAIServingModels, - openai_serving_render: OpenAIServingRender, + *args, supported_tasks: tuple[SupportedTask, ...], - *, - request_logger: RequestLogger | None, - chat_template: str | None, - chat_template_content_format: ChatTemplateContentFormatOption, - trust_request_chat_template: bool = False, - ) -> None: - super().__init__( - engine_client=engine_client, - models=models, - request_logger=request_logger, - ) + **kwargs, + ): + super().__init__(*args, **kwargs) + self.supported_tasks = supported_tasks self.pooling_task = self.model_config.get_pooling_task(supported_tasks) - self.openai_serving_render = openai_serving_render - self.chat_template = chat_template - self.chat_template_content_format: Final = chat_template_content_format - self.trust_request_chat_template = trust_request_chat_template + self.io_processors = init_pooling_io_processors( + supported_tasks=supported_tasks, + vllm_config=self.vllm_config, + renderer=self.renderer, + chat_template_config=self.chat_template_config, + ) - async def create_pooling( + async def __call__( self, - request: PoolingRequest, + request: AnyPoolingRequest, raw_request: Request | None = None, - ) -> PoolingResponse | IOProcessorResponse | PoolingBytesResponse | ErrorResponse: - """ - See https://platform.openai.com/docs/api-reference/embeddings/create - for the API specification. This API mimics the OpenAI Embedding API. - """ - error_check_ret = await self._check_model(request) - if error_check_ret is not None: - return error_check_ret - - model_name = self.models.model_name() - - request_id = f"pool-{self._base_request_id(raw_request)}" - created_time = int(time.time()) - - lora_request = self._maybe_get_adapters(request) - - if request.task is None: - request.task = self.pooling_task - - if getattr(request, "dimensions", None) is not None: - return self.create_error_response("dimensions is currently not supported") - - # plugin task uses io_processor.parse_request to verify inputs - if request.task != "plugin" and request.task != self.pooling_task: - if request.task not in self.supported_tasks: - raise ValueError( - f"Unsupported task: {request.task!r} " - f"Supported tasks: {self.supported_tasks}" - ) - else: - logger.warning_once( - "Pooling multitask support is deprecated and will be removed " - "in v0.20. When the default pooling task is not what you want, you " - 'need to manually specify it via --pooler-config.task "%s". ', - request.task, - ) - - engine_inputs: Sequence[EngineInput] - if use_io_processor := isinstance(request, IOProcessorRequest): - if self.io_processor is None: - raise ValueError( - "No IOProcessor plugin installed. Please refer " - "to the documentation and to the " - "'prithvi_geospatial_mae_io_processor' " - "offline inference example for more details." - ) - - validated_prompt = self.io_processor.parse_data(request.data) - - raw_prompts = await self.io_processor.pre_process_async( - prompt=validated_prompt, request_id=request_id - ) - engine_inputs = await self.openai_serving_render.preprocess_cmpl( - request, - prompt_to_seq(raw_prompts), - ) - elif isinstance(request, PoolingChatRequest): - error_check_ret = self.openai_serving_render.validate_chat_template( - request_chat_template=request.chat_template, - chat_template_kwargs=request.chat_template_kwargs, - trust_request_chat_template=self.trust_request_chat_template, - ) - if error_check_ret is not None: - return error_check_ret - - _, engine_inputs = await self.openai_serving_render.preprocess_chat( - request, - request.messages, - default_template=self.chat_template, - default_template_content_format=self.chat_template_content_format, - default_template_kwargs=None, - ) - elif isinstance(request, PoolingCompletionRequest): - engine_inputs = await self.openai_serving_render.preprocess_completion( - request, - prompt_input=request.input, - prompt_embeds=None, - ) - else: - raise ValueError(f"Unsupported request of type {type(request)}") - - # Schedule the request and get the result generator. - generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] - if use_io_processor: - assert self.io_processor is not None + ) -> Response: + ctx = await self._init_ctx(request, raw_request) + await self.io_processor.pre_process_online_async(ctx) - pooling_params = self.io_processor.merge_pooling_params() - if pooling_params.task is None: - pooling_params.task = "plugin" - else: - pooling_params = request.to_pooling_params() # type: ignore + if ctx.pooling_params is None: + ctx.pooling_params = self.io_processor.create_pooling_params(request) - for i, engine_input in enumerate(engine_inputs): - request_id_item = f"{request_id}-{i}" + await self._prepare_generators(ctx) + await self._collect_batch(ctx) + await self.io_processor.post_process_online_async(ctx) + return await self._build_response(ctx) - self._log_inputs( - request_id_item, - engine_input, - params=pooling_params, - lora_request=lora_request, - ) - - trace_headers = ( - None - if raw_request is None - else await self._get_trace_headers(raw_request.headers) - ) - - generator = self.engine_client.encode( - engine_input, - pooling_params, - request_id_item, - lora_request=lora_request, - trace_headers=trace_headers, - priority=request.priority, - ) - - generators.append(generator) - - result_generator = merge_async_iterators(*generators) + async def _build_response( + self, + ctx: PoolingServeContext, + ) -> Response: + encoding_format = ctx.request.encoding_format + embed_dtype = ctx.request.embed_dtype + endianness = ctx.request.endianness - if use_io_processor: - assert self.io_processor is not None - output = await self.io_processor.post_process_async( - result_generator, - request_id=request_id, + if encoding_format == "float" or encoding_format == "base64": + return self.request_output_to_pooling_json_response( + ctx.final_res_batch, + ctx.request_id, + ctx.created_time, + ctx.model_name, + encoding_format, + embed_dtype, + endianness, ) - if callable( - output_to_response := getattr( - self.io_processor, "output_to_response", None - ) - ): - logger.warning_once( - "`IOProcessor.output_to_response` is deprecated. To ensure " - "consistency between offline and online APIs, " - "`IOProcessorResponse` will become a transparent wrapper " - "around output data from v0.19 onwards.", - ) - - if hasattr(output, "request_id") and output.request_id is None: - output.request_id = request_id # type: ignore - - return output_to_response(output) # type: ignore - - return IOProcessorResponse(request_id=request_id, data=output) - - assert isinstance(request, (PoolingCompletionRequest, PoolingChatRequest)) - num_prompts = len(engine_inputs) - - # Non-streaming response - final_res_batch: list[PoolingRequestOutput | None] - final_res_batch = [None] * num_prompts - try: - async for i, res in result_generator: - final_res_batch[i] = res - - assert all(final_res is not None for final_res in final_res_batch) - - final_res_batch_checked = cast(list[PoolingRequestOutput], final_res_batch) - - response = self.request_output_to_pooling_response( - final_res_batch_checked, - request_id, - created_time, - model_name, - request.encoding_format, - request.embed_dtype, - request.endianness, + if encoding_format == "bytes" or encoding_format == "bytes_only": + return self.request_output_to_pooling_bytes_response( + ctx.final_res_batch, + ctx.request_id, + ctx.created_time, + ctx.model_name, + encoding_format, + embed_dtype, + endianness, ) - except asyncio.CancelledError: - return self.create_error_response("Client disconnected") - return response + assert_never(encoding_format) def request_output_to_pooling_json_response( self, @@ -330,37 +183,3 @@ def request_output_to_pooling_bytes_response( ) return PoolingBytesResponse(content=content, headers=headers) - - def request_output_to_pooling_response( - self, - final_res_batch: list[PoolingRequestOutput], - request_id: str, - created_time: int, - model_name: str, - encoding_format: EncodingFormat, - embed_dtype: EmbedDType, - endianness: Endianness, - ) -> PoolingResponse | PoolingBytesResponse: - if encoding_format == "float" or encoding_format == "base64": - return self.request_output_to_pooling_json_response( - final_res_batch, - request_id, - created_time, - model_name, - encoding_format, - embed_dtype, - endianness, - ) - - if encoding_format == "bytes" or encoding_format == "bytes_only": - return self.request_output_to_pooling_bytes_response( - final_res_batch, - request_id, - created_time, - model_name, - encoding_format, - embed_dtype, - endianness, - ) - - assert_never(encoding_format) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 7ff324f120d9..daefd09fe4e8 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -133,11 +133,6 @@ def __init__( ) self.renderer = renderer = renderer_from_config(self.vllm_config) - self.io_processor = get_io_processor( - self.vllm_config, - self.renderer, - self.model_config.io_processor_plugin, - ) # Convert EngineInput --> EngineCoreRequest. self.input_processor = InputProcessor(self.vllm_config, renderer) From 60cef4211742cd71ce8ab4301999f50b941c4caa Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 15:52:20 +0800 Subject: [PATCH 05/20] refine Signed-off-by: wang.yuqi --- vllm/engine/protocol.py | 1 - vllm/entrypoints/llm.py | 4 +- vllm/entrypoints/openai/api_server.py | 8 +--- vllm/entrypoints/openai/models/serving.py | 1 - vllm/entrypoints/pooling/base/io_processor.py | 6 +-- .../entrypoints/pooling/pooling/api_router.py | 30 ++----------- vllm/entrypoints/pooling/pooling/serving.py | 44 +++++++++++++++++-- vllm/entrypoints/serve/render/serving.py | 2 - vllm/v1/engine/async_llm.py | 1 - vllm/v1/engine/llm_engine.py | 1 - 10 files changed, 49 insertions(+), 49 deletions(-) diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 1cedb86bd17d..50013a060a8f 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -14,7 +14,6 @@ from vllm.inputs import EngineInput, PromptType from vllm.lora.request import LoRARequest from vllm.outputs import PoolingRequestOutput, RequestOutput -from vllm.plugins.io_processors import IOProcessor from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 3fb8648ceda0..661ce23d5a2f 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -49,9 +49,7 @@ load_chat_template, ) from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors -from vllm.entrypoints.pooling.scoring.io_processor import ( - ScoringIOProcessor, -) +from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessor from vllm.entrypoints.pooling.scoring.typing import ScoreInput from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext from vllm.entrypoints.utils import log_non_default_args diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index bf23c7e2e96b..48eec918ee41 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -370,7 +370,6 @@ async def init_app_state( state.openai_serving_render = OpenAIServingRender( model_config=engine_client.model_config, renderer=engine_client.renderer, - io_processor=engine_client.io_processor, model_registry=state.openai_serving_models.registry, request_logger=request_logger, chat_template=resolved_chat_template, @@ -440,13 +439,12 @@ async def init_render_app_state( Unlike :func:`init_app_state` this function does not require an :class:`~vllm.engine.protocol.EngineClient`; it bootstraps the - preprocessing pipeline (renderer, io_processor, input_processor) + preprocessing pipeline (renderer, input_processor) directly from the :class:`~vllm.config.VllmConfig`. """ from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.openai.models.serving import OpenAIModelRegistry from vllm.entrypoints.serve.render.serving import OpenAIServingRender - from vllm.plugins.io_processors import get_io_processor from vllm.renderers import renderer_from_config served_model_names = args.served_model_name or [args.model] @@ -464,15 +462,11 @@ async def init_render_app_state( request_logger = None renderer = renderer_from_config(vllm_config) - io_processor = get_io_processor( - vllm_config, renderer, vllm_config.model_config.io_processor_plugin - ) resolved_chat_template = load_chat_template(args.chat_template) state.openai_serving_render = OpenAIServingRender( model_config=vllm_config.model_config, renderer=renderer, - io_processor=io_processor, model_registry=model_registry, request_logger=request_logger, chat_template=resolved_chat_template, diff --git a/vllm/entrypoints/openai/models/serving.py b/vllm/entrypoints/openai/models/serving.py index dd7a8687f2b5..ba1902d5b772 100644 --- a/vllm/entrypoints/openai/models/serving.py +++ b/vllm/entrypoints/openai/models/serving.py @@ -112,7 +112,6 @@ def __init__( self.model_config = self.engine_client.model_config self.renderer = self.engine_client.renderer - self.io_processor = self.engine_client.io_processor self.input_processor = self.engine_client.input_processor async def init_static_loras(self): diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index f5de5c295170..c0237d2ac344 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -49,12 +49,12 @@ def __init__( chat_template_config.trust_request_chat_template ) - def create_pooling_params(self, request): - return request.to_pooling_params() - ####################################### # online APIs + def create_pooling_params(self, request): + return request.to_pooling_params() + def pre_process_online(self, ctx: PoolingServeContext): request = ctx.request diff --git a/vllm/entrypoints/pooling/pooling/api_router.py b/vllm/entrypoints/pooling/pooling/api_router.py index f63a8edf6ca8..a08570038c3a 100644 --- a/vllm/entrypoints/pooling/pooling/api_router.py +++ b/vllm/entrypoints/pooling/pooling/api_router.py @@ -3,24 +3,17 @@ from http import HTTPStatus from fastapi import APIRouter, Depends, Request -from fastapi.responses import JSONResponse, StreamingResponse -from typing_extensions import assert_never from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.pooling.pooling.protocol import ( - IOProcessorResponse, - PoolingBytesResponse, - PoolingRequest, - PoolingResponse, -) -from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling +from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest +from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.utils import load_aware_call, with_cancellation router = APIRouter() -def pooling(request: Request) -> OpenAIServingPooling | None: +def pooling(request: Request) -> ServingPooling | None: return request.app.state.serving_pooling @@ -39,19 +32,4 @@ async def create_pooling(request: PoolingRequest, raw_request: Request): if handler is None: raise NotImplementedError("The model does not support Pooling API") - generator = await handler.create_pooling(request, raw_request) - - if isinstance(generator, ErrorResponse): - return JSONResponse( - content=generator.model_dump(), status_code=generator.error.code - ) - elif isinstance(generator, (PoolingResponse, IOProcessorResponse)): - return JSONResponse(content=generator.model_dump()) - elif isinstance(generator, PoolingBytesResponse): - return StreamingResponse( - content=generator.content, - headers=generator.headers, - media_type=generator.media_type, - ) - - assert_never(generator) + return await handler(request, raw_request) diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index a552d7dd08e4..15b1283eae4d 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -15,6 +15,7 @@ from vllm.entrypoints.pooling.base.serving import PoolingServingBase from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors from vllm.entrypoints.pooling.pooling.protocol import ( + IOProcessorRequest, PoolingBytesResponse, PoolingResponse, PoolingResponseData, @@ -59,14 +60,49 @@ async def __call__( raw_request: Request | None = None, ) -> Response: ctx = await self._init_ctx(request, raw_request) - await self.io_processor.pre_process_online_async(ctx) - if ctx.pooling_params is None: - ctx.pooling_params = self.io_processor.create_pooling_params(request) + if request.task is None: + request.task = self.pooling_task + + if getattr(request, "dimensions", None) is not None: + raise ValueError("dimensions is currently not supported") + + # plugin task uses io_processor.parse_request to verify inputs + if request.task != "plugin" and request.task != self.pooling_task: + if request.task not in self.supported_tasks: + raise ValueError( + f"Unsupported task: {request.task!r} " + f"Supported tasks: {self.supported_tasks}" + ) + else: + logger.warning_once( + "Pooling multitask support is deprecated and will be removed " + "in v0.20. When the default pooling task is not what you want, you " + "need to manually specify it via --pooler-config.task %s. ", + request.task, + ) + + if isinstance(request, IOProcessorRequest): + if "plugin" not in self.io_processors: + raise ValueError( + "No IOProcessor plugin installed. Please refer " + "to the documentation and to the " + "'prithvi_geospatial_mae_io_processor' " + "offline inference example for more details." + ) + + io_processor = self.io_processors["plugin"] + else: + io_processor = self.io_processors[request.task] + await io_processor.pre_process_online_async(ctx) + + if ctx.pooling_params is None: + ctx.pooling_params = io_processor.create_pooling_params(request) await self._prepare_generators(ctx) await self._collect_batch(ctx) - await self.io_processor.post_process_online_async(ctx) + + await io_processor.post_process_online_async(ctx) return await self._build_response(ctx) async def _build_response( diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 43ea5127b3be..30422a5ca631 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -64,7 +64,6 @@ def __init__( self, model_config: ModelConfig, renderer: BaseRenderer, - io_processor: Any, model_registry: OpenAIModelRegistry, *, request_logger: RequestLogger | None, @@ -79,7 +78,6 @@ def __init__( ) -> None: self.model_config = model_config self.renderer = renderer - self.io_processor = io_processor self.model_registry = model_registry self.request_logger = request_logger self.chat_template = chat_template diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index daefd09fe4e8..1c87d9ec0944 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -26,7 +26,6 @@ from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput -from vllm.plugins.io_processors import get_io_processor from vllm.pooling_params import PoolingParams from vllm.renderers import renderer_from_config from vllm.renderers.inputs.preprocess import extract_prompt_components diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 55ccdde91339..d0545651b966 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -19,7 +19,6 @@ from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.outputs import PoolingRequestOutput, RequestOutput -from vllm.plugins.io_processors import get_io_processor from vllm.pooling_params import PoolingParams from vllm.renderers import renderer_from_config from vllm.renderers.inputs.preprocess import extract_prompt_components From 1078300a67c44698e525bd2412256a2e5091f329 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 16:41:01 +0800 Subject: [PATCH 06/20] refine Signed-off-by: wang.yuqi --- .../pooling/classify/test_offline.py | 7 +++++-- .../pooling/classify/test_online.py | 6 +++++- .../entrypoints/pooling/embed/test_offline.py | 7 +++++-- .../entrypoints/pooling/embed/test_online.py | 6 +++++- .../pooling/scoring/test_bi_encoder_online.py | 6 +++++- .../scoring/test_cross_encoder_online.py | 6 +++++- .../pooling/token_classify/test_offline.py | 7 +++++-- .../pooling/token_classify/test_online.py | 9 +++++---- .../pooling/token_embed/test_offline.py | 7 +++++-- .../pooling/token_embed/test_online.py | 9 +++++---- vllm/entrypoints/llm.py | 20 +++++++++---------- vllm/entrypoints/pooling/pooling/serving.py | 18 ++++++++--------- vllm/entrypoints/pooling/scoring/serving.py | 2 +- vllm/entrypoints/pooling/typing.py | 4 ++-- 14 files changed, 72 insertions(+), 42 deletions(-) diff --git a/tests/entrypoints/pooling/classify/test_offline.py b/tests/entrypoints/pooling/classify/test_offline.py index f556dd579e64..828443ebc419 100644 --- a/tests/entrypoints/pooling/classify/test_offline.py +++ b/tests/entrypoints/pooling/classify/test_offline.py @@ -110,8 +110,11 @@ def test_score_api(llm: LLM): llm.score("ping", "pong", use_tqdm=False) -@pytest.mark.parametrize("task", ["embed", "token_embed"]) +@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask): - err_msg = "Embedding API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Embedding API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/classify/test_online.py b/tests/entrypoints/pooling/classify/test_online.py index ed295a09b88c..848ce4083c50 100644 --- a/tests/entrypoints/pooling/classify/test_online.py +++ b/tests/entrypoints/pooling/classify/test_online.py @@ -469,4 +469,8 @@ async def test_pooling_not_supported( }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/embed/test_offline.py b/tests/entrypoints/pooling/embed/test_offline.py index e8d84ed45e0d..1ffeb027b489 100644 --- a/tests/entrypoints/pooling/embed/test_offline.py +++ b/tests/entrypoints/pooling/embed/test_offline.py @@ -107,8 +107,11 @@ def get_outputs(normalize): ) -@pytest.mark.parametrize("task", ["token_classify", "classify"]) +@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask): - err_msg = "Classification API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Classification API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index dc61244c9445..3032645c7b17 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -767,4 +767,8 @@ async def test_pooling_not_supported( }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py b/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py index 38146084e379..392514056645 100644 --- a/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py +++ b/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py @@ -411,4 +411,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str): }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py index ebb339263ced..059a32d7ca38 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py @@ -484,4 +484,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str): }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/token_classify/test_offline.py b/tests/entrypoints/pooling/token_classify/test_offline.py index f7a74675463a..d36761466dfe 100644 --- a/tests/entrypoints/pooling/token_classify/test_offline.py +++ b/tests/entrypoints/pooling/token_classify/test_offline.py @@ -65,14 +65,17 @@ def test_score_api(llm: LLM): llm.score("ping", "pong", use_tqdm=False) -@pytest.mark.parametrize("task", ["classify", "embed", "token_embed"]) +@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm): if task == "classify": with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"): llm.encode(prompt, pooling_task=task, use_tqdm=False) assert "deprecated" in caplog_vllm.text else: - err_msg = "Embedding API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Embedding API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/token_classify/test_online.py b/tests/entrypoints/pooling/token_classify/test_online.py index e91d0bc9a396..39fd378336c1 100644 --- a/tests/entrypoints/pooling/token_classify/test_online.py +++ b/tests/entrypoints/pooling/token_classify/test_online.py @@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"]) +@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"]) async def test_pooling_not_supported( server: RemoteOpenAIServer, model_name: str, task: str ): @@ -64,7 +64,8 @@ async def test_pooling_not_supported( }, ) - if task != "classify": - assert response.json()["error"]["type"] == "BadRequestError" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: err_msg = f"Unsupported task: {task!r}" - assert response.json()["error"]["message"].startswith(err_msg) + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/token_embed/test_offline.py b/tests/entrypoints/pooling/token_embed/test_offline.py index 697f4f81a11b..d2e87fbf23e5 100644 --- a/tests/entrypoints/pooling/token_embed/test_offline.py +++ b/tests/entrypoints/pooling/token_embed/test_offline.py @@ -62,14 +62,17 @@ def test_token_ids_prompts(llm: LLM): assert outputs[0].outputs.data.shape == (11, 384) -@pytest.mark.parametrize("task", ["embed", "classify", "token_classify"]) +@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm): if task == "embed": with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"): llm.encode(prompt, pooling_task=task, use_tqdm=False) assert "deprecated" in caplog_vllm.text else: - err_msg = "Classification API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Classification API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/token_embed/test_online.py b/tests/entrypoints/pooling/token_embed/test_online.py index 922c624e98ee..048491dac15d 100644 --- a/tests/entrypoints/pooling/token_embed/test_online.py +++ b/tests/entrypoints/pooling/token_embed/test_online.py @@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"]) +@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"]) async def test_pooling_not_supported( server: RemoteOpenAIServer, model_name: str, task: str ): @@ -87,7 +87,8 @@ async def test_pooling_not_supported( }, ) - if task != "embed": - assert response.json()["error"]["type"] == "BadRequestError" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: err_msg = f"Unsupported task: {task!r}" - assert response.json()["error"]["message"].startswith(err_msg) + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 661ce23d5a2f..b1651e4dea04 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1078,23 +1078,15 @@ def encode( pooled hidden states in the same order as the input prompts. """ - self._verify_pooling_task(pooling_task) - if isinstance(prompts, dict) and "data" in prompts: if pooling_task != "plugin": raise ValueError( "The 'data' field is only supported for the 'plugin' pooling task." ) - if "plugin" not in self.pooling_io_processors: - raise ValueError( - "No IOProcessor plugin installed. Please refer " - "to the documentation and to the " - "'prithvi_geospatial_mae_io_processor' " - "offline inference example for more details." - ) - + self._verify_pooling_task(pooling_task) assert pooling_task is not None and pooling_task in self.pooling_io_processors + io_processor = self.pooling_io_processors[pooling_task] ctx = OfflineInputsContext( @@ -1197,6 +1189,14 @@ def _verify_pooling_task(self, pooling_task: PoolingTask | None): pooling_task, ) + if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors: + raise ValueError( + "No IOProcessor plugin installed. Please refer " + "to the documentation and to the " + "'prithvi_geospatial_mae_io_processor' " + "offline inference example for more details." + ) + def embed( self, prompts: PromptType | Sequence[PromptType], diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 15b1283eae4d..3642f1b1b7f2 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -17,10 +17,11 @@ from vllm.entrypoints.pooling.pooling.protocol import ( IOProcessorRequest, PoolingBytesResponse, + PoolingRequest, PoolingResponse, PoolingResponseData, ) -from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext +from vllm.entrypoints.pooling.typing import PoolingServeContext from vllm.entrypoints.pooling.utils import ( encode_pooling_bytes, encode_pooling_output_base64, @@ -56,7 +57,7 @@ def __init__( async def __call__( self, - request: AnyPoolingRequest, + request: PoolingRequest, raw_request: Request | None = None, ) -> Response: ctx = await self._init_ctx(request, raw_request) @@ -82,7 +83,7 @@ async def __call__( request.task, ) - if isinstance(request, IOProcessorRequest): + if request.task == "plugin" or isinstance(request, IOProcessorRequest): if "plugin" not in self.io_processors: raise ValueError( "No IOProcessor plugin installed. Please refer " @@ -90,14 +91,13 @@ async def __call__( "'prithvi_geospatial_mae_io_processor' " "offline inference example for more details." ) + request.task = "plugin" - io_processor = self.io_processors["plugin"] - else: - io_processor = self.io_processors[request.task] - await io_processor.pre_process_online_async(ctx) + io_processor = self.io_processors[request.task] + await io_processor.pre_process_online_async(ctx) - if ctx.pooling_params is None: - ctx.pooling_params = io_processor.create_pooling_params(request) + if ctx.pooling_params is None: + ctx.pooling_params = io_processor.create_pooling_params(request) await self._prepare_generators(ctx) await self._collect_batch(ctx) diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index 0cde1daf6b85..fc52075611e6 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -60,7 +60,7 @@ def init_io_processor( assert score_type in ScoringIOProcessors processor_cls = ScoringIOProcessors[score_type] - return processor_cls(vllm_config=vllm_config, *args, **kwargs) + return processor_cls(vllm_config, *args, **kwargs) async def __call__(self, *args, **kwargs) -> Response: if not self.enable_flash_late_interaction: diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 66dd9dd4d2b4..ffe4a7d972e8 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -30,7 +30,7 @@ ) from vllm.entrypoints.pooling.scoring.protocol import ScoringRequest, ScoringResponse from vllm.entrypoints.pooling.scoring.typing import ScoringData -from vllm.inputs import EngineInput +from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest PoolingCompletionLikeRequest: TypeAlias = ( @@ -89,7 +89,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): @dataclass class OfflineInputsContext: - prompts: PromptType | Sequence[PromptType] | ScoringData + prompts: PromptType | Sequence[PromptType] | DataPrompt | ScoringData pooling_params: PoolingParams | list[PoolingParams] | None = None tokenization_kwargs: dict[str, Any] | None = None chat_template: str | None = None From 16028362f741192bc7d8332ccf9ad54259b1d5bd Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 17:14:47 +0800 Subject: [PATCH 07/20] refine Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/__init__.py | 22 +++++++++++-------- vllm/entrypoints/pooling/base/io_processor.py | 4 +++- vllm/entrypoints/pooling/base/serving.py | 11 ++++++++-- .../pooling/pooling/io_processor.py | 8 ++++++- vllm/entrypoints/pooling/pooling/serving.py | 6 +++-- vllm/entrypoints/sagemaker/api_router.py | 4 ++-- 6 files changed, 38 insertions(+), 17 deletions(-) diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index cc8368813612..a30f18558aac 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -73,11 +73,7 @@ def init_pooling_state( from vllm.tasks import POOLING_TASKS model_config = engine_client.model_config - chat_template_config = ChatTemplateConfig( - chat_template=load_chat_template(args.chat_template), - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, - ) + resolved_chat_template = load_chat_template(args.chat_template) state.serving_pooling = ( ( @@ -86,7 +82,9 @@ def init_pooling_state( state.openai_serving_models, supported_tasks=supported_tasks, request_logger=request_logger, - chat_template_config=chat_template_config, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, ) ) if any(t in supported_tasks for t in POOLING_TASKS) @@ -97,7 +95,9 @@ def init_pooling_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template_config=chat_template_config, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, ) if "embed" in supported_tasks else None @@ -107,7 +107,9 @@ def init_pooling_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template_config=chat_template_config, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, ) if "classify" in supported_tasks else None @@ -117,7 +119,9 @@ def init_pooling_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template_config=chat_template_config, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, enable_flash_late_interaction=getattr( args, "enable_flash_late_interaction", True ), diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index c0237d2ac344..79f350382ddd 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -101,7 +101,9 @@ async def post_process_online_async( # offline APIs def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: - assert not isinstance(ctx.prompts, ScoringData) + assert not isinstance(ctx.prompts, ScoringData) and not ( + isinstance(ctx.prompts, dict) and "data" in ctx.prompts + ) prompts_seq = prompt_to_seq(ctx.prompts) tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index ab9dc573b763..24138b58c353 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -13,6 +13,7 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ( ChatTemplateConfig, + ChatTemplateContentFormatOption, ) from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.engine.protocol import ErrorResponse @@ -42,8 +43,10 @@ def __init__( engine_client: EngineClient, models: OpenAIServingModels, *, - chat_template_config: ChatTemplateConfig, request_logger: RequestLogger | None, + chat_template: str | None = None, + chat_template_content_format: ChatTemplateContentFormatOption = "auto", + trust_request_chat_template: bool = False, return_tokens_as_token_ids: bool = False, log_error_stack: bool = False, ): @@ -56,7 +59,11 @@ def __init__( self.request_logger = request_logger self.return_tokens_as_token_ids = return_tokens_as_token_ids self.log_error_stack = log_error_stack - self.chat_template_config = chat_template_config + self.chat_template_config = ChatTemplateConfig( + chat_template=chat_template, + chat_template_content_format=chat_template_content_format, + trust_request_chat_template=trust_request_chat_template, + ) async def __call__( self, diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index f6ef1e7b8135..cd24577bbfcf 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -23,16 +23,22 @@ class PluginIOProcessor(PoolingIOProcessor): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.io_processor = get_io_processor( + io_processor = get_io_processor( self.vllm_config, self.renderer, self.model_config.io_processor_plugin, ) + assert io_processor is not None + self.io_processor = io_processor + ####################################### # offline APIs def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: + assert isinstance(ctx.prompts, dict) and "data" in ctx.prompts + assert ctx.pooling_params is not None + # Validate the request data is valid for the loaded plugin prompt_data = ctx.prompts.get("data") if prompt_data is None: diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 3642f1b1b7f2..3f95672e4451 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -21,7 +21,7 @@ PoolingResponse, PoolingResponseData, ) -from vllm.entrypoints.pooling.typing import PoolingServeContext +from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext from vllm.entrypoints.pooling.utils import ( encode_pooling_bytes, encode_pooling_output_base64, @@ -57,9 +57,11 @@ def __init__( async def __call__( self, - request: PoolingRequest, + request: AnyPoolingRequest, raw_request: Request | None = None, ) -> Response: + assert isinstance(request, PoolingRequest) + ctx = await self._init_ctx(request, raw_request) if request.task is None: diff --git a/vllm/entrypoints/sagemaker/api_router.py b/vllm/entrypoints/sagemaker/api_router.py index 45f5613bf5e8..1d63793df398 100644 --- a/vllm/entrypoints/sagemaker/api_router.py +++ b/vllm/entrypoints/sagemaker/api_router.py @@ -14,7 +14,7 @@ from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.pooling.base.serving import PoolingServing +from vllm.entrypoints.pooling.base.serving import PoolingServingBase from vllm.entrypoints.pooling.utils import enable_scoring_api from vllm.entrypoints.serve.instrumentator.basic import base from vllm.entrypoints.serve.instrumentator.health import health @@ -23,7 +23,7 @@ # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers # (requires typing_extensions >= 4.13) RequestType = Any -GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServing | None] +GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServingBase | None] EndpointFn = Callable[[RequestType, Request], Awaitable[Any]] From 7a65d08d09ee0c2d2d16f587ff0dc9d60d5debe8 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 17:33:01 +0800 Subject: [PATCH 08/20] mypy Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 9 ++++----- vllm/entrypoints/pooling/__init__.py | 1 - vllm/entrypoints/pooling/pooling/serving.py | 5 ++++- vllm/entrypoints/pooling/typing.py | 2 +- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index b1651e4dea04..3bba854b760b 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1078,11 +1078,10 @@ def encode( pooled hidden states in the same order as the input prompts. """ - if isinstance(prompts, dict) and "data" in prompts: - if pooling_task != "plugin": - raise ValueError( - "The 'data' field is only supported for the 'plugin' pooling task." - ) + if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin": + raise ValueError( + "The 'data' field is only supported for the 'plugin' pooling task." + ) self._verify_pooling_task(pooling_task) assert pooling_task is not None and pooling_task in self.pooling_io_processors diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index a30f18558aac..1980750ec20f 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -6,7 +6,6 @@ from fastapi import FastAPI from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.pooling.utils import enable_scoring_api from vllm.logger import init_logger diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 3f95672e4451..ecb4afa18647 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -95,7 +95,10 @@ async def __call__( ) request.task = "plugin" - io_processor = self.io_processors[request.task] + pooling_task = request.task + assert pooling_task is not None + + io_processor = self.io_processors[pooling_task] await io_processor.pre_process_online_async(ctx) if ctx.pooling_params is None: diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index ffe4a7d972e8..4d237f472492 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -90,7 +90,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): @dataclass class OfflineInputsContext: prompts: PromptType | Sequence[PromptType] | DataPrompt | ScoringData - pooling_params: PoolingParams | list[PoolingParams] | None = None + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None tokenization_kwargs: dict[str, Any] | None = None chat_template: str | None = None From 44e9f9628033212fd452768bdddb29072b2b4b2c Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 18:01:35 +0800 Subject: [PATCH 09/20] refine Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/embed/serving.py | 10 +++++++--- vllm/entrypoints/pooling/pooling/serving.py | 19 ++++++++++++++----- .../pooling/scoring/io_processor.py | 2 +- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index 8a477825e70f..33352288cc38 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -35,7 +35,6 @@ logger = init_logger(__name__) -JSONResponseCLS = get_json_response_cls() EmbeddingServeContext: TypeAlias = PoolingServeContext[EmbeddingRequest] @@ -46,6 +45,11 @@ class ServingEmbedding(PoolingServing): request_id_prefix = "embd" io_processor: EmbedIOProcessor + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.JSONResponseCLS = get_json_response_cls() + def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor: return EmbedIOProcessor(*args, **kwargs) @@ -137,7 +141,7 @@ def _openai_json_response( data=items, usage=usage, ) - return JSONResponseCLS(content=response.model_dump()) + return self.JSONResponseCLS(content=response.model_dump()) def _openai_bytes_response( self, @@ -206,4 +210,4 @@ def _build_cohere_response_from_ctx( ), ), ) - return JSONResponse(content=response.model_dump(exclude_none=True)) + return self.JSONResponse(content=response.model_dump(exclude_none=True)) diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index ecb4afa18647..c231efdf3af3 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -8,7 +8,7 @@ from typing import Literal, cast from fastapi import Request -from fastapi.responses import Response +from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never from vllm.entrypoints.openai.engine.protocol import UsageInfo @@ -26,6 +26,7 @@ encode_pooling_bytes, encode_pooling_output_base64, encode_pooling_output_float, + get_json_response_cls, ) from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput @@ -54,6 +55,7 @@ def __init__( renderer=self.renderer, chat_template_config=self.chat_template_config, ) + self.JSONResponseCLS = get_json_response_cls() async def __call__( self, @@ -151,7 +153,7 @@ def request_output_to_pooling_json_response( encoding_format: Literal["float", "base64"], embed_dtype: EmbedDType, endianness: Endianness, - ) -> PoolingResponse: + ) -> JSONResponse: encode_fn = cast( Callable[[PoolingRequestOutput], list[float] | str], ( @@ -183,13 +185,14 @@ def request_output_to_pooling_json_response( total_tokens=num_prompt_tokens, ) - return PoolingResponse( + response = PoolingResponse( id=request_id, created=created_time, model=model_name, data=items, usage=usage, ) + return self.JSONResponseCLS(content=response.model_dump()) def request_output_to_pooling_bytes_response( self, @@ -200,7 +203,7 @@ def request_output_to_pooling_bytes_response( encoding_format: Literal["bytes", "bytes_only"], embed_dtype: EmbedDType, endianness: Endianness, - ) -> PoolingBytesResponse: + ) -> StreamingResponse: content, items, usage = encode_pooling_bytes( pooling_outputs=final_res_batch, embed_dtype=embed_dtype, @@ -223,4 +226,10 @@ def request_output_to_pooling_bytes_response( } ) - return PoolingBytesResponse(content=content, headers=headers) + response = PoolingBytesResponse(content=content, headers=headers) + + return StreamingResponse( + content=response.content, + headers=response.headers, + media_type=response.media_type, + ) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index c520eb5ceb3d..dd505c79cf18 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -278,7 +278,7 @@ def pre_process_online(self, ctx: ScoringServeContext): def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: assert isinstance(ctx.prompts, ScoringData) - assert not isinstance(ctx.pooling_params, list) + assert not isinstance(ctx.pooling_params, Sequence) tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) From bc0916ce45cf7e1059533dfe004a12ba63835632 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Tue, 7 Apr 2026 18:11:23 +0800 Subject: [PATCH 10/20] refine Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/embed/serving.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index 33352288cc38..a1c45e2bf0e3 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -182,8 +182,8 @@ def _openai_bytes_response( media_type=response.media_type, ) - @staticmethod def _build_cohere_response_from_ctx( + self, ctx: PoolingServeContext, ) -> JSONResponse: request = ctx.request @@ -210,4 +210,4 @@ def _build_cohere_response_from_ctx( ), ), ) - return self.JSONResponse(content=response.model_dump(exclude_none=True)) + return self.JSONResponseCLS(content=response.model_dump(exclude_none=True)) From 3e8f11ba6b2889642d1dfbcd357d3beea1df813d Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 12:10:37 +0800 Subject: [PATCH 11/20] refine Signed-off-by: wang.yuqi --- .../openai/chat_completion/test_chat_error.py | 3 - .../chat_completion/test_serving_chat.py | 15 ---- .../completion/test_completion_error.py | 3 - .../openai/completion/test_lora_resolvers.py | 2 - .../test_generative_scoring.py | 1 - .../responses/test_serving_responses.py | 4 -- .../serve/disagg/test_generate_stream.py | 2 - .../serve/lora/test_serving_models.py | 1 - tests/v1/engine/test_async_llm.py | 1 - vllm/entrypoints/pooling/base/serving.py | 9 ++- vllm/entrypoints/pooling/embed/serving.py | 6 +- .../pooling/pooling/io_processor.py | 70 ++++++++++++++++++- vllm/entrypoints/pooling/pooling/serving.py | 8 ++- vllm/entrypoints/pooling/typing.py | 3 + 14 files changed, 88 insertions(+), 40 deletions(-) diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index 46070e4810be..f1fb7c7518b6 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -87,7 +87,6 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat: serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, @@ -123,7 +122,6 @@ async def test_chat_error_non_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -173,7 +171,6 @@ async def test_chat_error_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index cb356e0e1986..39d59d28f854 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -567,7 +567,6 @@ def _build_serving_render( return OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=model_registry, request_logger=None, chat_template=CHAT_TEMPLATE, @@ -599,7 +598,6 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat: class MockEngine: model_config: MockModelConfig = field(default_factory=MockModelConfig) input_processor: MagicMock = field(default_factory=MagicMock) - io_processor: MagicMock = field(default_factory=MagicMock) renderer: MagicMock = field(default_factory=MagicMock) @@ -632,7 +630,6 @@ async def test_serving_chat_returns_correct_model_name(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -662,7 +659,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -693,7 +689,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Initialize the serving chat @@ -737,7 +732,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -779,7 +773,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Initialize the serving chat @@ -823,7 +816,6 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_tokenizer = MagicMock(spec=MistralTokenizer) mock_renderer = MistralRenderer( @@ -863,7 +855,6 @@ async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected(): mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_tokenizer = MagicMock(spec=MistralTokenizer) mock_renderer = MistralRenderer( @@ -906,7 +897,6 @@ async def test_serving_chat_could_load_correct_generation_config(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Initialize the serving chat @@ -952,7 +942,6 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -1003,7 +992,6 @@ async def test_serving_chat_data_parallel_rank_extraction(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Mock the generate method to return an async generator @@ -1095,7 +1083,6 @@ def mock_engine(self) -> AsyncLLM: mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) return mock_engine @@ -1732,7 +1719,6 @@ async def test_tool_choice_validation_without_parser(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) models = OpenAIServingModels( @@ -1802,7 +1788,6 @@ async def test_streaming_n_gt1_independent_tool_parsers(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) models = OpenAIServingModels( diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 46eb02e3c599..3349f4126bc8 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -79,7 +79,6 @@ def _build_serving_completion(engine: AsyncLLM) -> OpenAIServingCompletion: serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, @@ -107,7 +106,6 @@ async def test_completion_error_non_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_completion = _build_serving_completion(mock_engine) @@ -157,7 +155,6 @@ async def test_completion_error_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_completion = _build_serving_completion(mock_engine) diff --git a/tests/entrypoints/openai/completion/test_lora_resolvers.py b/tests/entrypoints/openai/completion/test_lora_resolvers.py index 8d5283de5cf4..6a0bec92516d 100644 --- a/tests/entrypoints/openai/completion/test_lora_resolvers.py +++ b/tests/entrypoints/openai/completion/test_lora_resolvers.py @@ -137,7 +137,6 @@ async def mock_generate(*args, **kwargs): mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) models = OpenAIServingModels( @@ -148,7 +147,6 @@ async def mock_generate(*args, **kwargs): serving_render = OpenAIServingRender( model_config=mock_engine.model_config, renderer=mock_engine.renderer, - io_processor=mock_engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, diff --git a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py b/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py index a260027af0fc..632c4bcc90ae 100644 --- a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py @@ -77,7 +77,6 @@ def _create_mock_engine(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() # renderer is accessed by OpenAIServing.__init__ and serving.py mock_renderer = MagicMock() diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 39429cb9bf96..cc6179bb5233 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -148,7 +148,6 @@ async def serving_responses_instance(self): engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() models = MagicMock() @@ -237,7 +236,6 @@ async def serving_responses_instance(self): engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() models = MagicMock() @@ -299,7 +297,6 @@ def get_vocab(self): model_config.get_diff_sampling_param.return_value = {} engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() tokenizer = FakeTokenizer() @@ -602,7 +599,6 @@ def _make_serving_instance_with_reasoning(): model_config.get_diff_sampling_param.return_value = {} engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() models = MagicMock() diff --git a/tests/entrypoints/serve/disagg/test_generate_stream.py b/tests/entrypoints/serve/disagg/test_generate_stream.py index a9ca026306fa..76a9df22f69b 100644 --- a/tests/entrypoints/serve/disagg/test_generate_stream.py +++ b/tests/entrypoints/serve/disagg/test_generate_stream.py @@ -86,7 +86,6 @@ def _build_serving_tokens(engine: AsyncLLM, **kwargs) -> ServingTokens: serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, @@ -148,7 +147,6 @@ def _mock_engine() -> MagicMock: engine.errored = False engine.model_config = MockModelConfig() engine.input_processor = MagicMock() - engine.io_processor = MagicMock() engine.renderer = _build_renderer(engine.model_config) return engine diff --git a/tests/entrypoints/serve/lora/test_serving_models.py b/tests/entrypoints/serve/lora/test_serving_models.py index f6755f489343..ce9fdcc2bfb2 100644 --- a/tests/entrypoints/serve/lora/test_serving_models.py +++ b/tests/entrypoints/serve/lora/test_serving_models.py @@ -34,7 +34,6 @@ async def _async_serving_models_init() -> OpenAIServingModels: mock_model_config.max_model_len = 2048 mock_engine_client.model_config = mock_model_config mock_engine_client.input_processor = MagicMock() - mock_engine_client.io_processor = MagicMock() mock_engine_client.renderer = MagicMock() serving_models = OpenAIServingModels( diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 69a1c38a453d..21a651c62ab3 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -514,7 +514,6 @@ async def test_header_dp_rank_argument(): serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 24138b58c353..10f5cd6b5942 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod from collections.abc import AsyncGenerator, Mapping from http import HTTPStatus from typing import ClassVar @@ -35,7 +37,7 @@ from .io_processor import PoolingIOProcessor -class PoolingServingBase: +class PoolingServingBase(ABC): request_id_prefix: ClassVar[str] def __init__( @@ -65,6 +67,7 @@ def __init__( trust_request_chat_template=trust_request_chat_template, ) + @abstractmethod async def __call__( self, request: AnyPoolingRequest, @@ -171,6 +174,7 @@ async def _collect_batch( ctx.final_res_batch = [res for res in final_res_batch if res is not None] + @abstractmethod async def _build_response( self, ctx: PoolingServeContext, @@ -338,7 +342,7 @@ def _log_inputs( ) -class PoolingServing(PoolingServingBase): +class PoolingServing(PoolingServingBase, ABC): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -348,6 +352,7 @@ def __init__(self, *args, **kwargs): chat_template_config=self.chat_template_config, ) + @abstractmethod def init_io_processor( self, vllm_config: VllmConfig, diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index a1c45e2bf0e3..3abf0c7f3fd4 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -48,7 +48,7 @@ class ServingEmbedding(PoolingServing): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.JSONResponseCLS = get_json_response_cls() + self.json_response_cls = get_json_response_cls() def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor: return EmbedIOProcessor(*args, **kwargs) @@ -141,7 +141,7 @@ def _openai_json_response( data=items, usage=usage, ) - return self.JSONResponseCLS(content=response.model_dump()) + return self.json_response_cls(content=response.model_dump()) def _openai_bytes_response( self, @@ -210,4 +210,4 @@ def _build_cohere_response_from_ctx( ), ), ) - return self.JSONResponseCLS(content=response.model_dump(exclude_none=True)) + return self.json_response_cls(content=response.model_dump(exclude_none=True)) diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index cd24577bbfcf..11b0b7ebca5b 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -5,14 +5,19 @@ from vllm import PoolingParams, PoolingRequestOutput from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.entrypoints.pooling.pooling.protocol import ( + IOProcessorRequest, + IOProcessorResponse, +) from vllm.entrypoints.pooling.typing import ( OfflineInputsContext, OfflineOutputsContext, + PoolingServeContext, ) from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor -from vllm.renderers.inputs.preprocess import prompt_to_seq +from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq logger = init_logger(__name__) @@ -32,6 +37,69 @@ def __init__(self, *args, **kwargs): assert io_processor is not None self.io_processor = io_processor + ####################################### + # online APIs + + def pre_process_online(self, ctx: PoolingServeContext): + assert isinstance(ctx.request, IOProcessorRequest) + + validated_prompt = self.io_processor.parse_data(ctx.request.data) + + raw_prompts = self.io_processor.pre_process( + prompt=validated_prompt, request_id=ctx.request_id + ) + + parsed_prompts = [ + ( + prompt + if isinstance(prompt, bytes) + else parse_model_prompt(self.model_config, prompt) + ) + for prompt in prompt_to_seq(raw_prompts) + ] + + tok_params = ctx.request.build_tok_params(self.model_config) + + ctx.engine_inputs = self.renderer.render_cmpl( + parsed_prompts, + tok_params, + prompt_extras={ + k: v + for k in ("mm_processor_kwargs", "cache_salt") + if (v := getattr(ctx.request, k, None)) is not None + }, + ) + + pooling_params = self.io_processor.merge_pooling_params() + if pooling_params.task is None: + pooling_params.task = "plugin" + ctx.pooling_params = pooling_params + + def post_process_online( + self, + ctx: PoolingServeContext, + ): + output = self.io_processor.post_process( + ctx.final_res_batch, + request_id=ctx.request_id, + ) + + if callable( + output_to_response := getattr(self.io_processor, "output_to_response", None) + ): + logger.warning_once( + "`IOProcessor.output_to_response` is deprecated. To ensure " + "consistency between offline and online APIs, " + "`IOProcessorResponse` will become a transparent wrapper " + "around output data from v0.19 onwards.", + ) + + if hasattr(output, "request_id") and output.request_id is None: + output.request_id = request_id # type: ignore + + ctx.response = output_to_response(output) # type: ignore + ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) + ####################################### # offline APIs diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index c231efdf3af3..db2624772120 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -55,7 +55,7 @@ def __init__( renderer=self.renderer, chat_template_config=self.chat_template_config, ) - self.JSONResponseCLS = get_json_response_cls() + self.json_response_cls = get_json_response_cls() async def __call__( self, @@ -116,6 +116,10 @@ async def _build_response( self, ctx: PoolingServeContext, ) -> Response: + if ctx.response is not None: + # for IOProcessorResponse + return self.json_response_cls(content=ctx.response.model_dump()) + encoding_format = ctx.request.encoding_format embed_dtype = ctx.request.embed_dtype endianness = ctx.request.endianness @@ -192,7 +196,7 @@ def request_output_to_pooling_json_response( data=items, usage=usage, ) - return self.JSONResponseCLS(content=response.model_dump()) + return self.json_response_cls(content=response.model_dump()) def request_output_to_pooling_bytes_response( self, diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 4d237f472492..eb5d9af1f170 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -86,6 +86,9 @@ class PoolingServeContext(Generic[PoolingRequestT]): ## for bi-encoder & late-interaction n_queries: int | None = None + ## for IOProcessorResponse + response: Any | None = None + @dataclass class OfflineInputsContext: From d69a94a0d1d6f3a7ecd3ab2363ad41b3528f55cb Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 13:17:12 +0800 Subject: [PATCH 12/20] mypy Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/pooling/io_processor.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index 11b0b7ebca5b..b190ba8b934f 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -23,6 +23,9 @@ class PluginIOProcessor(PoolingIOProcessor): + """IO Processor plugins are a feature that allows pre- and post-processing + of the model input and output for pooling models.""" + name = "plugin" def __init__(self, *args, **kwargs): @@ -95,7 +98,7 @@ def post_process_online( ) if hasattr(output, "request_id") and output.request_id is None: - output.request_id = request_id # type: ignore + output.request_id = ctx.request_id # type: ignore ctx.response = output_to_response(output) # type: ignore ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) From 22bf2fad8b3067e621fb6d8791895c04222663e1 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 14:02:05 +0800 Subject: [PATCH 13/20] fix Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 3bba854b760b..f3d0b150b9de 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1088,6 +1088,9 @@ def encode( io_processor = self.pooling_io_processors[pooling_task] + if pooling_params is None: + pooling_params = PoolingParams() + ctx = OfflineInputsContext( prompts=prompts, pooling_params=pooling_params, @@ -1097,9 +1100,6 @@ def encode( engine_inputs = io_processor.pre_process_offline(ctx) n_inputs = len(engine_inputs) - if ctx.pooling_params is None: - ctx.pooling_params = PoolingParams() - params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) for param in params_seq: From 02af147c1a7580d3c728925719651671ba399625 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 14:53:55 +0800 Subject: [PATCH 14/20] fix Terratorch Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 1 - .../pooling/io_processor_factories.py | 25 +++++++++++-------- .../pooling/pooling/io_processor.py | 18 ++++++------- vllm/entrypoints/pooling/typing.py | 2 +- 4 files changed, 23 insertions(+), 23 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index f3d0b150b9de..0e93403c3862 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1082,7 +1082,6 @@ def encode( raise ValueError( "The 'data' field is only supported for the 'plugin' pooling task." ) - self._verify_pooling_task(pooling_task) assert pooling_task is not None and pooling_task in self.pooling_io_processors diff --git a/vllm/entrypoints/pooling/io_processor_factories.py b/vllm/entrypoints/pooling/io_processor_factories.py index 49182a9f0c8c..86bb899bfecb 100644 --- a/vllm/entrypoints/pooling/io_processor_factories.py +++ b/vllm/entrypoints/pooling/io_processor_factories.py @@ -3,13 +3,13 @@ from vllm.config import VllmConfig from vllm.entrypoints.chat_utils import ChatTemplateConfig -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessors -from vllm.entrypoints.pooling.utils import enable_scoring_api from vllm.plugins.io_processors import has_io_processor from vllm.renderers import BaseRenderer from vllm.tasks import SupportedTask +from .base.io_processor import PoolingIOProcessor +from .utils import enable_scoring_api + def init_pooling_io_processors( supported_tasks: tuple[SupportedTask, ...], @@ -21,24 +21,22 @@ def init_pooling_io_processors( processors: dict[str, type[PoolingIOProcessor]] = {} if "classify" in supported_tasks: - from vllm.entrypoints.pooling.classify.io_processor import ClassifyIOProcessor + from .classify.io_processor import ClassifyIOProcessor processors["classify"] = ClassifyIOProcessor if "token_classify" in supported_tasks: - from vllm.entrypoints.pooling.classify.io_processor import ( - TokenClassifyIOProcessor, - ) + from .classify.io_processor import TokenClassifyIOProcessor processors["token_classify"] = TokenClassifyIOProcessor if "embed" in supported_tasks: - from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor + from .embed.io_processor import EmbedIOProcessor processors["embed"] = EmbedIOProcessor if "token_embed" in supported_tasks: - from vllm.entrypoints.pooling.embed.io_processor import TokenEmbedIOProcessor + from .embed.io_processor import TokenEmbedIOProcessor processors["token_embed"] = TokenEmbedIOProcessor @@ -46,12 +44,17 @@ def init_pooling_io_processors( vllm_config, model_config.io_processor_plugin, ): - from vllm.entrypoints.pooling.pooling.io_processor import PluginIOProcessor + from .pooling.io_processor import PluginWithIOProcessor + + processors["plugin"] = PluginWithIOProcessor + elif "plugin" in supported_tasks: + from .pooling.io_processor import PluginWithoutIOProcessor - processors["plugin"] = PluginIOProcessor + processors["plugin"] = PluginWithoutIOProcessor if enable_scoring_api(supported_tasks, model_config): score_type = model_config.score_type + from .scoring.io_processor import ScoringIOProcessors if score_type is not None and score_type in ScoringIOProcessors: processors[score_type] = ScoringIOProcessors[score_type] diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index b190ba8b934f..c0b76d269df2 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -5,24 +5,22 @@ from vllm import PoolingParams, PoolingRequestOutput from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.pooling.protocol import ( - IOProcessorRequest, - IOProcessorResponse, -) -from vllm.entrypoints.pooling.typing import ( - OfflineInputsContext, - OfflineOutputsContext, - PoolingServeContext, -) from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq +from ..typing import OfflineInputsContext, OfflineOutputsContext, PoolingServeContext +from .protocol import IOProcessorRequest, IOProcessorResponse + logger = init_logger(__name__) -class PluginIOProcessor(PoolingIOProcessor): +class PluginWithoutIOProcessor(PoolingIOProcessor): + name = "plugin" + + +class PluginWithIOProcessor(PoolingIOProcessor): """IO Processor plugins are a feature that allows pre- and post-processing of the model input and output for pooling models.""" diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index eb5d9af1f170..ead772251f16 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -93,7 +93,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): @dataclass class OfflineInputsContext: prompts: PromptType | Sequence[PromptType] | DataPrompt | ScoringData - pooling_params: PoolingParams | Sequence[PoolingParams] | None = None + pooling_params: PoolingParams | Sequence[PoolingParams] tokenization_kwargs: dict[str, Any] | None = None chat_template: str | None = None From f622f6831aa6deeaf11538787eac5f61742a686e Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 15:14:50 +0800 Subject: [PATCH 15/20] mypy Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 1 + vllm/entrypoints/pooling/io_processor_factories.py | 8 ++++---- vllm/entrypoints/pooling/pooling/io_processor.py | 4 ++-- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 0e93403c3862..28dd6278ecc6 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1098,6 +1098,7 @@ def encode( engine_inputs = io_processor.pre_process_offline(ctx) n_inputs = len(engine_inputs) + assert ctx.pooling_params is not None params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) diff --git a/vllm/entrypoints/pooling/io_processor_factories.py b/vllm/entrypoints/pooling/io_processor_factories.py index 86bb899bfecb..de60a746a350 100644 --- a/vllm/entrypoints/pooling/io_processor_factories.py +++ b/vllm/entrypoints/pooling/io_processor_factories.py @@ -44,13 +44,13 @@ def init_pooling_io_processors( vllm_config, model_config.io_processor_plugin, ): - from .pooling.io_processor import PluginWithIOProcessor + from .pooling.io_processor import PluginWithIOProcessorPlugins - processors["plugin"] = PluginWithIOProcessor + processors["plugin"] = PluginWithIOProcessorPlugins elif "plugin" in supported_tasks: - from .pooling.io_processor import PluginWithoutIOProcessor + from .pooling.io_processor import PluginWithoutIOProcessorPlugins - processors["plugin"] = PluginWithoutIOProcessor + processors["plugin"] = PluginWithoutIOProcessorPlugins if enable_scoring_api(supported_tasks, model_config): score_type = model_config.score_type diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index c0b76d269df2..435cc988dfab 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -16,11 +16,11 @@ logger = init_logger(__name__) -class PluginWithoutIOProcessor(PoolingIOProcessor): +class PluginWithoutIOProcessorPlugins(PoolingIOProcessor): name = "plugin" -class PluginWithIOProcessor(PoolingIOProcessor): +class PluginWithIOProcessorPlugins(PoolingIOProcessor): """IO Processor plugins are a feature that allows pre- and post-processing of the model input and output for pooling models.""" From 124d33cc2bdead78c5b86d128e5202816bf34c61 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 15:37:15 +0800 Subject: [PATCH 16/20] mypy Signed-off-by: wang.yuqi --- vllm/entrypoints/llm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 28dd6278ecc6..d296e84d0411 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -1400,6 +1400,9 @@ def score( scoring_data = io_processor.valid_inputs(data_1, data_2) n_queries = len(scoring_data.data_1) + if pooling_params is None: + pooling_params = PoolingParams() + ctx = OfflineInputsContext( prompts=scoring_data, pooling_params=pooling_params, @@ -1412,9 +1415,6 @@ def score( n_inputs = len(engine_inputs) seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs) - - if ctx.pooling_params is None: - ctx.pooling_params = PoolingParams() params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) for param in params_seq: From 01582557ab5742184da222b2efd545cd203ddd4c Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 15:52:46 +0800 Subject: [PATCH 17/20] Update vllm/entrypoints/pooling/pooling/io_processor.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/pooling/io_processor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index 435cc988dfab..31f860144ea3 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -99,7 +99,8 @@ def post_process_online( output.request_id = ctx.request_id # type: ignore ctx.response = output_to_response(output) # type: ignore - ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) + else: + ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) ####################################### # offline APIs From a76447fedc3f744fc2f10d6d16281eea98e05bc0 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 15:58:33 +0800 Subject: [PATCH 18/20] refine Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/pooling/serving.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index db2624772120..ee00b6cfcba4 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -72,11 +72,14 @@ async def __call__( if getattr(request, "dimensions", None) is not None: raise ValueError("dimensions is currently not supported") + assert request.task is not None + pooling_task = request.task + # plugin task uses io_processor.parse_request to verify inputs - if request.task != "plugin" and request.task != self.pooling_task: - if request.task not in self.supported_tasks: + if pooling_task != "plugin" and pooling_task != self.pooling_task: + if pooling_task not in self.io_processors: raise ValueError( - f"Unsupported task: {request.task!r} " + f"Unsupported task: {pooling_task!r} " f"Supported tasks: {self.supported_tasks}" ) else: @@ -84,10 +87,10 @@ async def __call__( "Pooling multitask support is deprecated and will be removed " "in v0.20. When the default pooling task is not what you want, you " "need to manually specify it via --pooler-config.task %s. ", - request.task, + pooling_task, ) - if request.task == "plugin" or isinstance(request, IOProcessorRequest): + if pooling_task == "plugin" or isinstance(request, IOProcessorRequest): if "plugin" not in self.io_processors: raise ValueError( "No IOProcessor plugin installed. Please refer " @@ -95,10 +98,7 @@ async def __call__( "'prithvi_geospatial_mae_io_processor' " "offline inference example for more details." ) - request.task = "plugin" - - pooling_task = request.task - assert pooling_task is not None + pooling_task = "plugin" io_processor = self.io_processors[pooling_task] await io_processor.pre_process_online_async(ctx) From 03938fad0d2e8370a9752f9eb64031abaf3b154f Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 16:07:46 +0800 Subject: [PATCH 19/20] refine Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/pooling/serving.py | 25 +++++++++++---------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index ee00b6cfcba4..6df118d00aab 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -66,15 +66,26 @@ async def __call__( ctx = await self._init_ctx(request, raw_request) + if getattr(request, "dimensions", None) is not None: + raise ValueError("dimensions is currently not supported") + if request.task is None: request.task = self.pooling_task - if getattr(request, "dimensions", None) is not None: - raise ValueError("dimensions is currently not supported") + if isinstance(request, IOProcessorRequest): + request.task = "plugin" assert request.task is not None pooling_task = request.task + if pooling_task == "plugin" and "plugin" not in self.io_processors: + raise ValueError( + "No IOProcessor plugin installed. Please refer " + "to the documentation and to the " + "'prithvi_geospatial_mae_io_processor' " + "offline inference example for more details." + ) + # plugin task uses io_processor.parse_request to verify inputs if pooling_task != "plugin" and pooling_task != self.pooling_task: if pooling_task not in self.io_processors: @@ -90,16 +101,6 @@ async def __call__( pooling_task, ) - if pooling_task == "plugin" or isinstance(request, IOProcessorRequest): - if "plugin" not in self.io_processors: - raise ValueError( - "No IOProcessor plugin installed. Please refer " - "to the documentation and to the " - "'prithvi_geospatial_mae_io_processor' " - "offline inference example for more details." - ) - pooling_task = "plugin" - io_processor = self.io_processors[pooling_task] await io_processor.pre_process_online_async(ctx) From 3dd33260a9b353f0b27df9fcd6de0153529f33c4 Mon Sep 17 00:00:00 2001 From: "wang.yuqi" Date: Thu, 9 Apr 2026 16:16:08 +0800 Subject: [PATCH 20/20] refine Signed-off-by: wang.yuqi --- vllm/entrypoints/pooling/pooling/serving.py | 40 +++++++++++---------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 6df118d00aab..3d16c1f2cc9d 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -63,9 +63,23 @@ async def __call__( raw_request: Request | None = None, ) -> Response: assert isinstance(request, PoolingRequest) + pooling_task = self._verify_pooling_task(request) + io_processor = self.io_processors[pooling_task] ctx = await self._init_ctx(request, raw_request) + await io_processor.pre_process_online_async(ctx) + + if ctx.pooling_params is None: + ctx.pooling_params = io_processor.create_pooling_params(request) + + await self._prepare_generators(ctx) + await self._collect_batch(ctx) + + await io_processor.post_process_online_async(ctx) + return await self._build_response(ctx) + + def _verify_pooling_task(self, request: PoolingRequest) -> str: if getattr(request, "dimensions", None) is not None: raise ValueError("dimensions is currently not supported") @@ -78,14 +92,6 @@ async def __call__( assert request.task is not None pooling_task = request.task - if pooling_task == "plugin" and "plugin" not in self.io_processors: - raise ValueError( - "No IOProcessor plugin installed. Please refer " - "to the documentation and to the " - "'prithvi_geospatial_mae_io_processor' " - "offline inference example for more details." - ) - # plugin task uses io_processor.parse_request to verify inputs if pooling_task != "plugin" and pooling_task != self.pooling_task: if pooling_task not in self.io_processors: @@ -101,17 +107,15 @@ async def __call__( pooling_task, ) - io_processor = self.io_processors[pooling_task] - await io_processor.pre_process_online_async(ctx) - - if ctx.pooling_params is None: - ctx.pooling_params = io_processor.create_pooling_params(request) - - await self._prepare_generators(ctx) - await self._collect_batch(ctx) + if pooling_task == "plugin" and "plugin" not in self.io_processors: + raise ValueError( + "No IOProcessor plugin installed. Please refer " + "to the documentation and to the " + "'prithvi_geospatial_mae_io_processor' " + "offline inference example for more details." + ) - await io_processor.post_process_online_async(ctx) - return await self._build_response(ctx) + return pooling_task async def _build_response( self,