@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 592c6f3fbc

[Bugfix][Pooling] Handle multimodal cache misses without crashing (#58975)

Signed-off-by: linchaoying6 <linchaoying6@jd.com> Co-authored-by: linchaoying6 <linchaoying6@jd.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Chaoying committed Oct 2, 2026 at 22:14 UTC 592c6f3fbcefeb3a4e0c0b461531329bf4ee0489
21 files changed +525 -17
tests/entrypoints/pooling/scoring/test_io_processor_unit.py
+62 -1
@@ -1,20 +1,25 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Unit tests for ScoringIOProcessor post-tokenization helpers."""
3 +"""Unit tests for ScoringIOProcessor helpers."""
4
5 from dataclasses import dataclass
6 from types import SimpleNamespace
7
8 import pytest
9 +import torch
10 from tokenizers import Tokenizer, models, pre_tokenizers, processors
11 from transformers import TokenizersBackend
12
13 from vllm import TokensPrompt
14 from vllm.entrypoints.pooling.scoring.io_processor import (
15 + BiEncoderIOProcessor,
16 CrossEncoderIOProcessor,
17 + JinaRankingIOProcessor,
18 + LateInteractionIOProcessor,
19 _apply_post_tokenization_to_token_type_ids,
20 )
21 from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids
22 +from vllm.outputs import PoolingOutput, PoolingRequestOutput, RequestError
23 from vllm.renderers import TokenizeParams
24
25 pytestmark = pytest.mark.skip_global_cleanup
@@ -113,3 +118,59 @@ def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt():
118 assert first_doc == 10
119 assert prompt_token_ids[:first_doc] == list(range(10, num_query))
120 assert prompt_token_ids[first_doc:40] == list(range(num_query, 50))
121 +
122 +
123 +def _pooling_output(
124 + request_id: str,
125 + data: torch.Tensor,
126 + error: RequestError | None = None,
127 +) -> PoolingRequestOutput:
128 + return PoolingRequestOutput(
129 + request_id=request_id,
130 + outputs=PoolingOutput(data),
131 + prompt_token_ids=[1],
132 + num_cached_tokens=0,
133 + finished=True,
134 + error=error,
135 + )
136 +
137 +
138 +@pytest.mark.parametrize(
139 + "processor_type", [BiEncoderIOProcessor, LateInteractionIOProcessor]
140 +)
141 +@pytest.mark.parametrize("failed_index", [0, 1])
142 +def test_paired_scoring_preserves_request_error(processor_type, failed_index):
143 + processor = processor_type.__new__(processor_type)
144 + processor.pad_token_id = None
145 + error = RequestError(
146 + code="multimodal_cache_miss",
147 + message="Multi-modal processor cache miss.",
148 + retryable=True,
149 + )
150 + outputs = [
151 + _pooling_output("query", torch.tensor([1.0, 2.0])),
152 + _pooling_output("document", torch.tensor([3.0, 4.0])),
153 + ]
154 + outputs[failed_index] = _pooling_output(
155 + outputs[failed_index].request_id,
156 + torch.empty(0),
157 + error,
158 + )
159 +
160 + (result,) = processor._post_process(outputs, n_queries=1)
161 +
162 + assert result.request_id == "query_document"
163 + assert result.error is error
164 + assert result.outputs is outputs[failed_index].outputs
165 +
166 +
167 +def test_jina_scoring_preserves_request_error_before_indexing_output():
168 + processor = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor)
169 + error = RequestError(
170 + code="multimodal_cache_miss",
171 + message="Multi-modal processor cache miss.",
172 + retryable=True,
173 + )
174 + failed = _pooling_output("request", torch.empty(0), error)
175 +
176 + assert processor._post_process([failed], n_queries=1) == [failed]
tests/entrypoints/pooling/scoring/test_late_interaction_serving.py
+66 -1
@@ -49,10 +49,60 @@ def _query_key(context: PoolingServeContext) -> str:
49
50
51 @pytest.mark.asyncio
52 -async def test_colliding_request_ids_use_distinct_query_cache_keys(
52 +@pytest.mark.skip_global_cleanup
53 +@pytest.mark.parametrize("failed_stage", ["query", "doc"])
54 +async def test_flash_late_interaction_cleans_up_queries_on_failure(
55 + failed_stage: str,
56 monkeypatch: pytest.MonkeyPatch,
57 ):
58 serving = object.__new__(ServingScores)
59 + serving.io_processor = Mock()
60 + serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock())
61 +
62 + context = _make_context()
63 + monkeypatch.setattr(serving, "_init_ctx", AsyncMock(return_value=context))
64 + monkeypatch.setattr(serving, "_preprocessing", AsyncMock())
65 +
66 + async def encode_queries(ctx):
67 + ctx.late_interaction_query_keys = ["query-key"]
68 + if failed_stage == "query":
69 + raise RuntimeError("query failed")
70 +
71 + async def encode_docs(ctx):
72 + ctx.late_interaction_doc_keys = ["doc-key"]
73 + if failed_stage == "doc":
74 + raise RuntimeError("doc failed")
75 +
76 + monkeypatch.setattr(
77 + serving,
78 + "_flash_late_interaction_encode_queries",
79 + AsyncMock(side_effect=encode_queries),
80 + )
81 + monkeypatch.setattr(
82 + serving,
83 + "_flash_late_interaction_encode_docs",
84 + AsyncMock(side_effect=encode_docs),
85 + )
86 +
87 + with pytest.raises(RuntimeError, match=f"{failed_stage} failed"):
88 + await serving.flash_late_interaction()
89 +
90 + expected_doc_keys = [] if failed_stage == "query" else ["doc-key"]
91 + serving.engine_client.abort.assert_awaited_once_with(
92 + ["query-key", *expected_doc_keys]
93 + )
94 + serving.engine_client.collective_rpc.assert_awaited_once_with(
95 + "release_late_interaction_query_cache",
96 + args=(["query-key"],),
97 + )
98 +
99 +
100 +@pytest.mark.asyncio
101 +async def test_colliding_request_ids_use_distinct_late_interaction_keys(
102 + monkeypatch: pytest.MonkeyPatch,
103 +):
104 + serving = object.__new__(ServingScores)
105 + serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock())
106 prepare_generators = AsyncMock()
107 monkeypatch.setattr(serving, "_prepare_generators", prepare_generators)
108 monkeypatch.setattr(serving, "_collect_batch", AsyncMock())
@@ -68,9 +118,24 @@ async def test_colliding_request_ids_use_distinct_query_cache_keys(
118 first_doc_key = _query_key(prepared_contexts[1])
119 second_query_key = _query_key(prepared_contexts[2])
120 second_doc_key = _query_key(prepared_contexts[3])
121 + first_doc_request_ids = prepared_contexts[1].prompt_request_ids
122 + second_doc_request_ids = prepared_contexts[3].prompt_request_ids
123
124 assert first_query_key == first_doc_key
125 assert second_query_key == second_doc_key
126 assert first_query_key != second_query_key
127 assert "caller-controlled-id" not in first_query_key
128 assert "caller-controlled-id" not in second_query_key
129 + assert first_doc_request_ids is not None
130 + assert second_doc_request_ids is not None
131 + assert first.late_interaction_doc_keys == first_doc_request_ids
132 + assert second.late_interaction_doc_keys == second_doc_request_ids
133 + assert first_doc_request_ids != second_doc_request_ids
134 + assert all("caller-controlled-id" not in key for key in first_doc_request_ids)
135 + assert all("caller-controlled-id" not in key for key in second_doc_request_ids)
136 +
137 + assert first.late_interaction_query_keys is not None
138 + await serving._cleanup_flash_late_interaction(first)
139 + serving.engine_client.abort.assert_awaited_once_with(
140 + [*first.late_interaction_query_keys, *first_doc_request_ids]
141 + )
tests/test_outputs.py
+43 -1
@@ -2,8 +2,17 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 +import torch
6
6 -from vllm.outputs import RequestOutput
7 +from vllm.outputs import (
8 + ClassificationRequestOutput,
9 + EmbeddingRequestOutput,
10 + PoolingOutput,
11 + PoolingRequestOutput,
12 + RequestError,
13 + RequestOutput,
14 + ScoringRequestOutput,
15 +)
16
17 pytestmark = pytest.mark.cpu_test
18
@@ -19,3 +28,36 @@ def test_request_output_forward_compatible():
28 example_arg_added_in_new_version="some_value",
29 )
30 assert output is not None
31 +
32 +
33 +@pytest.mark.parametrize(
34 + ("output_type", "empty_value"),
35 + [
36 + (EmbeddingRequestOutput, []),
37 + (ClassificationRequestOutput, []),
38 + (ScoringRequestOutput, 0.0),
39 + ],
40 +)
41 +@pytest.mark.skip_global_cleanup
42 +def test_specialized_pooling_outputs_preserve_request_error(output_type, empty_value):
43 + error = RequestError(
44 + code="multimodal_cache_miss",
45 + message="Multi-modal processor cache miss.",
46 + retryable=True,
47 + )
48 + base_output = PoolingRequestOutput(
49 + request_id="request",
50 + outputs=PoolingOutput(torch.empty(0)),
51 + prompt_token_ids=[1, 2],
52 + num_cached_tokens=0,
53 + finished=True,
54 + error=error,
55 + )
56 +
57 + output = output_type.from_base(base_output)
58 +
59 + assert output.error is error
60 + value = getattr(output.outputs, "embedding", None)
61 + value = getattr(output.outputs, "probs", value)
62 + value = getattr(output.outputs, "score", value)
63 + assert value == empty_value
tests/v1/engine/test_llm_engine.py
+34
@@ -2,11 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import random
4 from typing import TYPE_CHECKING
5 +from unittest.mock import MagicMock, call
6
7 import pytest
8
9 from vllm import LLM
10 from vllm.sampling_params import SamplingParams, StructuredOutputsParams
11 +from vllm.v1.engine.llm_engine import LLMEngine
12 from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector
13
14 if TYPE_CHECKING:
@@ -18,6 +20,38 @@ MODEL = "facebook/opt-125m"
20 DTYPE = "half"
21
22
23 +@pytest.mark.skip_global_cleanup
24 +def test_step_invalidates_mm_processor_cache_misses():
25 + engine = object.__new__(LLMEngine)
26 + engine.should_execute_dummy_batch = False
27 + engine.log_stats = False
28 + engine.logger_manager = None
29 +
30 + mm_processor_cache = MagicMock()
31 + engine.renderer = MagicMock(mm_processor_cache=mm_processor_cache)
32 +
33 + engine_core_output = MagicMock(
34 + mm_cache_miss_hashes=["missing-hash-1", "missing-hash-2"]
35 + )
36 + outputs = MagicMock(
37 + outputs=[engine_core_output],
38 + timestamp=0.0,
39 + scheduler_stats=None,
40 + )
41 + engine.engine_core = MagicMock()
42 + engine.engine_core.get_output.return_value = outputs
43 +
44 + processed_outputs = MagicMock(reqs_to_abort=[], request_outputs=[])
45 + engine.output_processor = MagicMock()
46 + engine.output_processor.process_outputs.return_value = processed_outputs
47 +
48 + assert engine.step() == []
49 + assert mm_processor_cache.invalidate.call_args_list == [
50 + call("missing-hash-1"),
51 + call("missing-hash-2"),
52 + ]
53 +
54 +
55 def _vllm_model(
56 apc: bool,
57 vllm_runner: type[VllmRunner],
tests/v1/engine/test_output_processor.py
+96 -1
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, Mock
7
8 import numpy as np
9 import pytest
10 +import torch
11
12 from tests.v1.engine.utils import (
13 NUM_PROMPT_LOGPROBS_UNDER_TEST,
@@ -18,7 +19,12 @@ from tests.v1.engine.utils import (
19 from vllm import PoolingParams
20 from vllm.logprobs import FlatLogprobs, Logprob, PromptLogprobs, SampleLogprobs
21 from vllm.lora.request import LoRARequest
21 -from vllm.outputs import CompletionOutput, RequestOutput, SamplingMask
22 +from vllm.outputs import (
23 + CompletionOutput,
24 + PoolingRequestOutput,
25 + RequestOutput,
26 + SamplingMask,
27 +)
28 from vllm.sampling_params import RequestOutputKind, SamplingParams
29 from vllm.tokenizers import TokenizerLike
30 from vllm.v1.engine import (
@@ -1630,3 +1636,92 @@ def test_request_output_add_merges_delta_sampling_masks():
1636
1637 assert merged.outputs[0].token_ids == [10, 20]
1638 assert merged.outputs[0].sampling_mask.token_ids == [[10, 11], [20]]
1639 +
1640 +
1641 +@pytest.mark.skip_global_cleanup
1642 +def test_pooling_cache_miss_is_a_request_error_and_batch_continues():
1643 + output_processor = OutputProcessor(tokenizer=None, log_stats=False)
1644 + requests = [
1645 + EngineCoreRequest(
1646 + request_id=f"request-{idx}-int",
1647 + external_req_id=f"request-{idx}",
1648 + prompt_token_ids=[1, 2],
1649 + mm_features=None,
1650 + arrival_time=0,
1651 + lora_request=None,
1652 + cache_salt=None,
1653 + data_parallel_rank=None,
1654 + sampling_params=None,
1655 + pooling_params=PoolingParams(task="embed"),
1656 + )
1657 + for idx in range(2)
1658 + ]
1659 + for request in requests:
1660 + output_processor.add_request(request, prompt=None)
1661 +
1662 + processed = output_processor.process_outputs(
1663 + [
1664 + EngineCoreOutput(
1665 + request_id=requests[0].request_id,
1666 + new_token_ids=[],
1667 + finish_reason=FinishReason.ERROR,
1668 + mm_cache_miss_hashes=["missing-hash"],
1669 + ),
1670 + EngineCoreOutput(
1671 + request_id=requests[1].request_id,
1672 + new_token_ids=[],
1673 + pooling_output=torch.tensor([1.0, 2.0]),
1674 + finish_reason=FinishReason.STOP,
1675 + ),
1676 + ]
1677 + )
1678 +
1679 + assert len(processed.request_outputs) == 2
1680 + failed, succeeded = processed.request_outputs
1681 + assert isinstance(failed, PoolingRequestOutput)
1682 + assert failed.error is not None
1683 + assert failed.error.code == "multimodal_cache_miss"
1684 + assert failed.error.retryable
1685 + assert isinstance(succeeded, PoolingRequestOutput)
1686 + assert succeeded.error is None
1687 + assert torch.equal(succeeded.outputs.data, torch.tensor([1.0, 2.0]))
1688 +
1689 +
1690 +@pytest.mark.skip_global_cleanup
1691 +def test_async_pooling_cache_miss_does_not_fail_output_processor():
1692 + output_processor = OutputProcessor(tokenizer=None, log_stats=False)
1693 + request = EngineCoreRequest(
1694 + request_id="request-int",
1695 + external_req_id="request",
1696 + prompt_token_ids=[1, 2],
1697 + mm_features=None,
1698 + arrival_time=0,
1699 + lora_request=None,
1700 + cache_salt=None,
1701 + data_parallel_rank=None,
1702 + sampling_params=None,
1703 + pooling_params=PoolingParams(task="embed"),
1704 + )
1705 + queue = RequestOutputCollector(
1706 + output_kind=RequestOutputKind.FINAL_ONLY,
1707 + request_id=request.request_id,
1708 + )
1709 + output_processor.add_request(request, prompt=None, queue=queue)
1710 +
1711 + processed = output_processor.process_outputs(
1712 + [
1713 + EngineCoreOutput(
1714 + request_id=request.request_id,
1715 + new_token_ids=[],
1716 + finish_reason=FinishReason.ERROR,
1717 + mm_cache_miss_hashes=["missing-hash"],
1718 + )
1719 + ]
1720 + )
1721 +
1722 + assert processed.request_outputs == []
1723 + output = queue.get_nowait()
1724 + assert isinstance(output, PoolingRequestOutput)
1725 + assert output.error is not None
1726 + assert output.error.retryable
1727 + assert not output_processor.has_unfinished_requests()
tests/v1/worker/test_late_interaction_runner.py
+46
@@ -137,6 +137,52 @@ def test_finished_request_releases_unscored_doc_use():
137 )
138
139
140 +@pytest.mark.skip_global_cleanup
141 +def test_release_queries_is_targeted_and_idempotent():
142 + runner = LateInteractionRunner()
143 + query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
144 +
145 + for query_key in ("query-release", "query-keep"):
146 + query_params = _make_pooling_params(
147 + build_late_interaction_query_params(query_key=query_key, query_uses=1)
148 + )
149 + runner.postprocess_pooler_output(
150 + raw_pooler_output=[query_emb],
151 + pooling_params=[query_params],
152 + req_ids=[f"{query_key}-req"],
153 + finished_mask=[True],
154 + )
155 +
156 + released_doc_params = _make_pooling_params(
157 + build_late_interaction_doc_params(query_key="query-release")
158 + )
159 + kept_doc_params = _make_pooling_params(
160 + build_late_interaction_doc_params(query_key="query-keep")
161 + )
162 + runner.register_request("released-doc", released_doc_params)
163 + runner.register_request("kept-doc", kept_doc_params)
164 +
165 + runner.release_queries(["query-release"])
166 + runner.release_queries(["query-release"])
167 +
168 + with pytest.raises(ValueError, match="query cache miss"):
169 + runner.postprocess_pooler_output(
170 + raw_pooler_output=[query_emb],
171 + pooling_params=[released_doc_params],
172 + req_ids=["released-doc"],
173 + finished_mask=[True],
174 + )
175 +
176 + kept_output = runner.postprocess_pooler_output(
177 + raw_pooler_output=[query_emb],
178 + pooling_params=[kept_doc_params],
179 + req_ids=["kept-doc"],
180 + finished_mask=[True],
181 + )
182 + assert isinstance(kept_output, list)
183 + assert kept_output[0] is not None
184 +
185 +
186 def test_invalid_query_uses_raises():
187 runner = LateInteractionRunner()
188 bad_meta = LateInteractionParams(
vllm/entrypoints/pooling/base/serving.py
+5
@@ -21,6 +21,7 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
21 from vllm.entrypoints.serve.engine.serving import BaseServing
22 from vllm.entrypoints.serve.engine.typing import AnyRequest
23 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
24 +from vllm.exceptions import GenerationError, RetryableRequestError
25 from vllm.lora.request import LoRARequest
26 from vllm.renderers.base import BaseRenderer
27 from vllm.tracing import (
@@ -206,6 +207,10 @@ class PoolingBaseServing(ABC, BaseServing):
207 final_res_batch = [None] * num_inputs
208
209 async for i, res in ctx.result_generator:
210 + if res.error is not None:
211 + if res.error.retryable:
212 + raise RetryableRequestError(res.error.message)
213 + raise GenerationError(res.error.message)
214 final_res_batch[i] = res
215
216 if None in final_res_batch:
vllm/entrypoints/pooling/scoring/io_processor.py
+33
@@ -191,6 +191,27 @@ class ScoringIOProcessor(PoolingIOProcessor):
191 scoring_data = self.valid_inputs(data_1, data_2)
192 return scoring_data
193
194 + def _pair_error_output(
195 + self,
196 + output_1: PoolingRequestOutput,
197 + output_2: PoolingRequestOutput,
198 + ) -> PoolingRequestOutput | None:
199 + failed = output_1 if output_1.error is not None else output_2
200 + if failed.error is None:
201 + return None
202 +
203 + padding = [self.pad_token_id] if self.pad_token_id is not None else []
204 + return PoolingRequestOutput(
205 + request_id=f"{output_1.request_id}_{output_2.request_id}",
206 + outputs=failed.outputs,
207 + prompt_token_ids=(
208 + output_1.prompt_token_ids + padding + output_2.prompt_token_ids
209 + ),
210 + num_cached_tokens=(output_1.num_cached_tokens + output_2.num_cached_tokens),
211 + finished=failed.finished,
212 + error=failed.error,
213 + )
214 +
215
216 class BiEncoderIOProcessor(ScoringIOProcessor):
217 name = "bi-encoder"
@@ -310,6 +331,10 @@ class BiEncoderIOProcessor(ScoringIOProcessor):
331
332 final_res_batch: list[PoolingRequestOutput] = []
333 for emb_1, emb_2 in zip(emb_data_1, emb_data_2):
334 + if error_output := self._pair_error_output(emb_1, emb_2):
335 + final_res_batch.append(error_output)
336 + continue
337 +
338 pair_score = F.cosine_similarity(
339 emb_1.outputs.data.float(), emb_2.outputs.data.float(), dim=0
340 )
@@ -352,6 +377,10 @@ class LateInteractionIOProcessor(BiEncoderIOProcessor):
377
378 # Compute MaxSim scores
379 for emb_1, emb_2 in zip(emb_data_1, emb_data_2):
380 + if error_output := self._pair_error_output(emb_1, emb_2):
381 + final_res_batch.append(error_output)
382 + continue
383 +
384 # emb_1.outputs.data: [query_len, dim]
385 # emb_2.outputs.data: [doc_len, dim]
386 q_emb = emb_1.outputs.data
@@ -907,6 +936,10 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM
936 final_res_batch: list[PoolingRequestOutput] = []
937
938 for i in range(len(outputs)):
939 + if outputs[i].error is not None:
940 + final_res_batch.append(outputs[i])
941 + continue
942 +
943 embeds = outputs[i].outputs.data.float()
944
945 # The JinaForRanking model concatenates docs first, then query.
vllm/entrypoints/pooling/scoring/serving.py
+30 -5
@@ -195,13 +195,36 @@ class ServingScores(PoolingServing):
195 ctx = await self._init_ctx(self.io_processor, *args, **kwargs)
196 await self._preprocessing(self.io_processor, ctx)
197
198 - # stage 1: encode queries and cache token embeddings on workers.
199 - await self._flash_late_interaction_encode_queries(ctx)
200 - # stage 2: encode docs and return scalar scores from workers.
201 - await self._flash_late_interaction_encode_docs(ctx)
198 + try:
199 + # stage 1: encode queries and cache token embeddings on workers.
200 + await self._flash_late_interaction_encode_queries(ctx)
201 + # stage 2: encode docs and return scalar scores from workers.
202 + await self._flash_late_interaction_encode_docs(ctx)
203 + except BaseException:
204 + try:
205 + await self._cleanup_flash_late_interaction(ctx)
206 + except Exception:
207 + logger.exception("Failed to clean up late-interaction query cache.")
208 + raise
209
210 return await self._postprocessing_async(self.io_processor, ctx)
211
212 + async def _cleanup_flash_late_interaction(self, ctx: ScoringServeContext) -> None:
213 + query_keys = ctx.late_interaction_query_keys
214 + if not query_keys:
215 + return
216 +
217 + doc_keys = ctx.late_interaction_doc_keys or []
218 +
219 + try:
220 + # Stop documents before removing the query tensors they reference.
221 + await self.engine_client.abort([*query_keys, *doc_keys])
222 + finally:
223 + await self.engine_client.collective_rpc(
224 + "release_late_interaction_query_cache",
225 + args=(query_keys,),
226 + )
227 +
228 async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext):
229 assert ctx.n_queries is not None
230 assert ctx.engine_inputs is not None
@@ -259,7 +282,9 @@ class ServingScores(PoolingServing):
282 query_keys = ctx.late_interaction_query_keys
283 if query_keys is None:
284 raise RuntimeError("Late-interaction query keys were not initialized.")
262 - doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)]
285 + doc_namespace = random_uuid()
286 + doc_keys = [f"late-interaction-{doc_namespace}-doc-{i}" for i in range(n_docs)]
287 + ctx.late_interaction_doc_keys = doc_keys
288
289 for i in range(n_docs):
290 query_idx = 0 if n_queries == 1 else i
vllm/entrypoints/pooling/typing.py
+1
@@ -115,6 +115,7 @@ class PoolingServeContext(Generic[PoolingRequestT]):
115 ## for flash-late-interaction
116 query_final_res_batch: list[PoolingRequestOutput] | None = None
117 late_interaction_query_keys: list[str] | None = None
118 + late_interaction_doc_keys: list[str] | None = None
119
120
121 @dataclass
vllm/entrypoints/serve/exception_handling/error_response.py
+6 -2
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from http import HTTPStatus
4
5 -from vllm.exceptions import GenerationError
5 +from vllm.exceptions import GenerationError, RetryableRequestError
6 from vllm.logger import init_logger
7
8 from ..engine.protocol import ErrorInfo, ErrorResponse
@@ -56,7 +56,11 @@ def create_error_response(
56 status_code = HTTPStatus.BAD_REQUEST
57 param = None
58 elif isinstance(exc, GenerationError):
59 - err_type = "InternalServerError"
59 + err_type = (
60 + "ServiceUnavailableError"
61 + if isinstance(exc, RetryableRequestError)
62 + else "InternalServerError"
63 + )
64 status_code = exc.status_code
65 param = None
66 elif isinstance(exc, VLLMServerError):
vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py
+2 -2
@@ -70,8 +70,8 @@ async def generation_error_handler(req: Request, exc: GenerationError):
70 """Handle GenerationError without logging stack traces.
71
72 GenerationError is a known, expected error (e.g. KV cache load failure)
73 - that should be returned to the client as a 500 response without polluting
74 - server logs with stack traces.
73 + that should be returned to the client without polluting server logs with
74 + stack traces.
75 """
76 err = create_error_response(exc)
77 return JSONResponse(err.model_dump(), status_code=err.error.code)
vllm/exceptions.py
+8
@@ -130,6 +130,14 @@ class GenerationError(VLLMServerError):
130 self.status_code = HTTPStatus.INTERNAL_SERVER_ERROR
131
132
133 +class RetryableRequestError(GenerationError):
134 + """A transient request-level failure that is safe for clients to retry."""
135 +
136 + def __init__(self, message: str = "Request failed; please retry"):
137 + super().__init__(message)
138 + self.status_code = HTTPStatus.SERVICE_UNAVAILABLE
139 +
140 +
141 class GracefulHTTPError(VLLMError):
142 """Exception that should be translated into an HTTP error response.
143
vllm/outputs.py
+37 -4
@@ -18,6 +18,15 @@ from vllm.v1.metrics.stats import RequestSpecDecodeMetrics, RequestStateStats
18 logger = init_logger(__name__)
19
20
21 +@dataclass(frozen=True)
22 +class RequestError:
23 + """A request-level error that does not mark the engine as unhealthy."""
24 +
25 + code: str
26 + message: str
27 + retryable: bool = False
28 +
29 +
30 @dataclass
31 class SamplingMask:
32 """Per-token sampling support sets aligned with completion token IDs.
@@ -259,6 +268,7 @@ class PoolingRequestOutput(Generic[_O]):
268 prompt_token_ids (list[int]): A list of token IDs used in the prompt.
269 num_cached_tokens: The number of tokens with prefix cache hit.
270 finished (bool): A flag indicating whether the pooling is completed.
271 + error: Structured request-level error information, if the request failed.
272
273 """
274
@@ -269,20 +279,25 @@ class PoolingRequestOutput(Generic[_O]):
279 prompt_token_ids: list[int],
280 num_cached_tokens: int,
281 finished: bool,
282 + *,
283 + error: RequestError | None = None,
284 ):
285 self.request_id = request_id
286 self.prompt_token_ids = prompt_token_ids
287 self.num_cached_tokens = num_cached_tokens
288 self.finished = finished
289 self.outputs = outputs
290 + self.error = error
291
292 def __repr__(self) -> str:
293 + error = f", error={self.error!r}" if self.error is not None else ""
294 return (
295 f"{type(self).__name__}(request_id={self.request_id!r}, "
296 f"outputs={self.outputs!r}, "
297 f"prompt_token_ids={self.prompt_token_ids}, "
298 f"num_cached_tokens={self.num_cached_tokens}, "
285 - f"finished={self.finished})"
299 + f"finished={self.finished}"
300 + f"{error})"
301 )
302
303
@@ -319,12 +334,18 @@ class EmbeddingRequestOutput(PoolingRequestOutput[EmbeddingOutput]):
334 def from_base(
335 request_output: PoolingRequestOutput,
336 ) -> "EmbeddingRequestOutput":
337 + outputs = (
338 + EmbeddingOutput([])
339 + if request_output.error is not None
340 + else EmbeddingOutput.from_base(request_output.outputs)
341 + )
342 return EmbeddingRequestOutput(
343 request_id=request_output.request_id,
324 - outputs=EmbeddingOutput.from_base(request_output.outputs),
344 + outputs=outputs,
345 prompt_token_ids=request_output.prompt_token_ids,
346 num_cached_tokens=request_output.num_cached_tokens,
347 finished=request_output.finished,
348 + error=request_output.error,
349 )
350
351
@@ -362,12 +383,18 @@ class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]):
383 def from_base(
384 request_output: PoolingRequestOutput,
385 ) -> "ClassificationRequestOutput":
386 + outputs = (
387 + ClassificationOutput([])
388 + if request_output.error is not None
389 + else ClassificationOutput.from_base(request_output.outputs)
390 + )
391 return ClassificationRequestOutput(
392 request_id=request_output.request_id,
367 - outputs=ClassificationOutput.from_base(request_output.outputs),
393 + outputs=outputs,
394 prompt_token_ids=request_output.prompt_token_ids,
395 num_cached_tokens=request_output.num_cached_tokens,
396 finished=request_output.finished,
397 + error=request_output.error,
398 )
399
400
@@ -402,10 +429,16 @@ class ScoringRequestOutput(PoolingRequestOutput[ScoringOutput]):
429 def from_base(
430 request_output: PoolingRequestOutput,
431 ) -> "ScoringRequestOutput":
432 + outputs = (
433 + ScoringOutput(0.0)
434 + if request_output.error is not None
435 + else ScoringOutput.from_base(request_output.outputs)
436 + )
437 return ScoringRequestOutput(
438 request_id=request_output.request_id,
407 - outputs=ScoringOutput.from_base(request_output.outputs),
439 + outputs=outputs,
440 prompt_token_ids=request_output.prompt_token_ids,
441 num_cached_tokens=request_output.num_cached_tokens,
442 finished=request_output.finished,
443 + error=request_output.error,
444 )
vllm/v1/engine/llm_engine.py
+8
@@ -323,6 +323,14 @@ class LLMEngine:
323 engine_core_timestamp=outputs.timestamp,
324 iteration_stats=iteration_stats,
325 )
326 +
327 + mm_processor_cache = self.renderer.mm_processor_cache
328 + if mm_processor_cache is not None:
329 + for engine_core_output in outputs.outputs:
330 + if engine_core_output.mm_cache_miss_hashes:
331 + for mm_hash in engine_core_output.mm_cache_miss_hashes:
332 + mm_processor_cache.invalidate(mm_hash)
333 +
334 self.output_processor.update_scheduler_stats(outputs.scheduler_stats)
335
336 # 3) Abort any reqs that finished due to stop strings.
vllm/v1/engine/output_processor.py
+19
@@ -16,6 +16,7 @@ from vllm.outputs import (
16 CompletionOutput,
17 PoolingOutput,
18 PoolingRequestOutput,
19 + RequestError,
20 RequestOutput,
21 SamplingMask,
22 )
@@ -304,6 +305,7 @@ class RequestState:
305 stop_reason: int | str | None,
306 kv_transfer_params: dict[str, Any] | None = None,
307 ec_transfer_params: dict[str, Any] | None = None,
308 + error: RequestError | None = None,
309 ) -> RequestOutput | PoolingRequestOutput | None:
310 finished = finish_reason is not None
311 final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
@@ -342,6 +344,7 @@ class RequestState:
344 external_req_id,
345 [self._new_pooling_output(pooling_output)],
346 finished,
347 + error=error,
348 )
349
350 output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
@@ -369,6 +372,7 @@ class RequestState:
372 finished: bool,
373 kv_transfer_params: dict[str, Any] | None = None,
374 ec_transfer_params: dict[str, Any] | None = None,
375 + error: RequestError | None = None,
376 ) -> RequestOutput | PoolingRequestOutput:
377 # If prompt embeds were used, put placeholder prompt token ids
378 prompt_token_ids = self.prompt_token_ids
@@ -385,6 +389,7 @@ class RequestState:
389 num_cached_tokens=self.num_cached_tokens,
390 prompt_token_ids=prompt_token_ids,
391 finished=finished,
392 + error=error,
393 )
394 assert self.logprobs_processor is not None
395 if self.output_kind == RequestOutputKind.DELTA:
@@ -715,6 +720,19 @@ class OutputProcessor:
720 if engine_core_output.spec_decode_metrics is not None:
721 req_state.spec_decode_metrics = engine_core_output.spec_decode_metrics
722
723 + request_error = None
724 + if (
725 + req_state.detokenizer is None
726 + and engine_core_output.mm_cache_miss_hashes
727 + and finish_reason == FinishReason.ERROR
728 + ):
729 + request_error = RequestError(
730 + code="multimodal_cache_miss",
731 + message="Multi-modal processor cache miss.",
732 + retryable=True,
733 + )
734 + pooling_output = EMPTY_CPU_TENSOR
735 +
736 if pooling_output is None:
737 assert req_state.detokenizer is not None
738 assert req_state.logprobs_processor is not None
@@ -742,6 +760,7 @@ class OutputProcessor:
760 stop_reason,
761 kv_transfer_params,
762 ec_transfer_params,
763 + error=request_error,
764 ):
765 if req_state.streaming_input:
766 request_output.finished = False
vllm/v1/pool/late_interaction_runner.py
+12
@@ -42,6 +42,18 @@ class LateInteractionRunner:
42 self._query_uses.clear()
43 self._doc_query_keys.clear()
44
45 + def release_queries(self, query_keys: Iterable[str]) -> None:
46 + query_keys = set(query_keys)
47 + for query_key in query_keys:
48 + self._query_cache.pop(query_key, None)
49 + self._query_uses.pop(query_key, None)
50 +
51 + self._doc_query_keys = {
52 + req_id: query_key
53 + for req_id, query_key in self._doc_query_keys.items()
54 + if query_key not in query_keys
55 + }
56 +
57 def register_request(
58 self, req_id: str, pooling_params: PoolingParams | None
59 ) -> None:
vllm/v1/worker/gpu/model_runner.py
+4
@@ -1012,6 +1012,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1012 if self.pooling_runner is not None:
1013 self.pooling_runner.clear()
1014
1015 + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None:
1016 + if self.pooling_runner is not None:
1017 + self.pooling_runner.release_late_interaction_query_cache(query_keys)
1018 +
1019 @torch.inference_mode()
1020 def profile_cudagraph_memory(self) -> int:
1021 """Estimate the GPU memory required to capture CUDA graphs."""
vllm/v1/worker/gpu/pool/pooling_runner.py
+3
@@ -98,6 +98,9 @@ class PoolingRunner:
98 def clear(self) -> None:
99 self.late_interaction_runner.clear()
100
101 + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None:
102 + self.late_interaction_runner.release_queries(query_keys)
103 +
104 def _get_pooling_metadata(
105 self,
106 input_batch: InputBatch,
vllm/v1/worker/gpu_model_runner.py
+3
@@ -993,6 +993,9 @@ class GPUModelRunner(
993 self.encoder_cache.clear()
994 self.late_interaction_runner.clear()
995
996 + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None:
997 + self.late_interaction_runner.release_queries(query_keys)
998 +
999 def _get_positions(self, num_tokens: Any):
1000 if isinstance(num_tokens, int):
1001 if self.uses_mrope:
vllm/v1/worker/worker_base.py
+7
@@ -143,6 +143,13 @@ class WorkerBase:
143 if callable(reset_fn):
144 reset_fn()
145
146 + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None:
147 + release_fn = getattr(
148 + self.model_runner, "release_late_interaction_query_cache", None
149 + )
150 + if callable(release_fn):
151 + release_fn(query_keys)
152 +
153 def get_model(self) -> nn.Module:
154 raise NotImplementedError
155