@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c6843d397c

[Bugfix][Frontend] Validate reused prompt token ids; render messages for media and echo (#55771)

Signed-off-by: Shijie Lyu <sjlyu@amazon.com> Co-authored-by: Shijie Lyu <sjlyu@amazon.com> Co-authored-by: Claude <noreply@anthropic.com>

Shijie Lyu committed Sep 28, 2026 at 19:36 UTC c6843d397cce339ee60ee134e484ce936fb5338d

6 files changed +239 -2

docs/features/disagg_prefill.md

@@ -77,6 +77,8 @@ decode = client.chat.completions.create(

77 )

78 ```

80 +If `messages` has non-text content or `echo` is set, the ids are ignored and `messages` is rendered instead, so it must match the prefill request. Otherwise `kv_transfer_params["prompt_token_ids"]` must be a non-empty list of non-negative integers, or the request fails with HTTP 400, as it always does on `/v1/chat/completions/batch`.

81 +

82 ## Development

84 We implement disaggregated prefilling by running 2 vLLM instances. One for prefill (we call it prefill instance) and one for decode (we call it decode instance), and then use a connector to transfer the prefill KV caches and results from prefill instance to decode instance.

tests/entrypoints/openai/chat_completion/test_batched_chat_completions.py

+17

@@ -26,6 +26,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (

26 BatchChatCompletionRequest,

27 )

28 from vllm.entrypoints.openai.models.serving import OpenAIServingModels

29 +from vllm.exceptions import VLLMValidationError

30 from vllm.outputs import CompletionOutput, RequestOutput

31 from vllm.renderers.online_renderer import OnlineRenderer

32 from vllm.v1.engine.async_llm import AsyncLLM

@@ -385,3 +386,19 @@ async def test_batched_harmony_response_format_uses_structural_tag() -> None:

386 assert len(calls) == 2

387 for call in calls:

388 assert call.args[1].structured_outputs.structural_tag is not None

389 +

390 +

391 +@pytest.mark.skip_global_cleanup

392 +def test_batch_rejects_kv_transfer_prompt_token_ids():

393 + """One pre-tokenized prompt cannot stand in for every conversation."""

394 + with pytest.raises(

395 + VLLMValidationError, match=r"parameter=kv_transfer_params\.prompt_token_ids"

396 + ):

397 + BatchChatCompletionRequest(

398 + model="test-model",

399 + messages=[

400 + [{"role": "user", "content": "first"}],

401 + [{"role": "user", "content": "second"}],

402 + ],

403 + kv_transfer_params={"prompt_token_ids": [10, 20, 30]},

404 + )

tests/entrypoints/openai/chat_completion/test_serving_chat.py

+138

@@ -595,6 +595,15 @@ def _build_online_renderer(

595 )

596

597

598 +def _build_mock_engine() -> MagicMock:

599 + mock_engine = MagicMock(spec=AsyncLLM)

600 + mock_engine.errored = False

601 + mock_engine.model_config = MockModelConfig()

602 + mock_engine.input_processor = MagicMock()

603 + mock_engine.renderer = _build_renderer(mock_engine.model_config)

604 + return mock_engine

605 +

606 +

607 def _build_serving_chat(

608 engine: AsyncLLM,

609 *,

@@ -2817,3 +2826,132 @@ def test_make_request_with_harmony_reuses_kv_transfer_prompt_token_ids():

2826 assert engine_input["prompt_token_ids"] == [10, 20, 30]

2827 # The reuse key is consumed and other kv_transfer_params are preserved.

2828 assert request.kv_transfer_params == {"do_remote_prefill": True}

2829 +

2830 +

2831 +@pytest.mark.parametrize("ids", [[-1], [1.5]])

2832 +def test_make_request_with_harmony_rejects_invalid_kv_transfer_prompt_token_ids(ids):

2833 + engine = MockEngine()

2834 + engine.model_config.hf_config = MockHFConfig(model_type="gpt_oss")

2835 + models = OpenAIServingModels(engine, BASE_MODEL_PATHS)

2836 + online_renderer = _build_online_renderer(engine, models.registry)

2837 +

2838 + request = ChatCompletionRequest(

2839 + model=MODEL_NAME,

2840 + messages=[{"role": "user", "content": "hi"}],

2841 + kv_transfer_params={"prompt_token_ids": ids},

2842 + )

2843 + with pytest.raises(VLLMValidationError, match="non-negative integers"):

2844 + online_renderer._make_request_with_harmony(request)

2845 +

2846 +

2847 +@pytest.mark.asyncio

2848 +@pytest.mark.parametrize("ids", [[], [-1], [1.5], [1.0], ["1"], [True], "abc", 5])

2849 +async def test_chat_kv_transfer_prompt_token_ids_rejects_invalid_ids(ids):

2850 + """``kv_transfer_params`` is untyped, so the ids are checked on reuse."""

2851 + serving_chat = _build_serving_chat(_build_mock_engine())

2852 +

2853 + request = ChatCompletionRequest(

2854 + model=MODEL_NAME,

2855 + messages=[{"role": "user", "content": "hi"}],

2856 + kv_transfer_params={"prompt_token_ids": ids},

2857 + )

2858 + with pytest.raises(VLLMValidationError, match="non-negative integers"):

2859 + await serving_chat.render_chat_request(request)

2860 +

2861 +

2862 +@pytest.mark.asyncio

2863 +async def test_chat_kv_transfer_prompt_token_ids_rejection_notifies_kv_connector():

2864 + """A render-time 400 on a decode request notifies the KV connector."""

2865 + mock_engine = _build_mock_engine()

2866 + mock_engine.notify_kv_transfer_request_rejected = AsyncMock()

2867 + serving_chat = _build_serving_chat(mock_engine)

2868 + serving_chat.has_kv_connector = True

2869 +

2870 + request = ChatCompletionRequest(

2871 + model=MODEL_NAME,

2872 + messages=[{"role": "user", "content": "hi"}],

2873 + kv_transfer_params={"prompt_token_ids": [1.5], "do_remote_prefill": True},

2874 + )

2875 + with pytest.raises(VLLMValidationError, match="non-negative integers"):

2876 + await serving_chat.create_chat_completion(request)

2877 +

2878 + mock_engine.notify_kv_transfer_request_rejected.assert_awaited_once_with(

2879 + request.request_id, {"do_remote_prefill": True}, data_parallel_rank=None

2880 + )

2881 +

2882 +

2883 +@pytest.mark.asyncio

2884 +@pytest.mark.parametrize("ids", [[10, 20, 30], [1.5]])

2885 +async def test_chat_kv_transfer_prompt_token_ids_ignored_with_echo(ids):

2886 + """``echo`` needs the rendered conversation, so the ids are not used."""

2887 + serving_chat = _build_serving_chat(_build_mock_engine())

2888 +

2889 + request = ChatCompletionRequest(

2890 + model=MODEL_NAME,

2891 + messages=[{"role": "user", "content": "hi"}],

2892 + kv_transfer_params={"prompt_token_ids": ids, "do_remote_prefill": True},

2893 + echo=True,

2894 + )

2895 + result = await serving_chat.render_chat_request(request)

2896 + assert not isinstance(result, ErrorResponse)

2897 +

2898 + conversation, engine_inputs = result

2899 + assert [msg["role"] for msg in conversation] == ["user"]

2900 + assert engine_inputs[0]["prompt_token_ids"] != ids

2901 + assert request.kv_transfer_params == {"do_remote_prefill": True}

2902 +

2903 +

2904 +@pytest.mark.parametrize(

2905 + "content",

2906 + [

2907 + [

2908 + {"type": "text", "text": "what is in this image?"},

2909 + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},

2910 + ],

2911 + # A part written without a "type" is identified by its media key.

2912 + [{"image_url": "https://example.com/a.png"}],

2913 + # A media key counts even when the part claims to be text, which is

2914 + # how chat_utils reads a part carrying a "uuid".

2915 + [

2916 + {

2917 + "type": "text",

2918 + "text": "look",

2919 + "uuid": "u1",

2920 + "image_url": "https://example.com/a.png",

2921 + }

2922 + ],

2923 + [{"type": "hologram", "hologram": "https://a/b.holo"}],

2924 + [{"type": ["image_url"]}],

2925 + ],

2926 +)

2927 +def test_chat_kv_transfer_prompt_token_ids_ignored_with_non_text_content(content):

2928 + """The ids are dropped so that ``messages`` is rendered with its media."""

2929 + request = ChatCompletionRequest(

2930 + model=MODEL_NAME,

2931 + messages=[{"role": "user", "content": content}],

2932 + kv_transfer_params={

2933 + "prompt_token_ids": [10, 20, 30],

2934 + "do_remote_prefill": True,

2935 + },

2936 + )

2937 + assert request.kv_transfer_params == {"do_remote_prefill": True}

2938 +

2939 +

2940 +def test_chat_kv_transfer_prompt_token_ids_allows_text_only_parts():

2941 + """Text-bearing part types are not multimodal, so the ids are kept."""

2942 + request = ChatCompletionRequest(

2943 + model=MODEL_NAME,

2944 + messages=[

2945 + {"role": "user", "content": [{"type": "text", "text": "hi"}]},

2946 + {

2947 + "role": "assistant",

2948 + "content": [

2949 + {"type": "thinking", "thinking": "let me think"},

2950 + {"type": "refusal", "refusal": "no"},

2951 + ],

2952 + },

2953 + {"role": "user", "content": "plain string"},

2954 + ],

2955 + kv_transfer_params={"prompt_token_ids": [10, 20, 30]},

2956 + )

2957 + assert request.kv_transfer_params == {"prompt_token_ids": [10, 20, 30]}

vllm/entrypoints/chat_utils.py

+6 -1

@@ -1810,6 +1810,11 @@ PART_TYPES_TO_SKIP_NONE_CONTENT = (

1810 "refusal",

1811 )

1812

1813 +# Content part types parsed as text rather than multimodal data.

1814 +TEXT_PART_TYPES = frozenset(

1815 + {"text", "input_text", "output_text", "refusal", "thinking"}

1816 +)

1817 +

1818

1819 def _parse_chat_message_content_parts(

1820 role: str,

@@ -1904,7 +1909,7 @@ def _parse_chat_message_content_part(

1909 )

1910 return None

1911

1907 - if part_type in ("text", "input_text", "output_text", "refusal", "thinking"):

1912 + if part_type in TEXT_PART_TYPES:

1913 str_content = cast(str, content)

1914 _reject_reserved_placeholder_in_text(str_content, mm_parser.model_config)

1915 if wrap_dicts:

vllm/entrypoints/openai/chat_completion/protocol.py

+51

@@ -20,6 +20,8 @@ from pydantic import (

21 from vllm.config import ModelConfig

22 from vllm.entrypoints.chat_utils import (

23 + MM_PARSER_MAP,

24 + TEXT_PART_TYPES,

25 ChatCompletionMessageParam,

26 ChatTemplateContentFormatOption,

27 )

@@ -58,6 +60,11 @@ logger = init_logger(__name__)

60 _INT64_MIN = -(2**63)

61 _INT64_MAX = 2**63 - 1

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 +

69 class ChatMessage(OpenAIBaseModel):

70 role: str

@@ -976,6 +983,39 @@ class ChatCompletionRequest(OpenAIBaseModel):

983 )

984 return data

985

986 + @model_validator(mode="before")

987 + @classmethod

988 + def drop_prompt_token_ids_with_media(cls, data):

989 + # The forwarded ids would drop media in ``messages``, so ignore them. Runs

990 + # before validation, which can turn content into a one-shot iterator.

991 + if not isinstance(data, dict):

992 + return data

993 + kv_transfer_params = data.get("kv_transfer_params")

994 + messages = data.get("messages")

995 + 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)

999 + ):

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

1017 + return data

1018 +

1019 @model_validator(mode="before")

1020 @classmethod

1021 def check_system_message_content_type(cls, data):

@@ -1117,6 +1157,17 @@ class BatchChatCompletionRequest(OpenAIBaseModel):

1157 "when using `logprob_token_ids`, `logprobs` must be set to true.",

1158 parameter="logprob_token_ids",

1159 )

1160 + kv_transfer_params = data.get("kv_transfer_params")

1161 + if (

1162 + isinstance(kv_transfer_params, dict)

1163 + and kv_transfer_params.get("prompt_token_ids") is not None

1164 + ):

1165 + raise VLLMValidationError(

1166 + "Batch chat completions do not support "

1167 + "`kv_transfer_params['prompt_token_ids']`: one pre-tokenized "

1168 + "prompt cannot serve several conversations.",

1169 + parameter="kv_transfer_params.prompt_token_ids",

1170 + )

1171 response_format = data.get("response_format")

1172 if response_format is not None:

1173 rf_type = (

vllm/renderers/online_renderer.py

+25 -1

@@ -107,11 +107,35 @@ def _reused_prompt_token_ids(request: Any) -> list[int] | None:

107 Disaggregated serving carries the prefill stage's ids in

108 ``kv_transfer_params`` so the decode stage can skip re-tokenizing. Removing

109 the key keeps the id list out of the engine's sampling metadata.

110 +

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 """

115 kv = getattr(request, "kv_transfer_params", None)

116 if not isinstance(kv, dict):

117 return None

114 - return kv.pop("prompt_token_ids", None) or None

118 + ids = kv.pop("prompt_token_ids", None)

119 + if ids is None:

120 + return None

121 + if getattr(request, "echo", False):

122 + logger.debug(

123 + "Ignoring kv_transfer_params['prompt_token_ids']: "

124 + "echo is set, so messages are rendered instead."

125 + )

126 + return None

127 + # bool is an int subclass, hence the exact type check.

128 + if (

129 + not isinstance(ids, list)

130 + or not ids

131 + or any(type(x) is not int or x < 0 for x in ids)

132 + ):

133 + raise VLLMValidationError(

134 + "`kv_transfer_params['prompt_token_ids']` must be a non-empty list "

135 + "of non-negative integers.",

136 + parameter="kv_transfer_params.prompt_token_ids",

137 + )

138 + return ids

139

140

141 class OnlineRenderer: