@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 479611c0c6

[Frontend][EPD] Use JSON arrays for multimodal metadata (#56090)

Signed-off-by: Tianyu Guo <guoty@inferact.ai>

Tianyu Guo committed Sep 10, 2026 at 13:04 UTC 479611c0c6118f95d6028a337e69b4713c529843

6 files changed +167 -62

examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py

+9 -16

@@ -22,7 +22,6 @@ from __future__ import annotations

22 import argparse

23 import asyncio

24 import hashlib

25 -import io

25 import itertools

26 import json

27 import logging

@@ -34,7 +33,6 @@ from collections.abc import AsyncIterator

33 from typing import Any

35 import aiohttp

37 -import pybase64 as base64

36 import uvicorn

37 from fastapi import FastAPI, HTTPException, Request

38 from fastapi.responses import JSONResponse, StreamingResponse

@@ -125,19 +123,6 @@ def content_uuid(item: dict) -> str:

123 return hashlib.sha256(payload.encode()).hexdigest()

124

125

128 -def _b64_tensor(values: list) -> str:

129 - import torch

130 -

131 - buf = io.BytesIO()

132 - flat = [v for item in values for v in (item if isinstance(item, list) else [item])]

133 - # Floats stay float64 so timestamp strings format exactly as the

134 - # encoder computed them.

135 - dtype = torch.float64 if any(isinstance(v, float) for v in flat) else None

136 - # Downstream stacks per item, so hand over a flat vector.

137 - torch.save(torch.tensor(flat, dtype=dtype), buf)

138 - return base64.b64encode(buf.getvalue()).decode()

139 -

140 -

126 def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:

127 """Replace each media item with a metadata-only reference for the decoder.

128

@@ -172,7 +157,15 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:

157 # Whatever keys the encoder reported are the metadata its model

158 # declared as needed to size the placeholder range; the proxy does

159 # not need to know their names.

175 - metadata = {k: _b64_tensor(v) for k, v in meta.items()}

160 + # Downstream stacks per item; keep the existing per-item vector shape.

161 + metadata = {

162 + k: [

163 + x

164 + for item in v

165 + for x in (item if isinstance(item, list) else [item])

166 + ]

167 + for k, v in meta.items()

168 + }

169 if not metadata or not item_uuid:

170 # Nothing to size the placeholder range with. A processor cache

171 # hit is not a cause on its own: with the default `lru` type the

tests/entrypoints/unit_tests/test_chat_utils.py

+46 -1

@@ -18,6 +18,8 @@ from vllm.entrypoints.chat_utils import (

18 MEDIA_CONNECTOR_REGISTRY,

19 AsyncMultiModalItemTracker,

20 ConversationMessage,

21 + _load_embeds_dict,

22 + _parse_metadata_array,

23 _postprocess_messages,

24 parse_chat_messages,

25 parse_chat_messages_async,

@@ -39,6 +41,45 @@ QWEN25OMNI_MODEL_ID = "Qwen/Qwen2.5-Omni-7B"

41 MISTRAL_MODEL_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"

44 +@pytest.mark.parametrize("values", [[1, 32, 48], [0.0, 0.12345678912345678]])

45 +def test_json_metadata_preserves_numeric_precision(values):

46 + tensor = _parse_metadata_array("grid", values, {"grid"})

47 + assert tensor.tolist() == values

48 + assert tensor.dtype == (

49 + torch.float64 if isinstance(values[0], float) else torch.int64

50 + )

51 +

52 +

53 +@pytest.mark.parametrize(

54 + "key,values",

55 + [

56 + ("image_embeds", [1, 2]),

57 + ("grid", [True]),

58 + ("grid", [float("nan")]),

59 + ("grid", [float("inf")]),

60 + ("grid", [[1, 2]]),

61 + ("grid", [2**64]),

62 + ],

63 +)

64 +def test_json_arrays_are_only_valid_numeric_metadata(key, values):

65 + with pytest.raises(VLLMValidationError):

66 + _parse_metadata_array(key, values, {"grid"})

67 +

68 +

69 +@pytest.mark.asyncio

70 +async def test_json_metadata_bypasses_tensor_deserialization():

71 + from unittest.mock import AsyncMock

72 +

73 + tensor = torch.ones(2, 3)

74 + fetch = AsyncMock(return_value=tensor)

75 + result = await _load_embeds_dict(

76 + {"image_embeds": "legacy-base64", "image_grid_thw": [1, 2, 2]}, fetch

77 + )

78 + fetch.assert_awaited_once_with("legacy-base64")

79 + assert result["image_embeds"] is tensor

80 + assert result["image_grid_thw"] == [1, 2, 2]

81 +

82 +

83 @pytest.fixture(scope="function")

84 def kimi_k2_5_model_config():

85 return ModelConfig(

@@ -1355,8 +1396,10 @@ def test_parse_chat_messages_empty_dict_image_embeds(

1396 _assert_mm_uuids(mm_uuids, 1, expected_uuids=[None])

1397

1398

1399 +@pytest.mark.parametrize("json_metadata", [False, True])

1400 def test_parse_chat_messages_multiple_dict_image_embeds(

1401 qwen25omni_model_config_image_embeds,

1402 + json_metadata,

1403 ):

1404 """Test that multiple dictionaries for image_embeds is handled without errors."""

1405 # Create two sample image embedding tensors

@@ -1374,7 +1417,9 @@ def test_parse_chat_messages_multiple_dict_image_embeds(

1417 "type": "image_embeds",

1418 "image_embeds": {

1419 "image_embeds": tensor2base64(embeds),

1377 - "image_grid_thw": tensor2base64(grid_thw),

1420 + "image_grid_thw": grid_thw.tolist()

1421 + if json_metadata

1422 + else tensor2base64(grid_thw),

1423 },

1424 }

1425 for embeds, grid_thw in zip(

tests/utils_/test_collection_utils.py

+25 -1

@@ -2,7 +2,31 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import pytest

5 -from vllm.utils.collection_utils import common_prefix, swap_dict_values

5 +from vllm.utils.collection_utils import (

6 + common_prefix,

7 + is_list_of_numbers,

8 + swap_dict_values,

9 +)

10 +

11 +

12 +@pytest.mark.parametrize(

13 + "value,expected",

14 + [

15 + ([], True),

16 + ([1, -2, 2**1024], True),

17 + ([1, 0.5], True),

18 + ([1, True], False),

19 + ([1, float("nan")], False),

20 + ([1, float("inf")], False),

21 + ([1, -float("inf")], False),

22 + ([1, "2"], False),

23 + ([[1]], False),

24 + ((1, 2), False),

25 + (None, False),

26 + ],

27 +)

28 +def test_is_list_of_numbers_checks_all_finite_non_boolean_items(value, expected):

29 + assert is_list_of_numbers(value) is expected

32 @pytest.mark.parametrize(

tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py

@@ -123,6 +123,9 @@ def test_decode_rewrite_preserves_engine_reported_ec_hash(proxy):

123 )

124

125 assert rewritten["messages"][0]["content"][0]["uuid"] == "proxy-uuid"

126 + assert rewritten["messages"][0]["content"][0]["image_embeds"] == {

127 + "image_grid_thw": [1, 2, 3]

128 + }

129 assert rewritten["ec_transfer_params"]["ec_items"] == [

130 {"mm_hash": "engine-derived-hash", "transfer_id": "transfer"}

131 ]

vllm/entrypoints/chat_utils.py

+76 -44

@@ -73,7 +73,7 @@ from vllm.renderers.embed_utils import (

73 )

74 from vllm.transformers_utils.processor import get_video_processor_cls_name

75 from vllm.utils import random_uuid

76 -from vllm.utils.collection_utils import is_list_of

76 +from vllm.utils.collection_utils import is_list_of, is_list_of_numbers

77 from vllm.utils.import_utils import LazyLoader

79 if TYPE_CHECKING:

@@ -146,12 +146,15 @@ class ChatCompletionContentPartAudioParam(TypedDict, total=False):

146 """The type of the content part."""

147

148

149 +MultiModalEmbedsPayload: TypeAlias = str | dict[str, str | list[int | float]]

150 +

151 +

152 class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):

150 - image_embeds: str | dict[str, str] | None

153 + image_embeds: MultiModalEmbedsPayload | None

154 """

155 The image embeddings. It can be either:

156 - A single base64 string.

154 - - A dictionary where each value is a base64 string.

157 + - A dictionary of base64 tensors or numeric JSON metadata arrays.

158 """

159 type: Required[Literal["image_embeds"]]

160 """The type of the content part."""

@@ -163,11 +166,11 @@ class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):

166

167

168 class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False):

166 - audio_embeds: str | dict[str, str] | None

169 + audio_embeds: MultiModalEmbedsPayload | None

170 """

171 The audio embeddings. It can be either:

172 - A single base64 string representing a serialized torch tensor.

170 - - A dictionary where each value is a base64 string.

173 + - A dictionary of base64 tensors or numeric JSON metadata arrays.

174 """

175 type: Required[Literal["audio_embeds"]]

176 """The type of the content part."""

@@ -179,11 +182,11 @@ class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False):

182

183

184 class ChatCompletionContentPartVideoEmbedsParam(TypedDict, total=False):

182 - video_embeds: str | dict[str, str] | None

185 + video_embeds: MultiModalEmbedsPayload | None

186 """

187 The video embeddings. It can be either:

188 - A single base64 string representing a serialized torch tensor.

186 - - A dictionary where each value is a base64 string.

189 + - A dictionary of base64 tensors or numeric JSON metadata arrays.

190 """

191 type: Required[Literal["video_embeds"]]

192 """The type of the content part."""

@@ -527,6 +530,27 @@ def _merge_embeds(

530 return data_merged

531

532

533 +async def _load_embeds_dict(

534 + data: dict[str, str | list[int | float]],

535 + fetch: Callable[[str], Awaitable["torch.Tensor"]],

536 +) -> dict[str, Any]:

537 + encoded = {key: value for key, value in data.items() if isinstance(value, str)}

538 + tensors = await asyncio.gather(*(fetch(value) for value in encoded.values()))

539 + return {**data, **dict(zip(encoded, tensors))}

540 +

541 +

542 +def _parse_metadata_array(key: str, value: list, metadata_fields: set[str]):

543 + if key not in metadata_fields:

544 + raise VLLMValidationError(f"JSON arrays are only supported for metadata: {key}")

545 + if not is_list_of_numbers(value):

546 + raise VLLMValidationError(f"Metadata {key} must be a finite numeric array.")

547 + dtype = torch.float64 if any(isinstance(v, float) for v in value) else torch.int64

548 + try:

549 + return torch.tensor(value, dtype=dtype)

550 + except (ValueError, TypeError, OverflowError, RuntimeError) as error:

551 + raise VLLMValidationError(f"Invalid metadata array: {key}") from error

552 +

553 +

554 def _get_embeds_data(

555 modality: str,

556 data_items: list[Any],

@@ -544,6 +568,18 @@ def _get_embeds_data(

568 return _merge_embeds(dict_items, mm_processor)[embeds_key]

569

570 if is_list_of(data_items, dict):

571 + metadata_fields = mm_processor.info.data_parser.placeholder_metadata_fields(

572 + modality

573 + )

574 + data_items = [

575 + {

576 + key: _parse_metadata_array(key, value, metadata_fields)

577 + if isinstance(value, list)

578 + else value

579 + for key, value in item.items()

580 + }

581 + for item in data_items

582 + ]

583 return _merge_embeds(data_items, mm_processor)

584

585 raise NotImplementedError(type(data_items))

@@ -946,7 +982,7 @@ class BaseMultiModalContentParser(ABC):

982 @abstractmethod

983 def parse_image_embeds(

984 self,

949 - image_embeds: str | dict[str, str] | None,

985 + image_embeds: MultiModalEmbedsPayload | None,

986 uuid: str | None = None,

987 ) -> None:

988 raise NotImplementedError

@@ -970,7 +1006,7 @@ class BaseMultiModalContentParser(ABC):

1006 @abstractmethod

1007 def parse_audio_embeds(

1008 self,

973 - audio_embeds: str | dict[str, str] | None,

1009 + audio_embeds: MultiModalEmbedsPayload | None,

1010 uuid: str | None = None,

1011 ) -> None:

1012 raise NotImplementedError

@@ -986,7 +1022,7 @@ class BaseMultiModalContentParser(ABC):

1022 @abstractmethod

1023 def parse_video_embeds(

1024 self,

989 - video_embeds: str | dict[str, str] | None,

1025 + video_embeds: MultiModalEmbedsPayload | None,

1026 uuid: str | None = None,

1027 ) -> None:

1028 raise NotImplementedError

@@ -1043,7 +1079,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

1079

1080 def parse_image_embeds(

1081 self,

1046 - image_embeds: str | dict[str, str] | None,

1082 + image_embeds: MultiModalEmbedsPayload | None,

1083 uuid: str | None = None,

1084 ) -> None:

1085 mm_config = self.model_config.get_multimodal_config()

@@ -1055,7 +1091,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

1091

1092 if isinstance(image_embeds, dict):

1093 embeds = {

1058 - k: self._connector.fetch_image_embedding(v)

1094 + k: self._connector.fetch_image_embedding(v) if isinstance(v, str) else v

1095 for k, v in image_embeds.items()

1096 }

1097 placeholder = self._tracker.add("image_embeds", (embeds, uuid))

@@ -1071,7 +1107,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

1107

1108 def parse_audio_embeds(

1109 self,

1074 - audio_embeds: str | dict[str, str] | None,

1110 + audio_embeds: MultiModalEmbedsPayload | None,

1111 uuid: str | None = None,

1112 ) -> None:

1113 mm_config = self.model_config.get_multimodal_config()

@@ -1083,7 +1119,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

1119

1120 if isinstance(audio_embeds, dict):

1121 embeds = {

1086 - k: self._connector.fetch_audio_embedding(v)

1122 + k: self._connector.fetch_audio_embedding(v) if isinstance(v, str) else v

1123 for k, v in audio_embeds.items()

1124 }

1125 placeholder = self._tracker.add("audio_embeds", (embeds, uuid))

@@ -1148,7 +1184,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

1184

1185 def parse_video_embeds(

1186 self,

1151 - video_embeds: str | dict[str, str] | None,

1187 + video_embeds: MultiModalEmbedsPayload | None,

1188 uuid: str | None = None,

1189 ) -> None:

1190 mm_config = self.model_config.get_multimodal_config()

@@ -1160,7 +1196,7 @@ class MultiModalContentParser(BaseMultiModalContentParser):

1196

1197 if isinstance(video_embeds, dict):

1198 embeds = {

1163 - k: self._connector.fetch_video_embedding(v)

1199 + k: self._connector.fetch_video_embedding(v) if isinstance(v, str) else v

1200 for k, v in video_embeds.items()

1201 }

1202 placeholder = self._tracker.add("video_embeds", (embeds, uuid))

@@ -1242,7 +1278,7 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

1278

1279 def parse_image_embeds(

1280 self,

1245 - image_embeds: str | dict[str, str] | None,

1281 + image_embeds: MultiModalEmbedsPayload | None,

1282 uuid: str | None = None,

1283 ) -> None:

1284 mm_config = self.model_config.get_multimodal_config()

@@ -1260,17 +1296,13 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

1296

1297 async def _image_embeds_with_uuid_async(

1298 self,

1263 - image_embeds: str | dict[str, str] | None,

1299 + image_embeds: MultiModalEmbedsPayload | None,

1300 uuid: str | None,

1301 ):

1302 if isinstance(image_embeds, dict):

1267 - tensors = await asyncio.gather(

1268 - *(

1269 - self._connector.fetch_image_embedding_async(v)

1270 - for v in image_embeds.values()

1271 - )

1303 + embeds = await _load_embeds_dict(

1304 + image_embeds, self._connector.fetch_image_embedding_async

1305 )

1273 - embeds = dict(zip(image_embeds, tensors))

1306 elif isinstance(image_embeds, str):

1307 embeds = await self._connector.fetch_image_embedding_async(image_embeds)

1308 else:

@@ -1279,7 +1311,7 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

1311

1312 def parse_audio_embeds(

1313 self,

1282 - audio_embeds: str | dict[str, str] | None,

1314 + audio_embeds: MultiModalEmbedsPayload | None,

1315 uuid: str | None = None,

1316 ) -> None:

1317 mm_config = self.model_config.get_multimodal_config()

@@ -1297,17 +1329,13 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

1329

1330 async def _audio_embeds_with_uuid_async(

1331 self,

1300 - audio_embeds: str | dict[str, str] | None,

1332 + audio_embeds: MultiModalEmbedsPayload | None,

1333 uuid: str | None,

1334 ):

1335 if isinstance(audio_embeds, dict):

1304 - tensors = await asyncio.gather(

1305 - *(

1306 - self._connector.fetch_audio_embedding_async(v)

1307 - for v in audio_embeds.values()

1308 - )

1336 + embeds = await _load_embeds_dict(

1337 + audio_embeds, self._connector.fetch_audio_embedding_async

1338 )

1310 - embeds = dict(zip(audio_embeds, tensors))

1339 elif isinstance(audio_embeds, str):

1340 embeds = await self._connector.fetch_audio_embedding_async(audio_embeds)

1341 else:

@@ -1382,7 +1410,7 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

1410

1411 def parse_video_embeds(

1412 self,

1385 - video_embeds: str | dict[str, str] | None,

1413 + video_embeds: MultiModalEmbedsPayload | None,

1414 uuid: str | None = None,

1415 ) -> None:

1416 mm_config = self.model_config.get_multimodal_config()

@@ -1400,17 +1428,13 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

1428

1429 async def _video_embeds_with_uuid_async(

1430 self,

1403 - video_embeds: str | dict[str, str] | None,

1431 + video_embeds: MultiModalEmbedsPayload | None,

1432 uuid: str | None,

1433 ):

1434 if isinstance(video_embeds, dict):

1407 - tensors = await asyncio.gather(

1408 - *(

1409 - self._connector.fetch_video_embedding_async(v)

1410 - for v in video_embeds.values()

1411 - )

1435 + embeds = await _load_embeds_dict(

1436 + video_embeds, self._connector.fetch_video_embedding_async

1437 )

1413 - embeds = dict(zip(video_embeds, tensors))

1438 elif isinstance(video_embeds, str):

1439 embeds = await self._connector.fetch_video_embedding_async(video_embeds)

1440 else:

@@ -1611,7 +1635,9 @@ _AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python

1635 _VideoParser = TypeAdapter(ChatCompletionContentPartVideoParam).validate_python

1636

1637 _ResponsesInputImageParser = TypeAdapter(ResponseInputImageParam).validate_python

1614 -_ContentPart: TypeAlias = str | dict[str, str] | InputAudio | PILImage

1638 +_ContentPart: TypeAlias = (

1639 + MultiModalEmbedsPayload | dict[str, str] | InputAudio | PILImage

1640 +)

1641

1642 # Define a mapping from part types to their corresponding parsing functions.

1643 MM_PARSER_MAP: dict[

@@ -1901,15 +1927,21 @@ def _parse_chat_message_content_part(

1927 mm_parser.parse_image(str_content, uuid)

1928 modality = "image"

1929 elif part_type == "image_embeds":

1904 - content = cast(str | dict[str, str], content) if content is not None else None

1930 + content = (

1931 + cast(MultiModalEmbedsPayload, content) if content is not None else None

1932 + )

1933 mm_parser.parse_image_embeds(content, uuid)

1934 modality = "image"

1935 elif part_type == "audio_embeds":

1908 - content = cast(str | dict[str, str], content) if content is not None else None

1936 + content = (

1937 + cast(MultiModalEmbedsPayload, content) if content is not None else None

1938 + )

1939 mm_parser.parse_audio_embeds(content, uuid)

1940 modality = "audio"

1941 elif part_type == "video_embeds":

1912 - content = cast(str | dict[str, str], content) if content is not None else None

1942 + content = (

1943 + cast(MultiModalEmbedsPayload, content) if content is not None else None

1944 + )

1945 mm_parser.parse_video_embeds(content, uuid)

1946 modality = "video"

1947 elif part_type == "prompt_embeds":

vllm/utils/collection_utils.py

@@ -6,6 +6,7 @@ Contains helpers that are applied to collections.

6 This is similar in concept to the `collections` module.

7 """

9 +import math

10 from collections import defaultdict

11 from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence

12 from typing import Generic, Literal, TypeVar

@@ -68,6 +69,13 @@ def is_list_of(

69 assert_never(check)

72 +def is_list_of_numbers(value: object) -> bool:

73 + """Check every item is an int or finite float, excluding booleans."""

74 + return isinstance(value, list) and all(

75 + type(v) is int or (type(v) is float and math.isfinite(v)) for v in value

76 + )

77 +

78 +

79 @overload

80 def common_prefix(items: Sequence[str]) -> str: ...