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, 16:00 UTC
[Frontend] Re-enable running MaxSim on GPU (#38620)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Apr 3, 2026 at 00:03 UTC
a9b4f07ba259f6fa95b61a6aa64ea80ea7681493
12 files changed
+207
-54
tests/entrypoints/pooling/scoring/test_late_interaction_online.py
+7
-2
@@ -26,13 +26,18 @@ TEXTS_2 = [
26
]
27
28
29
-@pytest.fixture(scope="module")
30
-def server():
29
+@pytest.fixture(scope="module", params=[True, False])
30
+def server(request):
31
args = [
32
"--max-model-len",
33
str(MAX_MODEL_LEN),
34
]
35
36
+ # Test run pooling score MaxSim on worker side (GPU)
37
+ # aka flash-late-interaction
38
+ if not request.param:
39
+ args += ["--no-enable-flash-late-interaction"]
40
+
41
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
42
yield remote_server
43
tests/v1/worker/test_late_interaction_runner.py
+1
-1
@@ -4,12 +4,12 @@
4
import pytest
5
import torch
6
7
+from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
8
from vllm.pooling_params import LateInteractionParams, PoolingParams
9
from vllm.v1.pool.late_interaction import (
10
LATE_INTERACTION_MODE_CACHE_QUERY,
11
build_late_interaction_doc_params,
12
build_late_interaction_query_params,
12
- compute_maxsim_score,
13
)
14
from vllm.v1.worker.gpu.pool.late_interaction_runner import LateInteractionRunner
15
vllm/entrypoints/llm.py
+3
-3
@@ -1449,14 +1449,14 @@ class LLM:
1449
1450
pooling_task = io_processor.pooling_task
1451
scoring_data = io_processor.valid_inputs(data_1, data_2)
1452
- offset = len(scoring_data.data_1)
1452
+ n_queries = len(scoring_data.data_1)
1453
1454
ctx = OfflineInputsContext(
1455
prompts=scoring_data,
1456
pooling_params=pooling_params,
1457
tokenization_kwargs=tokenization_kwargs,
1458
chat_template=chat_template,
1459
- offset=offset,
1459
+ n_queries=n_queries,
1460
)
1461
1462
processor_inputs = io_processor.pre_process_offline(ctx)
@@ -1487,7 +1487,7 @@ class LLM:
1487
1488
outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
1489
outputs = io_processor.post_process_offline(
1490
- ctx=OfflineOutputsContext(outputs=outputs, offset=offset),
1490
+ ctx=OfflineOutputsContext(outputs=outputs, n_queries=n_queries),
1491
)
1492
1493
return [ScoringRequestOutput.from_base(item) for item in outputs]
vllm/entrypoints/openai/cli_args.py
+3
@@ -278,6 +278,9 @@ class FrontendArgs(BaseFrontendArgs):
278
Enable offline FastAPI documentation for air-gapped environments.
279
Uses vendored static assets bundled with vLLM.
280
"""
281
+ enable_flash_late_interaction: bool = True
282
+ """If set, run pooling score MaxSim on GPU in the API server process.
283
+ Can significantly improve late-interaction scoring performance."""
284
285
@classmethod
286
def _customize_cli_kwargs(
vllm/entrypoints/pooling/__init__.py
+3
@@ -123,6 +123,9 @@ def init_pooling_state(
123
chat_template=resolved_chat_template,
124
chat_template_content_format=args.chat_template_content_format,
125
trust_request_chat_template=args.trust_request_chat_template,
126
+ enable_flash_late_interaction=getattr(
127
+ args, "enable_flash_late_interaction", True
128
+ ),
129
)
130
if enable_scoring_api(supported_tasks, model_config)
131
else None
vllm/entrypoints/pooling/base/serving.py
+13
-6
@@ -82,9 +82,20 @@ class PoolingServing:
82
request: AnyPoolingRequest,
83
raw_request: Request | None = None,
84
) -> Response:
85
+ ctx = await self._init_ctx(request, raw_request)
86
+ await self.io_processor.pre_process_online_async(ctx)
87
+ await self._prepare_generators(ctx)
88
+ await self._collect_batch(ctx)
89
+ await self.io_processor.post_process_online_async(ctx)
90
+ return await self._build_response(ctx)
91
+
92
+ async def _init_ctx(
93
+ self,
94
+ request: AnyPoolingRequest,
95
+ raw_request: Request | None = None,
96
+ ):
97
model_name = self.models.model_name()
98
request_id = f"{self.request_id_prefix}-{self._base_request_id(raw_request)}"
87
-
99
await self._check_model(request)
100
101
ctx = PoolingServeContext(
@@ -96,11 +107,7 @@ class PoolingServing:
107
108
self._validate_request(ctx)
109
self._maybe_get_adapters(ctx)
99
- await self.io_processor.pre_process_online_async(ctx)
100
- await self._prepare_generators(ctx)
101
- await self._collect_batch(ctx)
102
- await self.io_processor.post_process_online_async(ctx)
103
- return await self._build_response(ctx)
110
+ return ctx
111
112
async def _prepare_generators(
113
self,
vllm/entrypoints/pooling/scoring/io_processor.py
+33
-25
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import time
4
from collections.abc import Sequence
5
-from typing import Any, TypeAlias, cast
5
+from typing import Any, TypeAlias
6
7
import torch.nn.functional as F
8
@@ -16,7 +16,7 @@ from vllm.entrypoints.pooling.typing import (
16
from vllm.inputs import EngineInput
17
from vllm.renderers import TokenizeParams
18
from vllm.renderers.hf import safe_apply_chat_template
19
-from vllm.tasks import PoolingTask, ScoreType
19
+from vllm.tasks import PoolingTask
20
from vllm.utils.mistral import is_mistral_tokenizer
21
22
from ...chat_utils import ChatTemplateResolutionError
@@ -34,7 +34,7 @@ ScoringServeContext: TypeAlias = PoolingServeContext[ScoringRequest]
34
35
36
class ScoringIOProcessor(PoolingIOProcessor):
37
- name: ScoreType
37
+ name: str
38
pooling_task: PoolingTask
39
40
def __init__(self, *args, **kwargs):
@@ -63,7 +63,7 @@ class ScoringIOProcessor(PoolingIOProcessor):
63
64
65
class BiEncoderIOProcessor(ScoringIOProcessor):
66
- name: ScoreType = "bi-encoder"
66
+ name = "bi-encoder"
67
pooling_task: PoolingTask = "embed"
68
69
#######################################
@@ -94,20 +94,17 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
94
)
95
96
ctx.engine_inputs = engine_inputs
97
- ctx.intermediates = len(scoring_data.data_1)
97
+ ctx.n_queries = len(scoring_data.data_1)
98
99
def post_process_online(
100
self,
101
ctx: ScoringServeContext,
102
):
103
- if ctx.final_res_batch is None:
104
- raise ValueError("Final response batch not available")
105
-
106
- if ctx.intermediates is None:
107
- raise ValueError("data_1 len not available")
103
+ assert ctx.final_res_batch is not None
104
+ assert isinstance(ctx.n_queries, int)
105
106
ctx.final_res_batch = self._post_process(
110
- outputs=ctx.final_res_batch, offset=cast(int, ctx.intermediates)
107
+ outputs=ctx.final_res_batch, n_queries=ctx.n_queries
108
)
109
110
#######################################
@@ -124,8 +121,8 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
121
self,
122
ctx: OfflineOutputsContext,
123
) -> list[PoolingRequestOutput]:
127
- assert ctx.offset is not None
128
- return self._post_process(outputs=ctx.outputs, offset=ctx.offset)
124
+ assert ctx.n_queries is not None
125
+ return self._post_process(outputs=ctx.outputs, n_queries=ctx.n_queries)
126
127
#######################################
128
# helpers
@@ -145,9 +142,9 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
142
prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras
143
)
144
148
- def _post_process(self, outputs: list[PoolingRequestOutput], offset: int):
149
- emb_data_1 = outputs[:offset]
150
- emb_data_2 = outputs[offset:]
145
+ def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
146
+ emb_data_1 = outputs[:n_queries]
147
+ emb_data_2 = outputs[n_queries:]
148
149
if len(emb_data_1) == 1:
150
emb_data_1 = emb_data_1 * len(emb_data_2)
@@ -177,13 +174,13 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
174
175
176
class LateInteractionIOProcessor(BiEncoderIOProcessor):
180
- name: ScoreType = "late-interaction"
177
+ name = "late-interaction"
178
pooling_task: PoolingTask = "token_embed"
179
183
- def _post_process(self, outputs: list[PoolingRequestOutput], offset: int):
180
+ def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
181
# Split into query and document embeddings
185
- emb_data_1 = outputs[:offset]
186
- emb_data_2 = outputs[offset:]
182
+ emb_data_1 = outputs[:n_queries]
183
+ emb_data_2 = outputs[n_queries:]
184
185
# Expand queries if 1:N scoring
186
if len(emb_data_1) == 1:
@@ -217,8 +214,15 @@ class LateInteractionIOProcessor(BiEncoderIOProcessor):
214
return final_res_batch
215
216
217
+class FlashLateInteractionIOProcessor(LateInteractionIOProcessor):
218
+ name = "flash-late-interaction"
219
+
220
+ def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int):
221
+ return outputs
222
+
223
+
224
class CrossEncoderIOProcessor(ScoringIOProcessor):
221
- name: ScoreType = "cross-encoder"
225
+ name = "cross-encoder"
226
pooling_task: PoolingTask = "classify"
227
228
def __init__(self, *args, **kwargs):
@@ -412,8 +416,12 @@ class CrossEncoderIOProcessor(ScoringIOProcessor):
416
return full_prompt, engine_prompt
417
418
415
-ScoringIOProcessors: dict[ScoreType, type[ScoringIOProcessor]] = {
416
- "bi-encoder": BiEncoderIOProcessor,
417
- "late-interaction": LateInteractionIOProcessor,
418
- "cross-encoder": CrossEncoderIOProcessor,
419
+ScoringIOProcessors: dict[str, type[ScoringIOProcessor]] = {
420
+ p.name: p
421
+ for p in [
422
+ BiEncoderIOProcessor,
423
+ LateInteractionIOProcessor,
424
+ FlashLateInteractionIOProcessor,
425
+ CrossEncoderIOProcessor,
426
+ ]
427
}
vllm/entrypoints/pooling/scoring/serving.py
+134
-2
@@ -1,9 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from fastapi.responses import JSONResponse
4
+from fastapi.responses import JSONResponse, Response
5
6
+from vllm import PoolingParams
7
from vllm.config import ModelConfig
8
+from vllm.engine.protocol import EngineClient
9
from vllm.entrypoints.chat_utils import ChatTemplateConfig
10
from vllm.entrypoints.openai.engine.protocol import UsageInfo
11
from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
@@ -11,6 +13,10 @@ from vllm.entrypoints.pooling.base.serving import PoolingServing
13
from vllm.logger import init_logger
14
from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
15
from vllm.renderers import BaseRenderer
16
+from vllm.v1.pool.late_interaction import (
17
+ build_late_interaction_doc_params,
18
+ build_late_interaction_query_params,
19
+)
20
21
from .io_processor import ScoringIOProcessors, ScoringServeContext
22
from .protocol import (
@@ -31,13 +37,30 @@ logger = init_logger(__name__)
37
class ServingScores(PoolingServing):
38
request_id_prefix = "score"
39
40
+ def __init__(
41
+ self,
42
+ engine_client: EngineClient,
43
+ *args,
44
+ enable_flash_late_interaction: bool = True,
45
+ **kwargs,
46
+ ):
47
+ self.score_type = engine_client.model_config.score_type
48
+ self.enable_flash_late_interaction = (
49
+ self.score_type == "late-interaction" and enable_flash_late_interaction
50
+ )
51
+
52
+ super().__init__(engine_client, *args, **kwargs)
53
+
54
def init_io_processor(
55
self,
56
model_config: ModelConfig,
57
renderer: BaseRenderer,
58
chat_template_config: ChatTemplateConfig,
59
) -> PoolingIOProcessor:
40
- score_type = model_config.score_type
60
+ score_type: str = model_config.score_type
61
+ if self.enable_flash_late_interaction:
62
+ score_type = "flash-late-interaction"
63
+
64
assert score_type in ScoringIOProcessors
65
processor_cls = ScoringIOProcessors[score_type]
66
return processor_cls(
@@ -46,6 +69,12 @@ class ServingScores(PoolingServing):
69
chat_template_config=chat_template_config,
70
)
71
72
+ async def __call__(self, *args, **kwargs) -> Response:
73
+ if not self.enable_flash_late_interaction:
74
+ return await super().__call__(*args, **kwargs)
75
+
76
+ return await self.flash_late_interaction(*args, **kwargs)
77
+
78
async def _build_response(
79
self,
80
ctx: ScoringServeContext,
@@ -158,3 +187,106 @@ class ServingScores(PoolingServing):
187
)
188
189
return JSONResponse(content=response.model_dump())
190
+
191
+ ###################################################################################
192
+ ### Run pooling score MaxSim on worker side (GPU) in the API server process
193
+ ### Can significantly improve late-interaction scoring performance.
194
+
195
+ async def flash_late_interaction(self, *args, **kwargs) -> Response:
196
+ ctx = await self._init_ctx(*args, **kwargs)
197
+ ctx.pooling_params = self.io_processor.create_pooling_params(ctx.request)
198
+ await self.io_processor.pre_process_online_async(ctx)
199
+
200
+ # stage 1: encode queries and cache token embeddings on workers.
201
+ await self._flash_late_interaction_encode_queries(ctx)
202
+ # stage 2: encode docs and return scalar scores from workers.
203
+ await self._flash_late_interaction_encode_docs(ctx)
204
+
205
+ await self.io_processor.post_process_online_async(ctx)
206
+ return await self._build_response(ctx)
207
+
208
+ async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext):
209
+ assert ctx.n_queries is not None
210
+ assert ctx.engine_inputs is not None
211
+ assert isinstance(ctx.pooling_params, PoolingParams)
212
+
213
+ n_queries = ctx.n_queries
214
+ n_docs = len(ctx.engine_inputs) - n_queries
215
+ query_engine_inputs = ctx.engine_inputs[:n_queries]
216
+
217
+ query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
218
+ query_uses = [n_docs if n_queries == 1 else 1] * n_queries
219
+
220
+ query_pooling_params_list = []
221
+ for i in range(n_queries):
222
+ pooling_params = ctx.pooling_params.clone()
223
+ pooling_params.late_interaction_params = (
224
+ build_late_interaction_query_params(
225
+ query_key=query_keys[i],
226
+ query_uses=query_uses[i],
227
+ )
228
+ )
229
+ query_pooling_params_list.append(pooling_params)
230
+
231
+ assert (
232
+ n_queries
233
+ == len(query_pooling_params_list)
234
+ == len(query_engine_inputs)
235
+ == len(query_keys)
236
+ )
237
+
238
+ query_ctx = ScoringServeContext(
239
+ request=ctx.request,
240
+ raw_request=ctx.raw_request,
241
+ model_name=ctx.model_name,
242
+ request_id=ctx.request_id,
243
+ pooling_params=query_pooling_params_list,
244
+ prompt_request_ids=query_keys,
245
+ engine_inputs=query_engine_inputs,
246
+ )
247
+
248
+ await self._prepare_generators(query_ctx)
249
+ await self._collect_batch(query_ctx)
250
+
251
+ async def _flash_late_interaction_encode_docs(self, ctx: ScoringServeContext):
252
+ assert ctx.n_queries is not None
253
+ assert ctx.engine_inputs is not None
254
+ assert isinstance(ctx.pooling_params, PoolingParams)
255
+
256
+ n_queries = ctx.n_queries
257
+ n_docs = len(ctx.engine_inputs) - n_queries
258
+ doc_engine_inputs = ctx.engine_inputs[n_queries:]
259
+
260
+ query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
261
+ doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)]
262
+
263
+ doc_pooling_params_list = []
264
+ for i in range(n_docs):
265
+ query_idx = 0 if n_queries == 1 else i
266
+ pooling_params = ctx.pooling_params.clone()
267
+ pooling_params.late_interaction_params = build_late_interaction_doc_params(
268
+ query_key=query_keys[query_idx]
269
+ )
270
+ doc_pooling_params_list.append(pooling_params)
271
+
272
+ assert (
273
+ n_docs
274
+ == len(doc_pooling_params_list)
275
+ == len(doc_engine_inputs)
276
+ == len(doc_keys)
277
+ )
278
+
279
+ doc_ctx = ScoringServeContext(
280
+ request=ctx.request,
281
+ raw_request=ctx.raw_request,
282
+ model_name=ctx.model_name,
283
+ request_id=ctx.request_id,
284
+ pooling_params=doc_pooling_params_list,
285
+ prompt_request_ids=doc_keys,
286
+ engine_inputs=doc_engine_inputs,
287
+ )
288
+
289
+ await self._prepare_generators(doc_ctx)
290
+ await self._collect_batch(doc_ctx)
291
+
292
+ ctx.final_res_batch = doc_ctx.final_res_batch
vllm/entrypoints/pooling/scoring/utils.py
+2
-1
@@ -36,8 +36,9 @@ def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tens
36
Returns:
37
MaxSim score (sum over query tokens of max similarity to any doc token)
38
"""
39
+ # compute in float32 for numerical stability
40
# [query_len, doc_len]
40
- token_scores = torch.matmul(q_emb, d_emb.T)
41
+ token_scores = torch.matmul(q_emb.float(), d_emb.float().T)
42
# Max over document tokens, sum over query tokens
43
return token_scores.amax(dim=-1).sum()
44
vllm/entrypoints/pooling/typing.py
+5
-2
@@ -83,6 +83,9 @@ class PoolingServeContext(Generic[PoolingRequestT]):
83
84
model_config = ConfigDict(arbitrary_types_allowed=True)
85
86
+ ## for bi-encoder & late-interaction
87
+ n_queries: int | None = None
88
+
89
90
@dataclass
91
class OfflineInputsContext:
@@ -92,7 +95,7 @@ class OfflineInputsContext:
95
chat_template: str | None = None
96
97
## for bi-encoder & late-interaction
95
- offset: int | None = None
98
+ n_queries: int | None = None
99
100
101
@dataclass
@@ -100,4 +103,4 @@ class OfflineOutputsContext:
103
outputs: list[PoolingRequestOutput]
104
105
## for bi-encoder & late-interaction
103
- offset: int | None = None
106
+ n_queries: int | None = None
vllm/v1/pool/late_interaction.py
+1
-10
@@ -56,16 +56,7 @@ def build_late_interaction_doc_params(
56
)
57
58
59
-def compute_maxsim_score(
60
- q_emb: torch.Tensor,
61
- d_emb: torch.Tensor,
62
-) -> torch.Tensor:
63
- # compute in float32 for numerical stability
64
- token_scores = torch.matmul(q_emb.float(), d_emb.float().T)
65
- return token_scores.amax(dim=-1).sum()
66
-
67
-
68
-def compute_maxsim_scores(
59
+def compute_maxsim_score_batched(
60
q_embs: Sequence[torch.Tensor],
61
d_embs: Sequence[torch.Tensor],
62
max_batch_size: int = 64,
vllm/v1/worker/gpu/pool/late_interaction_runner.py
+2
-2
@@ -9,7 +9,7 @@ from vllm.v1.outputs import PoolerOutput
9
from vllm.v1.pool.late_interaction import (
10
LATE_INTERACTION_MODE_CACHE_QUERY,
11
LATE_INTERACTION_MODE_SCORE_DOC,
12
- compute_maxsim_scores,
12
+ compute_maxsim_score_batched,
13
)
14
15
@@ -116,7 +116,7 @@ class LateInteractionRunner:
116
raise ValueError(f"Unsupported late-interaction mode: {mode!r}")
117
118
if score_indices:
119
- score_values = compute_maxsim_scores(score_queries, score_docs)
119
+ score_values = compute_maxsim_score_batched(score_queries, score_docs)
120
for i, req_id, query_key, score in zip(
121
score_indices, score_req_ids, score_query_keys, score_values
122
):