[Bugfix][Frontend] Ignore reused prompt token ids for media in /v1/responses (#59173)
Signed-off-by: Shijie Lyu <lshjhf@gmail.com>
Shijie Lyu committed
Sep 29, 2026 at 22:21 UTC
ac68c3087215e0a4f3cdfa218508c6aada57235d
4 files changed
+146
-33
tests/entrypoints/openai/responses/test_serving_responses.py
+97
@@ -34,6 +34,7 @@ from openai_harmony import Message as OpenAIHarmonyMessage
34
from openai_harmony import Role
35
36
import vllm.envs as envs
37
+from tests.entrypoints.openai.chat_completion.test_serving_chat import MockModelConfig
38
from vllm.entrypoints.generate.base.protocol import (
39
DeltaFunctionCall,
40
DeltaMessage,
@@ -294,6 +295,102 @@ async def test_online_renderer_renders_non_harmony_responses_with_explicit_histo
295
assert preprocess_kwargs["default_template_kwargs"]["request_default"] == "kept"
296
297
298
+_REUSED_IDS = [10, 20, 30]
299
+_IMAGE_PART = {
300
+ "type": "input_image",
301
+ "image_url": "https://example.com/a.png",
302
+ "detail": "auto",
303
+}
304
+
305
+
306
+async def _render_responses_with_reuse(request, previous_messages=None):
307
+ online_renderer = OnlineRenderer(
308
+ model_config=MockModelConfig(),
309
+ renderer=MagicMock(),
310
+ request_logger=None,
311
+ chat_template=None,
312
+ chat_template_content_format="auto",
313
+ )
314
+ render_chat = AsyncMock(return_value=([[]], [tokens_input([1, 2])]))
315
+ online_renderer.renderer.render_chat_async = render_chat
316
+
317
+ result = await online_renderer.render_responses(
318
+ request, previous_messages=previous_messages
319
+ )
320
+
321
+ assert not isinstance(result, ErrorResponse)
322
+ return render_chat, result.engine_input
323
+
324
+
325
+@pytest.mark.asyncio
326
+@pytest.mark.parametrize("ids", [_REUSED_IDS, [1.5]])
327
+@pytest.mark.parametrize(
328
+ "request_input",
329
+ [
330
+ [
331
+ {
332
+ "role": "user",
333
+ "content": [{"type": "input_text", "text": "look"}, _IMAGE_PART],
334
+ }
335
+ ],
336
+ [
337
+ {
338
+ "type": "function_call",
339
+ "call_id": "call_1",
340
+ "name": "shot",
341
+ "arguments": "{}",
342
+ },
343
+ {
344
+ "type": "function_call_output",
345
+ "call_id": "call_1",
346
+ "output": [_IMAGE_PART],
347
+ },
348
+ ],
349
+ ],
350
+ ids=["input_image", "tool_output_image"],
351
+)
352
+async def test_responses_kv_transfer_prompt_token_ids_ignored_with_non_text_input(
353
+ request_input, ids
354
+):
355
+ request = ResponsesRequest(
356
+ input=request_input,
357
+ kv_transfer_params={"do_remote_prefill": True, "prompt_token_ids": ids},
358
+ )
359
+
360
+ render_chat, _ = await _render_responses_with_reuse(request)
361
+
362
+ render_chat.assert_awaited_once()
363
+ assert request.kv_transfer_params == {"do_remote_prefill": True}
364
+
365
+
366
+@pytest.mark.asyncio
367
+async def test_responses_kv_transfer_prompt_token_ids_ignored_with_media_in_history():
368
+ request = ResponsesRequest(
369
+ input="next", kv_transfer_params={"prompt_token_ids": _REUSED_IDS}
370
+ )
371
+
372
+ render_chat, _ = await _render_responses_with_reuse(
373
+ request, previous_messages=[{"role": "user", "content": [_IMAGE_PART]}]
374
+ )
375
+
376
+ render_chat.assert_awaited_once()
377
+
378
+
379
+@pytest.mark.asyncio
380
+async def test_responses_kv_transfer_prompt_token_ids_kept_with_text_input():
381
+ request = ResponsesRequest(
382
+ input=[
383
+ {"role": "user", "content": [{"type": "input_text", "text": "hi"}]},
384
+ {"role": "assistant", "content": "ok"},
385
+ ],
386
+ kv_transfer_params={"prompt_token_ids": _REUSED_IDS},
387
+ )
388
+
389
+ _, engine_input = await _render_responses_with_reuse(request)
390
+
391
+ assert engine_input["prompt_token_ids"] == _REUSED_IDS
392
+
393
+
394
@pytest.mark.asyncio
395
async def test_online_renderer_rejects_untrusted_responses_chat_template():
396
renderer = _new_online_renderer()
vllm/entrypoints/chat_utils.py
+26
@@ -1814,6 +1814,32 @@ PART_TYPES_TO_SKIP_NONE_CONTENT = (
1814
TEXT_PART_TYPES = frozenset(
1815
{"text", "input_text", "output_text", "refusal", "thinking"}
1816
)
1817
+# Content part types that carry no multimodal data.
1818
+_TEXT_CONTENT_PART_TYPES = TEXT_PART_TYPES | {"tool_reference"}
1819
+# Keys that mark a content part as multimodal, whatever its ``type``.
1820
+_MEDIA_CONTENT_PART_KEYS = frozenset(MM_PARSER_MAP) - _TEXT_CONTENT_PART_TYPES
1821
+
1822
+
1823
+def has_non_text_content(messages: Any) -> bool:
1824
+ """Whether any message in ``messages`` has a non-text content part.
1825
+
1826
+ Only list content is inspected, so validated chat content that is a
1827
+ one-shot iterator is never consumed.
1828
+ """
1829
+ if not isinstance(messages, list):
1830
+ return False
1831
+ for msg in messages:
1832
+ content = msg.get("content") if isinstance(msg, dict) else None
1833
+ if not isinstance(content, list):
1834
+ continue
1835
+ for part in content:
1836
+ if isinstance(part, dict) and (
1837
+ any(key in part for key in _MEDIA_CONTENT_PART_KEYS)
1838
+ or not isinstance(part_type := part.get("type", "text"), str)
1839
+ or part_type not in _TEXT_CONTENT_PART_TYPES
1840
+ ):
1841
+ return True
1842
+ return False
1843
1844
1845
def _parse_chat_message_content_parts(
vllm/entrypoints/openai/chat_completion/protocol.py
+9
-28
@@ -20,10 +20,9 @@ from pydantic import (
20
21
from vllm.config import ModelConfig
22
from vllm.entrypoints.chat_utils import (
23
- MM_PARSER_MAP,
24
- TEXT_PART_TYPES,
23
ChatCompletionMessageParam,
24
ChatTemplateContentFormatOption,
25
+ has_non_text_content,
26
)
27
from vllm.entrypoints.generate.base.protocol import (
28
AnyResponseFormat,
@@ -60,11 +59,6 @@ logger = init_logger(__name__)
59
_INT64_MIN = -(2**63)
60
_INT64_MAX = 2**63 - 1
61
63
-# Content part types that carry no multimodal data.
64
-_TEXT_CONTENT_PART_TYPES = TEXT_PART_TYPES | {"tool_reference"}
65
-# Keys that mark a content part as multimodal, whatever its ``type``.
66
-_MEDIA_CONTENT_PART_KEYS = frozenset(MM_PARSER_MAP) - _TEXT_CONTENT_PART_TYPES
67
-
62
63
class ChatMessage(OpenAIBaseModel):
64
role: str
@@ -991,29 +985,16 @@ class ChatCompletionRequest(OpenAIBaseModel):
985
if not isinstance(data, dict):
986
return data
987
kv_transfer_params = data.get("kv_transfer_params")
994
- messages = data.get("messages")
988
if (
996
- not isinstance(kv_transfer_params, dict)
997
- or kv_transfer_params.get("prompt_token_ids") is None
998
- or not isinstance(messages, list)
989
+ isinstance(kv_transfer_params, dict)
990
+ and kv_transfer_params.get("prompt_token_ids") is not None
991
+ and has_non_text_content(data.get("messages"))
992
):
1000
- return data
1001
- for msg in messages:
1002
- content = msg.get("content") if isinstance(msg, dict) else None
1003
- if not isinstance(content, list):
1004
- continue
1005
- for part in content:
1006
- if isinstance(part, dict) and (
1007
- any(key in part for key in _MEDIA_CONTENT_PART_KEYS)
1008
- or not isinstance(part_type := part.get("type", "text"), str)
1009
- or part_type not in _TEXT_CONTENT_PART_TYPES
1010
- ):
1011
- logger.debug(
1012
- "Ignoring kv_transfer_params['prompt_token_ids']: "
1013
- "messages have non-text content and are rendered instead."
1014
- )
1015
- kv_transfer_params.pop("prompt_token_ids")
1016
- return data
993
+ logger.debug(
994
+ "Ignoring kv_transfer_params['prompt_token_ids']: "
995
+ "messages have non-text content and are rendered instead."
996
+ )
997
+ kv_transfer_params.pop("prompt_token_ids")
998
return data
999
1000
@model_validator(mode="before")
vllm/renderers/online_renderer.py
+14
-5
@@ -16,6 +16,7 @@ from vllm.entrypoints.chat_utils import (
16
ChatCompletionMessageParam,
17
ChatTemplateContentFormatOption,
18
ConversationMessage,
19
+ has_non_text_content,
20
)
21
from vllm.entrypoints.generate.base.protocol import validate_request_mm_kwargs
22
from vllm.entrypoints.openai.chat_completion.protocol import (
@@ -101,16 +102,18 @@ def _extract_allowed_tools_from_mcp_requests(
102
return allowed_tools_map
103
104
104
-def _reused_prompt_token_ids(request: Any) -> list[int] | None:
105
+def _reused_prompt_token_ids(
106
+ request: Any, messages: list[Any] | None = None
107
+) -> list[int] | None:
108
"""Pop prompt token ids forwarded for decode-side reuse, if any.
109
110
Disaggregated serving carries the prefill stage's ids in
111
``kv_transfer_params`` so the decode stage can skip re-tokenizing. Removing
112
the key keeps the id list out of the engine's sampling metadata.
113
111
- Returns None without checking the ids when ``echo`` is set, since echo
112
- needs ``messages`` to be rendered. Otherwise raises VLLMValidationError if
113
- the ids are malformed.
114
+ Returns None without checking the ids when ``echo`` is set or ``messages``
115
+ has non-text content, since both need ``messages`` to be rendered.
116
+ Otherwise raises VLLMValidationError if the ids are malformed.
117
"""
118
kv = getattr(request, "kv_transfer_params", None)
119
if not isinstance(kv, dict):
@@ -124,6 +127,12 @@ def _reused_prompt_token_ids(request: Any) -> list[int] | None:
127
"echo is set, so messages are rendered instead."
128
)
129
return None
130
+ if has_non_text_content(messages):
131
+ logger.debug(
132
+ "Ignoring kv_transfer_params['prompt_token_ids']: "
133
+ "messages have non-text content and are rendered instead."
134
+ )
135
+ return None
136
# bool is an int subclass, hence the exact type check.
137
if (
138
not isinstance(ids, list)
@@ -772,7 +781,7 @@ class OnlineRenderer:
781
default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
782
)
783
775
- reuse_ids = _reused_prompt_token_ids(request)
784
+ reuse_ids = _reused_prompt_token_ids(request, messages)
785
if reuse_ids:
786
# Decode-side token reuse: feed the forwarded ids straight to the
787
# engine, skipping templating and tokenization. ``messages`` are not