@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 61c9ef986a

[Frontend] Parallelize preprocessing within the same request for pooling models online serving. (#49153)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

wang.yuqi committed Jul 22, 2026 at 18:56 UTC 61c9ef986a807aa3b9c6ccd25bb223b8f4116ac7
17 files changed +530 -812
tests/entrypoints/pooling/embed/test_io_processor.py
+23 -231
@@ -9,8 +9,6 @@ from pydantic import TypeAdapter, ValidationError
9 from vllm import PoolingParams
10 from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
11 from vllm.entrypoints.pooling.embed.protocol import (
12 - CohereEmbedContent,
13 - CohereEmbedInput,
12 CohereEmbedRequest,
13 EmbeddingBatchChatInputRequest,
14 EmbeddingBatchChatRequest,
@@ -19,7 +17,10 @@ from vllm.entrypoints.pooling.embed.protocol import (
17 EmbeddingCompletionRequest,
18 EmbeddingRequest,
19 )
22 -from vllm.entrypoints.pooling.typing import PoolingServeContext
20 +from vllm.entrypoints.pooling.typing import (
21 + PoolingEngineInput,
22 + PoolingServeContext,
23 +)
24 from vllm.outputs import PoolingOutput, PoolingRequestOutput
25
26
@@ -410,6 +411,7 @@ class TestChunkedEmbeddingProcessing:
411 def _make_handler(cls):
412 handler = object.__new__(EmbedIOProcessor)
413 handler.model_config = cls._FakeModelConfig()
414 + handler.enable_chunked_processing = True
415 return handler
416
417 @staticmethod
@@ -421,15 +423,29 @@ class TestChunkedEmbeddingProcessing:
423 }
424 )
425 assert isinstance(request, EmbeddingCompletionRequest)
426 + pooling_params = PoolingParams()
427 return PoolingServeContext(
428 request=request,
426 - pooling_params=PoolingParams(),
429 + pooling_params=pooling_params,
430 model_name="test",
431 request_id="embd-client-prompt-999-chunk-888",
432 engine_inputs=[
430 - {"prompt_token_ids": [0, 1, 2, 3, 4]},
431 - {"prompt_token_ids": [10, 11]},
433 + PoolingEngineInput(
434 + prompts={"prompt_token_ids": [0, 1, 2, 3, 4]},
435 + params=pooling_params,
436 + lora_requests=None,
437 + priorities=0,
438 + ),
439 + PoolingEngineInput(
440 + prompts={"prompt_token_ids": [10, 11]},
441 + params=pooling_params,
442 + lora_requests=None,
443 + priorities=0,
444 + ),
445 ],
446 + lora_request=None,
447 + priorities=0,
448 + prompt_extras=None,
449 )
450
451 @staticmethod
@@ -450,7 +466,7 @@ class TestChunkedEmbeddingProcessing:
466 handler = self._make_handler()
467 ctx = self._make_context()
468
453 - handler._pre_process_chunked(ctx)
469 + handler.maybe_pre_process_chunked(ctx)
470
471 assert ctx.prompt_request_ids == [
472 "embd-client-prompt-999-chunk-888-prompt-0-chunk-0",
@@ -488,227 +504,3 @@ class TestChunkedEmbeddingProcessing:
504 ctx.final_res_batch[1].outputs.data,
505 torch.tensor([9.0, 9.0]),
506 )
491 -
492 -
493 -class TestPreProcessCohereOnline:
494 - """Unit tests for EmbedIOProcessor._pre_process_cohere_online."""
495 -
496 - @staticmethod
497 - def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
498 - return PoolingServeContext(
499 - request=CohereEmbedRequest(model="test", **request_kwargs),
500 - pooling_params=PoolingParams(),
501 - model_name="test",
502 - request_id="embd-test",
503 - )
504 -
505 - @staticmethod
506 - def _make_handler():
507 - handler = object.__new__(EmbedIOProcessor)
508 - handler._validate_input_type = lambda _input_type: None
509 - return handler
510 -
511 - def test_text_only_without_task_prefix_uses_completion_path(self):
512 - handler = self._make_handler()
513 - ctx = self._make_context(texts=["hello"])
514 - calls: list[tuple[str, object]] = []
515 -
516 - def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
517 - calls.append(("completion", prompt_input))
518 - return ["completion"]
519 -
520 - handler._get_task_instruction_prefix = lambda _input_type: None
521 - handler._has_chat_template = lambda: False
522 - handler._preprocess_cmpl_online = preprocess_cmpl_online
523 - handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
524 - "text-only request should not require chat rendering"
525 - )
526 -
527 - handler._pre_process_cohere_online(ctx)
528 -
529 - assert ctx.engine_inputs == ["completion"]
530 - assert calls == [("completion", ["hello"])]
531 -
532 - def test_text_only_falls_back_to_prefixed_completion_without_template(self):
533 - handler = self._make_handler()
534 - ctx = self._make_context(texts=["hello"], input_type="query")
535 - calls: list[tuple[str, object]] = []
536 -
537 - def preprocess_cmpl(request, prompt_input, prompt_embeds):
538 - calls.append(("completion", prompt_input))
539 - return ["fallback"]
540 -
541 - handler._get_task_instruction_prefix = lambda _input_type: "query: "
542 - handler._has_chat_template = lambda: False
543 - handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
544 - "chat rendering should be skipped without a template"
545 - )
546 - handler._preprocess_cmpl_online = preprocess_cmpl
547 -
548 - handler._pre_process_cohere_online(ctx)
549 -
550 - assert ctx.engine_inputs == ["fallback"]
551 - assert calls == [("completion", ["query: hello"])]
552 -
553 - def test_text_only_with_template_uses_chat_path(self):
554 - handler = self._make_handler()
555 - ctx = self._make_context(texts=["hello"], input_type="query")
556 - calls: list[tuple[str, object]] = []
557 -
558 - def batch_render_chat(
559 - request,
560 - all_messages,
561 - truncate_prompt_tokens,
562 - truncation_side,
563 - ):
564 - calls.append(
565 - (
566 - "chat",
567 - {
568 - "request": request,
569 - "all_messages": all_messages,
570 - "truncate_prompt_tokens": truncate_prompt_tokens,
571 - "truncation_side": truncation_side,
572 - },
573 - )
574 - )
575 - return ["chat"]
576 -
577 - handler._get_task_instruction_prefix = lambda _input_type: "query: "
578 - handler._has_chat_template = lambda: True
579 - handler._batch_render_chat = batch_render_chat
580 - handler._preprocess_cmpl_online = lambda *_args, **_kwargs: pytest.fail(
581 - "completion path should be skipped when a template exists"
582 - )
583 -
584 - handler._pre_process_cohere_online(ctx)
585 -
586 - assert ctx.engine_inputs == ["chat"]
587 - assert calls == [
588 - (
589 - "chat",
590 - {
591 - "request": ctx.request,
592 - "all_messages": [
593 - handler._mixed_input_to_messages(
594 - CohereEmbedInput(
595 - content=[CohereEmbedContent(type="text", text="hello")]
596 - ),
597 - task_prefix="query: ",
598 - )
599 - ],
600 - "truncate_prompt_tokens": -1,
601 - "truncation_side": None,
602 - },
603 - )
604 - ]
605 -
606 -
607 -class TestPreProcessOpenAIEmbeddingChatOnline:
608 - """Unit tests for OpenAI embedding chat preprocessing."""
609 -
610 - class _FakeModelConfig:
611 - max_model_len = 128
612 - encoder_config: dict[str, object] = {}
613 - pooler_config = None
614 - multimodal_config = None
615 - is_encoder_decoder = False
616 -
617 - class _FakeRenderer:
618 - tokenizer = object()
619 -
620 - def __init__(self):
621 - self.calls = []
622 -
623 - def render_chat(
624 - self,
625 - all_messages,
626 - chat_params,
627 - tok_params,
628 - prompt_extras=None,
629 - ):
630 - self.calls.append(
631 - {
632 - "all_messages": all_messages,
633 - "chat_params": chat_params,
634 - "tok_params": tok_params,
635 - "prompt_extras": prompt_extras,
636 - }
637 - )
638 - return all_messages, [
639 - {"prompt_token_ids": [index]} for index, _ in enumerate(all_messages)
640 - ]
641 -
642 - @classmethod
643 - def _make_handler(cls, renderer):
644 - handler = object.__new__(EmbedIOProcessor)
645 - handler.renderer = renderer
646 - handler.model_config = cls._FakeModelConfig()
647 - handler.chat_template = "template"
648 - handler.chat_template_content_format = "auto"
649 - handler.trust_request_chat_template = False
650 - handler.enable_chunked_processing = False
651 - return handler
652 -
653 - @staticmethod
654 - def _make_context(
655 - request: (
656 - EmbeddingChatRequest
657 - | EmbeddingBatchChatRequest
658 - | EmbeddingChatInputRequest
659 - | EmbeddingBatchChatInputRequest
660 - ),
661 - ) -> PoolingServeContext[
662 - EmbeddingChatRequest
663 - | EmbeddingBatchChatRequest
664 - | EmbeddingChatInputRequest
665 - | EmbeddingBatchChatInputRequest
666 - ]:
667 - return PoolingServeContext(
668 - request=request,
669 - pooling_params=PoolingParams(),
670 - model_name="test",
671 - request_id="embd-test",
672 - )
673 -
674 - def test_chat_template_kwargs_forwarded_for_batched_input_messages(self):
675 - request = TypeAdapter(EmbeddingRequest).validate_python(
676 - {
677 - "model": "test",
678 - "input": [
679 - [{"role": "user", "content": "hello"}],
680 - [{"role": "user", "content": "goodbye"}],
681 - ],
682 - "add_generation_prompt": True,
683 - "chat_template_kwargs": {"instruction": "Represent the query: "},
684 - "mm_processor_kwargs": {"max_pixels": 1},
685 - "cache_salt": "salt",
686 - }
687 - )
688 - assert isinstance(request, EmbeddingBatchChatInputRequest)
689 -
690 - renderer = self._FakeRenderer()
691 - handler = self._make_handler(renderer)
692 - ctx = self._make_context(request)
693 -
694 - handler.pre_process_online(ctx)
695 -
696 - assert ctx.engine_inputs == [
697 - {"prompt_token_ids": [0]},
698 - {"prompt_token_ids": [1]},
699 - ]
700 - assert len(renderer.calls) == 1
701 -
702 - call = renderer.calls[0]
703 - assert call["all_messages"] == request.messages
704 - assert call["prompt_extras"] == {
705 - "mm_processor_kwargs": {"max_pixels": 1},
706 - "cache_salt": "salt",
707 - }
708 -
709 - chat_template_kwargs = call["chat_params"].chat_template_kwargs
710 - assert chat_template_kwargs["instruction"] == "Represent the query: "
711 - assert chat_template_kwargs["add_generation_prompt"] is True
712 - assert chat_template_kwargs["continue_final_message"] is False
713 - assert "tools" not in chat_template_kwargs
714 - assert chat_template_kwargs["tokenize"] is False
tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py
-43
@@ -2,7 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import weakref
5 -from types import SimpleNamespace
5
6 import pytest
7 import torch
@@ -10,10 +9,7 @@ import torch
9 from tests.models.utils import softmax
10 from vllm import LLM, PoolingParams
11 from vllm.distributed import cleanup_dist_env_and_memory
13 -from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
14 -from vllm.entrypoints.pooling.scoring.typing import ScoringData
12 from vllm.platforms import current_platform
16 -from vllm.renderers import TokenizeParams
13
14 MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
15 PROMPT = "The chef prepared a delicious meal."
@@ -145,45 +141,6 @@ def test_max_tokens_per_doc(llm: LLM):
141 assert with_limit_tokens < no_limit_tokens
142
143
148 -def test_token_type_ids_follow_post_tokenization():
149 - processor = object.__new__(CrossEncoderIOProcessor)
150 - processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1)
151 - processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt)
152 - processor.model_config = None
153 - processor.get_score_prompt = lambda **_: (
154 - "",
155 - {
156 - "prompt_token_ids": list(range(32)),
157 - "token_type_ids": [0] * 16 + [1] * 16,
158 - },
159 - )
160 -
161 - engine_inputs, pooling_params = processor._pre_process(
162 - ScoringData(data_1=["query"], data_2=["document"]),
163 - TokenizeParams(
164 - max_total_tokens=None,
165 - truncate_prompt_tokens=16,
166 - truncation_side="left",
167 - ),
168 - PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}),
169 - )
170 -
171 - assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32))
172 - assert pooling_params[0].extra_kwargs == {
173 - "cache_salt": "salt",
174 - "compressed_token_type_ids": 0,
175 - }
176 -
177 - engine_inputs, pooling_params = processor._pre_process(
178 - ScoringData(data_1=["query"], data_2=["document"]),
179 - TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40),
180 - PoolingParams(task="classify"),
181 - )
182 -
183 - assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8
184 - assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16}
185 -
186 -
144 def test_pooling_params(llm: LLM):
145 def get_outputs(use_activation):
146 outputs = llm.score(
tests/entrypoints/serve/lora/test_serving_models.py
+8 -6
@@ -161,7 +161,9 @@ def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing:
161 return serving
162
163
164 -def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
164 +def _make_pooling_ctx(
165 + model_name: str, serving: PoolingBaseServing
166 +) -> PoolingServeContext:
167 mock_request = MagicMock()
168 mock_request.model = model_name
169 return PoolingServeContext(
@@ -169,6 +171,9 @@ def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
171 model_name=MODEL_NAME,
172 request_id="test-id",
173 pooling_params=PoolingParams(),
174 + lora_request=serving._maybe_get_adapters(mock_request),
175 + priorities=0,
176 + prompt_extras=None,
177 )
178
179
@@ -176,9 +181,7 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
181 """LoRA adapter name must populate ctx.lora_request without raising."""
182 lora_name = "bot-embed-lora"
183 serving = _make_pooling_serving(lora_name)
179 - ctx = _make_pooling_ctx(lora_name)
180 -
181 - ctx.lora_request = serving._maybe_get_adapters(ctx.request)
184 + ctx = _make_pooling_ctx(lora_name, serving)
185
186 assert ctx.lora_request is not None
187 assert ctx.lora_request.lora_name == lora_name
@@ -187,7 +190,6 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
190 def test_pooling_maybe_get_adapters_unknown_model_raises():
191 """An unrecognised model name must still raise VLLMNotFoundError."""
192 serving = _make_pooling_serving("some-lora")
190 - ctx = _make_pooling_ctx("unknown-model")
193
194 with pytest.raises(VLLMNotFoundError):
193 - serving._maybe_get_adapters(ctx.request)
195 + _make_pooling_ctx("unknown-model", serving)
vllm/entrypoints/pooling/base/io_processor.py
+96 -100
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Sequence
5 +from concurrent.futures import Executor
6 from typing import Any, Final, cast
7
8 from vllm import (
@@ -10,21 +11,17 @@ from vllm import (
11 )
12 from vllm.config import VllmConfig
13 from vllm.entrypoints.chat_utils import (
13 - ChatCompletionMessageParam,
14 ChatTemplateConfig,
15 - ChatTemplateContentFormatOption,
16 - ConversationMessage,
15 )
18 -from vllm.entrypoints.serve.engine.typing import RendererChatRequest, RendererRequest
19 -from vllm.inputs import EngineInput, SingletonPrompt
16 from vllm.lora.request import LoRARequest
17 from vllm.renderers import BaseRenderer, merge_kwargs
18 from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
23 -from vllm.tool_parsers import ToolParser
19 +from vllm.utils.async_utils import make_async
20 from vllm.utils.mistral import is_mistral_tokenizer
21
22 from ..typing import (
27 - ALLOfflineInputsContext,
23 + AnyOfflineInputsContext,
24 + AnyRenderParam,
25 EncodeChatRenderParams,
26 EncodeCMPLRenderParams,
27 OfflineEncodeInputsContext,
@@ -66,14 +63,25 @@ class PoolingIOProcessor:
63 chat_template_config.trust_request_chat_template
64 )
65
66 + self.template_kwargs = None
67 + self.tool_dicts = None
68 +
69 + # Shared thread pool executor for preprocessing
70 + self._executor: Executor = self.renderer._executor
71 + self.render_async = make_async(self.render, executor=self._executor)
72 +
73 #######################################
74 # online APIs
75
76 def create_pooling_params(self, request):
77 return request.to_pooling_params()
78
75 - def pre_process_online(self, ctx: PoolingServeContext):
79 + def get_request_factory_online(
80 + self, ctx: PoolingServeContext
81 + ) -> Sequence[AnyRenderParam]:
82 request = ctx.request
83 + renderer = self.renderer
84 + requests: Sequence[AnyRenderParam]
85
86 if isinstance(request, PoolingChatLikeRequest):
87 self._validate_chat_template(
@@ -81,24 +89,87 @@ class PoolingIOProcessor:
89 chat_template_kwargs=request.chat_template_kwargs,
90 trust_request_chat_template=self.trust_request_chat_template,
91 )
84 - _, engine_inputs = self._preprocess_chat_online(
85 - request,
86 - request.messages,
87 - default_template=self.chat_template,
88 - default_template_content_format=self.chat_template_content_format,
89 - default_template_kwargs=None,
92 +
93 + num_requests = 1
94 + default_template_kwargs = merge_kwargs(
95 + self.template_kwargs,
96 + dict(
97 + tools=self.tool_dicts,
98 + tokenize=is_mistral_tokenizer(renderer.tokenizer),
99 + ),
100 + )
101 +
102 + mm_config = self.model_config.multimodal_config
103 + tok_params = request.build_tok_params(self.model_config)
104 + chat_params = request.build_chat_params(
105 + self.chat_template, self.chat_template_content_format
106 + ).with_defaults(
107 + default_template_kwargs,
108 + default_media_io_kwargs=(
109 + mm_config.media_io_kwargs if mm_config else None
110 + ),
111 + )
112 +
113 + params_seq = self._params_to_seq(ctx.pooling_params, num_requests)
114 + seq_lora_requests = self._lora_request_to_seq(
115 + ctx.lora_request, num_requests
116 )
117 + seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
118 +
119 + requests = [
120 + EncodeChatRenderParams(
121 + conversations=request.messages,
122 + chat_params=chat_params,
123 + tok_params=tok_params,
124 + prompt_extras=ctx.prompt_extras,
125 + skip_mm_cache=False,
126 + params=params_seq[i],
127 + lora_requests=seq_lora_requests[i],
128 + priorities=seq_priority[i],
129 + )
130 + for i in range(num_requests)
131 + ]
132 +
133 + return requests
134 +
135 elif isinstance(request, PoolingCompletionLikeRequest):
92 - engine_inputs = self._preprocess_cmpl_online(
93 - request,
94 - prompt_input=request.input,
95 - prompt_embeds=None,
136 + model_config = self.model_config
137 + prompts_seq = prompt_to_seq(request.input)
138 + num_requests = len(prompts_seq)
139 +
140 + parsed_prompts = [
141 + (
142 + prompt
143 + if isinstance(prompt, bytes)
144 + else parse_model_prompt(model_config, prompt)
145 + )
146 + for prompt in prompts_seq
147 + ]
148 + tok_params = request.build_tok_params(model_config)
149 +
150 + params_seq = self._params_to_seq(ctx.pooling_params, num_requests)
151 + seq_lora_requests = self._lora_request_to_seq(
152 + ctx.lora_request, num_requests
153 )
154 + seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
155 +
156 + requests = [
157 + EncodeCMPLRenderParams(
158 + prompts=parsed_prompts[i],
159 + tok_params=tok_params,
160 + prompt_extras=ctx.prompt_extras,
161 + skip_mm_cache=False,
162 + params=params_seq[i],
163 + lora_requests=seq_lora_requests[i],
164 + priorities=seq_priority[i],
165 + )
166 + for i in range(num_requests)
167 + ]
168 +
169 + return requests
170 else:
171 raise ValueError(f"Invalid {self.name} request type")
172
100 - ctx.engine_inputs = engine_inputs
101 -
173 def post_process_online(
174 self,
175 ctx: PoolingServeContext,
@@ -109,7 +180,7 @@ class PoolingIOProcessor:
180 # offline APIs
181
182 def get_request_factory_offline(
112 - self, ctx: ALLOfflineInputsContext
183 + self, ctx: AnyOfflineInputsContext
184 ) -> tuple[RequestFactory, int]:
185 assert isinstance(ctx, OfflineEncodeInputsContext)
186
@@ -209,84 +280,6 @@ class PoolingIOProcessor:
280 priorities=render_params["priorities"],
281 )
282
212 - def _preprocess_cmpl_online(
213 - self,
214 - request: RendererRequest,
215 - prompt_input: str | list[str] | list[int] | list[list[int]] | None,
216 - prompt_embeds: bytes | list[bytes] | None,
217 - ) -> list[EngineInput]:
218 - renderer = self.renderer
219 - model_config = self.model_config
220 -
221 - prompts = list[SingletonPrompt | bytes]()
222 - if prompt_embeds is not None: # embeds take higher priority
223 - prompts.extend(prompt_to_seq(prompt_embeds))
224 - if prompt_input is not None:
225 - prompts.extend(prompt_to_seq(prompt_input))
226 -
227 - parsed_prompts = [
228 - (
229 - prompt
230 - if isinstance(prompt, bytes)
231 - else parse_model_prompt(model_config, prompt)
232 - )
233 - for prompt in prompts
234 - ]
235 - tok_params = request.build_tok_params(model_config)
236 -
237 - return renderer.render_cmpl(
238 - parsed_prompts,
239 - tok_params,
240 - prompt_extras={
241 - k: v
242 - for k in ("mm_processor_kwargs", "cache_salt")
243 - if (v := getattr(request, k, None)) is not None
244 - },
245 - )
246 -
247 - def _preprocess_chat_online(
248 - self,
249 - request: RendererChatRequest,
250 - messages: list[ChatCompletionMessageParam],
251 - default_template: str | None,
252 - default_template_content_format: ChatTemplateContentFormatOption,
253 - default_template_kwargs: dict[str, Any] | None,
254 - tool_dicts: list[dict[str, Any]] | None = None,
255 - tool_parser: type[ToolParser] | None = None,
256 - ) -> tuple[list[ConversationMessage], list[EngineInput]]:
257 - renderer = self.renderer
258 -
259 - default_template_kwargs = merge_kwargs(
260 - default_template_kwargs,
261 - dict(
262 - tools=tool_dicts,
263 - tokenize=is_mistral_tokenizer(renderer.tokenizer),
264 - ),
265 - )
266 -
267 - mm_config = self.model_config.multimodal_config
268 -
269 - tok_params = request.build_tok_params(self.model_config)
270 - chat_params = request.build_chat_params(
271 - default_template, default_template_content_format
272 - ).with_defaults(
273 - default_template_kwargs,
274 - default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
275 - )
276 -
277 - (conversation,), (engine_input,) = renderer.render_chat(
278 - [messages],
279 - chat_params,
280 - tok_params,
281 - prompt_extras={
282 - k: v
283 - for k in ("mm_processor_kwargs", "cache_salt")
284 - if (v := getattr(request, k, None)) is not None
285 - },
286 - )
287 -
288 - return conversation, [engine_input]
289 -
283 def _validate_chat_template(
284 self,
285 request_chat_template: str | None,
@@ -341,10 +334,13 @@ class PoolingIOProcessor:
334
335 def _priority_to_seq(
336 self,
344 - priority: Sequence[int] | None,
337 + priority: int | Sequence[int] | None,
338 num_requests: int,
339 ) -> Sequence[int]:
340 if priority is not None:
341 + if isinstance(priority, int):
342 + return [priority] * num_requests
343 +
344 if len(priority) != num_requests:
345 raise ValueError(
346 f"The lengths of prompts ({num_requests}) "
vllm/entrypoints/pooling/base/serving.py
+31 -29
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 +import asyncio
4 from abc import ABC, abstractmethod
5 from collections.abc import AsyncGenerator, Mapping
6 from concurrent.futures import Executor
@@ -63,9 +63,6 @@ class PoolingBaseServing(ABC, BaseServing):
63
64 # Shared thread pool executor for preprocessing and postprocessing.
65 self._executor: Executor = self.renderer._executor
66 - self._preprocessing_async = make_async(
67 - self._preprocessing, executor=self._executor
68 - )
66 self._postprocessing_async = make_async(
67 self._postprocessing, executor=self._executor
68 )
@@ -77,7 +74,7 @@ class PoolingBaseServing(ABC, BaseServing):
74 ) -> Response:
75 io_processor = self.get_io_processor(request)
76 ctx = await self._init_ctx(io_processor, request, raw_request)
80 - await self._preprocessing_async(io_processor, ctx)
77 + await self._preprocessing(io_processor, ctx)
78 await self._prepare_generators(ctx)
79 await self._collect_batch(ctx)
80 return await self._postprocessing_async(io_processor, ctx)
@@ -86,11 +83,17 @@ class PoolingBaseServing(ABC, BaseServing):
83 def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor:
84 raise NotImplementedError
85
89 - @torch.inference_mode()
90 - def _preprocessing(
86 + async def _preprocessing(
87 self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext
88 ):
93 - return io_processor.pre_process_online(ctx)
89 + requests = io_processor.get_request_factory_online(ctx)
90 +
91 + if len(requests) == 0:
92 + raise ValueError("You must pass at least one prompt")
93 +
94 + ctx.engine_inputs = await asyncio.gather(
95 + *[io_processor.render_async(request) for request in requests]
96 + )
97
98 @torch.inference_mode()
99 def _postprocessing(
@@ -110,16 +113,26 @@ class PoolingBaseServing(ABC, BaseServing):
113 await self._check_model(request)
114
115 pooling_params = io_processor.create_pooling_params(request)
116 + lora_request = self._maybe_get_adapters(request)
117 + priorities = getattr(request, "priority", 0)
118 + prompt_extras = {
119 + k: v
120 + for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs")
121 + if (v := getattr(request, k, None)) is not None
122 + }
123 +
124 ctx = PoolingServeContext(
125 request=request,
126 raw_request=raw_request,
127 model_name=model_name,
128 pooling_params=pooling_params,
129 request_id=request_id,
130 + lora_request=lora_request,
131 + priorities=priorities,
132 + prompt_extras=prompt_extras,
133 )
134
135 self._validate_request(ctx)
122 - ctx.lora_request = self._maybe_get_adapters(ctx.request)
136 return ctx
137
138 async def _prepare_generators(
@@ -127,7 +140,7 @@ class PoolingBaseServing(ABC, BaseServing):
140 ctx: PoolingServeContext,
141 ):
142 if ctx.engine_inputs is None:
130 - raise ValueError("Engine prompts not available")
143 + raise ValueError("Engine inputs not available")
144
145 generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
146
@@ -139,12 +152,7 @@ class PoolingBaseServing(ABC, BaseServing):
152
153 assert ctx.pooling_params is not None
154 pooling_params = ctx.pooling_params
142 -
143 - if isinstance(pooling_params, list):
144 - for params in pooling_params:
145 - params.verify(self.model_config)
146 - else:
147 - pooling_params.verify(self.model_config)
155 + pooling_params.verify(self.model_config)
156
157 for i, engine_input in enumerate(ctx.engine_inputs):
158 prompt_request_id = (
@@ -153,26 +161,20 @@ class PoolingBaseServing(ABC, BaseServing):
161 else ctx.prompt_request_ids[i]
162 )
163
156 - params = (
157 - pooling_params[i]
158 - if isinstance(pooling_params, list)
159 - else pooling_params
160 - )
161 -
164 self._log_inputs(
165 prompt_request_id,
164 - engine_input,
165 - params=params,
166 + engine_input["prompts"],
167 + params=engine_input["params"],
168 lora_request=ctx.lora_request,
169 )
170
171 generator = self.engine_client.encode(
170 - engine_input,
171 - params,
172 - prompt_request_id,
173 - lora_request=ctx.lora_request,
172 + request_id=prompt_request_id,
173 + prompt=engine_input["prompts"],
174 + pooling_params=engine_input["params"],
175 + lora_request=engine_input["lora_requests"],
176 + priority=engine_input["priorities"],
177 trace_headers=trace_headers,
175 - priority=getattr(ctx.request, "priority", 0),
178 )
179
180 generators.append(generator)
vllm/entrypoints/pooling/classify/api_router.py
+5 -5
@@ -16,8 +16,11 @@ from .serving import ServingClassification
16 router = APIRouter()
17
18
19 -def classify(request: Request) -> ServingClassification | None:
20 - return request.app.state.serving_classification
19 +def classify(request: Request) -> ServingClassification:
20 + handler = getattr(request.app.state, "serving_classification", None)
21 + if handler is None:
22 + raise NotImplementedError("The model does not support Classification API")
23 + return handler
24
25
26 @router.post("/classify", dependencies=[Depends(validate_json_request)])
@@ -27,7 +30,4 @@ async def create_classify(
30 request: ClassificationRequest, raw_request: Request
31 ) -> Response:
32 handler = classify(raw_request)
30 - if handler is None:
31 - raise NotImplementedError("The model does not support Classification API")
32 -
33 return await handler(request, raw_request)
vllm/entrypoints/pooling/embed/api_router.py
+5 -8
@@ -18,8 +18,11 @@ from .serving import ServingEmbedding
18 router = APIRouter()
19
20
21 -def embedding(request: Request) -> ServingEmbedding | None:
22 - return request.app.state.serving_embedding
21 +def embedding(request: Request) -> ServingEmbedding:
22 + handler = getattr(request.app.state, "serving_embedding", None)
23 + if handler is None:
24 + raise NotImplementedError("The model does not support Embeddings API")
25 + return handler
26
27
28 @router.post(
@@ -37,9 +40,6 @@ async def create_embedding(
40 raw_request: Request,
41 ):
42 handler = embedding(raw_request)
40 - if handler is None:
41 - raise NotImplementedError("The model does not support Embeddings API")
42 -
43 return await handler(request, raw_request)
44
45
@@ -58,7 +58,4 @@ async def create_cohere_embedding(
58 raw_request: Request,
59 ):
60 handler = embedding(raw_request)
61 - if handler is None:
62 - raise NotImplementedError("The model does not support Embeddings API")
63 -
61 return await handler(request, raw_request)
vllm/entrypoints/pooling/embed/io_processor.py
+125 -131
@@ -17,7 +17,7 @@ from vllm.entrypoints.chat_utils import (
17 ChatCompletionMessageParam,
18 CustomChatCompletionMessageParam,
19 )
20 -from vllm.inputs import EngineInput, tokens_input
20 +from vllm.inputs import tokens_input
21 from vllm.logger import init_logger
22 from vllm.outputs import PoolingOutput, PoolingRequestOutput
23 from vllm.renderers import merge_kwargs
@@ -28,11 +28,14 @@ from vllm.utils.mistral import is_mistral_tokenizer
28 from ..base.io_processor import PoolingIOProcessor
29 from ..scoring.io_processor import JinaRankingIOProcessorMixin
30 from ..typing import (
31 - ALLOfflineInputsContext,
31 + AnyOfflineInputsContext,
32 + AnyRenderParam,
33 ChunkedEmbeddingMetadata,
34 + EncodeChatRenderParams,
35 OfflineEncodeInputsContext,
36 PoolingChatLikeRequest,
37 PoolingCompletionLikeRequest,
38 + PoolingEngineInput,
39 PoolingServeContext,
40 RequestFactory,
41 )
@@ -76,9 +79,11 @@ class EmbedIOProcessor(PoolingIOProcessor):
79 list(self.task_instructions.keys()),
80 )
81
79 - def pre_process_online(self, ctx: PoolingServeContext):
82 + def get_request_factory_online(
83 + self, ctx: PoolingServeContext
84 + ) -> Sequence[AnyRenderParam]:
85 if isinstance(ctx.request, CohereEmbedRequest):
81 - self._pre_process_cohere_online(ctx)
86 + requests = self._get_request_factory_cohere_online(ctx)
87 elif isinstance(
88 ctx.request,
89 (
@@ -88,12 +93,11 @@ class EmbedIOProcessor(PoolingIOProcessor):
93 EmbeddingBatchChatInputRequest,
94 ),
95 ):
91 - self._pre_process_openai_chat_online(ctx)
96 + requests = self._get_request_factory_chat_input_online(ctx)
97 else:
93 - super().pre_process_online(ctx)
98 + requests = super().get_request_factory_online(ctx)
99
95 - if self.enable_chunked_processing:
96 - self._pre_process_chunked(ctx)
100 + return requests
101
102 def post_process_online(
103 self,
@@ -114,18 +118,21 @@ class EmbedIOProcessor(PoolingIOProcessor):
118 # PTAL: examples/pooling/embed/openai_embedding_long_text
119 #################################################################
120
117 - def _pre_process_chunked(self, ctx: PoolingServeContext) -> None:
121 + def maybe_pre_process_chunked(self, ctx: PoolingServeContext) -> None:
122 + if not self.enable_chunked_processing:
123 + return None
124 +
125 if ctx.engine_inputs is None:
126 raise ValueError("Engine prompts not available")
127
128 ctx.original_engine_inputs = ctx.engine_inputs
129 request_id = ctx.request_id
130 max_model_len = self.model_config.max_model_len
124 - chunked_engine_inputs: list[EngineInput] = []
131 + chunked_engine_inputs: list[PoolingEngineInput] = []
132 prompt_request_ids: list[str] = []
133 chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] = []
134 for prompt_idx, engine_input in enumerate(ctx.engine_inputs):
128 - token_ids = engine_input.get("prompt_token_ids", None)
135 + token_ids = engine_input["prompts"].get("prompt_token_ids", None)
136 if token_ids is None:
137 raise NotImplementedError(
138 "Long Text Embedding with Chunked Processing does "
@@ -138,7 +145,12 @@ class EmbedIOProcessor(PoolingIOProcessor):
145 chunk_list(prompt_token_ids, max_model_len)
146 ):
147 chunked_engine_inputs.append(
141 - tokens_input(prompt_token_ids=chunk_tokens)
148 + PoolingEngineInput(
149 + prompts=tokens_input(prompt_token_ids=chunk_tokens),
150 + params=engine_input["params"],
151 + lora_requests=engine_input["lora_requests"],
152 + priorities=engine_input["priorities"],
153 + )
154 )
155 prompt_request_ids.append(
156 f"{request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}"
@@ -234,7 +246,7 @@ class EmbedIOProcessor(PoolingIOProcessor):
246
247 # Get original prompt token IDs for this prompt
248 original_prompt = original_engine_inputs[prompt_idx]
237 - token_ids = original_prompt.get("prompt_token_ids", None)
249 + token_ids = original_prompt["prompts"].get("prompt_token_ids", None)
250 if token_ids is None:
251 raise NotImplementedError(
252 "Long Text Embedding with Chunked Processing does "
@@ -262,6 +274,68 @@ class EmbedIOProcessor(PoolingIOProcessor):
274
275 return None
276
277 + #################################################################
278 + # Chat input Request Preprocessing & Postprocessing
279 + #################################################################
280 +
281 + def _get_request_factory_chat_input_online(
282 + self,
283 + ctx: PoolingServeContext,
284 + ) -> Sequence[AnyRenderParam]:
285 + request = ctx.request
286 + renderer = self.renderer
287 +
288 + self._validate_chat_template(
289 + request_chat_template=request.chat_template,
290 + chat_template_kwargs=request.chat_template_kwargs,
291 + trust_request_chat_template=self.trust_request_chat_template,
292 + )
293 +
294 + if isinstance(
295 + request, (EmbeddingBatchChatRequest, EmbeddingBatchChatInputRequest)
296 + ):
297 + all_messages = request.messages
298 + else:
299 + all_messages = [request.messages]
300 + num_requests = len(all_messages)
301 +
302 + default_template_kwargs = merge_kwargs(
303 + self.template_kwargs,
304 + dict(
305 + tools=self.tool_dicts,
306 + tokenize=is_mistral_tokenizer(renderer.tokenizer),
307 + ),
308 + )
309 +
310 + mm_config = self.model_config.multimodal_config
311 + tok_params = request.build_tok_params(self.model_config)
312 + chat_params = request.build_chat_params(
313 + self.chat_template, self.chat_template_content_format
314 + ).with_defaults(
315 + default_template_kwargs,
316 + default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
317 + )
318 +
319 + params_seq = self._params_to_seq(ctx.pooling_params, num_requests)
320 + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
321 + seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
322 +
323 + requests = [
324 + EncodeChatRenderParams(
325 + conversations=all_messages[i],
326 + chat_params=chat_params,
327 + tok_params=tok_params,
328 + prompt_extras=ctx.prompt_extras,
329 + skip_mm_cache=False,
330 + params=params_seq[i],
331 + lora_requests=seq_lora_requests[i],
332 + priorities=seq_priority[i],
333 + )
334 + for i in range(num_requests)
335 + ]
336 +
337 + return requests
338 +
339 #################################################################
340 # Cohere Request Preprocessing & Postprocessing
341 #################################################################
@@ -378,71 +452,9 @@ class EmbedIOProcessor(PoolingIOProcessor):
452 )
453 return super().create_pooling_params(request)
454
381 - def _pre_process_openai_chat_online(
382 - self,
383 - ctx: PoolingServeContext[
384 - EmbeddingChatRequest
385 - | EmbeddingBatchChatRequest
386 - | EmbeddingChatInputRequest
387 - | EmbeddingBatchChatInputRequest
388 - ],
389 - ) -> None:
390 - request = ctx.request
391 - self._validate_chat_template(
392 - request_chat_template=request.chat_template,
393 - chat_template_kwargs=request.chat_template_kwargs,
394 - trust_request_chat_template=self.trust_request_chat_template,
395 - )
396 -
397 - if isinstance(
398 - request, (EmbeddingBatchChatRequest, EmbeddingBatchChatInputRequest)
399 - ):
400 - all_messages = request.messages
401 - else:
402 - all_messages = [request.messages]
403 - ctx.engine_inputs = self._batch_render_openai_chat(request, all_messages)
404 -
405 - def _batch_render_openai_chat(
406 - self,
407 - request: (
408 - EmbeddingChatRequest
409 - | EmbeddingBatchChatRequest
410 - | EmbeddingChatInputRequest
411 - | EmbeddingBatchChatInputRequest
412 - ),
413 - all_messages: Sequence[list[ChatCompletionMessageParam]],
414 - ) -> list[EngineInput]:
415 - renderer = self.renderer
416 - mm_config = self.model_config.multimodal_config
417 -
418 - tok_params = request.build_tok_params(self.model_config)
419 - chat_params = request.build_chat_params(
420 - self.chat_template,
421 - self.chat_template_content_format,
422 - ).with_defaults(
423 - merge_kwargs(
424 - None,
425 - dict(
426 - tools=None,
427 - tokenize=is_mistral_tokenizer(renderer.tokenizer),
428 - ),
429 - ),
430 - default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
431 - )
432 -
433 - _, engine_inputs = renderer.render_chat(
434 - all_messages,
435 - chat_params,
436 - tok_params,
437 - prompt_extras={
438 - k: v
439 - for k in ("mm_processor_kwargs", "cache_salt")
440 - if (v := getattr(request, k, None)) is not None
441 - },
442 - )
443 - return engine_inputs
444 -
445 - def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None:
455 + def _get_request_factory_cohere_online(
456 + self, ctx: PoolingServeContext
457 + ) -> Sequence[AnyRenderParam]:
458 """Convert a ``CohereEmbedRequest`` into engine prompts.
459
460 If a model has a chat template the task instruction are rendered
@@ -479,13 +491,12 @@ class EmbedIOProcessor(PoolingIOProcessor):
491 task_prefix = self._get_task_instruction_prefix(input_type)
492
493 if task_prefix is None:
482 - ctx.engine_inputs = self._preprocess_cohere_text_completion(
483 - request,
494 + return self._get_request_factory_cohere_text_completion(
495 + ctx,
496 texts,
497 truncate_prompt_tokens,
498 truncation_side,
499 )
488 - return
500
501 all_messages = [
502 self._mixed_input_to_messages(
@@ -497,27 +508,30 @@ class EmbedIOProcessor(PoolingIOProcessor):
508 for text in texts
509 ]
510 if self._has_chat_template():
500 - ctx.engine_inputs = self._batch_render_chat(
501 - request,
511 + return self._batch_render_chat(
512 + ctx,
513 all_messages,
514 truncate_prompt_tokens,
515 truncation_side,
516 )
517 else:
507 - ctx.engine_inputs = self._preprocess_cohere_text_completion(
508 - request,
518 + return self._get_request_factory_cohere_text_completion(
519 + ctx,
520 self._apply_task_instruction(texts, input_type),
521 truncate_prompt_tokens,
522 truncation_side,
523 )
513 - return
524
525 task_prefix = self._get_task_instruction_prefix(input_type)
526 all_messages = [
527 self._mixed_input_to_messages(inp, task_prefix=task_prefix) for inp in input
528 ]
519 - ctx.engine_inputs = self._batch_render_chat(
520 - request, all_messages, truncate_prompt_tokens, truncation_side
529 +
530 + return self._batch_render_chat(
531 + ctx,
532 + all_messages,
533 + truncate_prompt_tokens,
534 + truncation_side,
535 )
536
537 def _has_chat_template(self) -> bool:
@@ -531,13 +545,14 @@ class EmbedIOProcessor(PoolingIOProcessor):
545 is not None
546 )
547
534 - def _preprocess_cohere_text_completion(
548 + def _get_request_factory_cohere_text_completion(
549 self,
536 - request: CohereEmbedRequest,
550 + ctx: PoolingServeContext,
551 texts: list[str],
552 truncate_prompt_tokens: int | None,
553 truncation_side: Literal["left", "right"] | None,
540 - ) -> list[EngineInput]:
554 + ) -> Sequence[AnyRenderParam]:
555 + request = ctx.request
556 proxy = EmbeddingCompletionRequest(
557 model=request.model,
558 input=texts,
@@ -546,50 +561,32 @@ class EmbedIOProcessor(PoolingIOProcessor):
561 truncate_prompt_tokens=truncate_prompt_tokens,
562 truncation_side=truncation_side,
563 )
549 - return self._preprocess_cmpl_online(
550 - proxy, prompt_input=proxy.input, prompt_embeds=None
551 - )
564 + ctx.request = proxy
565 + requests = super().get_request_factory_online(ctx)
566 + ctx.request = request
567 + return requests
568
569 def _batch_render_chat(
570 self,
555 - request: CohereEmbedRequest,
571 + ctx: PoolingServeContext,
572 all_messages: Sequence[list[ChatCompletionMessageParam]],
573 truncate_prompt_tokens: int | None,
574 truncation_side: Literal["left", "right"] | None,
559 - ) -> list[EngineInput]:
575 + ) -> Sequence[AnyRenderParam]:
576 """Batch-render multiple conversations through the chat template."""
561 - if not all_messages:
562 - return []
563 -
564 - proxy = EmbeddingChatRequest(
577 + request = ctx.request
578 + proxy = EmbeddingBatchChatRequest(
579 model=request.model,
566 - messages=list(all_messages[0]),
580 + messages=all_messages,
581 dimensions=request.output_dimension,
582 encoding_format="float",
583 truncate_prompt_tokens=truncate_prompt_tokens,
584 truncation_side=truncation_side,
585 )
572 -
573 - renderer = self.renderer
574 - mm_config = self.model_config.multimodal_config
575 -
576 - tok_params = proxy.build_tok_params(self.model_config)
577 - chat_params = proxy.build_chat_params(
578 - self.chat_template,
579 - self.chat_template_content_format,
580 - ).with_defaults(
581 - merge_kwargs(
582 - None,
583 - dict(
584 - tools=None,
585 - tokenize=is_mistral_tokenizer(renderer.tokenizer),
586 - ),
587 - ),
588 - default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
589 - )
590 -
591 - _, engine_inputs = renderer.render_chat(all_messages, chat_params, tok_params)
592 - return engine_inputs
586 + ctx.request = proxy
587 + requests = self._get_request_factory_chat_input_online(ctx)
588 + ctx.request = request
589 + return requests
590
591 def _validate_input_type(self, input_type: str | None) -> None:
592 """Raise if *input_type* is not supported by this model."""
@@ -640,7 +637,9 @@ class TokenEmbedIOProcessor(PoolingIOProcessor):
637 class JinaRankingTokenEmbedIOProcessor(
638 TokenEmbedIOProcessor, JinaRankingIOProcessorMixin
639 ):
643 - def pre_process_online(self, ctx: PoolingServeContext):
640 + def get_request_factory_online(
641 + self, ctx: PoolingServeContext
642 + ) -> Sequence[AnyRenderParam]:
643 request = ctx.request
644 if isinstance(request, PoolingCompletionLikeRequest):
645 prompts = request.input
@@ -655,20 +654,15 @@ class JinaRankingTokenEmbedIOProcessor(
654 query=text_prompts[-1], docs=text_prompts[:-1]
655 )
656
658 - engine_inputs = self._preprocess_cmpl_online(
659 - request,
660 - prompt_input=prompt_input,
661 - prompt_embeds=None,
662 - )
657 + request.input = prompt_input
658 + return super().get_request_factory_online(ctx)
659 elif isinstance(request, PoolingChatLikeRequest):
660 raise ValueError("The JinaForRanking does not support chat Request.")
661 else:
662 raise ValueError(f"Invalid {self.name} request type")
663
668 - ctx.engine_inputs = engine_inputs
669 -
664 def get_request_factory_offline(
671 - self, ctx: ALLOfflineInputsContext
665 + self, ctx: AnyOfflineInputsContext
666 ) -> tuple[RequestFactory, int]:
667 assert isinstance(ctx, OfflineEncodeInputsContext)
668 if not isinstance(ctx.prompts, Sequence) or len(ctx.prompts) < 2:
vllm/entrypoints/pooling/embed/protocol.py
+6 -6
@@ -93,9 +93,9 @@ class EmbeddingBatchChatRequest(
93 ``messages`` instead of introducing a separate batch-specific field.
94 """
95
96 - messages: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = (
97 - Field(..., min_length=1)
98 - )
96 + messages: Sequence[
97 + Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
98 + ] = Field(..., min_length=1)
99
100 def to_pooling_params(self):
101 return PoolingParams(
@@ -133,9 +133,9 @@ class EmbeddingChatInputRequest(
133 class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest):
134 """OpenAI embeddings request with batched chat conversations in ``input``."""
135
136 - input: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = (
137 - Field(..., min_length=1)
138 - )
136 + input: Sequence[
137 + Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
138 + ] = Field(..., min_length=1)
139
140 @model_validator(mode="before")
141 @classmethod
vllm/entrypoints/pooling/embed/serving.py
+7
@@ -9,6 +9,7 @@ from vllm.logger import init_logger
9 from vllm.outputs import PoolingRequestOutput
10 from vllm.utils.serial_utils import EmbedDType, Endianness
11
12 +from ..base.io_processor import PoolingIOProcessor
13 from ..base.serving import PoolingServing
14 from ..typing import PoolingServeContext
15 from ..utils import (
@@ -53,6 +54,12 @@ class ServingEmbedding(PoolingServing):
54 def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor:
55 return EmbedIOProcessor(*args, **kwargs)
56
57 + async def _preprocessing(
58 + self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext
59 + ):
60 + await super()._preprocessing(io_processor, ctx)
61 + self.io_processor.maybe_pre_process_chunked(ctx)
62 +
63 def _build_response(
64 self,
65 ctx: PoolingServeContext,
vllm/entrypoints/pooling/offline.py
+5 -2
@@ -25,7 +25,7 @@ from .factories import init_pooling_io_processors
25 from .scoring.io_processor import ScoringIOProcessor
26 from .scoring.typing import ScoreInput
27 from .typing import (
28 - ALLOfflineInputsContext,
28 + AnyOfflineInputsContext,
29 OfflineEncodeInputsContext,
30 OfflineOutputsContext,
31 OfflinePluginInputsContext,
@@ -104,7 +104,7 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
104
105 io_processor = self.pooling_io_processors[pooling_task]
106
107 - ctx: ALLOfflineInputsContext
107 + ctx: AnyOfflineInputsContext
108 if isinstance(prompts, dict) and "data" in prompts:
109 ctx = OfflinePluginInputsContext(
110 pooling_task=pooling_task,
@@ -399,6 +399,9 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
399 num_requests: int,
400 use_tqdm: bool | Callable[..., tqdm] = True,
401 ):
402 + if num_requests == 0:
403 + raise ValueError("You must pass at least one prompt")
404 +
405 # Keeping max_num_seqs * 2 requests in the core can already saturate the core.
406 # Therefore, keep most requests waiting outside the core.
407 max_requests_in_core = (
vllm/entrypoints/pooling/pooling/api_router.py
+5 -5
@@ -17,8 +17,11 @@ from .serving import ServingPooling
17 router = APIRouter()
18
19
20 -def pooling(request: Request) -> ServingPooling | None:
21 - return request.app.state.serving_pooling
20 +def pooling(request: Request) -> ServingPooling:
21 + handler = getattr(request.app.state, "serving_pooling", None)
22 + if handler is None:
23 + raise NotImplementedError("The model does not support Pooling API")
24 + return handler
25
26
27 @router.post(
@@ -33,7 +36,4 @@ def pooling(request: Request) -> ServingPooling | None:
36 @load_aware_call
37 async def create_pooling(request: PoolingRequest, raw_request: Request):
38 handler = pooling(raw_request)
36 - if handler is None:
37 - raise NotImplementedError("The model does not support Pooling API")
38 -
39 return await handler(request, raw_request)
vllm/entrypoints/pooling/pooling/io_processor.py
+26 -13
@@ -10,7 +10,9 @@ from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
10
11 from ..base.io_processor import PoolingIOProcessor
12 from ..typing import (
13 - ALLOfflineInputsContext,
13 + AnyOfflineInputsContext,
14 + AnyRenderParam,
15 + EncodeCMPLRenderParams,
16 OfflineEncodeInputsContext,
17 OfflineOutputsContext,
18 OfflinePluginInputsContext,
@@ -49,7 +51,9 @@ class PluginWithIOProcessorPlugins(PoolingIOProcessor):
51 #######################################
52 # online APIs
53
52 - def pre_process_online(self, ctx: PoolingServeContext):
54 + def get_request_factory_online(
55 + self, ctx: PoolingServeContext
56 + ) -> Sequence[AnyRenderParam]:
57 assert isinstance(ctx.request, IOProcessorRequest)
58
59 validated_prompt = self.io_processor.parse_data(ctx.request.data)
@@ -66,24 +70,33 @@ class PluginWithIOProcessorPlugins(PoolingIOProcessor):
70 )
71 for prompt in prompt_to_seq(raw_prompts)
72 ]
73 + num_requests = len(parsed_prompts)
74
75 tok_params = ctx.request.build_tok_params(self.model_config)
76
72 - ctx.engine_inputs = self.renderer.render_cmpl(
73 - parsed_prompts,
74 - tok_params,
75 - prompt_extras={
76 - k: v
77 - for k in ("mm_processor_kwargs", "cache_salt")
78 - if (v := getattr(ctx.request, k, None)) is not None
79 - },
80 - )
81 -
77 pooling_params = self.io_processor.merge_pooling_params()
78 if pooling_params.task is None:
79 pooling_params.task = "plugin"
80 ctx.pooling_params = pooling_params
81
82 + params_seq = self._params_to_seq(ctx.pooling_params, num_requests)
83 + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
84 + seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
85 +
86 + requests = [
87 + EncodeCMPLRenderParams(
88 + prompts=parsed_prompts[i],
89 + tok_params=tok_params,
90 + prompt_extras=ctx.prompt_extras,
91 + skip_mm_cache=False,
92 + params=params_seq[i],
93 + lora_requests=seq_lora_requests[i],
94 + priorities=seq_priority[i],
95 + )
96 + for i in range(num_requests)
97 + ]
98 + return requests
99 +
100 def post_process_online(
101 self,
102 ctx: PoolingServeContext,
@@ -114,7 +127,7 @@ class PluginWithIOProcessorPlugins(PoolingIOProcessor):
127 # offline APIs
128
129 def get_request_factory_offline(
117 - self, ctx: ALLOfflineInputsContext
130 + self, ctx: AnyOfflineInputsContext
131 ) -> tuple[RequestFactory, int]:
132 assert isinstance(ctx, OfflinePluginInputsContext)
133 assert isinstance(ctx.prompts, dict) and "data" in ctx.prompts
vllm/entrypoints/pooling/scoring/api_router.py
+10 -10
@@ -20,12 +20,18 @@ router = APIRouter()
20 logger = init_logger(__name__)
21
22
23 -def score(request: Request) -> ServingScores | None:
24 - return request.app.state.serving_scores
23 +def score(request: Request) -> ServingScores:
24 + handler = getattr(request.app.state, "serving_scores", None)
25 + if handler is None:
26 + raise NotImplementedError("The model does not support Score API")
27 + return handler
28
29
27 -def rerank(request: Request) -> ServingScores | None:
28 - return request.app.state.serving_scores
30 +def rerank(request: Request) -> ServingScores:
31 + handler = getattr(request.app.state, "serving_scores", None)
32 + if handler is None:
33 + raise NotImplementedError("The model does not support Rerank (Score) API")
34 + return handler
35
36
37 @router.post(
@@ -40,9 +46,6 @@ def rerank(request: Request) -> ServingScores | None:
46 @load_aware_call
47 async def create_score(request: ScoreRequest, raw_request: Request):
48 handler = score(raw_request)
43 - if handler is None:
44 - raise NotImplementedError("The model does not support Score API")
45 -
49 return await handler(request, raw_request)
50
51
@@ -77,9 +80,6 @@ async def create_score_v1(request: ScoreRequest, raw_request: Request):
80 @load_aware_call
81 async def do_rerank(request: RerankRequest, raw_request: Request):
82 handler = rerank(raw_request)
80 - if handler is None:
81 - raise NotImplementedError("The model does not support Rerank (Score) API")
82 -
83 return await handler(request, raw_request)
84
85
vllm/entrypoints/pooling/scoring/io_processor.py
+147 -192
@@ -6,8 +6,7 @@ from typing import Any, TypeAlias, cast
6
7 import torch.nn.functional as F
8
9 -from vllm import PoolingParams, PoolingRequestOutput, PromptType, TokensPrompt
10 -from vllm.inputs import EngineInput
9 +from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt
10 from vllm.renderers import TokenizeParams
11 from vllm.renderers.hf import safe_apply_chat_template
12 from vllm.renderers.inputs.preprocess import (
@@ -20,8 +19,11 @@ from vllm.utils.mistral import is_mistral_tokenizer
19
20 from ...chat_utils import ChatTemplateResolutionError
21 from ..base.io_processor import PoolingIOProcessor
22 +from ..pooling.protocol import PoolingCompletionRequest
23 from ..typing import (
24 - ALLOfflineInputsContext,
24 + AnyOfflineInputsContext,
25 + AnyPoolingRequest,
26 + AnyRenderParam,
27 EncodeChatRenderParams,
28 EncodeCMPLRenderParams,
29 OfflineEncodeInputsContext,
@@ -167,17 +169,7 @@ class ScoringIOProcessor(PoolingIOProcessor):
169 )
170 return scoring_data
171
170 -
171 -class BiEncoderIOProcessor(ScoringIOProcessor):
172 - name = "bi-encoder"
173 - pooling_task: PoolingTask = "embed"
174 -
175 - #######################################
176 - # online APIs
177 -
178 - def pre_process_online(self, ctx: ScoringServeContext):
179 - request = ctx.request
180 -
172 + def valid_inputs_online(self, request: AnyPoolingRequest):
173 if isinstance(request, ScoreRequest):
174 data_1 = request.data_1
175 data_2 = request.data_2
@@ -185,9 +177,24 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
177 data_1 = request.query
178 data_2 = request.documents
179 else:
188 - raise ValueError(f"Invalid {self.name} request type")
180 + raise ValueError(f"Invalid {request.__class__.__name__} request type")
181
182 scoring_data = self.valid_inputs(data_1, data_2)
183 + return scoring_data
184 +
185 +
186 +class BiEncoderIOProcessor(ScoringIOProcessor):
187 + name = "bi-encoder"
188 + pooling_task: PoolingTask = "embed"
189 +
190 + #######################################
191 + # online APIs
192 +
193 + def get_request_factory_online(
194 + self, ctx: PoolingServeContext
195 + ) -> Sequence[AnyRenderParam]:
196 + request = ctx.request
197 + scoring_data = self.valid_inputs_online(request)
198
199 max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
200 request=request
@@ -197,19 +204,37 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
204 scoring_data, max_tokens_per_query, max_tokens_per_doc
205 )
206
200 - tok_params = request.build_tok_params(self.model_config)
201 - engine_inputs = self._pre_process(
202 - scoring_data,
203 - tok_params,
204 - prompt_extras={
205 - k: v
206 - for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs")
207 - if (v := getattr(request, k, None)) is not None
208 - },
207 + data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config)
208 + data_2 = score_data_to_prompts(
209 + scoring_data.data_2, "document", self.model_config
210 )
211 + prompts = data_1 + data_2
212 + ctx.n_queries = len(data_1)
213
211 - ctx.engine_inputs = engine_inputs
212 - ctx.n_queries = len(scoring_data.data_1)
214 + prompts_seq = prompt_to_seq(prompts)
215 + parsed_prompts = [
216 + parse_model_prompt(self.model_config, prompt) for prompt in prompts_seq
217 + ]
218 + num_requests = len(parsed_prompts)
219 +
220 + tok_params = request.build_tok_params(self.model_config)
221 + params_seq = self._params_to_seq(ctx.pooling_params, num_requests)
222 + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
223 + seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
224 +
225 + requests = [
226 + EncodeCMPLRenderParams(
227 + prompts=parsed_prompts[i],
228 + tok_params=tok_params,
229 + prompt_extras=ctx.prompt_extras,
230 + skip_mm_cache=False,
231 + params=params_seq[i],
232 + lora_requests=seq_lora_requests[i],
233 + priorities=seq_priority[i],
234 + )
235 + for i in range(num_requests)
236 + ]
237 + return requests
238
239 def post_process_online(
240 self,
@@ -226,7 +251,7 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
251 # offline APIs
252
253 def get_request_factory_offline(
229 - self, ctx: ALLOfflineInputsContext
254 + self, ctx: AnyOfflineInputsContext
255 ) -> tuple[RequestFactory, int]:
256 assert isinstance(ctx, OfflineScoringInputsContext)
257
@@ -267,41 +292,6 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
292 #######################################
293 # helpers
294
270 - def _pre_process(
271 - self,
272 - scoring_data: ScoringData,
273 - tok_params: TokenizeParams,
274 - prompt_extras: dict[str, Any] | None = None,
275 - ) -> Sequence[EngineInput]:
276 - data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config)
277 - data_2 = score_data_to_prompts(
278 - scoring_data.data_2, "document", self.model_config
279 - )
280 -
281 - return self._preprocess_cmpl_offline(
282 - prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras
283 - )
284 -
285 - def _preprocess_cmpl_offline(
286 - self,
287 - prompts: PromptType | Sequence[PromptType],
288 - tok_params: TokenizeParams,
289 - prompt_extras: dict[str, Any] | None = None,
290 - ) -> Sequence[EngineInput]:
291 - prompts = prompt_to_seq(prompts)
292 - parsed_prompts = [
293 - (
294 - prompt
295 - if isinstance(prompt, bytes)
296 - else parse_model_prompt(self.model_config, prompt)
297 - )
298 - for prompt in prompts
299 - ]
300 -
301 - return self.renderer.render_cmpl(
302 - parsed_prompts, tok_params, prompt_extras=prompt_extras
303 - )
304 -
295 def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
296 emb_data_1 = outputs[:n_queries]
297 emb_data_2 = outputs[n_queries:]
@@ -432,49 +422,50 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
422 #######################################
423 # online APIs
424
435 - def pre_process_online(self, ctx: ScoringServeContext):
425 + def get_request_factory_online(
426 + self, ctx: PoolingServeContext
427 + ) -> Sequence[AnyRenderParam]:
428 request = ctx.request
429 + scoring_data = self.valid_inputs_online(request)
430 + data_1 = scoring_data.data_1
431 + data_2 = scoring_data.data_2
432 + num_requests = len(data_2)
433
438 - if isinstance(request, ScoreRequest):
439 - data_1 = request.data_1
440 - data_2 = request.data_2
441 - elif isinstance(request, RerankRequest):
442 - data_1 = request.query
443 - data_2 = request.documents
444 - else:
445 - raise ValueError(f"Invalid {self.name} request type")
446 -
447 - scoring_data = self.valid_inputs(data_1, data_2)
434 + if len(data_1) == 1:
435 + data_1 = data_1 * num_requests
436
437 max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
438 request=request
439 )
440
441 tok_params = request.build_tok_params(self.model_config)
454 - pooling_params = self.create_pooling_params(request)
455 -
456 - engine_inputs, pooling_params_list = self._pre_process(
457 - scoring_data,
458 - tok_params,
459 - pooling_params,
460 - chat_template=self.chat_template,
461 - max_tokens_per_query=max_tokens_per_query,
462 - max_tokens_per_doc=max_tokens_per_doc,
463 - prompt_extras={
464 - k: v
465 - for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs")
466 - if (v := getattr(request, k, None)) is not None
467 - },
468 - )
442 + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
443 + seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
444 +
445 + requests = [
446 + ScoringRenderParams(
447 + data_1=data_1[i],
448 + data_2=data_2[i],
449 + chat_template=self.chat_template,
450 + max_tokens_per_query=max_tokens_per_query,
451 + max_tokens_per_doc=max_tokens_per_doc,
452 + tok_params=tok_params,
453 + prompt_extras=ctx.prompt_extras,
454 + skip_mm_cache=False,
455 + params=ctx.pooling_params,
456 + lora_requests=seq_lora_requests[i],
457 + priorities=seq_priority[i],
458 + )
459 + for i in range(num_requests)
460 + ]
461
470 - ctx.engine_inputs = engine_inputs
471 - ctx.pooling_params = pooling_params_list
462 + return requests
463
464 #######################################
465 # offline APIs
466
467 def get_request_factory_offline(
477 - self, ctx: ALLOfflineInputsContext
468 + self, ctx: AnyOfflineInputsContext
469 ) -> tuple[RequestFactory, int]:
470 assert isinstance(ctx, OfflineScoringInputsContext)
471
@@ -485,10 +476,12 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
476 if len(data_1) == 1:
477 data_1 = data_1 * num_requests
478
479 + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
480 + pooling_params=ctx.pooling_params
481 + )
482 tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
483 **(ctx.tokenization_kwargs or {})
484 )
491 -
485 prompt_extras = ctx.pooling_params.extra_kwargs
486
487 seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
@@ -500,6 +493,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
493 data_1=data_1[i],
494 data_2=data_2[i],
495 chat_template=ctx.chat_template,
496 + max_tokens_per_query=max_tokens_per_query,
497 + max_tokens_per_doc=max_tokens_per_doc,
498 tok_params=tok_params,
499 prompt_extras=prompt_extras,
500 skip_mm_cache=False,
@@ -531,17 +526,13 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
526 params = render_params["params"]
527 prompt_extras = render_params["prompt_extras"]
528
534 - max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
535 - pooling_params=params
536 - )
537 -
529 _, engine_prompt = self.get_score_prompt(
530 data_1=render_params["data_1"],
531 data_2=render_params["data_2"],
532 encode_kwargs=tok_params.get_encode_kwargs(),
533 chat_template=render_params["chat_template"],
543 - max_tokens_per_query=max_tokens_per_query,
544 - max_tokens_per_doc=max_tokens_per_doc,
534 + max_tokens_per_query=render_params["max_tokens_per_query"],
535 + max_tokens_per_doc=render_params["max_tokens_per_doc"],
536 chat_template_kwargs=prompt_extras.get("chat_template_kwargs")
537 if prompt_extras
538 else None,
@@ -551,29 +542,16 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
542
543 if token_type_ids := engine_prompt.pop("token_type_ids", None):
544 params = params.clone()
554 - compressed = compress_token_type_ids(token_type_ids)
555 - params.extra_kwargs = {"compressed_token_type_ids": compressed}
556 -
557 - engine_input = self.renderer.process_for_engine(engine_prompt, arrival_time)
558 -
559 - return PoolingEngineInput(
560 - prompts=engine_input,
561 - params=params,
562 - lora_requests=render_params["lora_requests"],
563 - priorities=render_params["priorities"],
564 - )
545 + compressed = compress_token_type_ids(
546 + _apply_post_tokenization_to_token_type_ids(
547 + self.tokenizer, tok_params, token_type_ids
548 + )
549 + )
550 + params.extra_kwargs = {
551 + **(params.extra_kwargs or {}),
552 + "compressed_token_type_ids": compressed,
553 + }
554
566 - def _pre_process(
567 - self,
568 - scoring_data: ScoringData,
569 - tok_params: TokenizeParams,
570 - pooling_params: PoolingParams | None,
571 - chat_template: str | None = None,
572 - max_tokens_per_query: int = 0,
573 - max_tokens_per_doc: int = 0,
574 - prompt_extras: dict[str, Any] | None = None,
575 - ) -> tuple[Sequence[EngineInput], list[PoolingParams]]:
576 - arrival_time = time.time()
555 engine_prompt_extras = (
556 {
557 k: v
@@ -584,55 +562,18 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
562 else None
563 )
564
587 - data_1 = scoring_data.data_1
588 - data_2 = scoring_data.data_2
565 + if engine_prompt_extras:
566 + target_prompt = extract_target_prompt(self.model_config, engine_prompt)
567 + target_prompt.update(engine_prompt_extras)
568
590 - if len(data_1) == 1:
591 - data_1 = data_1 * len(data_2)
592 -
593 - if pooling_params is None:
594 - pooling_params = PoolingParams(task="classify")
595 -
596 - pooling_params_list = list[PoolingParams]()
597 - engine_inputs = list[EngineInput]()
598 - for q, d in zip(data_1, data_2):
599 - _, engine_prompt = self.get_score_prompt(
600 - data_1=q,
601 - data_2=d,
602 - encode_kwargs=tok_params.get_encode_kwargs(),
603 - chat_template=chat_template,
604 - max_tokens_per_query=max_tokens_per_query,
605 - max_tokens_per_doc=max_tokens_per_doc,
606 - chat_template_kwargs=prompt_extras.get("chat_template_kwargs")
607 - if prompt_extras
608 - else None,
609 - )
610 -
611 - token_type_ids = engine_prompt.pop("token_type_ids", None)
612 - tok_params.apply_post_tokenization(self.tokenizer, engine_prompt)
613 -
614 - if token_type_ids is not None:
615 - params = pooling_params.clone()
616 - compressed = compress_token_type_ids(
617 - _apply_post_tokenization_to_token_type_ids(
618 - self.tokenizer, tok_params, token_type_ids
619 - )
620 - )
621 - params.extra_kwargs = {
622 - **(params.extra_kwargs or {}),
623 - "compressed_token_type_ids": compressed,
624 - }
625 - pooling_params_list.append(params)
626 - else:
627 - pooling_params_list.append(pooling_params)
569 + engine_input = self.renderer.process_for_engine(engine_prompt, arrival_time)
570
629 - if engine_prompt_extras:
630 - target_prompt = extract_target_prompt(self.model_config, engine_prompt)
631 - target_prompt.update(engine_prompt_extras)
632 - engine_inputs.append(
633 - self.renderer.process_for_engine(engine_prompt, arrival_time)
634 - )
635 - return engine_inputs, pooling_params_list
571 + return PoolingEngineInput(
572 + prompts=engine_input,
573 + params=params,
574 + lora_requests=render_params["lora_requests"],
575 + priorities=render_params["priorities"],
576 + )
577
578 def get_score_prompt(
579 self,
@@ -837,16 +778,27 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
778 name = "jina-reranking-scoring"
779 pooling_task: PoolingTask = "token_embed"
780
840 - def get_request_factory_offline(
841 - self, ctx: ALLOfflineInputsContext
842 - ) -> tuple[RequestFactory, int]:
843 - assert isinstance(ctx, OfflineScoringInputsContext)
781 + def get_request_factory_online(
782 + self, ctx: PoolingServeContext
783 + ) -> Sequence[AnyRenderParam]:
784 + request = ctx.request
785 + ctx.n_queries = 1
786
845 - scoring_data = ctx.scoring_data
846 - prompt_extras = ctx.pooling_params.extra_kwargs
787 + prompt_extras = ctx.prompt_extras
788 + scoring_data = self.valid_inputs_online(request)
789 +
790 + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
791 + request=request
792 + )
793 +
794 + if max_tokens_per_query > 0 or max_tokens_per_doc > 0:
795 + scoring_data = self._truncate_scoring_data(
796 + scoring_data, max_tokens_per_query, max_tokens_per_doc
797 + )
798
799 queries = self.ensure_str(scoring_data.data_1)
800 docs = self.ensure_str(scoring_data.data_2)
801 +
802 chat_template_kwargs = (
803 prompt_extras.get("chat_template_kwargs") if prompt_extras else None
804 )
@@ -868,24 +820,19 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
820 for q, d in zip(queries, docs)
821 ]
822
871 - return PoolingIOProcessor.get_request_factory_offline(
872 - self,
873 - OfflineEncodeInputsContext(
874 - pooling_task=self.pooling_task,
875 - prompts=prompts,
876 - tokenization_kwargs=ctx.tokenization_kwargs,
877 - pooling_params=ctx.pooling_params,
878 - lora_request=ctx.lora_request,
879 - priorities=ctx.priorities,
880 - ),
881 - )
823 + ctx.request = PoolingCompletionRequest(task="token_embed", input=prompts)
824 + requests = PoolingIOProcessor.get_request_factory_online(self, ctx)
825 + ctx.request = request
826 + return requests
827 +
828 + def get_request_factory_offline(
829 + self, ctx: AnyOfflineInputsContext
830 + ) -> tuple[RequestFactory, int]:
831 + assert isinstance(ctx, OfflineScoringInputsContext)
832 +
833 + scoring_data = ctx.scoring_data
834 + prompt_extras = ctx.pooling_params.extra_kwargs
835
883 - def _pre_process(
884 - self,
885 - scoring_data: ScoringData,
886 - tok_params: TokenizeParams,
887 - prompt_extras: dict[str, Any] | None = None,
888 - ) -> Sequence[EngineInput]:
836 queries = self.ensure_str(scoring_data.data_1)
837 docs = self.ensure_str(scoring_data.data_2)
838 chat_template_kwargs = (
@@ -909,8 +856,16 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
856 for q, d in zip(queries, docs)
857 ]
858
912 - return self._preprocess_cmpl_offline(
913 - prompts=prompts, tok_params=tok_params, prompt_extras=prompt_extras
859 + return PoolingIOProcessor.get_request_factory_offline(
860 + self,
861 + OfflineEncodeInputsContext(
862 + pooling_task=self.pooling_task,
863 + prompts=prompts,
864 + tokenization_kwargs=ctx.tokenization_kwargs,
865 + pooling_params=ctx.pooling_params,
866 + lora_request=ctx.lora_request,
867 + priorities=ctx.priorities,
868 + ),
869 )
870
871 def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
vllm/entrypoints/pooling/scoring/serving.py
+13 -19
@@ -190,7 +190,7 @@ class ServingScores(PoolingServing):
190
191 async def flash_late_interaction(self, *args, **kwargs) -> Response:
192 ctx = await self._init_ctx(self.io_processor, *args, **kwargs)
193 - await self._preprocessing_async(self.io_processor, ctx)
193 + await self._preprocessing(self.io_processor, ctx)
194
195 # stage 1: encode queries and cache token embeddings on workers.
196 await self._flash_late_interaction_encode_queries(ctx)
@@ -211,7 +211,6 @@ class ServingScores(PoolingServing):
211 query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
212 query_uses = [n_docs if n_queries == 1 else 1] * n_queries
213
214 - query_pooling_params_list = []
214 for i in range(n_queries):
215 pooling_params = ctx.pooling_params.clone()
216 pooling_params.late_interaction_params = (
@@ -220,23 +219,21 @@ class ServingScores(PoolingServing):
219 query_uses=query_uses[i],
220 )
221 )
223 - query_pooling_params_list.append(pooling_params)
222 + query_engine_inputs[i]["params"] = pooling_params
223
225 - assert (
226 - n_queries
227 - == len(query_pooling_params_list)
228 - == len(query_engine_inputs)
229 - == len(query_keys)
230 - )
224 + assert n_queries == len(query_engine_inputs) == len(query_keys)
225
226 query_ctx = ScoringServeContext(
227 request=ctx.request,
228 raw_request=ctx.raw_request,
229 model_name=ctx.model_name,
230 request_id=ctx.request_id,
237 - pooling_params=query_pooling_params_list,
231 + pooling_params=ctx.pooling_params,
232 prompt_request_ids=query_keys,
233 engine_inputs=query_engine_inputs,
234 + lora_request=ctx.lora_request,
235 + priorities=ctx.priorities,
236 + prompt_extras=ctx.prompt_extras,
237 )
238
239 await self._prepare_generators(query_ctx)
@@ -255,30 +252,27 @@ class ServingScores(PoolingServing):
252 query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
253 doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)]
254
258 - doc_pooling_params_list = []
255 for i in range(n_docs):
256 query_idx = 0 if n_queries == 1 else i
257 pooling_params = ctx.pooling_params.clone()
258 pooling_params.late_interaction_params = build_late_interaction_doc_params(
259 query_key=query_keys[query_idx]
260 )
265 - doc_pooling_params_list.append(pooling_params)
261 + doc_engine_inputs[i]["params"] = pooling_params
262
267 - assert (
268 - n_docs
269 - == len(doc_pooling_params_list)
270 - == len(doc_engine_inputs)
271 - == len(doc_keys)
272 - )
263 + assert n_docs == len(doc_engine_inputs) == len(doc_keys)
264
265 doc_ctx = ScoringServeContext(
266 request=ctx.request,
267 raw_request=ctx.raw_request,
268 model_name=ctx.model_name,
269 request_id=ctx.request_id,
279 - pooling_params=doc_pooling_params_list,
270 + pooling_params=ctx.pooling_params,
271 prompt_request_ids=doc_keys,
272 engine_inputs=doc_engine_inputs,
273 + lora_request=ctx.lora_request,
274 + priorities=ctx.priorities,
275 + prompt_extras=ctx.prompt_extras,
276 )
277
278 await self._prepare_generators(doc_ctx)
vllm/entrypoints/pooling/typing.py
+18 -12
@@ -88,10 +88,13 @@ class PoolingServeContext(Generic[PoolingRequestT]):
88 raw_request: Request | None = None
89 model_name: str
90 request_id: str
91 - pooling_params: PoolingParams | list[PoolingParams]
91 + pooling_params: PoolingParams
92 + lora_request: LoRARequest | None
93 + priorities: int | Sequence[int] | None
94 + prompt_extras: dict[str, Any] | None
95 +
96 created_time: int = field(default_factory=lambda: int(time.time()))
93 - lora_request: LoRARequest | None = None
94 - engine_inputs: Sequence[EngineInput] | None = None
97 + engine_inputs: Sequence["PoolingEngineInput"] | None = None
98 prompt_request_ids: list[str] | None = None
99
100 result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = (
@@ -100,7 +103,7 @@ class PoolingServeContext(Generic[PoolingRequestT]):
103 final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)
104
105 ## for Long Text Embedding with Chunked Processing
103 - original_engine_inputs: Sequence[EngineInput] | None = None
106 + original_engine_inputs: Sequence["PoolingEngineInput"] | None = None
107 chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] | None = None
108
109 ## for bi-encoder & late-interaction
@@ -118,7 +121,7 @@ class OfflineInputsContext:
121 pooling_task: PoolingTask
122 tokenization_kwargs: dict[str, Any] | None
123 lora_request: Sequence[LoRARequest | None] | None
121 - priorities: Sequence[int] | None
124 + priorities: int | Sequence[int] | None
125
126
127 @dataclass
@@ -140,7 +143,7 @@ class OfflinePluginInputsContext(OfflineInputsContext):
143 pooling_params: PoolingParams | Sequence[PoolingParams] | None
144
145
143 -ALLOfflineInputsContext: TypeAlias = (
146 +AnyOfflineInputsContext: TypeAlias = (
147 OfflineEncodeInputsContext
148 | OfflineScoringInputsContext
149 | OfflinePluginInputsContext
@@ -178,6 +181,15 @@ class ScoringRenderParams(RenderParams):
181 data_1: ScoreData
182 data_2: ScoreData
183 chat_template: str | None
184 + max_tokens_per_query: int
185 + max_tokens_per_doc: int
186 +
187 +
188 +AnyRenderParam: TypeAlias = (
189 + EncodeCMPLRenderParams | EncodeChatRenderParams | ScoringRenderParams
190 +)
191 +RequestGenerator: TypeAlias = Generator[AnyRenderParam]
192 +RequestFactory: TypeAlias = Callable[[], RequestGenerator]
193
194
195 class PoolingEngineInput(TypedDict):
@@ -185,9 +197,3 @@ class PoolingEngineInput(TypedDict):
197 params: PoolingParams
198 lora_requests: LoRARequest | None
199 priorities: int
188 -
189 -
190 -RequestGenerator: TypeAlias = Generator[
191 - EncodeCMPLRenderParams | EncodeChatRenderParams | ScoringRenderParams
192 -]
193 -RequestFactory: TypeAlias = Callable[[], RequestGenerator]