@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 52358e6e19

[Frontend] Expose multimodal metadata for disaggregated prefill (#54659)

Signed-off-by: zhouyou9505 <zhouyou9505@gmail.com> Co-authored-by: Cursor Grok 4.6 <noreply@cursor.com>

zhouyou committed Sep 6, 2026 at 22:46 UTC 52358e6e192aeae73dd3764046c98fc4b156ea83
7 files changed +605 -100
docs/serving/online_serving/renderer.md
+144
@@ -27,3 +27,147 @@ VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 vllm serve <model>
27 - Render chat completions
28
29 For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md).
30 +
31 +## Multimodal Render Features
32 +
33 +Multimodal render responses include a `features` object with per-modality
34 +hashes, placeholder ranges, and serialized processor data. When the model
35 +exposes placeholder-metadata or `keep_on_cpu` fields (for example
36 +`image_grid_thw`), the response also includes `mm_metadata`. Each
37 +`mm_metadata` entry is a base64-encoded `MultiModalKwargsItem` containing
38 +only those fields, not encoder inputs such as `pixel_values`.
39 +
40 +The arrays in `mm_hashes`, `mm_placeholders`, `kwargs_data`, and
41 +`mm_metadata` use the same per-modality item order. Downstream workers
42 +should split those fields:
43 +
44 +- Encode requests keep `kwargs_data`.
45 +- Prefill requests may omit `kwargs_data` and send `mm_metadata` only when
46 + `ec_transfer_params` is also set, so embeddings are loaded by the EC
47 + connector. Omitting `kwargs_data` without `ec_transfer_params` is rejected.
48 +- Legacy clients that ignore `mm_metadata` and keep sending `kwargs_data`
49 + continue to work.
50 +
51 +## Example
52 +
53 +The example below shows how a disaggregated encode / prefill coordinator can
54 +split a multimodal render response. The render step returns both
55 +`kwargs_data` (encoder tensors plus metadata) and `mm_metadata` (metadata
56 +only). Encode keeps the full payload; prefill drops `kwargs_data` after the
57 +EC connector has published embeddings.
58 +
59 +```python
60 +import httpx
61 +
62 +MODEL = "Qwen/Qwen3-VL-2B-Instruct"
63 +RENDER = "http://localhost:8100" # vllm launch render ...
64 +ENCODE = "http://localhost:8200" # encode worker
65 +PREFILL = "http://localhost:8300" # prefill worker
66 +
67 +chat_request = {
68 + "model": MODEL,
69 + "messages": [
70 + {
71 + "role": "user",
72 + "content": [
73 + {"type": "image_url", "image_url": {"url": "<data-url>"}},
74 + {"type": "text", "text": "Describe this image."},
75 + ],
76 + }
77 + ],
78 +}
79 +
80 +with httpx.Client(timeout=120.0) as client:
81 + # 1. Render: preprocess into token IDs and multimodal features.
82 + render_response = client.post(
83 + f"{RENDER}/v1/chat/completions/render", json=chat_request
84 + ).json()
85 +
86 + features = render_response["features"]
87 + # features["kwargs_data"]["image"][0] -> pixel_values + image_grid_thw
88 + # features["mm_metadata"]["image"][0] -> image_grid_thw only
89 +
90 + # 2. Encode: send full kwargs_data so the encoder can run vision towers.
91 + encode_response = client.post(
92 + f"{ENCODE}/inference/v1/generate",
93 + json={
94 + "token_ids": render_response["token_ids"],
95 + "features": {
96 + "mm_hashes": features["mm_hashes"],
97 + "mm_placeholders": features["mm_placeholders"],
98 + "kwargs_data": features["kwargs_data"],
99 + },
100 + "sampling_params": {"max_tokens": 1},
101 + },
102 + ).json()
103 + ec_transfer_params = encode_response["ec_transfer_params"]
104 +
105 + # 3. Prefill: omit kwargs_data; load embeddings via EC connector.
106 + prefill_response = client.post(
107 + f"{PREFILL}/inference/v1/generate",
108 + json={
109 + "token_ids": render_response["token_ids"],
110 + "features": {
111 + "mm_hashes": features["mm_hashes"],
112 + "mm_placeholders": features["mm_placeholders"],
113 + "mm_metadata": features["mm_metadata"],
114 + },
115 + "ec_transfer_params": ec_transfer_params,
116 + "sampling_params": {"max_tokens": 64},
117 + },
118 + ).json()
119 +
120 +print(prefill_response["choices"][0]["token_ids"])
121 +```
122 +
123 +Single-process clients can keep passing the full render response to
124 +`/inference/v1/generate` unchanged; `mm_metadata` is optional and ignored
125 +when `kwargs_data` is present.
126 +
127 +### Payload shape
128 +
129 +#### Render response
130 +
131 +`/v1/chat/completions/render` returns both `kwargs_data` and `mm_metadata`.
132 +The arrays share the same per-modality item order. Base64 blobs are truncated
133 +below for readability.
134 +
135 +```json
136 +{
137 + "token_ids": [151644, 872],
138 + "features": {
139 + "mm_hashes": {"image": ["abc123..."]},
140 + "mm_placeholders": {"image": [{"offset": 0, "length": 256}]},
141 + "kwargs_data": {
142 + "image": ["<base64 MultiModalKwargsItem: pixel_values + image_grid_thw>"]
143 + },
144 + "mm_metadata": {
145 + "image": ["<base64 MultiModalKwargsItem: image_grid_thw only>"]
146 + }
147 + }
148 +}
149 +```
150 +
151 +Forward `kwargs_data` to the encode worker. Keep `mm_metadata` for prefill.
152 +
153 +#### Prefill request
154 +
155 +Prefill omits `kwargs_data` and sends `mm_metadata` with
156 +`ec_transfer_params` from the encode response:
157 +
158 +```json
159 +{
160 + "token_ids": [151644, 872],
161 + "features": {
162 + "mm_hashes": {"image": ["abc123..."]},
163 + "mm_placeholders": {"image": [{"offset": 0, "length": 256}]},
164 + "mm_metadata": {
165 + "image": ["<base64 MultiModalKwargsItem: image_grid_thw only>"]
166 + }
167 + },
168 + "ec_transfer_params": {
169 + "ec_items": [{"mm_hash": "abc123...", "peer_host": "10.0.0.1"}]
170 + },
171 + "sampling_params": {"max_tokens": 64}
172 +}
173 +```
tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py
+214
@@ -6,21 +6,30 @@ token_in_token_out generate endpoint.
6 """
7
8 import torch
9 +from pydantic import ValidationError
10
11 +from vllm.entrypoints.scale_out.token_in_token_out.mm_features import (
12 + extract_mm_features,
13 + merge_mm_kwargs_items,
14 + mm_kwargs_from_features,
15 +)
16 from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import (
17 decode_mm_kwargs_item,
18 encode_mm_kwargs_item,
19 )
20 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
21 + GenerateRequest,
22 MultiModalFeatures,
23 PlaceholderRangeInfo,
24 )
25 +from vllm.inputs import mm_input
26 from vllm.multimodal.inputs import (
27 MultiModalBatchedField,
28 MultiModalFieldElem,
29 MultiModalFlatField,
30 MultiModalKwargsItem,
31 MultiModalSharedField,
32 + PlaceholderRange,
33 )
34
35
@@ -112,3 +121,208 @@ def test_mm_features_with_kwargs_data():
121
122 decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
123 assert torch.equal(elem.data, decoded["pixel_values"].data)
124 +
125 +
126 +def _image_engine_input(
127 + *,
128 + pixel_keep_on_cpu: bool = False,
129 + grid_keep_on_cpu: bool = True,
130 +) -> tuple[object, MultiModalFieldElem, MultiModalFieldElem]:
131 + pixel_values = MultiModalFieldElem(
132 + data=torch.randn(5, 3, dtype=torch.float32),
133 + field=MultiModalBatchedField(keep_on_cpu=pixel_keep_on_cpu),
134 + )
135 + image_grid_thw = MultiModalFieldElem(
136 + data=torch.tensor([[1, 24, 24]], dtype=torch.int32),
137 + field=MultiModalBatchedField(keep_on_cpu=grid_keep_on_cpu),
138 + )
139 + item = MultiModalKwargsItem(
140 + {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}
141 + )
142 + engine_input = mm_input(
143 + prompt_token_ids=[1, 2, 3],
144 + mm_kwargs={"image": [item]},
145 + mm_hashes={"image": ["abc123"]},
146 + mm_placeholders={"image": [PlaceholderRange(offset=0, length=1)]},
147 + )
148 + return engine_input, pixel_values, image_grid_thw
149 +
150 +
151 +def test_render_extracts_metadata_fields_separately():
152 + """Render metadata contains CPU fields but excludes encoder inputs."""
153 + engine_input, pixel_values, image_grid_thw = _image_engine_input()
154 +
155 + features = extract_mm_features(engine_input)
156 +
157 + assert features is not None
158 + assert features.mm_metadata is not None
159 + features = MultiModalFeatures.model_validate_json(features.model_dump_json())
160 + assert features.mm_metadata is not None
161 + metadata = decode_mm_kwargs_item(features.mm_metadata["image"][0])
162 + assert set(metadata) == {"image_grid_thw"}
163 + assert torch.equal(metadata["image_grid_thw"].data, image_grid_thw.data)
164 +
165 + assert features.kwargs_data is not None
166 + full_data = decode_mm_kwargs_item(features.kwargs_data["image"][0])
167 + assert set(full_data) == {"pixel_values", "image_grid_thw"}
168 + assert torch.equal(full_data["pixel_values"].data, pixel_values.data)
169 +
170 +
171 +def test_extract_includes_declared_placeholder_metadata_fields():
172 + """EC placeholder metadata is kept even when keep_on_cpu is unset."""
173 + engine_input, _pixel_values, image_grid_thw = _image_engine_input(
174 + grid_keep_on_cpu=False
175 + )
176 +
177 + features = extract_mm_features(
178 + engine_input,
179 + metadata_fields_for=lambda modality: (
180 + {"image_grid_thw"} if modality == "image" else set()
181 + ),
182 + )
183 +
184 + assert features is not None
185 + assert features.mm_metadata is not None
186 + metadata = decode_mm_kwargs_item(features.mm_metadata["image"][0])
187 + assert set(metadata) == {"image_grid_thw"}
188 + assert torch.equal(metadata["image_grid_thw"].data, image_grid_thw.data)
189 +
190 +
191 +def test_legacy_kwargs_only_generate_keeps_full_payload():
192 + """Old clients that omit mm_metadata still reconstruct full kwargs."""
193 + engine_input, pixel_values, image_grid_thw = _image_engine_input()
194 + rendered = extract_mm_features(engine_input)
195 + assert rendered is not None
196 + assert rendered.kwargs_data is not None
197 +
198 + features = MultiModalFeatures(
199 + mm_hashes=rendered.mm_hashes,
200 + mm_placeholders=rendered.mm_placeholders,
201 + kwargs_data=rendered.kwargs_data,
202 + )
203 + merged = mm_kwargs_from_features(features)
204 + item = merged["image"][0]
205 + assert item is not None
206 + assert set(item) == {"pixel_values", "image_grid_thw"}
207 + assert torch.equal(item["pixel_values"].data, pixel_values.data)
208 + assert torch.equal(item["image_grid_thw"].data, image_grid_thw.data)
209 +
210 +
211 +def test_metadata_only_generate_requires_ec_transfer_params():
212 + engine_input, _, _ = _image_engine_input()
213 + rendered = extract_mm_features(engine_input)
214 + assert rendered is not None
215 + payload = {
216 + "token_ids": [1, 2, 3],
217 + "sampling_params": {"max_tokens": 1},
218 + "features": {
219 + "mm_hashes": rendered.mm_hashes,
220 + "mm_placeholders": {
221 + modality: [p.model_dump() for p in ranges]
222 + for modality, ranges in rendered.mm_placeholders.items()
223 + },
224 + "mm_metadata": rendered.mm_metadata,
225 + },
226 + }
227 +
228 + try:
229 + GenerateRequest.model_validate(payload)
230 + raise AssertionError("expected metadata-only request without EC to fail")
231 + except ValidationError as exc:
232 + assert "ec_transfer_params" in str(exc)
233 +
234 + payload["ec_transfer_params"] = {"image-hash": {"peer_host": "10.0.0.1"}}
235 + request = GenerateRequest.model_validate(payload)
236 + assert request.features is not None
237 + assert request.features.kwargs_data is None
238 + merged = mm_kwargs_from_features(request.features)
239 + item = merged["image"][0]
240 + assert item is not None
241 + assert set(item) == {"image_grid_thw"}
242 +
243 +
244 +def test_kwargs_and_metadata_generate_does_not_require_ec():
245 + engine_input, pixel_values, image_grid_thw = _image_engine_input()
246 + rendered = extract_mm_features(engine_input)
247 + assert rendered is not None
248 + request = GenerateRequest.model_validate(
249 + {
250 + "token_ids": [1, 2, 3],
251 + "sampling_params": {"max_tokens": 1},
252 + "features": rendered.model_dump(),
253 + }
254 + )
255 + assert request.features is not None
256 + merged = mm_kwargs_from_features(request.features)
257 + item = merged["image"][0]
258 + assert item is not None
259 + assert set(item) == {"pixel_values", "image_grid_thw"}
260 + assert torch.equal(item["pixel_values"].data, pixel_values.data)
261 + assert torch.equal(item["image_grid_thw"].data, image_grid_thw.data)
262 +
263 +
264 +def test_cache_hit_null_kwargs_keeps_metadata_slot():
265 + pixel_values = MultiModalFieldElem(
266 + data=torch.randn(5, 3, dtype=torch.float32),
267 + field=MultiModalBatchedField(),
268 + )
269 + image_grid_thw = MultiModalFieldElem(
270 + data=torch.tensor([[1, 24, 24]], dtype=torch.int32),
271 + field=MultiModalBatchedField(keep_on_cpu=True),
272 + )
273 + item = MultiModalKwargsItem(
274 + {"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}
275 + )
276 + engine_input = mm_input(
277 + prompt_token_ids=[1, 2, 3, 4],
278 + mm_kwargs={"image": [item, None]},
279 + mm_hashes={"image": ["hit", "cached"]},
280 + mm_placeholders={
281 + "image": [
282 + PlaceholderRange(offset=0, length=1),
283 + PlaceholderRange(offset=2, length=1),
284 + ]
285 + },
286 + )
287 +
288 + features = extract_mm_features(engine_input)
289 + assert features is not None
290 + assert features.kwargs_data is not None
291 + assert features.mm_metadata is not None
292 + assert features.kwargs_data["image"][1] is None
293 + assert features.mm_metadata["image"][1] is None
294 + assert len(features.mm_metadata["image"]) == 2
295 + metadata = decode_mm_kwargs_item(features.mm_metadata["image"][0])
296 + assert torch.equal(metadata["image_grid_thw"].data, image_grid_thw.data)
297 +
298 +
299 +def test_features_alignment_rejects_length_mismatch():
300 + try:
301 + MultiModalFeatures(
302 + mm_hashes={"image": ["a", "b"]},
303 + mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=1)]},
304 + )
305 + except ValidationError:
306 + return
307 + raise AssertionError("expected placeholder length mismatch to fail")
308 +
309 +
310 +def test_metadata_can_replace_or_extend_full_mm_data():
311 + """Metadata-only requests retain the fields required by the model."""
312 + metadata_elem = MultiModalFieldElem(
313 + data=torch.tensor([[1, 24, 24]], dtype=torch.int32),
314 + field=MultiModalBatchedField(keep_on_cpu=True),
315 + )
316 + metadata_item = MultiModalKwargsItem({"image_grid_thw": metadata_elem})
317 + pixel_elem = MultiModalFieldElem(
318 + data=torch.randn(5, 3),
319 + field=MultiModalBatchedField(),
320 + )
321 + full_item = MultiModalKwargsItem(
322 + {"pixel_values": pixel_elem, "image_grid_thw": metadata_elem}
323 + )
324 +
325 + assert merge_mm_kwargs_items(None, metadata_item) == metadata_item
326 + merged = merge_mm_kwargs_items(full_item, metadata_item)
327 + assert merged is not None
328 + assert set(merged) == {"pixel_values", "image_grid_thw"}
vllm/entrypoints/scale_out/derender/serving.py
+3 -43
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import time
4 -from typing import cast
4
5 import vllm.envs as envs
6 from vllm.entrypoints.openai.chat_completion.protocol import (
@@ -19,16 +18,11 @@ from vllm.entrypoints.openai.models.serving import (
18 from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo
19 from vllm.entrypoints.serve.engine.serving import BaseServing
20 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
22 -from vllm.inputs import (
23 - EngineInput,
24 - MultiModalHashes,
25 - MultiModalInput,
26 - MultiModalPlaceholders,
27 -)
21 +from vllm.inputs import EngineInput
22 from vllm.logger import init_logger
23 from vllm.renderers.online_derenderer import OnlineDerenderer
24
31 -from ..token_in_token_out.mm_serde import encode_mm_kwargs_item
25 +from ..token_in_token_out.mm_features import extract_mm_features
26 from ..token_in_token_out.protocol import (
27 DerenderChatRequest,
28 DerenderChatStreamRequest,
@@ -37,7 +31,6 @@ from ..token_in_token_out.protocol import (
31 DerenderStreamState,
32 GenerateResponse,
33 MultiModalFeatures,
40 - PlaceholderRangeInfo,
34 )
35
36 logger = init_logger(__name__)
@@ -337,37 +330,4 @@ class ServingDerender(BaseServing):
330 def _extract_mm_features(
331 engine_input: EngineInput,
332 ) -> MultiModalFeatures | None:
340 - """Extract multimodal metadata from a rendered engine prompt.
341 -
342 - Returns ``None`` for text-only prompts.
343 - """
344 - if engine_input.get("type") != "multimodal":
345 - return None
346 -
347 - # At this point engine_input is a MultiModalInput TypedDict.
348 - mm_engine_input = cast(MultiModalInput, engine_input)
349 - mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
350 - raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
351 -
352 - mm_placeholders = {
353 - modality: [
354 - PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges
355 - ]
356 - for modality, ranges in raw_placeholders.items()
357 - }
358 -
359 - # Serialize tensor data per modality.
360 - kwargs_data: dict[str, list[str | None]] | None = None
361 - if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
362 - kwargs_data = {}
363 - for modality, items in raw_mm_kwargs.items():
364 - kwargs_data[modality] = [
365 - encode_mm_kwargs_item(item) if item is not None else None
366 - for item in items
367 - ]
368 -
369 - return MultiModalFeatures(
370 - mm_hashes=mm_hashes,
371 - mm_placeholders=mm_placeholders,
372 - kwargs_data=kwargs_data,
373 - )
333 + return extract_mm_features(engine_input)
vllm/entrypoints/scale_out/render/serving.py
+36 -44
@@ -1,7 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import cast
4 -
3 from vllm.entrypoints.anthropic.protocol import AnthropicMessagesRequest
4 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
5 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
@@ -10,23 +8,20 @@ from vllm.entrypoints.openai.models.serving import (
8 OpenAIModelRegistry,
9 OpenAIServingModels,
10 )
13 -from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import encode_mm_kwargs_item
11 +from vllm.entrypoints.scale_out.token_in_token_out.mm_features import (
12 + extract_mm_features,
13 +)
14 from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
15 GenerateRequest,
16 MultiModalFeatures,
17 - PlaceholderRangeInfo,
17 )
18 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
19 from vllm.entrypoints.serve.engine.serving import BaseServing
20 from vllm.entrypoints.serve.utils.api_utils import get_max_tokens
21 from vllm.entrypoints.serve.utils.request_logger import RequestLogger
23 -from vllm.inputs import (
24 - EngineInput,
25 - MultiModalHashes,
26 - MultiModalInput,
27 - MultiModalPlaceholders,
28 -)
22 +from vllm.inputs import EngineInput
23 from vllm.logger import init_logger
24 +from vllm.multimodal.parse import MultiModalDataParser
25 from vllm.renderers.inputs.preprocess import (
26 extract_prompt_components,
27 extract_prompt_len,
@@ -59,6 +54,9 @@ class ServingRender(BaseServing):
54 )
55 )
56
57 + self._placeholder_metadata_parser: MultiModalDataParser | None = None
58 + self._placeholder_metadata_parser_failed = False
59 +
60 self.default_sampling_params = (
61 online_renderer.model_config.get_diff_sampling_param()
62 )
@@ -231,41 +229,35 @@ class ServingRender(BaseServing):
229
230 return generate_requests
231
234 - @staticmethod
232 + def _placeholder_metadata_fields(self, modality: str) -> set[str]:
233 + """Fields the EC consumer still requires after embeddings transfer."""
234 + if self._placeholder_metadata_parser_failed:
235 + return set()
236 + if self._placeholder_metadata_parser is None:
237 + try:
238 + from vllm.multimodal import MULTIMODAL_REGISTRY
239 +
240 + self._placeholder_metadata_parser = (
241 + MULTIMODAL_REGISTRY.create_processor(
242 + self.model_config
243 + ).info.data_parser
244 + )
245 + except Exception:
246 + logger.debug(
247 + "Could not load placeholder metadata fields; "
248 + "mm_metadata will use keep_on_cpu fields only."
249 + )
250 + self._placeholder_metadata_parser_failed = True
251 + return set()
252 + return set(
253 + self._placeholder_metadata_parser.placeholder_metadata_fields(modality)
254 + )
255 +
256 def _extract_mm_features(
257 + self,
258 engine_input: EngineInput,
259 ) -> MultiModalFeatures | None:
238 - """Extract multimodal metadata from a rendered engine prompt.
239 -
240 - Returns ``None`` for text-only prompts.
241 - """
242 - if engine_input.get("type") != "multimodal":
243 - return None
244 -
245 - # At this point engine_input is a MultiModalInput TypedDict.
246 - mm_engine_input = cast(MultiModalInput, engine_input)
247 - mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
248 - raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
249 -
250 - mm_placeholders = {
251 - modality: [
252 - PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges
253 - ]
254 - for modality, ranges in raw_placeholders.items()
255 - }
256 -
257 - # Serialize tensor data per modality.
258 - kwargs_data: dict[str, list[str | None]] | None = None
259 - if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
260 - kwargs_data = {}
261 - for modality, items in raw_mm_kwargs.items():
262 - kwargs_data[modality] = [
263 - encode_mm_kwargs_item(item) if item is not None else None
264 - for item in items
265 - ]
266 -
267 - return MultiModalFeatures(
268 - mm_hashes=mm_hashes,
269 - mm_placeholders=mm_placeholders,
270 - kwargs_data=kwargs_data,
260 + return extract_mm_features(
261 + engine_input,
262 + metadata_fields_for=self._placeholder_metadata_fields,
263 )
vllm/entrypoints/scale_out/token_in_token_out/mm_features.py new
+159
@@ -0,0 +1,159 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Helpers for render-time multimodal feature extraction and generate input."""
4 +
5 +from __future__ import annotations
6 +
7 +from collections.abc import Callable, Collection, Sequence
8 +from typing import cast
9 +
10 +from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import (
11 + decode_mm_kwargs_item,
12 + encode_mm_kwargs_item,
13 +)
14 +from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
15 + MultiModalFeatures,
16 + PlaceholderRangeInfo,
17 +)
18 +from vllm.inputs import (
19 + EngineInput,
20 + MultiModalHashes,
21 + MultiModalInput,
22 + MultiModalPlaceholders,
23 +)
24 +from vllm.multimodal.inputs import (
25 + MultiModalKwargsItem,
26 + MultiModalKwargsOptionalItems,
27 +)
28 +
29 +
30 +def merge_mm_kwargs_items(
31 + kwargs_item: MultiModalKwargsItem | None,
32 + metadata_item: MultiModalKwargsItem | None,
33 +) -> MultiModalKwargsItem | None:
34 + """Merge full multimodal data with its metadata-only counterpart."""
35 + if kwargs_item is None:
36 + return metadata_item
37 + if metadata_item is None:
38 + return kwargs_item
39 + return MultiModalKwargsItem({**kwargs_item, **metadata_item})
40 +
41 +
42 +def _encode_metadata_items(
43 + items: Sequence[MultiModalKwargsItem | None],
44 + *,
45 + declared: Collection[str],
46 +) -> list[str | None]:
47 + """Serialize placeholder-metadata and ``keep_on_cpu`` fields per item."""
48 + metadata_items: list[str | None] = []
49 + for item in items:
50 + if item is None:
51 + metadata_items.append(None)
52 + continue
53 +
54 + metadata_item = MultiModalKwargsItem(
55 + {
56 + key: elem
57 + for key, elem in item.items()
58 + if elem.field.keep_on_cpu or key in declared
59 + }
60 + )
61 + metadata_items.append(
62 + encode_mm_kwargs_item(metadata_item) if metadata_item else None
63 + )
64 + return metadata_items
65 +
66 +
67 +def _encode_mm_kwargs_with_metadata(
68 + raw_mm_kwargs: MultiModalKwargsOptionalItems,
69 + *,
70 + metadata_fields_for: Callable[[str], Collection[str]] | None = None,
71 +) -> tuple[dict[str, list[str | None]], dict[str, list[str | None]] | None]:
72 + """Serialize full kwargs and their metadata-only subsets per modality."""
73 + kwargs_data: dict[str, list[str | None]] = {}
74 + metadata_by_modality: dict[str, list[str | None]] = {}
75 +
76 + for modality, items in raw_mm_kwargs.items():
77 + kwargs_data[modality] = [
78 + encode_mm_kwargs_item(item) if item is not None else None for item in items
79 + ]
80 +
81 + declared = (
82 + set(metadata_fields_for(modality))
83 + if metadata_fields_for is not None
84 + else set()
85 + )
86 + metadata_items = _encode_metadata_items(items, declared=declared)
87 + if any(item is not None for item in metadata_items):
88 + metadata_by_modality[modality] = metadata_items
89 +
90 + mm_metadata = metadata_by_modality or None
91 + return kwargs_data, mm_metadata
92 +
93 +
94 +def mm_kwargs_from_features(
95 + features: MultiModalFeatures,
96 +) -> dict[str, list[MultiModalKwargsItem | None]]:
97 + """Deserialize ``features`` into per-modality kwargs for ``mm_input``."""
98 + mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {}
99 + kwargs_data = features.kwargs_data or {}
100 + mm_metadata = features.mm_metadata or {}
101 + for modality, hashes in features.mm_hashes.items():
102 + n = len(hashes)
103 + kwargs_items = [
104 + decode_mm_kwargs_item(item) if item is not None else None
105 + for item in kwargs_data.get(modality, [None] * n)
106 + ]
107 + metadata_items = [
108 + decode_mm_kwargs_item(item) if item is not None else None
109 + for item in mm_metadata.get(modality, [None] * n)
110 + ]
111 + mm_kwargs[modality] = [
112 + merge_mm_kwargs_items(kwargs_item, metadata_item)
113 + for kwargs_item, metadata_item in zip(
114 + kwargs_items, metadata_items, strict=True
115 + )
116 + ]
117 + return mm_kwargs
118 +
119 +
120 +def extract_mm_features(
121 + engine_input: EngineInput,
122 + *,
123 + metadata_fields_for: Callable[[str], Collection[str]] | None = None,
124 +) -> MultiModalFeatures | None:
125 + """Extract multimodal features from a rendered engine prompt.
126 +
127 + Returns ``None`` for text-only prompts. ``mm_metadata`` keeps the
128 + intersection of processed kwargs that prefill needs after EC transfer:
129 + fields declared as embedding metadata, plus fields marked
130 + ``keep_on_cpu`` (for example M-RoPE grid dims).
131 + """
132 + if engine_input.get("type") != "multimodal":
133 + return None
134 +
135 + mm_engine_input = cast(MultiModalInput, engine_input)
136 + mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
137 + raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
138 +
139 + mm_placeholders = {
140 + modality: [
141 + PlaceholderRangeInfo(offset=p.offset, length=p.length) for p in ranges
142 + ]
143 + for modality, ranges in raw_placeholders.items()
144 + }
145 +
146 + kwargs_data: dict[str, list[str | None]] | None = None
147 + mm_metadata: dict[str, list[str | None]] | None = None
148 + if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
149 + kwargs_data, mm_metadata = _encode_mm_kwargs_with_metadata(
150 + raw_mm_kwargs,
151 + metadata_fields_for=metadata_fields_for,
152 + )
153 +
154 + return MultiModalFeatures(
155 + mm_hashes=mm_hashes,
156 + mm_placeholders=mm_placeholders,
157 + kwargs_data=kwargs_data,
158 + mm_metadata=mm_metadata,
159 + )
vllm/entrypoints/scale_out/token_in_token_out/protocol.py
+45
@@ -44,6 +44,14 @@ class PlaceholderRangeInfo(BaseModel):
44 # placeholder masks that cannot be recomputed from offset+length alone.
45
46
47 +def _has_serialized_mm_items(
48 + payload: dict[str, list[str | None]] | None,
49 +) -> bool:
50 + if not payload:
51 + return False
52 + return any(item is not None for items in payload.values() for item in items)
53 +
54 +
55 class MultiModalFeatures(BaseModel):
56 """Lightweight multimodal metadata produced by the render step.
57
@@ -67,6 +75,18 @@ class MultiModalFeatures(BaseModel):
75 ``None`` for metadata-only (cache-hit) responses.
76 """
77
78 + mm_metadata: dict[str, list[str | None]] | None = None
79 + """Per-modality serialized metadata for disaggregated prefill.
80 +
81 + Each value is a list parallel to ``mm_hashes[modality]``. A ``str``
82 + entry is a base64-encoded ``MultiModalKwargsItem`` containing only
83 + placeholder-metadata and ``keep_on_cpu`` fields. ``None`` means that
84 + the metadata is unavailable for that item. Prefill can use this
85 + instead of ``kwargs_data`` only when ``ec_transfer_params`` is also
86 + set, so embeddings arrive through the EC connector rather than from
87 + ``pixel_values``.
88 + """
89 +
90 @model_validator(mode="after")
91 def _validate_parallel_fields(self) -> "MultiModalFeatures":
92 modalities = set(self.mm_hashes)
@@ -76,6 +96,8 @@ class MultiModalFeatures(BaseModel):
96 )
97 if self.kwargs_data is not None and set(self.kwargs_data) != modalities:
98 raise ValueError("kwargs_data must use the same modalities as mm_hashes")
99 + if self.mm_metadata is not None and set(self.mm_metadata) != modalities:
100 + raise ValueError("mm_metadata must use the same modalities as mm_hashes")
101
102 flattened_ranges: list[tuple[int, int]] = []
103 for modality in modalities:
@@ -93,6 +115,13 @@ class MultiModalFeatures(BaseModel):
115 raise ValueError(
116 f"{modality} kwargs_data and mm_hashes must have the same length"
117 )
118 + if (
119 + self.mm_metadata is not None
120 + and len(self.mm_metadata[modality]) != num_hashes
121 + ):
122 + raise ValueError(
123 + f"{modality} mm_metadata and mm_hashes must have the same length"
124 + )
125 flattened_ranges.extend(
126 (placeholder.offset, placeholder.offset + placeholder.length)
127 for placeholder in self.mm_placeholders[modality]
@@ -156,6 +185,22 @@ class GenerateRequest(BaseModel):
185 raise ValueError("content_parts and features are mutually exclusive")
186 return self
187
188 + @model_validator(mode="after")
189 + def _require_ec_for_metadata_only(self) -> "GenerateRequest":
190 + features = self.features
191 + if features is None:
192 + return self
193 + if not _has_serialized_mm_items(features.mm_metadata):
194 + return self
195 + if _has_serialized_mm_items(features.kwargs_data):
196 + return self
197 + if self.ec_transfer_params:
198 + return self
199 + raise ValueError(
200 + "features.mm_metadata without features.kwargs_data requires "
201 + "ec_transfer_params so embeddings can be loaded by the EC connector"
202 + )
203 +
204 sampling_params: SamplingParams
205 """The sampling parameters for the model."""
206
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+4 -13
@@ -35,7 +35,6 @@ from vllm.inputs import EngineInput, TokensPrompt, mm_input
35 from vllm.logger import init_logger
36 from vllm.logprobs import Logprob
37 from vllm.multimodal.inputs import (
38 - MultiModalKwargsItem,
38 MultiModalKwargsItems,
39 PlaceholderRange,
40 )
@@ -45,7 +44,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams
44 from vllm.utils.collection_utils import as_list
45 from vllm.utils.serial_utils import numpy2base64
46
48 -from .mm_serde import decode_mm_kwargs_item
47 +from .mm_features import mm_kwargs_from_features
48 from .protocol import (
49 GenerateRequest,
50 GenerateResponse,
@@ -177,17 +176,9 @@ class ServingTokens(GenerateBaseServing):
176 for modality, ranges in features.mm_placeholders.items()
177 }
178
180 - # Deserialize tensor data when present; None → cache hit.
181 - mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {}
182 - if features.kwargs_data is not None:
183 - for modality, items in features.kwargs_data.items():
184 - mm_kwargs[modality] = [
185 - decode_mm_kwargs_item(item) if item is not None else None
186 - for item in items
187 - ]
188 - else:
189 - for modality, hashes in features.mm_hashes.items():
190 - mm_kwargs[modality] = [None] * len(hashes)
179 + # Deserialize full tensor data and optional metadata-only data.
180 + # Metadata-only items are valid when ec_transfer_params is set.
181 + mm_kwargs = mm_kwargs_from_features(features)
182
183 engine_input = mm_input(
184 prompt_token_ids=request.token_ids,