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 7, 00:15 UTC
[Perf] Compute maxsim in worker side, reducing redundant copies, 2.7% E2E throughput improvement (#36159)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 9, 2026 at 23:55 UTC
7279374f9108652296a8f38b6f9c7f0585a0cda4
11 files changed
+520
-60
tests/v1/engine/test_engine_core_client.py
+71
@@ -24,17 +24,23 @@ from vllm import SamplingParams
24
from vllm.distributed.kv_events import BlockStored, KVEventBatch, ZmqEventPublisher
25
from vllm.engine.arg_utils import EngineArgs
26
from vllm.platforms import current_platform
27
+from vllm.pooling_params import LateInteractionParams, PoolingParams
28
from vllm.usage.usage_lib import UsageContext
29
from vllm.utils.torch_utils import set_default_torch_num_threads
30
from vllm.v1.engine import EngineCoreRequest
31
from vllm.v1.engine.core import EngineCore
32
from vllm.v1.engine.core_client import (
33
AsyncMPClient,
34
+ DPLBAsyncMPClient,
35
EngineCoreClient,
36
SyncMPClient,
37
)
38
from vllm.v1.engine.utils import CoreEngineProcManager
39
from vllm.v1.executor.abstract import Executor
40
+from vllm.v1.pool.late_interaction import (
41
+ LATE_INTERACTION_MODE_CACHE_QUERY,
42
+ LATE_INTERACTION_MODE_SCORE_DOC,
43
+)
44
45
from ...distributed.conftest import MockSubscriber
46
from ...utils import create_new_process_for_each_test
@@ -164,6 +170,71 @@ def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch):
170
client.shutdown()
171
172
173
+def _make_pooling_request(
174
+ request_id: str, *, mode: str | None = None, query_key: str | None = None
175
+) -> EngineCoreRequest:
176
+ late_interaction_params = None
177
+ if mode is not None and query_key is not None:
178
+ late_interaction_params = LateInteractionParams(
179
+ mode=mode,
180
+ query_key=query_key,
181
+ )
182
+
183
+ return EngineCoreRequest(
184
+ request_id=request_id,
185
+ prompt_token_ids=[1, 2, 3],
186
+ mm_features=None,
187
+ sampling_params=None,
188
+ pooling_params=PoolingParams(
189
+ task="token_embed",
190
+ late_interaction_params=late_interaction_params,
191
+ ),
192
+ arrival_time=time.time(),
193
+ lora_request=None,
194
+ cache_salt=None,
195
+ data_parallel_rank=None,
196
+ )
197
+
198
+
199
+def test_dplb_late_interaction_sticky_routing():
200
+ client = object.__new__(DPLBAsyncMPClient)
201
+ client.client_count = 1
202
+ client.reqs_in_flight = {}
203
+ client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]
204
+ client.lb_engines = [[0, 0], [0, 0], [0, 0]]
205
+ client.eng_start_index = 0
206
+
207
+ query_key = "rerank-abc-query-0"
208
+ query_request = _make_pooling_request(
209
+ "query-req", mode=LATE_INTERACTION_MODE_CACHE_QUERY, query_key=query_key
210
+ )
211
+ doc_request = _make_pooling_request(
212
+ "doc-req", mode=LATE_INTERACTION_MODE_SCORE_DOC, query_key=query_key
213
+ )
214
+
215
+ query_engine = client.get_core_engine_for_request(query_request)
216
+ doc_engine = client.get_core_engine_for_request(doc_request)
217
+
218
+ assert query_engine == doc_engine
219
+ assert client.reqs_in_flight["query-req"] == query_engine
220
+ assert client.reqs_in_flight["doc-req"] == doc_engine
221
+
222
+
223
+def test_dplb_non_late_interaction_still_uses_lb():
224
+ client = object.__new__(DPLBAsyncMPClient)
225
+ client.client_count = 1
226
+ client.reqs_in_flight = {}
227
+ client.core_engines = [b"\x00\x00", b"\x01\x00", b"\x02\x00"]
228
+ client.lb_engines = [[2, 1], [0, 0], [1, 0]]
229
+ client.eng_start_index = 0
230
+
231
+ request = make_request(SamplingParams(max_tokens=1))
232
+ chosen_engine = client.get_core_engine_for_request(request)
233
+
234
+ assert chosen_engine == client.core_engines[1]
235
+ assert client.lb_engines[1][0] == 1
236
+
237
+
238
def loop_until_done(client: EngineCoreClient, outputs: dict):
239
while True:
240
engine_core_outputs = client.get_output().outputs
tests/v1/worker/test_late_interaction_runner.py
new
+113
@@ -0,0 +1,113 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.pooling_params import LateInteractionParams, PoolingParams
8
+from vllm.v1.pool.late_interaction import (
9
+ LATE_INTERACTION_MODE_CACHE_QUERY,
10
+ build_late_interaction_doc_params,
11
+ build_late_interaction_query_params,
12
+ compute_maxsim_score,
13
+)
14
+from vllm.v1.worker.gpu.pool.late_interaction_runner import LateInteractionRunner
15
+
16
+
17
+def _make_pooling_params(
18
+ late_interaction_params: LateInteractionParams,
19
+) -> PoolingParams:
20
+ return PoolingParams(
21
+ task="token_embed",
22
+ late_interaction_params=late_interaction_params,
23
+ )
24
+
25
+
26
+def test_postprocess_scores_and_releases_query_cache():
27
+ runner = LateInteractionRunner()
28
+ query_key = "query-0"
29
+ query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
30
+ doc_emb = torch.tensor([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]], dtype=torch.float32)
31
+
32
+ query_params = _make_pooling_params(
33
+ build_late_interaction_query_params(query_key=query_key, query_uses=1)
34
+ )
35
+ query_output = runner.postprocess_pooler_output(
36
+ raw_pooler_output=[query_emb],
37
+ pooling_params=[query_params],
38
+ req_ids=["query-req"],
39
+ finished_mask=[True],
40
+ )
41
+ assert isinstance(query_output, list)
42
+ assert query_output[0] is not None
43
+ assert query_output[0].shape == torch.Size([])
44
+
45
+ doc_params = _make_pooling_params(
46
+ build_late_interaction_doc_params(query_key=query_key)
47
+ )
48
+ doc_output = runner.postprocess_pooler_output(
49
+ raw_pooler_output=[doc_emb],
50
+ pooling_params=[doc_params],
51
+ req_ids=["doc-req"],
52
+ finished_mask=[True],
53
+ )
54
+ assert isinstance(doc_output, list)
55
+ assert doc_output[0] is not None
56
+ assert torch.allclose(doc_output[0], compute_maxsim_score(query_emb, doc_emb))
57
+
58
+ with pytest.raises(ValueError, match="query cache miss"):
59
+ runner.postprocess_pooler_output(
60
+ raw_pooler_output=[doc_emb],
61
+ pooling_params=[doc_params],
62
+ req_ids=["doc-req-2"],
63
+ finished_mask=[True],
64
+ )
65
+
66
+
67
+def test_finished_request_releases_unscored_doc_use():
68
+ runner = LateInteractionRunner()
69
+ query_key = "query-cancel"
70
+ query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
71
+ doc_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
72
+
73
+ query_params = _make_pooling_params(
74
+ build_late_interaction_query_params(query_key=query_key, query_uses=1)
75
+ )
76
+ runner.postprocess_pooler_output(
77
+ raw_pooler_output=[query_emb],
78
+ pooling_params=[query_params],
79
+ req_ids=["query-req"],
80
+ finished_mask=[True],
81
+ )
82
+
83
+ doc_params = _make_pooling_params(
84
+ build_late_interaction_doc_params(query_key=query_key)
85
+ )
86
+ runner.register_request("doc-req", doc_params)
87
+ runner.on_requests_finished({"doc-req"})
88
+
89
+ with pytest.raises(ValueError, match="query cache miss"):
90
+ runner.postprocess_pooler_output(
91
+ raw_pooler_output=[doc_emb],
92
+ pooling_params=[doc_params],
93
+ req_ids=["doc-req-retry"],
94
+ finished_mask=[True],
95
+ )
96
+
97
+
98
+def test_invalid_query_uses_raises():
99
+ runner = LateInteractionRunner()
100
+ bad_meta = LateInteractionParams(
101
+ mode=LATE_INTERACTION_MODE_CACHE_QUERY,
102
+ query_key="query-bad",
103
+ )
104
+ bad_meta.query_uses = "bad-int" # type: ignore[assignment]
105
+ bad_query_params = _make_pooling_params(bad_meta)
106
+
107
+ with pytest.raises(ValueError, match="must be an integer value"):
108
+ runner.postprocess_pooler_output(
109
+ raw_pooler_output=[torch.ones((2, 2), dtype=torch.float32)],
110
+ pooling_params=[bad_query_params],
111
+ req_ids=["query-req"],
112
+ finished_mask=[True],
113
+ )
vllm/entrypoints/cli/serve.py
-6
@@ -225,12 +225,6 @@ def run_multi_api_server(args: argparse.Namespace):
225
num_api_servers: int = args.api_server_count
226
assert num_api_servers > 0
227
228
- if num_api_servers > 1 and getattr(args, "use_gpu_for_pooling_score", False):
229
- # TODO(wentao): remove this once well tested
230
- raise ValueError(
231
- "--use-gpu-for-pooling-score cannot be used with api_server_count > 1 now"
232
- )
233
-
228
if num_api_servers > 1:
229
setup_multiprocess_prometheus()
230
vllm/entrypoints/openai/cli_args.py
-4
@@ -281,10 +281,6 @@ class FrontendArgs(BaseFrontendArgs):
281
Enable offline FastAPI documentation for air-gapped environments.
282
Uses vendored static assets bundled with vLLM.
283
"""
284
- use_gpu_for_pooling_score: bool = False
285
- """If set, run pooling score MaxSim on GPU in the API server process.
286
- Can significantly improve late-interaction scoring performance.
287
- https://github.com/vllm-project/vllm/pull/35330"""
284
285
@classmethod
286
def _customize_cli_kwargs(
vllm/entrypoints/pooling/__init__.py
+1
-1
@@ -111,7 +111,7 @@ def init_pooling_state(
111
state.openai_serving_models,
112
request_logger=request_logger,
113
score_template=resolved_chat_template,
114
- use_gpu_for_pooling_score=getattr(args, "use_gpu_for_pooling_score", False),
114
+ log_error_stack=args.log_error_stack,
115
)
116
if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
117
else None
vllm/entrypoints/pooling/score/serving.py
+78
-48
@@ -31,7 +31,6 @@ from vllm.entrypoints.pooling.score.utils import (
31
ScoreInputs,
32
_cosine_similarity,
33
compress_token_type_ids,
34
- compute_maxsim_scores,
34
get_score_prompt,
35
parse_score_data_single,
36
validate_score_input,
@@ -43,6 +42,10 @@ from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
42
from vllm.tokenizers import TokenizerLike
43
from vllm.utils.async_utils import make_async, merge_async_iterators
44
from vllm.utils.mistral import is_mistral_tokenizer
45
+from vllm.v1.pool.late_interaction import (
46
+ build_late_interaction_doc_params,
47
+ build_late_interaction_query_params,
48
+)
49
50
logger = init_logger(__name__)
51
@@ -56,7 +59,6 @@ class ServingScores(OpenAIServing):
59
request_logger: RequestLogger | None,
60
score_template: str | None = None,
61
log_error_stack: bool = False,
59
- use_gpu_for_pooling_score: bool = False,
62
) -> None:
63
super().__init__(
64
engine_client=engine_client,
@@ -64,7 +66,6 @@ class ServingScores(OpenAIServing):
66
request_logger=request_logger,
67
)
68
self.score_template = score_template
67
- self.use_gpu_for_pooling_score = use_gpu_for_pooling_score
69
70
self._tokenizer_executor = ThreadPoolExecutor(max_workers=1)
71
@@ -253,19 +254,30 @@ class ServingScores(OpenAIServing):
254
)
255
)
256
256
- input_texts: list[str] = []
257
- engine_prompts: list[TokensPrompt] = []
258
- for text, engine_prompt in preprocessed:
259
- input_texts.append(text)
260
- engine_prompts.append(engine_prompt)
257
+ query_prompts: list[TokensPrompt] = [
258
+ prompt for _, prompt in preprocessed[: len(data_1)]
259
+ ]
260
+ doc_prompts: list[TokensPrompt] = [
261
+ prompt for _, prompt in preprocessed[len(data_1) :]
262
+ ]
263
262
- # Schedule the request and get the result generator.
263
- generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
264
+ default_pooling_params = request.to_pooling_params("token_embed")
265
265
- pooling_params = request.to_pooling_params("token_embed")
266
-
267
- for i, engine_prompt in enumerate(engine_prompts):
268
- request_id_item = f"{request_id}-{i}"
266
+ # stage 1: encode queries and cache token embeddings on workers.
267
+ query_keys = [f"{request_id}-query-{i}" for i in range(len(query_prompts))]
268
+ query_uses = [len(doc_prompts) if len(query_prompts) == 1 else 1] * len(
269
+ query_prompts
270
+ )
271
+ query_generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
272
+ for i, engine_prompt in enumerate(query_prompts):
273
+ request_id_item = f"{request_id}-query-{i}"
274
+ pooling_params = default_pooling_params.clone()
275
+ pooling_params.late_interaction_params = (
276
+ build_late_interaction_query_params(
277
+ query_key=query_keys[i],
278
+ query_uses=query_uses[i],
279
+ )
280
+ )
281
282
self._log_inputs(
283
request_id_item,
@@ -274,7 +286,7 @@ class ServingScores(OpenAIServing):
286
lora_request=lora_request,
287
)
288
277
- generators.append(
289
+ query_generators.append(
290
self.engine_client.encode(
291
engine_prompt,
292
pooling_params,
@@ -285,53 +297,71 @@ class ServingScores(OpenAIServing):
297
)
298
)
299
288
- result_generator = merge_async_iterators(*generators)
289
-
290
- # Collect token embeddings
291
- embeddings: list[PoolingRequestOutput | None] = [None] * len(engine_prompts)
292
-
293
- async for i, res in result_generator:
294
- embeddings[i] = res
295
-
296
- # Split into query and document embeddings
297
- emb_data_1: list[PoolingRequestOutput] = []
298
- emb_data_2: list[PoolingRequestOutput] = []
299
-
300
- for i in range(0, len(data_1)):
301
- assert (emb := embeddings[i]) is not None
302
- emb_data_1.append(emb)
300
+ query_outputs: list[PoolingRequestOutput | None] = [None] * len(query_prompts)
301
+ if query_generators:
302
+ async for i, res in merge_async_iterators(*query_generators):
303
+ query_outputs[i] = res
304
+
305
+ assert all(res is not None for res in query_outputs)
306
+ query_results = [res for res in query_outputs if res is not None]
307
+
308
+ # stage 2: encode docs and return scalar scores from workers.
309
+ doc_generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
310
+ for i, engine_prompt in enumerate(doc_prompts):
311
+ request_id_item = f"{request_id}-doc-{i}"
312
+ query_idx = 0 if len(query_prompts) == 1 else i
313
+ pooling_params = default_pooling_params.clone()
314
+ pooling_params.late_interaction_params = build_late_interaction_doc_params(
315
+ query_key=query_keys[query_idx]
316
+ )
317
304
- for i in range(len(data_1), len(embeddings)):
305
- assert (emb := embeddings[i]) is not None
306
- emb_data_2.append(emb)
318
+ self._log_inputs(
319
+ request_id_item,
320
+ engine_prompt,
321
+ params=pooling_params,
322
+ lora_request=lora_request,
323
+ )
324
308
- # Expand queries if 1:N scoring
309
- if len(emb_data_1) == 1:
310
- emb_data_1 = emb_data_1 * len(emb_data_2)
325
+ doc_generators.append(
326
+ self.engine_client.encode(
327
+ engine_prompt,
328
+ pooling_params,
329
+ request_id_item,
330
+ lora_request=lora_request,
331
+ trace_headers=trace_headers,
332
+ priority=request.priority,
333
+ )
334
+ )
335
312
- # Compute MaxSim scores
313
- from vllm.outputs import PoolingOutput
336
+ doc_outputs: list[PoolingRequestOutput | None] = [None] * len(doc_prompts)
337
+ if doc_generators:
338
+ async for i, res in merge_async_iterators(*doc_generators):
339
+ doc_outputs[i] = res
340
315
- maxsim_scores = compute_maxsim_scores(
316
- [emb.outputs.data for emb in emb_data_1],
317
- [emb.outputs.data for emb in emb_data_2],
318
- use_gpu_for_pooling_score=self.use_gpu_for_pooling_score,
319
- )
341
+ assert all(res is not None for res in doc_outputs)
342
+ doc_results = [res for res in doc_outputs if res is not None]
343
344
scores: list[PoolingRequestOutput] = []
345
padding: list[int] = []
346
if (pad_token_id := tokenizer.pad_token_id) is not None:
347
padding = [pad_token_id]
348
326
- for emb_1, emb_2, maxsim_score in zip(emb_data_1, emb_data_2, maxsim_scores):
327
- tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
349
+ if len(query_results) == 1:
350
+ query_results = query_results * len(doc_results)
351
+
352
+ for query_result, doc_result in zip(query_results, doc_results):
353
+ tokens = (
354
+ query_result.prompt_token_ids + padding + doc_result.prompt_token_ids
355
+ )
356
357
scores.append(
358
PoolingRequestOutput(
331
- request_id=f"{emb_1.request_id}_{emb_2.request_id}",
332
- outputs=PoolingOutput(data=maxsim_score),
359
+ request_id=f"{query_result.request_id}_{doc_result.request_id}",
360
+ outputs=doc_result.outputs,
361
prompt_token_ids=tokens,
334
- num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
362
+ num_cached_tokens=(
363
+ query_result.num_cached_tokens + doc_result.num_cached_tokens
364
+ ),
365
finished=True,
366
)
367
)
vllm/pooling_params.py
+22
@@ -11,6 +11,26 @@ from vllm.sampling_params import RequestOutputKind
11
from vllm.tasks import PoolingTask
12
13
14
+class LateInteractionParams(
15
+ msgspec.Struct,
16
+ omit_defaults=True, # type: ignore[call-arg]
17
+ array_like=True,
18
+): # type: ignore[call-arg]
19
+ """Metadata for worker-side late-interaction scoring.
20
+
21
+ Attributes:
22
+ mode:
23
+ - "cache_query": cache query token embeddings
24
+ - "score_doc": score a document against a cached query.
25
+ query_key: stable key used for both DP routing and worker cache lookup.
26
+ query_uses: expected number of document requests
27
+ """
28
+
29
+ mode: str
30
+ query_key: str
31
+ query_uses: int | None = None
32
+
33
+
34
class PoolingParams(
35
msgspec.Struct,
36
omit_defaults=True, # type: ignore[call-arg]
@@ -46,6 +66,7 @@ class PoolingParams(
66
task: PoolingTask | None = None
67
requires_token_ids: bool = False
68
skip_reading_prefix_cache: bool | None = None
69
+ late_interaction_params: LateInteractionParams | None = None
70
extra_kwargs: dict[str, Any] | None = None
71
output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
72
@@ -193,6 +214,7 @@ class PoolingParams(
214
f"returned_token_ids={self.returned_token_ids}, "
215
f"requires_token_ids={self.requires_token_ids}, "
216
f"skip_reading_prefix_cache={self.skip_reading_prefix_cache}, "
217
+ f"late_interaction_params={self.late_interaction_params}, "
218
f"extra_kwargs={self.extra_kwargs})"
219
)
220
vllm/v1/engine/core_client.py
+6
-1
@@ -52,6 +52,7 @@ from vllm.v1.engine.utils import (
52
launch_core_engines,
53
)
54
from vllm.v1.executor import Executor
55
+from vllm.v1.pool.late_interaction import get_late_interaction_engine_index
56
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
57
58
logger = init_logger(__name__)
@@ -1360,7 +1361,11 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
1361
1362
def get_core_engine_for_request(self, request: EngineCoreRequest) -> EngineIdentity:
1363
# Engines are in rank order.
1363
- if (eng_index := request.data_parallel_rank) is None:
1364
+ if (eng_index := request.data_parallel_rank) is None and (
1365
+ eng_index := get_late_interaction_engine_index(
1366
+ request.pooling_params, len(self.core_engines)
1367
+ )
1368
+ ) is None:
1369
current_counts = self.lb_engines
1370
# TODO use P2C alg for larger DP sizes
1371
num_engines = len(current_counts)
vllm/v1/pool/late_interaction.py
new
+64
@@ -0,0 +1,64 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import zlib
5
+
6
+import torch
7
+
8
+from vllm.pooling_params import LateInteractionParams, PoolingParams
9
+
10
+LATE_INTERACTION_MODE_CACHE_QUERY = "cache_query"
11
+LATE_INTERACTION_MODE_SCORE_DOC = "score_doc"
12
+
13
+
14
+def get_late_interaction_engine_index(
15
+ pooling_params: PoolingParams | None,
16
+ num_engines: int,
17
+) -> int | None:
18
+ if pooling_params is None or pooling_params.late_interaction_params is None:
19
+ return None
20
+
21
+ late_interaction_params = pooling_params.late_interaction_params
22
+ mode = late_interaction_params.mode
23
+ if mode not in (
24
+ LATE_INTERACTION_MODE_CACHE_QUERY,
25
+ LATE_INTERACTION_MODE_SCORE_DOC,
26
+ ):
27
+ return None
28
+
29
+ query_key = late_interaction_params.query_key
30
+ if not isinstance(query_key, str) or not query_key:
31
+ return None
32
+
33
+ # query embeddings are cached in process-local worker memory,
34
+ # pin requests sharing the same query key to the same engine.
35
+ return zlib.crc32(query_key.encode("utf-8")) % num_engines
36
+
37
+
38
+def build_late_interaction_query_params(
39
+ query_key: str,
40
+ query_uses: int,
41
+) -> LateInteractionParams:
42
+ return LateInteractionParams(
43
+ mode=LATE_INTERACTION_MODE_CACHE_QUERY,
44
+ query_key=query_key,
45
+ query_uses=max(1, int(query_uses)),
46
+ )
47
+
48
+
49
+def build_late_interaction_doc_params(
50
+ query_key: str,
51
+) -> LateInteractionParams:
52
+ return LateInteractionParams(
53
+ mode=LATE_INTERACTION_MODE_SCORE_DOC,
54
+ query_key=query_key,
55
+ )
56
+
57
+
58
+def compute_maxsim_score(
59
+ q_emb: torch.Tensor,
60
+ d_emb: torch.Tensor,
61
+) -> torch.Tensor:
62
+ # compute in float32 for numerical stability
63
+ token_scores = torch.matmul(q_emb.float(), d_emb.float().T)
64
+ return token_scores.amax(dim=-1).sum()
vllm/v1/worker/gpu/pool/late_interaction_runner.py
new
+150
@@ -0,0 +1,150 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
+
5
+import torch
6
+
7
+from vllm.pooling_params import PoolingParams
8
+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_score,
13
+)
14
+
15
+
16
+class LateInteractionRunner:
17
+ """Worker-side state and postprocessing for late-interaction scoring."""
18
+
19
+ def __init__(self) -> None:
20
+ # query_key -> token embeddings for late-interaction scoring.
21
+ self._query_cache: dict[str, torch.Tensor] = {}
22
+ # query_key -> remaining number of docs that should use this query.
23
+ self._query_uses: dict[str, int] = {}
24
+ # doc request id -> query key.
25
+ self._doc_query_keys: dict[str, str] = {}
26
+
27
+ def clear(self) -> None:
28
+ self._query_cache.clear()
29
+ self._query_uses.clear()
30
+ self._doc_query_keys.clear()
31
+
32
+ def register_request(
33
+ self, req_id: str, pooling_params: PoolingParams | None
34
+ ) -> None:
35
+ mode, query_key, _ = self._parse_late_interaction_meta(pooling_params)
36
+ if mode == LATE_INTERACTION_MODE_SCORE_DOC and query_key is not None:
37
+ self._doc_query_keys[req_id] = query_key
38
+ else:
39
+ self._doc_query_keys.pop(req_id, None)
40
+
41
+ def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None:
42
+ for req_id in finished_req_ids:
43
+ query_key = self._doc_query_keys.pop(req_id, None)
44
+ if query_key is not None:
45
+ self._release_query_use(query_key)
46
+
47
+ def postprocess_pooler_output(
48
+ self,
49
+ raw_pooler_output: PoolerOutput,
50
+ pooling_params: list[PoolingParams],
51
+ req_ids: list[str],
52
+ finished_mask: list[bool],
53
+ ) -> PoolerOutput:
54
+ if not isinstance(raw_pooler_output, list):
55
+ return raw_pooler_output
56
+
57
+ num_reqs = len(pooling_params)
58
+ if len(raw_pooler_output) != num_reqs:
59
+ raise ValueError(
60
+ "raw_pooler_output and pooling_params must have the same length."
61
+ )
62
+ if len(req_ids) != num_reqs:
63
+ raise ValueError("req_ids and pooling_params must have the same length.")
64
+ if len(finished_mask) != num_reqs:
65
+ raise ValueError(
66
+ "finished_mask and pooling_params must have the same length."
67
+ )
68
+
69
+ if not any(finished_mask):
70
+ return raw_pooler_output
71
+ if not any(p.late_interaction_params is not None for p in pooling_params):
72
+ return raw_pooler_output
73
+
74
+ outputs: list[torch.Tensor | None] = list(raw_pooler_output)
75
+ for i, (req_id, output, params, finished) in enumerate(
76
+ zip(req_ids, outputs, pooling_params, finished_mask)
77
+ ):
78
+ if not finished or output is None:
79
+ continue
80
+
81
+ mode, query_key, query_uses = self._parse_late_interaction_meta(params)
82
+ if mode is None:
83
+ continue
84
+
85
+ assert query_key is not None
86
+ if mode == LATE_INTERACTION_MODE_CACHE_QUERY:
87
+ assert query_uses is not None
88
+ # `output` can be a view into the current step's hidden-states
89
+ # buffer, so clone it before storing across scheduling steps.
90
+ self._query_cache[query_key] = output.clone()
91
+ self._query_uses[query_key] = query_uses
92
+ outputs[i] = torch.zeros((), device=output.device, dtype=torch.float32)
93
+ continue
94
+
95
+ if mode == LATE_INTERACTION_MODE_SCORE_DOC:
96
+ query_output = self._query_cache.get(query_key)
97
+ if query_output is None:
98
+ raise ValueError(
99
+ "late-interaction query cache miss for key "
100
+ f"{query_key!r}. Ensure query requests are executed "
101
+ "before their paired document requests."
102
+ )
103
+
104
+ outputs[i] = compute_maxsim_score(query_output, output)
105
+ self._doc_query_keys.pop(req_id, None)
106
+ self._release_query_use(query_key)
107
+ continue
108
+
109
+ raise ValueError(f"Unsupported late-interaction mode: {mode!r}")
110
+
111
+ return outputs
112
+
113
+ def _release_query_use(self, query_key: str) -> None:
114
+ remaining = self._query_uses.get(query_key, 1) - 1
115
+ if remaining <= 0:
116
+ self._query_uses.pop(query_key, None)
117
+ self._query_cache.pop(query_key, None)
118
+ else:
119
+ self._query_uses[query_key] = remaining
120
+
121
+ @staticmethod
122
+ def _parse_late_interaction_meta(
123
+ pooling_params: PoolingParams | None,
124
+ ) -> tuple[str | None, str | None, int | None]:
125
+ if pooling_params is None or pooling_params.late_interaction_params is None:
126
+ return None, None, None
127
+
128
+ late_interaction_params = pooling_params.late_interaction_params
129
+ mode = late_interaction_params.mode
130
+
131
+ query_key = late_interaction_params.query_key
132
+ if not isinstance(query_key, str) or not query_key:
133
+ raise ValueError(
134
+ "late-interaction request is missing a valid query key in "
135
+ "pooling_params.late_interaction_params."
136
+ )
137
+
138
+ if mode == LATE_INTERACTION_MODE_CACHE_QUERY:
139
+ query_uses_raw = late_interaction_params.query_uses
140
+ if query_uses_raw is None:
141
+ query_uses_raw = 1
142
+ try:
143
+ query_uses = max(1, int(query_uses_raw))
144
+ except (TypeError, ValueError) as exc:
145
+ raise ValueError(
146
+ "late-interaction query uses must be an integer value."
147
+ ) from exc
148
+ return mode, query_key, query_uses
149
+
150
+ return mode, query_key, None
vllm/v1/worker/gpu_model_runner.py
+15
@@ -181,6 +181,7 @@ from vllm.v1.worker.cp_utils import (
181
)
182
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
183
from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
184
+from vllm.v1.worker.gpu.pool.late_interaction_runner import LateInteractionRunner
185
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
186
from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
187
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
@@ -491,6 +492,7 @@ class GPUModelRunner(
492
493
# mm_hash -> encoder_output
494
self.encoder_cache: dict[str, torch.Tensor] = {}
495
+ self.late_interaction_runner = LateInteractionRunner()
496
497
self.use_aux_hidden_state_outputs = False
498
# Set up speculative decoding.
@@ -831,6 +833,7 @@ class GPUModelRunner(
833
"""
834
if self.mm_budget:
835
self.mm_budget.reset_cache()
836
+ self.late_interaction_runner.clear()
837
838
def reset_encoder_cache(self) -> None:
839
"""Clear the GPU-side encoder cache storing vision embeddings.
@@ -839,6 +842,7 @@ class GPUModelRunner(
842
stale embeddings computed with old weights are not reused.
843
"""
844
self.encoder_cache.clear()
845
+ self.late_interaction_runner.clear()
846
847
@torch.inference_mode()
848
def init_fp8_kv_scales(self) -> None:
@@ -1002,6 +1006,9 @@ class GPUModelRunner(
1006
for req_id in scheduler_output.finished_req_ids:
1007
self.requests.pop(req_id, None)
1008
self.num_prompt_logprobs.pop(req_id, None)
1009
+ self.late_interaction_runner.on_requests_finished(
1010
+ scheduler_output.finished_req_ids
1011
+ )
1012
# Remove the finished requests from the persistent batch.
1013
# NOTE(woosuk): There could be an edge case where finished_req_ids and
1014
# scheduled_req_ids overlap. This happens when a request is aborted and
@@ -1089,6 +1096,7 @@ class GPUModelRunner(
1096
lora_request=new_req_data.lora_request,
1097
)
1098
self.requests[req_id] = req_state
1099
+ self.late_interaction_runner.register_request(req_id, pooling_params)
1100
1101
if sampling_params and sampling_params.prompt_logprobs is not None:
1102
self.num_prompt_logprobs[req_id] = (
@@ -1360,6 +1368,7 @@ class GPUModelRunner(
1368
req_state.prompt_embeds = new_req_data.prompt_embeds
1369
req_state.sampling_params = new_req_data.sampling_params
1370
req_state.pooling_params = new_req_data.pooling_params
1371
+ self.late_interaction_runner.register_request(req_id, req_state.pooling_params)
1372
req_state.block_ids = new_req_data.block_ids
1373
req_state.num_computed_tokens = new_req_data.num_computed_tokens
1374
req_state.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
@@ -2875,6 +2884,12 @@ class GPUModelRunner(
2884
seq_len == prompt_len
2885
for seq_len, prompt_len in zip(seq_lens_cpu, pooling_metadata.prompt_lens)
2886
]
2887
+ raw_pooler_output = self.late_interaction_runner.postprocess_pooler_output(
2888
+ raw_pooler_output=raw_pooler_output,
2889
+ pooling_params=pooling_metadata.pooling_params,
2890
+ req_ids=self.input_batch.req_ids,
2891
+ finished_mask=finished_mask,
2892
+ )
2893
2894
model_runner_output = ModelRunnerOutput(
2895
req_ids=self.input_batch.req_ids.copy(),