A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 04:45 UTC
[Frontend] Offload blocking preprocessing & postprocessing ops to thread pool for pooling entrypoints. (#39763)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
wang.yuqi committed
Apr 14, 2026 at 16:29 UTC
c0ecaed95030815c49699995bf6461209ccd6a0f
10 files changed
+68
-83
tests/entrypoints/pooling/embed/test_io_processor.py
+7
-5
@@ -4,6 +4,7 @@
4
5
import pytest
6
7
+from vllm import PoolingParams
8
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
9
from vllm.entrypoints.pooling.embed.protocol import (
10
CohereEmbedContent,
@@ -218,6 +219,7 @@ class TestPreProcessCohereOnline:
219
def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
220
return PoolingServeContext(
221
request=CohereEmbedRequest(model="test", **request_kwargs),
222
+ pooling_params=PoolingParams(),
223
model_name="test",
224
request_id="embd-test",
225
)
@@ -233,13 +235,13 @@ class TestPreProcessCohereOnline:
235
ctx = self._make_context(texts=["hello"])
236
calls: list[tuple[str, object]] = []
237
236
- def preprocess_completion(request, prompt_input, prompt_embeds):
238
+ def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
239
calls.append(("completion", prompt_input))
240
return ["completion"]
241
242
handler._get_task_instruction_prefix = lambda _input_type: None
243
handler._has_chat_template = lambda: False
242
- handler._preprocess_completion_online = preprocess_completion
244
+ handler._preprocess_cmpl_online = preprocess_cmpl_online
245
handler._batch_render_chat = lambda *_args, **_kwargs: (
246
pytest.fail("text-only request should not require chat rendering")
247
)
@@ -254,7 +256,7 @@ class TestPreProcessCohereOnline:
256
ctx = self._make_context(texts=["hello"], input_type="query")
257
calls: list[tuple[str, object]] = []
258
257
- def preprocess_completion(request, prompt_input, prompt_embeds):
259
+ def preprocess_cmpl(request, prompt_input, prompt_embeds):
260
calls.append(("completion", prompt_input))
261
return ["fallback"]
262
@@ -263,7 +265,7 @@ class TestPreProcessCohereOnline:
265
handler._batch_render_chat = lambda *_args, **_kwargs: (
266
pytest.fail("chat rendering should be skipped without a template")
267
)
266
- handler._preprocess_completion_online = preprocess_completion
268
+ handler._preprocess_cmpl_online = preprocess_cmpl
269
270
handler._pre_process_cohere_online(ctx)
271
@@ -297,7 +299,7 @@ class TestPreProcessCohereOnline:
299
handler._get_task_instruction_prefix = lambda _input_type: "query: "
300
handler._has_chat_template = lambda: True
301
handler._batch_render_chat = batch_render_chat
300
- handler._preprocess_completion_online = lambda *_args, **_kwargs: (
302
+ handler._preprocess_cmpl_online = lambda *_args, **_kwargs: (
303
pytest.fail("completion path should be skipped when a template exists")
304
)
305
vllm/entrypoints/pooling/base/io_processor.py
+4
-24
@@ -72,7 +72,7 @@ class PoolingIOProcessor:
72
default_template_kwargs=None,
73
)
74
elif isinstance(request, PoolingCompletionLikeRequest):
75
- engine_inputs = self._preprocess_completion_online(
75
+ engine_inputs = self._preprocess_cmpl_online(
76
request,
77
prompt_input=request.input,
78
prompt_embeds=None,
@@ -82,21 +82,12 @@ class PoolingIOProcessor:
82
83
ctx.engine_inputs = engine_inputs
84
85
- async def pre_process_online_async(self, ctx: PoolingServeContext):
86
- self.pre_process_online(ctx)
87
-
85
def post_process_online(
86
self,
87
ctx: PoolingServeContext,
88
):
89
pass
90
94
- async def post_process_online_async(
95
- self,
96
- ctx: PoolingServeContext,
97
- ):
98
- self.post_process_online(ctx)
99
-
91
#######################################
92
# offline APIs
93
@@ -109,12 +100,7 @@ class PoolingIOProcessor:
100
tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
101
**(ctx.tokenization_kwargs or {})
102
)
112
- return self._preprocess_completion_offline(
113
- prompts=prompts_seq, tok_params=tok_params
114
- )
115
-
116
- async def pre_process_offline_async(self, ctx: OfflineInputsContext):
117
- return self.pre_process_offline(ctx)
103
+ return self._preprocess_cmpl_offline(prompts=prompts_seq, tok_params=tok_params)
104
105
def post_process_offline(
106
self,
@@ -122,16 +108,10 @@ class PoolingIOProcessor:
108
) -> list[PoolingRequestOutput]:
109
return ctx.outputs
110
125
- async def post_process_offline_async(
126
- self,
127
- ctx: OfflineOutputsContext,
128
- ) -> list[PoolingRequestOutput]:
129
- return self.post_process_offline(ctx)
130
-
111
#######################################
112
# helpers
113
134
- def _preprocess_completion_online(
114
+ def _preprocess_cmpl_online(
115
self,
116
request: RendererRequest,
117
prompt_input: str | list[str] | list[int] | list[list[int]] | None,
@@ -209,7 +189,7 @@ class PoolingIOProcessor:
189
190
return conversation, [engine_input]
191
212
- def _preprocess_completion_offline(
192
+ def _preprocess_cmpl_offline(
193
self,
194
prompts: PromptType | Sequence[PromptType],
195
tok_params: TokenizeParams,
vllm/entrypoints/pooling/base/serving.py
+40
-18
@@ -3,9 +3,11 @@
3
4
from abc import ABC, abstractmethod
5
from collections.abc import AsyncGenerator, Mapping
6
+from concurrent.futures import Executor
7
from http import HTTPStatus
8
from typing import ClassVar
9
10
+import torch
11
from fastapi import Request
12
from fastapi.responses import Response
13
from starlette.datastructures import Headers
@@ -32,7 +34,7 @@ from vllm.tracing import (
34
log_tracing_disabled_warning,
35
)
36
from vllm.utils import random_uuid
35
-from vllm.utils.async_utils import merge_async_iterators
37
+from vllm.utils.async_utils import make_async, merge_async_iterators
38
39
from .io_processor import PoolingIOProcessor
40
@@ -67,16 +69,47 @@ class PoolingServingBase(ABC):
69
trust_request_chat_template=trust_request_chat_template,
70
)
71
70
- @abstractmethod
72
+ # Shared thread pool executor for preprocessing and postprocessing.
73
+ self._executor: Executor = models.renderer._executor
74
+ self._preprocessing_async = make_async(
75
+ self._preprocessing, executor=self._executor
76
+ )
77
+ self._postprocessing_async = make_async(
78
+ self._postprocessing, executor=self._executor
79
+ )
80
+
81
async def __call__(
82
self,
83
request: AnyPoolingRequest,
84
raw_request: Request | None = None,
85
) -> Response:
86
+ io_processor = self.get_io_processor(request)
87
+ ctx = await self._init_ctx(io_processor, request, raw_request)
88
+ await self._preprocessing_async(io_processor, ctx)
89
+ await self._prepare_generators(ctx)
90
+ await self._collect_batch(ctx)
91
+ return await self._postprocessing_async(io_processor, ctx)
92
+
93
+ @abstractmethod
94
+ def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor:
95
raise NotImplementedError
96
97
+ @torch.inference_mode()
98
+ def _preprocessing(
99
+ self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext
100
+ ):
101
+ return io_processor.pre_process_online(ctx)
102
+
103
+ @torch.inference_mode()
104
+ def _postprocessing(
105
+ self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext
106
+ ):
107
+ io_processor.post_process_online(ctx)
108
+ return self._build_response(ctx)
109
+
110
async def _init_ctx(
111
self,
112
+ io_processor: PoolingIOProcessor,
113
request: AnyPoolingRequest,
114
raw_request: Request | None = None,
115
):
@@ -84,10 +117,12 @@ class PoolingServingBase(ABC):
117
request_id = f"{self.request_id_prefix}-{self._base_request_id(raw_request)}"
118
await self._check_model(request)
119
120
+ pooling_params = io_processor.create_pooling_params(request)
121
ctx = PoolingServeContext(
122
request=request,
123
raw_request=raw_request,
124
model_name=model_name,
125
+ pooling_params=pooling_params,
126
request_id=request_id,
127
)
128
@@ -175,7 +210,7 @@ class PoolingServingBase(ABC):
210
ctx.final_res_batch = [res for res in final_res_batch if res is not None]
211
212
@abstractmethod
178
- async def _build_response(
213
+ def _build_response(
214
self,
215
ctx: PoolingServeContext,
216
) -> Response:
@@ -362,18 +397,5 @@ class PoolingServing(PoolingServingBase, ABC):
397
) -> PoolingIOProcessor:
398
raise NotImplementedError
399
365
- async def __call__(
366
- self,
367
- request: AnyPoolingRequest,
368
- raw_request: Request | None = None,
369
- ) -> Response:
370
- ctx = await self._init_ctx(request, raw_request)
371
- await self.io_processor.pre_process_online_async(ctx)
372
-
373
- if ctx.pooling_params is None:
374
- ctx.pooling_params = self.io_processor.create_pooling_params(request)
375
-
376
- await self._prepare_generators(ctx)
377
- await self._collect_batch(ctx)
378
- await self.io_processor.post_process_online_async(ctx)
379
- return await self._build_response(ctx)
400
+ def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor:
401
+ return self.io_processor
vllm/entrypoints/pooling/classify/serving.py
+1
-1
@@ -31,7 +31,7 @@ class ServingClassification(PoolingServing):
31
def init_io_processor(self, *args, **kwargs) -> ClassifyIOProcessor:
32
return ClassifyIOProcessor(*args, **kwargs)
33
34
- async def _build_response(
34
+ def _build_response(
35
self,
36
ctx: ClassificationServeContext,
37
) -> JSONResponse:
vllm/entrypoints/pooling/embed/io_processor.py
+2
-2
@@ -470,7 +470,7 @@ class EmbedIOProcessor(PoolingIOProcessor):
470
truncate_prompt_tokens=truncate_prompt_tokens,
471
truncation_side=truncation_side,
472
)
473
- return self._preprocess_completion_online(
473
+ return self._preprocess_cmpl_online(
474
proxy, prompt_input=proxy.input, prompt_embeds=None
475
)
476
@@ -579,7 +579,7 @@ class JinaRankingTokenEmbedIOProcessor(
579
query=text_prompts[-1], docs=text_prompts[:-1]
580
)
581
582
- engine_inputs = self._preprocess_completion_online(
582
+ engine_inputs = self._preprocess_cmpl_online(
583
request,
584
prompt_input=prompt_input,
585
prompt_embeds=None,
vllm/entrypoints/pooling/embed/serving.py
+3
-3
@@ -53,15 +53,15 @@ class ServingEmbedding(PoolingServing):
53
def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor:
54
return EmbedIOProcessor(*args, **kwargs)
55
56
- async def _build_response(
56
+ def _build_response(
57
self,
58
ctx: PoolingServeContext,
59
) -> Response:
60
if isinstance(ctx.request, CohereEmbedRequest):
61
return self._build_cohere_response_from_ctx(ctx)
62
- return await self._build_openai_response(ctx)
62
+ return self._build_openai_response(ctx)
63
64
- async def _build_openai_response(
64
+ def _build_openai_response(
65
self,
66
ctx: EmbeddingServeContext,
67
) -> JSONResponse | StreamingResponse:
vllm/entrypoints/pooling/pooling/serving.py
+4
-21
@@ -7,11 +7,11 @@ from collections.abc import Callable
7
from functools import partial
8
from typing import Literal, cast
9
10
-from fastapi import Request
10
from fastapi.responses import JSONResponse, Response, StreamingResponse
11
from typing_extensions import assert_never
12
13
from vllm.entrypoints.openai.engine.protocol import UsageInfo
14
+from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
15
from vllm.entrypoints.pooling.base.serving import PoolingServingBase
16
from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
17
from vllm.entrypoints.pooling.pooling.protocol import (
@@ -57,27 +57,10 @@ class ServingPooling(PoolingServingBase):
57
)
58
self.json_response_cls = get_json_response_cls()
59
60
- async def __call__(
61
- self,
62
- request: AnyPoolingRequest,
63
- raw_request: Request | None = None,
64
- ) -> Response:
60
+ def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor:
61
assert isinstance(request, PoolingRequest)
62
pooling_task = self._verify_pooling_task(request)
67
-
68
- io_processor = self.io_processors[pooling_task]
69
- ctx = await self._init_ctx(request, raw_request)
70
-
71
- await io_processor.pre_process_online_async(ctx)
72
-
73
- if ctx.pooling_params is None:
74
- ctx.pooling_params = io_processor.create_pooling_params(request)
75
-
76
- await self._prepare_generators(ctx)
77
- await self._collect_batch(ctx)
78
-
79
- await io_processor.post_process_online_async(ctx)
80
- return await self._build_response(ctx)
63
+ return self.io_processors[pooling_task]
64
65
def _verify_pooling_task(self, request: PoolingRequest) -> str:
66
if getattr(request, "dimensions", None) is not None:
@@ -117,7 +100,7 @@ class ServingPooling(PoolingServingBase):
100
101
return pooling_task
102
120
- async def _build_response(
103
+ def _build_response(
104
self,
105
ctx: PoolingServeContext,
106
) -> Response:
vllm/entrypoints/pooling/scoring/io_processor.py
+2
-2
@@ -220,7 +220,7 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
220
scoring_data.data_2, "document", self.model_config
221
)
222
223
- return self._preprocess_completion_offline(
223
+ return self._preprocess_cmpl_offline(
224
prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras
225
)
226
@@ -682,7 +682,7 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
682
for q, d in zip(queries, docs)
683
]
684
685
- return self._preprocess_completion_offline(
685
+ return self._preprocess_cmpl_offline(
686
prompts=prompts, tok_params=tok_params, prompt_extras=prompt_extras
687
)
688
vllm/entrypoints/pooling/scoring/serving.py
+4
-6
@@ -65,7 +65,7 @@ class ServingScores(PoolingServing):
65
66
return await self.flash_late_interaction(*args, **kwargs)
67
68
- async def _build_response(
68
+ def _build_response(
69
self,
70
ctx: ScoringServeContext,
71
) -> JSONResponse:
@@ -183,17 +183,15 @@ class ServingScores(PoolingServing):
183
### Can significantly improve late-interaction scoring performance.
184
185
async def flash_late_interaction(self, *args, **kwargs) -> Response:
186
- ctx = await self._init_ctx(*args, **kwargs)
187
- ctx.pooling_params = self.io_processor.create_pooling_params(ctx.request)
188
- await self.io_processor.pre_process_online_async(ctx)
186
+ ctx = await self._init_ctx(self.io_processor, *args, **kwargs)
187
+ await self._preprocessing_async(self.io_processor, ctx)
188
189
# stage 1: encode queries and cache token embeddings on workers.
190
await self._flash_late_interaction_encode_queries(ctx)
191
# stage 2: encode docs and return scalar scores from workers.
192
await self._flash_late_interaction_encode_docs(ctx)
193
195
- await self.io_processor.post_process_online_async(ctx)
196
- return await self._build_response(ctx)
194
+ return await self._postprocessing_async(self.io_processor, ctx)
195
196
async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext):
197
assert ctx.n_queries is not None
vllm/entrypoints/pooling/typing.py
+1
-1
@@ -69,9 +69,9 @@ class PoolingServeContext(Generic[PoolingRequestT]):
69
raw_request: Request | None = None
70
model_name: str
71
request_id: str
72
+ pooling_params: PoolingParams | list[PoolingParams]
73
created_time: int = field(default_factory=lambda: int(time.time()))
74
lora_request: LoRARequest | None = None
74
- pooling_params: PoolingParams | list[PoolingParams] | None = None
75
engine_inputs: Sequence[EngineInput] | None = None
76
prompt_request_ids: list[str] | None = None
77
intermediates: Any | None = None