[Frontend] Overlap preprocessing and computation for pooling models offline inference (#47699)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Jul 16, 2026 at 22:00 UTC
3e90d015ba5bbad6f229ddb4c1f467ca33101524
10 files changed
+673
-150
.buildkite/test_areas/models_language.yaml
+3
-3
@@ -103,10 +103,10 @@ steps:
103
commands:
104
- pytest -v -s models/language/generation_ppl_test
105
106
-- label: Language Models Test (Extended Pooling) # 36min
106
+- label: Language Models Test (Extended Pooling)
107
device: h200_35gb
108
key: language-models-test-extended-pooling
109
- timeout_in_minutes: 70
109
+ timeout_in_minutes: 120
110
optional: true
111
source_file_dependencies:
112
- vllm/
@@ -116,7 +116,7 @@ steps:
116
mirror:
117
amd:
118
device: mi325_1
119
- timeout_in_minutes: 100
119
+ timeout_in_minutes: 120
120
depends_on:
121
- image-build-amd
122
tests/entrypoints/pooling/basic/test_tiling_engine.py
new
+108
@@ -0,0 +1,108 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import weakref
5
+from unittest import mock
6
+
7
+import pytest
8
+
9
+from vllm import LLM, PoolingParams
10
+from vllm.distributed import cleanup_dist_env_and_memory
11
+
12
+MODEL_NAME = "intfloat/multilingual-e5-small"
13
+
14
+
15
+@pytest.fixture(scope="module")
16
+def llm():
17
+ llm = LLM(
18
+ model=MODEL_NAME,
19
+ max_num_seqs=2, # small to trigger tiling
20
+ tensor_parallel_size=1,
21
+ gpu_memory_utilization=0.75,
22
+ enforce_eager=True,
23
+ seed=0,
24
+ )
25
+
26
+ yield weakref.proxy(llm)
27
+
28
+ del llm
29
+
30
+ cleanup_dist_env_and_memory()
31
+
32
+
33
+@pytest.mark.skip_global_cleanup
34
+def test_tiling_engine_basic(llm):
35
+ """
36
+ Basic test with a small number of prompts (less than max_num_seqs).
37
+ No tiling should be triggered, but the engine still processes correctly.
38
+ """
39
+ prompts = ["Hello", "World"]
40
+ outputs = llm.encode(prompts, pooling_task="embed")
41
+ assert len(outputs) == len(prompts)
42
+
43
+
44
+@pytest.mark.skip_global_cleanup
45
+def test_tiling_engine_many_requests(llm):
46
+ """
47
+ Test with a large number of prompts that exceeds max_num_seqs.
48
+ This verifies that _run_tiling_engine correctly chunks requests,
49
+ processes all of them, and returns outputs in the correct order.
50
+ """
51
+ num_prompts = 10
52
+ prompts = [f"Prompt {i}" for i in range(num_prompts)]
53
+ outputs = llm.encode(prompts, pooling_task="embed")
54
+ assert len(outputs) == num_prompts
55
+
56
+
57
+@pytest.mark.skip_global_cleanup
58
+def test_tiling_engine_with_pooling_params(llm):
59
+ """
60
+ Test the tiling engine when different PoolingParams are provided.
61
+ The engine must handle a list of params that matches the number of prompts.
62
+ """
63
+ num_prompts = 10
64
+ prompts = [f"Prompt {i}" for i in range(num_prompts)]
65
+ pooling_params = [PoolingParams() for _ in range(num_prompts)]
66
+
67
+ outputs = llm.encode(prompts, pooling_params=pooling_params, pooling_task="embed")
68
+ assert len(outputs) == num_prompts
69
+
70
+ # Single PoolingParams shared across all prompts
71
+ single_param = PoolingParams()
72
+ outputs = llm.encode(prompts, pooling_params=single_param, pooling_task="embed")
73
+ assert len(outputs) == num_prompts
74
+
75
+ # None PoolingParams should fall back to default
76
+ outputs = llm.encode(prompts, pooling_params=None, pooling_task="embed")
77
+ assert len(outputs) == num_prompts
78
+
79
+
80
+@pytest.mark.skip_global_cleanup
81
+def test_tiling_engine_abort_on_exception(llm):
82
+ """
83
+ Test that abort_request IS called with the correct arguments when an
84
+ exception occurs inside the engine's step() loop.
85
+ """
86
+ prompts = ["Prompt 0", "Prompt 1", "Prompt 2"]
87
+
88
+ # Mock the step method to throw an exception on the second call
89
+ original_step = llm.llm_engine.step
90
+ call_count = 0
91
+
92
+ def mocked_step():
93
+ nonlocal call_count
94
+ call_count += 1
95
+ if call_count == 2:
96
+ raise RuntimeError("Simulated engine error")
97
+ return original_step()
98
+
99
+ with mock.patch.object(llm.llm_engine, "step", side_effect=mocked_step):
100
+ # We expect an exception to be raised from encode
101
+ with mock.patch.object(llm.llm_engine, "abort_request") as mock_abort: # noqa: SIM117
102
+ with pytest.raises(RuntimeError, match="Simulated engine error"):
103
+ llm.encode(prompts, pooling_task="embed")
104
+
105
+ args, kwargs = mock_abort.call_args
106
+ request_ids = args[0]
107
+ assert isinstance(request_ids, list)
108
+ assert len(request_ids) > 0
tests/models/language/pooling/test_multi_vector_retrieval.py
+3
-6
@@ -12,7 +12,7 @@ from vllm.config import PoolerConfig
12
"model",
13
["BAAI/bge-m3"],
14
)
15
-@pytest.mark.parametrize("dtype", ["half"])
15
+@pytest.mark.parametrize("dtype", ["float"])
16
@torch.inference_mode
17
def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype: str):
18
with vllm_runner(
@@ -20,13 +20,11 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
20
runner="pooling",
21
pooler_config=PoolerConfig(task="token_embed"),
22
max_model_len=None,
23
+ dtype=dtype,
24
) as vllm_model:
25
vllm_outputs = vllm_model.token_embed(example_prompts)
26
26
- with hf_runner(
27
- model,
28
- auto_cls=AutoModel,
29
- ) as hf_model:
27
+ with hf_runner(model, auto_cls=AutoModel, dtype=dtype) as hf_model:
28
tokenizer = hf_model.tokenizer
29
hf_outputs = []
30
for prompt in example_prompts:
@@ -34,7 +32,6 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
32
inputs = hf_model.wrap_device(inputs)
33
output = hf_model.model(**inputs)
34
embedding = output.last_hidden_state[0].float()
37
- # normal
35
hf_outputs.append(embedding.cpu())
36
37
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
vllm/entrypoints/llm.py
+1
-1
@@ -356,7 +356,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
356
# path; the synchronous `LLM` entrypoint runs multimodal
357
# preprocessing serially. Warn so the setting is not a silent
358
# no-op. See vllm-project/vllm#42901.
359
- if self.model_config.renderer_num_workers > 1:
359
+ if self.model_config.renderer_num_workers > 1 and self.runner_type != "pooling":
360
logger.warning_once(
361
"`renderer_num_workers=%d` was set, but the offline `LLM` "
362
"entrypoint uses the synchronous renderer path and runs "
vllm/entrypoints/pooling/base/io_processor.py
+133
-30
@@ -2,9 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Sequence
5
-from typing import Any, Final
5
+from typing import Any, Final, cast
6
7
-from vllm import PoolingParams, PoolingRequestOutput, PromptType
7
+from vllm import (
8
+ PoolingParams,
9
+ PoolingRequestOutput,
10
+)
11
from vllm.config import VllmConfig
12
from vllm.entrypoints.chat_utils import (
13
ChatCompletionMessageParam,
@@ -14,18 +17,25 @@ from vllm.entrypoints.chat_utils import (
17
)
18
from vllm.entrypoints.serve.engine.typing import RendererChatRequest, RendererRequest
19
from vllm.inputs import EngineInput, SingletonPrompt
17
-from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs
20
+from vllm.lora.request import LoRARequest
21
+from vllm.renderers import BaseRenderer, merge_kwargs
22
from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
23
from vllm.tool_parsers import ToolParser
24
from vllm.utils.mistral import is_mistral_tokenizer
25
22
-from ..scoring.typing import ScoringData
26
from ..typing import (
24
- OfflineInputsContext,
27
+ ALLOfflineInputsContext,
28
+ EncodeChatRenderParams,
29
+ EncodeCMPLRenderParams,
30
+ OfflineEncodeInputsContext,
31
OfflineOutputsContext,
32
PoolingChatLikeRequest,
33
PoolingCompletionLikeRequest,
34
+ PoolingEngineInput,
35
PoolingServeContext,
36
+ RequestFactory,
37
+ RequestGenerator,
38
+ ScoringRenderParams,
39
)
40
41
@@ -98,16 +108,62 @@ class PoolingIOProcessor:
108
#######################################
109
# offline APIs
110
101
- def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
102
- assert not isinstance(ctx.prompts, ScoringData) and not (
103
- isinstance(ctx.prompts, dict) and "data" in ctx.prompts
104
- )
111
+ def get_request_factory_offline(
112
+ self, ctx: ALLOfflineInputsContext
113
+ ) -> tuple[RequestFactory, int]:
114
+ assert isinstance(ctx, OfflineEncodeInputsContext)
115
116
prompts_seq = prompt_to_seq(ctx.prompts)
117
+ num_requests = len(prompts_seq)
118
+ pooling_task = ctx.pooling_task
119
+
120
+ parsed_prompts = [
121
+ (
122
+ prompt
123
+ if isinstance(prompt, bytes)
124
+ else parse_model_prompt(self.model_config, prompt)
125
+ )
126
+ for prompt in prompts_seq
127
+ ]
128
tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
129
**(ctx.tokenization_kwargs or {})
130
)
110
- return self._preprocess_cmpl_offline(prompts=prompts_seq, tok_params=tok_params)
131
+
132
+ pooling_params: PoolingParams | Sequence[PoolingParams]
133
+ if ctx.pooling_params is None:
134
+ pooling_params = PoolingParams()
135
+ else:
136
+ pooling_params = ctx.pooling_params
137
+
138
+ params_seq = self._params_to_seq(pooling_params, num_requests)
139
+
140
+ for param in params_seq:
141
+ if param.task is None:
142
+ param.task = pooling_task
143
+ elif pooling_task == "plugin":
144
+ # `plugin` task uses io_processor.parse_request to verify inputs.
145
+ # We actually allow plugin to overwrite pooling_task.
146
+ pass
147
+ elif param.task != pooling_task:
148
+ msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
149
+ raise ValueError(msg)
150
+
151
+ seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
152
+ seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
153
+
154
+ def request_factory() -> RequestGenerator:
155
+ for i in range(num_requests):
156
+ yield EncodeCMPLRenderParams(
157
+ prompts=parsed_prompts[i],
158
+ tok_params=tok_params,
159
+ prompt_extras=None,
160
+ skip_mm_cache=False,
161
+ params=params_seq[i],
162
+ lora_requests=seq_lora_requests[i],
163
+ priorities=seq_priority[i],
164
+ )
165
+
166
+ return request_factory, num_requests
167
168
def post_process_offline(
169
self,
@@ -118,6 +174,41 @@ class PoolingIOProcessor:
174
#######################################
175
# helpers
176
177
+ def render(
178
+ self,
179
+ render_params: EncodeCMPLRenderParams
180
+ | EncodeChatRenderParams
181
+ | ScoringRenderParams,
182
+ ) -> PoolingEngineInput:
183
+ if "conversations" in render_params:
184
+ render_params = cast(EncodeChatRenderParams, render_params)
185
+ (_,), engine_input = self.renderer.render_chat(
186
+ conversations=[render_params["conversations"]],
187
+ chat_params=render_params["chat_params"],
188
+ tok_params=render_params["tok_params"],
189
+ prompt_extras=render_params["prompt_extras"],
190
+ skip_mm_cache=render_params["skip_mm_cache"],
191
+ )
192
+ elif "prompts" in render_params:
193
+ render_params = cast(EncodeCMPLRenderParams, render_params)
194
+ engine_input = self.renderer.render_cmpl(
195
+ prompts=[render_params["prompts"]],
196
+ tok_params=render_params["tok_params"],
197
+ prompt_extras=render_params["prompt_extras"],
198
+ skip_mm_cache=render_params["skip_mm_cache"],
199
+ )
200
+ else:
201
+ raise ValueError(
202
+ f"Unsupported render_params type {render_params.__class__.__name__}"
203
+ )
204
+
205
+ return PoolingEngineInput(
206
+ prompts=engine_input[0],
207
+ params=render_params["params"],
208
+ lora_requests=render_params["lora_requests"],
209
+ priorities=render_params["priorities"],
210
+ )
211
+
212
def _preprocess_cmpl_online(
213
self,
214
request: RendererRequest,
@@ -196,26 +287,6 @@ class PoolingIOProcessor:
287
288
return conversation, [engine_input]
289
199
- def _preprocess_cmpl_offline(
200
- self,
201
- prompts: PromptType | Sequence[PromptType],
202
- tok_params: TokenizeParams,
203
- prompt_extras: dict[str, Any] | None = None,
204
- ) -> Sequence[EngineInput]:
205
- prompts = prompt_to_seq(prompts)
206
- parsed_prompts = [
207
- (
208
- prompt
209
- if isinstance(prompt, bytes)
210
- else parse_model_prompt(self.model_config, prompt)
211
- )
212
- for prompt in prompts
213
- ]
214
-
215
- return self.renderer.render_cmpl(
216
- parsed_prompts, tok_params, prompt_extras=prompt_extras
217
- )
218
-
290
def _validate_chat_template(
291
self,
292
request_chat_template: str | None,
@@ -251,3 +322,35 @@ class PoolingIOProcessor:
322
return params
323
324
return [params] * num_requests
325
+
326
+ def _lora_request_to_seq(
327
+ self,
328
+ lora_request: LoRARequest | None | Sequence[LoRARequest | None],
329
+ num_requests: int,
330
+ ) -> Sequence[LoRARequest | None]:
331
+ if isinstance(lora_request, Sequence):
332
+ if len(lora_request) != num_requests:
333
+ raise ValueError(
334
+ f"The lengths of prompts ({num_requests}) "
335
+ f"and lora_request ({len(lora_request)}) must be the same."
336
+ )
337
+
338
+ return lora_request
339
+
340
+ return [lora_request] * num_requests
341
+
342
+ def _priority_to_seq(
343
+ self,
344
+ priority: Sequence[int] | None,
345
+ num_requests: int,
346
+ ) -> Sequence[int]:
347
+ if priority is not None:
348
+ if len(priority) != num_requests:
349
+ raise ValueError(
350
+ f"The lengths of prompts ({num_requests}) "
351
+ f"and priority ({len(priority)}) must be the same."
352
+ )
353
+
354
+ return priority
355
+
356
+ return [0] * num_requests
vllm/entrypoints/pooling/embed/io_processor.py
+8
-3
@@ -28,11 +28,13 @@ 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,
32
ChunkedEmbeddingMetadata,
32
- OfflineInputsContext,
33
+ OfflineEncodeInputsContext,
34
PoolingChatLikeRequest,
35
PoolingCompletionLikeRequest,
36
PoolingServeContext,
37
+ RequestFactory,
38
)
39
from .protocol import (
40
CohereEmbedContent,
@@ -665,7 +667,10 @@ class JinaRankingTokenEmbedIOProcessor(
667
668
ctx.engine_inputs = engine_inputs
669
668
- def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
670
+ def get_request_factory_offline(
671
+ self, ctx: ALLOfflineInputsContext
672
+ ) -> tuple[RequestFactory, int]:
673
+ assert isinstance(ctx, OfflineEncodeInputsContext)
674
if not isinstance(ctx.prompts, Sequence) or len(ctx.prompts) < 2:
675
raise ValueError("The JinaForRanking model requires at least 2 inputs.")
676
@@ -677,4 +682,4 @@ class JinaRankingTokenEmbedIOProcessor(
682
query=text_prompts[-1], docs=text_prompts[:-1]
683
)
684
680
- return super().pre_process_offline(ctx)
685
+ return super().get_request_factory_offline(ctx)
vllm/entrypoints/pooling/offline.py
+134
-61
@@ -20,10 +20,18 @@ from vllm.outputs import (
20
from vllm.pooling_params import PoolingParams
21
from vllm.tasks import SCORE_TYPE_MAP, PoolingTask, SupportedTask
22
23
+from .base.io_processor import PoolingIOProcessor
24
from .factories import init_pooling_io_processors
25
from .scoring.io_processor import ScoringIOProcessor
26
from .scoring.typing import ScoreInput
26
-from .typing import OfflineInputsContext, OfflineOutputsContext
27
+from .typing import (
28
+ ALLOfflineInputsContext,
29
+ OfflineEncodeInputsContext,
30
+ OfflineOutputsContext,
31
+ OfflinePluginInputsContext,
32
+ OfflineScoringInputsContext,
33
+ RequestFactory,
34
+)
35
36
logger = init_logger(__name__)
37
@@ -48,6 +56,9 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
56
chat_template_config=self.chat_template_config,
57
)
58
59
+ # Use thread pool executor to accelerate preprocessing.
60
+ self._executor = self.renderer._executor
61
+
62
def encode(
63
self,
64
prompts: PromptType | Sequence[PromptType] | DataPrompt,
@@ -93,43 +104,30 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
104
105
io_processor = self.pooling_io_processors[pooling_task]
106
96
- if pooling_params is None:
97
- pooling_params = PoolingParams()
98
-
99
- ctx = OfflineInputsContext(
100
- prompts=prompts,
101
- pooling_params=pooling_params,
102
- tokenization_kwargs=tokenization_kwargs,
103
- )
107
+ ctx: ALLOfflineInputsContext
108
+ if isinstance(prompts, dict) and "data" in prompts:
109
+ ctx = OfflinePluginInputsContext(
110
+ pooling_task=pooling_task,
111
+ prompts=prompts, # type: ignore[arg-type]
112
+ tokenization_kwargs=tokenization_kwargs,
113
+ pooling_params=pooling_params,
114
+ lora_request=lora_request,
115
+ priorities=None,
116
+ )
117
+ else:
118
+ ctx = OfflineEncodeInputsContext(
119
+ pooling_task=pooling_task,
120
+ prompts=prompts,
121
+ tokenization_kwargs=tokenization_kwargs,
122
+ pooling_params=pooling_params,
123
+ lora_request=lora_request,
124
+ priorities=None,
125
+ )
126
105
- engine_inputs = io_processor.pre_process_offline(ctx)
106
- n_inputs = len(engine_inputs)
107
- assert ctx.pooling_params is not None
108
-
109
- params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)
110
-
111
- for param in params_seq:
112
- if param.task is None:
113
- param.task = pooling_task
114
- elif pooling_task == "plugin":
115
- # `plugin` task uses io_processor.parse_request to verify inputs.
116
- # We actually allow plugin to overwrite pooling_task.
117
- pass
118
- elif param.task != pooling_task:
119
- msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
120
- raise ValueError(msg)
121
-
122
- seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)
123
- seq_priority = self._priority_to_seq(None, n_inputs)
124
-
125
- self._render_and_add_requests(
126
- prompts=engine_inputs,
127
- params=params_seq,
128
- lora_requests=seq_lora_requests,
129
- priorities=seq_priority,
127
+ request_factory, num_requests = io_processor.get_request_factory_offline(ctx)
128
+ outputs = self._run_tiling_engine(
129
+ io_processor, request_factory, num_requests, use_tqdm=use_tqdm
130
)
131
-
132
- outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
131
outputs = io_processor.post_process_offline(
132
ctx=OfflineOutputsContext(outputs=outputs)
133
)
@@ -357,46 +355,121 @@ class PoolingOfflineMixin(OfflineInferenceMixin):
355
io_processor = self.pooling_io_processors[score_type]
356
assert isinstance(io_processor, ScoringIOProcessor)
357
360
- pooling_task = io_processor.pooling_task
358
scoring_data = io_processor.valid_inputs(data_1, data_2)
359
n_queries = len(scoring_data.data_1)
360
361
if pooling_params is None:
362
pooling_params = PoolingParams()
363
367
- ctx = OfflineInputsContext(
368
- prompts=scoring_data,
364
+ assert isinstance(pooling_params, PoolingParams)
365
+ pooling_task = io_processor.pooling_task
366
+ if pooling_params.task is None:
367
+ pooling_params.task = pooling_task
368
+ elif pooling_params.task != pooling_task:
369
+ msg = (
370
+ f"You cannot overwrite {pooling_params.task=!r} with {pooling_task=!r}!"
371
+ )
372
+ raise ValueError(msg)
373
+
374
+ ctx = OfflineScoringInputsContext(
375
+ pooling_task=pooling_task,
376
+ scoring_data=scoring_data,
377
pooling_params=pooling_params,
378
tokenization_kwargs=tokenization_kwargs,
379
+ lora_request=lora_request,
380
chat_template=chat_template,
372
- n_queries=n_queries,
381
+ priorities=None,
382
)
383
375
- engine_inputs = io_processor.pre_process_offline(ctx)
376
- n_inputs = len(engine_inputs)
377
-
378
- seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)
379
- params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)
380
-
381
- for param in params_seq:
382
- if param.task is None:
383
- param.task = pooling_task
384
- elif param.task != pooling_task:
385
- msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
386
- raise ValueError(msg)
387
-
388
- seq_priority = self._priority_to_seq(None, n_inputs)
389
-
390
- self._render_and_add_requests(
391
- prompts=engine_inputs,
392
- params=params_seq,
393
- lora_requests=seq_lora_requests,
394
- priorities=seq_priority,
384
+ request_factory, num_requests = io_processor.get_request_factory_offline(ctx)
385
+ outputs = self._run_tiling_engine(
386
+ io_processor, request_factory, num_requests, use_tqdm=use_tqdm
387
)
388
397
- outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
389
outputs = io_processor.post_process_offline(
390
ctx=OfflineOutputsContext(outputs=outputs, n_queries=n_queries),
391
)
392
393
return [ScoringRequestOutput.from_base(item) for item in outputs]
394
+
395
+ def _run_tiling_engine(
396
+ self,
397
+ io_processor: PoolingIOProcessor,
398
+ request_factory: RequestFactory,
399
+ num_requests: int,
400
+ use_tqdm: bool | Callable[..., tqdm] = True,
401
+ ):
402
+ # Keeping max_num_seqs * 2 requests in the core can already saturate the core.
403
+ # Therefore, keep most requests waiting outside the core.
404
+ max_requests_in_core = (
405
+ self.llm_engine.vllm_config.scheduler_config.max_num_seqs * 2
406
+ )
407
+ num_requests_in_core = 0
408
+ num_waited_requests = num_requests
409
+
410
+ if use_tqdm:
411
+ tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
412
+ pbar = tqdm_func(
413
+ total=num_requests,
414
+ desc="Processed prompts",
415
+ dynamic_ncols=True,
416
+ postfix=f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s",
417
+ )
418
+
419
+ outputs: list[PoolingRequestOutput] = []
420
+ added_request_ids: set[str] = set()
421
+
422
+ it = self._executor.map(io_processor.render, request_factory())
423
+
424
+ try:
425
+ while num_waited_requests or self.llm_engine.has_unfinished_requests():
426
+ requests = []
427
+ for _ in range(max_requests_in_core - num_requests_in_core):
428
+ if num_waited_requests == 0:
429
+ break
430
+ try:
431
+ request = next(it)
432
+ requests.append(request)
433
+ except StopIteration:
434
+ num_waited_requests = 0
435
+ break
436
+
437
+ num_waited_requests -= 1
438
+ num_requests_in_core += 1
439
+
440
+ if requests:
441
+ request_ids = self._render_and_add_requests(
442
+ prompts=[x["prompts"] for x in requests],
443
+ params=[x["params"] for x in requests],
444
+ lora_requests=[x["lora_requests"] for x in requests],
445
+ priorities=[x["priorities"] for x in requests],
446
+ )
447
+
448
+ for request_id in request_ids:
449
+ # undo assign_request_id
450
+ request_id = request_id.split("-", 1)[0]
451
+ added_request_ids.add(request_id)
452
+
453
+ step_outputs = self.llm_engine.step()
454
+ for output in step_outputs:
455
+ assert isinstance(output, PoolingRequestOutput)
456
+ assert output.finished
457
+ outputs.append(output)
458
+ added_request_ids.discard(output.request_id)
459
+ num_requests_in_core -= 1
460
+
461
+ if use_tqdm:
462
+ pbar.update(1)
463
+ if pbar.n == num_requests:
464
+ pbar.refresh()
465
+
466
+ except Exception:
467
+ if added_request_ids:
468
+ self.llm_engine.abort_request(list(added_request_ids))
469
+ raise
470
+
471
+ finally:
472
+ if use_tqdm:
473
+ pbar.close()
474
+
475
+ return sorted(outputs, key=lambda x: int(x.request_id))
vllm/entrypoints/pooling/pooling/io_processor.py
+32
-9
@@ -4,13 +4,19 @@ from collections.abc import Sequence
4
from typing import Any
5
6
from vllm import PoolingParams, PoolingRequestOutput
7
-from vllm.inputs import EngineInput
7
from vllm.logger import init_logger
8
from vllm.plugins.io_processors import get_io_processor
9
from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
10
11
from ..base.io_processor import PoolingIOProcessor
13
-from ..typing import OfflineInputsContext, OfflineOutputsContext, PoolingServeContext
12
+from ..typing import (
13
+ ALLOfflineInputsContext,
14
+ OfflineEncodeInputsContext,
15
+ OfflineOutputsContext,
16
+ OfflinePluginInputsContext,
17
+ PoolingServeContext,
18
+ RequestFactory,
19
+)
20
from .protocol import IOProcessorRequest, IOProcessorResponse
21
22
logger = init_logger(__name__)
@@ -107,9 +113,11 @@ class PluginWithIOProcessorPlugins(PoolingIOProcessor):
113
#######################################
114
# offline APIs
115
110
- def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
116
+ def get_request_factory_offline(
117
+ self, ctx: ALLOfflineInputsContext
118
+ ) -> tuple[RequestFactory, int]:
119
+ assert isinstance(ctx, OfflinePluginInputsContext)
120
assert isinstance(ctx.prompts, dict) and "data" in ctx.prompts
112
- assert ctx.pooling_params is not None
121
122
# Validate the request data is valid for the loaded plugin
123
prompt_data = ctx.prompts.get("data")
@@ -126,20 +134,35 @@ class PluginWithIOProcessorPlugins(PoolingIOProcessor):
134
prompts = self.io_processor.pre_process(prompt=validated_prompt)
135
prompts_seq = prompt_to_seq(prompts)
136
137
+ num_requests = len(prompts_seq)
138
+
139
+ pooling_params: PoolingParams | Sequence[PoolingParams]
140
+ if ctx.pooling_params is None:
141
+ pooling_params = PoolingParams()
142
+ else:
143
+ pooling_params = ctx.pooling_params
144
+
145
params_seq: list[PoolingParams] = [
146
self.io_processor.merge_pooling_params(param)
147
for param in self._params_to_seq(
132
- ctx.pooling_params,
133
- len(prompts_seq),
148
+ pooling_params,
149
+ num_requests,
150
)
151
]
152
for p in params_seq:
153
if p.task is None:
154
p.task = "plugin"
155
140
- ctx.pooling_params = params_seq
141
- ctx.prompts = prompts_seq
142
- return super().pre_process_offline(ctx)
156
+ return super().get_request_factory_offline(
157
+ OfflineEncodeInputsContext(
158
+ prompts=prompts_seq,
159
+ pooling_params=params_seq,
160
+ pooling_task="plugin",
161
+ tokenization_kwargs=ctx.tokenization_kwargs,
162
+ lora_request=ctx.lora_request,
163
+ priorities=ctx.priorities,
164
+ )
165
+ )
166
167
def post_process_offline(
168
self,
vllm/entrypoints/pooling/scoring/io_processor.py
+177
-28
@@ -2,24 +2,36 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
from collections.abc import Sequence
5
-from typing import Any, TypeAlias
5
+from typing import Any, TypeAlias, cast
6
7
import torch.nn.functional as F
8
9
-from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt
9
+from vllm import PoolingParams, PoolingRequestOutput, PromptType, TokensPrompt
10
from vllm.inputs import EngineInput
11
from vllm.renderers import TokenizeParams
12
from vllm.renderers.hf import safe_apply_chat_template
13
-from vllm.renderers.inputs.preprocess import extract_target_prompt
13
+from vllm.renderers.inputs.preprocess import (
14
+ extract_target_prompt,
15
+ parse_model_prompt,
16
+ prompt_to_seq,
17
+)
18
from vllm.tasks import PoolingTask
19
from vllm.utils.mistral import is_mistral_tokenizer
20
21
from ...chat_utils import ChatTemplateResolutionError
22
from ..base.io_processor import PoolingIOProcessor
23
from ..typing import (
20
- OfflineInputsContext,
24
+ ALLOfflineInputsContext,
25
+ EncodeChatRenderParams,
26
+ EncodeCMPLRenderParams,
27
+ OfflineEncodeInputsContext,
28
OfflineOutputsContext,
29
+ OfflineScoringInputsContext,
30
+ PoolingEngineInput,
31
PoolingServeContext,
32
+ RequestFactory,
33
+ RequestGenerator,
34
+ ScoringRenderParams,
35
)
36
from .protocol import RerankRequest, ScoreRequest, ScoringRequest
37
from .typing import ScoreData, ScoreInput, ScoringData
@@ -213,25 +225,37 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
225
#######################################
226
# offline APIs
227
216
- def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
217
- assert isinstance(ctx.prompts, ScoringData)
218
- assert not isinstance(ctx.pooling_params, Sequence)
219
-
220
- tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
221
- **(ctx.tokenization_kwargs or {})
222
- )
228
+ def get_request_factory_offline(
229
+ self, ctx: ALLOfflineInputsContext
230
+ ) -> tuple[RequestFactory, int]:
231
+ assert isinstance(ctx, OfflineScoringInputsContext)
232
233
max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
234
pooling_params=ctx.pooling_params
235
)
236
228
- scoring_data = ctx.prompts
237
+ scoring_data = ctx.scoring_data
238
if max_tokens_per_query > 0 or max_tokens_per_doc > 0:
239
scoring_data = self._truncate_scoring_data(
240
scoring_data, max_tokens_per_query, max_tokens_per_doc
241
)
242
234
- return self._pre_process(scoring_data, tok_params)
243
+ data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config)
244
+ data_2 = score_data_to_prompts(
245
+ scoring_data.data_2, "document", self.model_config
246
+ )
247
+ prompts = data_1 + data_2
248
+
249
+ return super().get_request_factory_offline(
250
+ OfflineEncodeInputsContext(
251
+ pooling_task=self.pooling_task,
252
+ prompts=prompts,
253
+ tokenization_kwargs=ctx.tokenization_kwargs,
254
+ pooling_params=ctx.pooling_params,
255
+ lora_request=ctx.lora_request,
256
+ priorities=ctx.priorities,
257
+ )
258
+ )
259
260
def post_process_offline(
261
self,
@@ -258,6 +282,26 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
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
+
305
def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
306
emb_data_1 = outputs[:n_queries]
307
emb_data_2 = outputs[n_queries:]
@@ -429,33 +473,95 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
473
#######################################
474
# offline APIs
475
432
- def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
433
- assert isinstance(ctx.prompts, ScoringData)
434
- assert not isinstance(ctx.pooling_params, Sequence)
476
+ def get_request_factory_offline(
477
+ self, ctx: ALLOfflineInputsContext
478
+ ) -> tuple[RequestFactory, int]:
479
+ assert isinstance(ctx, OfflineScoringInputsContext)
480
+
481
+ data_1 = ctx.scoring_data.data_1
482
+ data_2 = ctx.scoring_data.data_2
483
+ num_requests = len(data_2)
484
+
485
+ if len(data_1) == 1:
486
+ data_1 = data_1 * num_requests
487
488
tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
489
**(ctx.tokenization_kwargs or {})
490
)
491
492
+ prompt_extras = ctx.pooling_params.extra_kwargs
493
+
494
+ seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests)
495
+ seq_priority = self._priority_to_seq(ctx.priorities, num_requests)
496
+
497
+ def request_factory() -> RequestGenerator:
498
+ for i in range(num_requests):
499
+ yield ScoringRenderParams(
500
+ data_1=data_1[i],
501
+ data_2=data_2[i],
502
+ chat_template=ctx.chat_template,
503
+ tok_params=tok_params,
504
+ prompt_extras=prompt_extras,
505
+ skip_mm_cache=False,
506
+ params=ctx.pooling_params,
507
+ lora_requests=seq_lora_requests[i],
508
+ priorities=seq_priority[i],
509
+ )
510
+
511
+ return request_factory, num_requests
512
+
513
+ #######################################
514
+ # helpers
515
+
516
+ def render(
517
+ self,
518
+ render_params: EncodeCMPLRenderParams
519
+ | EncodeChatRenderParams
520
+ | ScoringRenderParams,
521
+ ) -> PoolingEngineInput:
522
+ if "data_1" not in render_params:
523
+ raise ValueError(
524
+ f"Unsupported render_params type {render_params.__class__.__name__}"
525
+ )
526
+ render_params = cast(ScoringRenderParams, render_params)
527
+
528
+ arrival_time = time.time()
529
+
530
+ tok_params = render_params["tok_params"]
531
+ params = render_params["params"]
532
+ prompt_extras = render_params["prompt_extras"]
533
+
534
max_tokens_per_query, max_tokens_per_doc = self._get_token_limits(
441
- pooling_params=ctx.pooling_params
535
+ pooling_params=params
536
)
537
444
- prompt_extras = ctx.pooling_params.extra_kwargs if ctx.pooling_params else None
445
- engine_inputs, pooling_params_list = self._pre_process(
446
- ctx.prompts,
447
- tok_params,
448
- ctx.pooling_params,
449
- ctx.chat_template,
538
+ _, engine_prompt = self.get_score_prompt(
539
+ data_1=render_params["data_1"],
540
+ data_2=render_params["data_2"],
541
+ encode_kwargs=tok_params.get_encode_kwargs(),
542
+ chat_template=render_params["chat_template"],
543
max_tokens_per_query=max_tokens_per_query,
544
max_tokens_per_doc=max_tokens_per_doc,
452
- prompt_extras=prompt_extras,
545
+ chat_template_kwargs=prompt_extras.get("chat_template_kwargs")
546
+ if prompt_extras
547
+ else None,
548
)
454
- ctx.pooling_params = pooling_params_list
455
- return engine_inputs
549
457
- #######################################
458
- # helpers
550
+ tok_params.apply_post_tokenization(self.tokenizer, engine_prompt)
551
+
552
+ if token_type_ids := engine_prompt.pop("token_type_ids", None):
553
+ 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
+ )
565
566
def _pre_process(
567
self,
@@ -731,6 +837,49 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
837
name = "jina-reranking-scoring"
838
pooling_task: PoolingTask = "token_embed"
839
840
+ def get_request_factory_offline(
841
+ self, ctx: ALLOfflineInputsContext
842
+ ) -> tuple[RequestFactory, int]:
843
+ assert isinstance(ctx, OfflineScoringInputsContext)
844
+
845
+ scoring_data = ctx.scoring_data
846
+ prompt_extras = ctx.pooling_params.extra_kwargs
847
+
848
+ queries = self.ensure_str(scoring_data.data_1)
849
+ docs = self.ensure_str(scoring_data.data_2)
850
+ chat_template_kwargs = (
851
+ prompt_extras.get("chat_template_kwargs") if prompt_extras else None
852
+ )
853
+ instruction = (
854
+ chat_template_kwargs.get("instruction") if chat_template_kwargs else None
855
+ )
856
+
857
+ if len(queries) == 1:
858
+ prompts = [
859
+ self.format_docs_prompts_func(
860
+ query=queries[0], docs=docs, instruction=instruction
861
+ )
862
+ ]
863
+ else:
864
+ prompts = [
865
+ self.format_docs_prompts_func(
866
+ query=q, docs=[d], instruction=instruction
867
+ )
868
+ for q, d in zip(queries, docs)
869
+ ]
870
+
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
+ )
882
+
883
def _pre_process(
884
self,
885
scoring_data: ScoringData,
vllm/entrypoints/pooling/typing.py
+74
-9
@@ -1,17 +1,21 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
-from collections.abc import AsyncGenerator, Sequence
4
+from collections.abc import AsyncGenerator, Callable, Generator, Sequence
5
from dataclasses import dataclass, field
6
-from typing import Any, Generic, TypeAlias, TypeVar
6
+from typing import Any, Generic, TypeAlias, TypedDict, TypeVar
7
8
from fastapi import Request
9
from pydantic import ConfigDict
10
11
from vllm import PoolingParams, PoolingRequestOutput, PromptType
12
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
13
from vllm.inputs import DataPrompt, EngineInput
14
from vllm.lora.request import LoRARequest
15
+from vllm.renderers import ChatParams, TokenizeParams
16
+from vllm.renderers.inputs import DictPrompt
17
18
+from ...tasks import PoolingTask
19
from .classify.protocol import (
20
ClassificationChatRequest,
21
ClassificationCompletionRequest,
@@ -34,7 +38,7 @@ from .pooling.protocol import (
38
PoolingResponse,
39
)
40
from .scoring.protocol import ScoringRequest, ScoringResponse
37
-from .scoring.typing import ScoringData
41
+from .scoring.typing import ScoreData, ScoringData
42
43
PoolingCompletionLikeRequest: TypeAlias = (
44
EmbeddingCompletionRequest
@@ -111,13 +115,36 @@ class PoolingServeContext(Generic[PoolingRequestT]):
115
116
@dataclass
117
class OfflineInputsContext:
114
- prompts: PromptType | Sequence[PromptType] | DataPrompt | ScoringData
115
- pooling_params: PoolingParams | Sequence[PoolingParams]
116
- tokenization_kwargs: dict[str, Any] | None = None
117
- chat_template: str | None = None
118
+ pooling_task: PoolingTask
119
+ tokenization_kwargs: dict[str, Any] | None
120
+ lora_request: Sequence[LoRARequest | None] | None
121
+ priorities: Sequence[int] | None
122
119
- ## for bi-encoder & late-interaction
120
- n_queries: int | None = None
123
+
124
+@dataclass
125
+class OfflineEncodeInputsContext(OfflineInputsContext):
126
+ prompts: PromptType | Sequence[PromptType]
127
+ pooling_params: PoolingParams | Sequence[PoolingParams] | None
128
+
129
+
130
+@dataclass
131
+class OfflineScoringInputsContext(OfflineInputsContext):
132
+ scoring_data: ScoringData
133
+ chat_template: str | None
134
+ pooling_params: PoolingParams
135
+
136
+
137
+@dataclass
138
+class OfflinePluginInputsContext(OfflineInputsContext):
139
+ prompts: DataPrompt
140
+ pooling_params: PoolingParams | Sequence[PoolingParams] | None
141
+
142
+
143
+ALLOfflineInputsContext: TypeAlias = (
144
+ OfflineEncodeInputsContext
145
+ | OfflineScoringInputsContext
146
+ | OfflinePluginInputsContext
147
+)
148
149
150
@dataclass
@@ -126,3 +153,41 @@ class OfflineOutputsContext:
153
154
## for bi-encoder & late-interaction
155
n_queries: int | None = None
156
+
157
+
158
+class RenderParams(TypedDict):
159
+ tok_params: TokenizeParams
160
+ prompt_extras: dict[str, Any] | None
161
+ skip_mm_cache: bool
162
+
163
+ params: PoolingParams
164
+ lora_requests: LoRARequest | None
165
+ priorities: int
166
+
167
+
168
+class EncodeCMPLRenderParams(RenderParams):
169
+ prompts: DictPrompt
170
+
171
+
172
+class EncodeChatRenderParams(RenderParams):
173
+ conversations: list["ChatCompletionMessageParam"]
174
+ chat_params: ChatParams
175
+
176
+
177
+class ScoringRenderParams(RenderParams):
178
+ data_1: ScoreData
179
+ data_2: ScoreData
180
+ chat_template: str | None
181
+
182
+
183
+class PoolingEngineInput(TypedDict):
184
+ prompts: EngineInput
185
+ params: PoolingParams
186
+ lora_requests: LoRARequest | None
187
+ priorities: int
188
+
189
+
190
+RequestGenerator: TypeAlias = Generator[
191
+ EncodeCMPLRenderParams | EncodeChatRenderParams | ScoringRenderParams
192
+]
193
+RequestFactory: TypeAlias = Callable[[], RequestGenerator]