Use server-generated keys for late-interaction query caches (#51445)
Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com>
Clinton Thomas committed
Sep 3, 2026 at 03:15 UTC
ee17d0d869203ef9a35ad73358a4987bba14b1fc
3 files changed
+75
-2
tests/entrypoints/pooling/scoring/test_late_interaction_serving.py
new
+65
@@ -0,0 +1,65 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from unittest.mock import AsyncMock, Mock
5
+
6
+import pytest
7
+
8
+from vllm.entrypoints.pooling.scoring.serving import ServingScores
9
+from vllm.entrypoints.pooling.typing import PoolingServeContext
10
+from vllm.inputs import tokens_input
11
+from vllm.pooling_params import PoolingParams
12
+
13
+
14
+def _make_context() -> PoolingServeContext:
15
+ return PoolingServeContext(
16
+ request=Mock(),
17
+ model_name="model",
18
+ request_id="score-caller-controlled-id",
19
+ pooling_params=PoolingParams(task="token_embed"),
20
+ lora_request=None,
21
+ priorities=None,
22
+ prompt_extras=None,
23
+ engine_inputs=[
24
+ tokens_input(prompt_token_ids=[1]),
25
+ tokens_input(prompt_token_ids=[2]),
26
+ ],
27
+ n_queries=1,
28
+ )
29
+
30
+
31
+def _query_key(context: PoolingServeContext) -> str:
32
+ assert context.engine_inputs is not None
33
+ pooling_params = context.engine_inputs[0]["params"]
34
+ late_interaction_params = pooling_params.late_interaction_params
35
+ assert late_interaction_params is not None
36
+ query_key = late_interaction_params.query_key
37
+ assert query_key is not None
38
+ return query_key
39
+
40
+
41
+@pytest.mark.asyncio
42
+async def test_colliding_request_ids_use_distinct_query_cache_keys():
43
+ serving = object.__new__(ServingScores)
44
+ serving._prepare_generators = AsyncMock()
45
+ serving._collect_batch = AsyncMock()
46
+
47
+ first = _make_context()
48
+ second = _make_context()
49
+ for context in (first, second):
50
+ await serving._flash_late_interaction_encode_queries(context)
51
+ await serving._flash_late_interaction_encode_docs(context)
52
+
53
+ prepared_contexts = [
54
+ call.args[0] for call in serving._prepare_generators.await_args_list
55
+ ]
56
+ first_query_key = _query_key(prepared_contexts[0])
57
+ first_doc_key = _query_key(prepared_contexts[1])
58
+ second_query_key = _query_key(prepared_contexts[2])
59
+ second_doc_key = _query_key(prepared_contexts[3])
60
+
61
+ assert first_query_key == first_doc_key
62
+ assert second_query_key == second_doc_key
63
+ assert first_query_key != second_query_key
64
+ assert "caller-controlled-id" not in first_query_key
65
+ assert "caller-controlled-id" not in second_query_key
vllm/entrypoints/pooling/scoring/serving.py
+9
-2
@@ -9,6 +9,7 @@ from vllm.entrypoints.serve.engine.protocol import UsageInfo
9
from vllm.logger import init_logger
10
from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
11
from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
12
+from vllm.utils import random_uuid
13
from vllm.v1.pool.late_interaction import (
14
build_late_interaction_doc_params,
15
build_late_interaction_query_params,
@@ -208,7 +209,11 @@ class ServingScores(PoolingServing):
209
n_docs = len(ctx.engine_inputs) - n_queries
210
query_engine_inputs = ctx.engine_inputs[:n_queries]
211
211
- query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
212
+ query_namespace = random_uuid()
213
+ query_keys = [
214
+ f"late-interaction-{query_namespace}-query-{i}" for i in range(n_queries)
215
+ ]
216
+ ctx.late_interaction_query_keys = query_keys
217
query_uses = [n_docs if n_queries == 1 else 1] * n_queries
218
219
for i in range(n_queries):
@@ -249,7 +254,9 @@ class ServingScores(PoolingServing):
254
n_docs = len(ctx.engine_inputs) - n_queries
255
doc_engine_inputs = ctx.engine_inputs[n_queries:]
256
252
- query_keys = [f"{ctx.request_id}-query-{i}" for i in range(n_queries)]
257
+ query_keys = ctx.late_interaction_query_keys
258
+ if query_keys is None:
259
+ raise RuntimeError("Late-interaction query keys were not initialized.")
260
doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)]
261
262
for i in range(n_docs):
vllm/entrypoints/pooling/typing.py
+1
@@ -114,6 +114,7 @@ class PoolingServeContext(Generic[PoolingRequestT]):
114
115
## for flash-late-interaction
116
query_final_res_batch: list[PoolingRequestOutput] | None = None
117
+ late_interaction_query_keys: list[str] | None = None
118
119
120
@dataclass