@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b09806e28f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 11, 12:30 UTC

[Bugfix] Dictionary MM embeddings for online chat (#30507)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Dec 13, 2025 at 15:48 UTC b09806e28ffcc3e63176d668b2b3e965b35c986c

3 files changed +193 -44

tests/entrypoints/test_chat_utils.py

+105 -5

@@ -796,9 +796,13 @@ def test_parse_chat_messages_empty_image_embeds_with_uuid(

796 "content": "<|image_1|>\nWhat's in this image?",

797 }

798 ]

799 +

800 assert mm_data is not None

801 assert "image" in mm_data

801 - assert mm_data["image"] is None

802 + assert isinstance(mm_data["image"], list)

803 + assert len(mm_data["image"]) == 1

804 + assert mm_data["image"][0] is None

805 +

806 _assert_mm_uuids(mm_uuids, 1, expected_uuids=[uuid])

807

808

@@ -825,10 +829,11 @@ def test_parse_chat_messages_empty_audio_embeds_with_uuid(

829 # Should have audio in mm_data as None (UUID provided)

830 assert mm_data is not None

831 assert "audio" in mm_data

828 - assert mm_data["audio"] is None

832 + assert isinstance(mm_data["audio"], list)

833 + assert len(mm_data["audio"]) == 1

834 + assert mm_data["audio"][0] is None

835 +

836 # UUID should be recorded

830 - assert mm_uuids is not None

831 - assert "audio" in mm_uuids

837 _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[uuid])

838

839

@@ -1121,10 +1126,105 @@ async def test_parse_chat_messages_empty_image_embeds_with_uuid_async(

1126 mm_data = await mm_future

1127 assert mm_data is not None

1128 assert "image" in mm_data

1124 - assert mm_data["image"] is None

1129 + assert isinstance(mm_data["image"], list)

1130 + assert len(mm_data["image"]) == 1

1131 + assert mm_data["image"][0] is None

1132 +

1133 _assert_mm_uuids(mm_uuids, 1, expected_uuids=[uuid])

1134

1135

1136 +def test_parse_chat_messages_empty_dict_image_embeds(

1137 + phi3v_model_config_image_embeds,

1138 +):

1139 + """Test that empty dictionary for image_embeds is handled without errors."""

1140 + conversation, mm_data, mm_uuids = parse_chat_messages(

1141 + [

1142 + {

1143 + "role": "user",

1144 + "content": [

1145 + {"type": "image_embeds", "image_embeds": {}},

1146 + {"type": "text", "text": "What's in this image?"},

1147 + ],

1148 + }

1149 + ],

1150 + phi3v_model_config_image_embeds,

1151 + content_format="string",

1152 + )

1153 +

1154 + # Verify conversation structure

1155 + assert conversation == [

1156 + {

1157 + "role": "user",

1158 + "content": "<|image_1|>\nWhat's in this image?",

1159 + }

1160 + ]

1161 +

1162 + # Verify mm_data contains an empty dictionary of embeddings

1163 + assert mm_data is not None

1164 + assert "image" in mm_data

1165 + assert isinstance(mm_data["image"], dict)

1166 + assert len(mm_data["image"]) == 0

1167 +

1168 + # Verify UUIDs (None since we didn't provide any)

1169 + _assert_mm_uuids(mm_uuids, 1, expected_uuids=[None])

1170 +

1171 +

1172 +def test_parse_chat_messages_multiple_dict_image_embeds(

1173 + phi3v_model_config_image_embeds,

1174 +):

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

1176 + # Create two sample image embedding tensors

1177 + batch_size = 2

1178 + image_embedding_1 = torch.randn(batch_size, 256, 1024)

1179 + image_embedding_2 = torch.randn(batch_size, 3)

1180 +

1181 + conversation, mm_data, mm_uuids = parse_chat_messages(

1182 + [

1183 + {

1184 + "role": "user",

1185 + "content": [

1186 + {

1187 + "type": "image_embeds",

1188 + "image_embeds": {

1189 + "image_embedding_1": tensor2base64(p),

1190 + "image_embedding_2": tensor2base64(i),

1191 + },

1192 + }

1193 + for p, i in zip(image_embedding_1, image_embedding_2)

1194 + ]

1195 + + [

1196 + {"type": "text", "text": "Describe these two images."},

1197 + ],

1198 + }

1199 + ],

1200 + phi3v_model_config_image_embeds,

1201 + content_format="string",

1202 + )

1203 +

1204 + # Verify conversation structure

1205 + assert conversation == [

1206 + {

1207 + "role": "user",

1208 + "content": "<|image_1|>\n<|image_2|>\nDescribe these two images.",

1209 + }

1210 + ]

1211 +

1212 + # Verify mm_data contains a dictionary of multi-embeddings

1213 + assert mm_data is not None

1214 + assert "image" in mm_data

1215 + assert isinstance(mm_data["image"], dict)

1216 + assert len(mm_data["image"]) == batch_size

1217 +

1218 + # Verify each embedding has the correct shape

1219 + assert isinstance(mm_data["image"]["image_embedding_1"], torch.Tensor)

1220 + assert mm_data["image"]["image_embedding_1"].shape == image_embedding_1.shape

1221 + assert isinstance(mm_data["image"]["image_embedding_2"], torch.Tensor)

1222 + assert mm_data["image"]["image_embedding_2"].shape == image_embedding_2.shape

1223 +

1224 + # Verify UUIDs (None since we didn't provide any)

1225 + _assert_mm_uuids(mm_uuids, batch_size, expected_uuids=[None, None])

1226 +

1227 +

1228 @pytest.mark.asyncio

1229 async def test_parse_chat_messages_multiple_images_async(

1230 phi3v_model_config,

vllm/entrypoints/chat_utils.py

+68 -29

@@ -9,7 +9,7 @@ from collections import Counter, defaultdict, deque

9 from collections.abc import Awaitable, Callable, Iterable

10 from functools import cached_property, lru_cache, partial

11 from pathlib import Path

12 -from typing import Any, Generic, Literal, TypeAlias, TypeVar, cast

12 +from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar, cast

14 import jinja2

15 import jinja2.ext

@@ -53,7 +53,14 @@ from vllm.tokenizers import MistralTokenizer, TokenizerLike

53 from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path

54 from vllm.transformers_utils.processor import cached_get_processor

55 from vllm.utils import random_uuid

56 +from vllm.utils.collection_utils import is_list_of

57 from vllm.utils.func_utils import supports_kw

58 +from vllm.utils.import_utils import LazyLoader

59 +

60 +if TYPE_CHECKING:

61 + import torch

62 +else:

63 + torch = LazyLoader("torch", globals(), "torch")

65 logger = init_logger(__name__)

@@ -620,6 +627,44 @@ ModalityStr = Literal["image", "audio", "video", "image_embeds", "audio_embeds"]

627 _T = TypeVar("_T")

628

629

630 +def _extract_embeds(tensors: list[torch.Tensor]):

631 + if len(tensors) == 0:

632 + return tensors

633 +

634 + if len(tensors) == 1:

635 + tensors[0]._is_single_item = True # type: ignore

636 + return tensors[0] # To keep backwards compatibility for single item input

637 +

638 + first_shape = tensors[0].shape

639 + if all(t.shape == first_shape for t in tensors):

640 + return torch.stack(tensors)

641 +

642 + return tensors

643 +

644 +

645 +def _get_embeds_data(items_by_modality: dict[str, list[Any]], modality: str):

646 + embeds_key = f"{modality}_embeds"

647 + embeds = items_by_modality[embeds_key]

648 +

649 + if len(embeds) == 0:

650 + return embeds

651 + if is_list_of(embeds, torch.Tensor):

652 + return _extract_embeds(embeds)

653 + if is_list_of(embeds, dict):

654 + if not embeds:

655 + return {}

656 +

657 + first_keys = set(embeds[0].keys())

658 + if any(set(item.keys()) != first_keys for item in embeds[1:]):

659 + raise ValueError(

660 + "All dictionaries in the list of embeddings must have the same keys."

661 + )

662 +

663 + return {k: _extract_embeds([item[k] for item in embeds]) for k in first_keys}

664 +

665 + return embeds

666 +

667 +

668 class BaseMultiModalItemTracker(ABC, Generic[_T]):

669 """

670 Tracks multi-modal items in a given request and ensures that the number

@@ -688,11 +733,14 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

733 def all_mm_uuids(self) -> MultiModalUUIDDict | None:

734 if not self._items_by_modality:

735 return None

691 - mm_uuids = {}

736 +

737 uuids_by_modality = dict(self._uuids_by_modality)

738 if "image" in uuids_by_modality and "image_embeds" in uuids_by_modality:

739 raise ValueError("Mixing raw image and embedding inputs is not allowed")

740 + if "audio" in uuids_by_modality and "audio_embeds" in uuids_by_modality:

741 + raise ValueError("Mixing raw audio and embedding inputs is not allowed")

742

743 + mm_uuids = {}

744 if "image_embeds" in uuids_by_modality:

745 mm_uuids["image"] = uuids_by_modality["image_embeds"]

746 if "image" in uuids_by_modality:

@@ -703,6 +751,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

751 mm_uuids["audio"] = uuids_by_modality["audio"] # UUIDs of audios

752 if "video" in uuids_by_modality:

753 mm_uuids["video"] = uuids_by_modality["video"] # UUIDs of videos

754 +

755 return mm_uuids

756

757 @abstractmethod

@@ -714,29 +763,25 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[object]):

763 def all_mm_data(self) -> MultiModalDataDict | None:

764 if not self._items_by_modality:

765 return None

717 - mm_inputs = {}

766 +

767 items_by_modality = dict(self._items_by_modality)

768 if "image" in items_by_modality and "image_embeds" in items_by_modality:

769 raise ValueError("Mixing raw image and embedding inputs is not allowed")

770 if "audio" in items_by_modality and "audio_embeds" in items_by_modality:

771 raise ValueError("Mixing raw audio and embedding inputs is not allowed")

772

773 + mm_inputs = {}

774 if "image_embeds" in items_by_modality:

725 - image_embeds_lst = items_by_modality["image_embeds"]

726 - mm_inputs["image"] = (

727 - image_embeds_lst if len(image_embeds_lst) != 1 else image_embeds_lst[0]

728 - )

775 + mm_inputs["image"] = _get_embeds_data(items_by_modality, "image")

776 if "image" in items_by_modality:

777 mm_inputs["image"] = items_by_modality["image"] # A list of images

778 if "audio_embeds" in items_by_modality:

732 - audio_embeds_lst = items_by_modality["audio_embeds"]

733 - mm_inputs["audio"] = (

734 - audio_embeds_lst if len(audio_embeds_lst) != 1 else audio_embeds_lst[0]

735 - )

779 + mm_inputs["audio"] = _get_embeds_data(items_by_modality, "audio")

780 if "audio" in items_by_modality:

781 mm_inputs["audio"] = items_by_modality["audio"] # A list of audios

782 if "video" in items_by_modality:

783 mm_inputs["video"] = items_by_modality["video"] # A list of videos

784 +

785 return mm_inputs

786

787 def create_parser(self) -> "BaseMultiModalContentParser":

@@ -747,38 +792,32 @@ class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):

792 async def all_mm_data(self) -> MultiModalDataDict | None:

793 if not self._items_by_modality:

794 return None

750 - mm_inputs = {}

751 - items_by_modality = {}

752 - for modality, items in self._items_by_modality.items():

753 - coros = []

754 - for item in items:

755 - if item is not None:

756 - coros.append(item)

757 - else:

758 - coros.append(asyncio.sleep(0))

759 - items_by_modality[modality] = await asyncio.gather(*coros)

795

796 + coros_by_modality = {

797 + modality: [item or asyncio.sleep(0) for item in items]

798 + for modality, items in self._items_by_modality.items()

799 + }

800 + items_by_modality: dict[str, list[object | None]] = {

801 + modality: await asyncio.gather(*coros)

802 + for modality, coros in coros_by_modality.items()

803 + }

804 if "image" in items_by_modality and "image_embeds" in items_by_modality:

805 raise ValueError("Mixing raw image and embedding inputs is not allowed")

806 if "audio" in items_by_modality and "audio_embeds" in items_by_modality:

807 raise ValueError("Mixing raw audio and embedding inputs is not allowed")

808

809 + mm_inputs = {}

810 if "image_embeds" in items_by_modality:

767 - image_embeds_lst = items_by_modality["image_embeds"]

768 - mm_inputs["image"] = (

769 - image_embeds_lst if len(image_embeds_lst) != 1 else image_embeds_lst[0]

770 - )

811 + mm_inputs["image"] = _get_embeds_data(items_by_modality, "image")

812 if "image" in items_by_modality:

813 mm_inputs["image"] = items_by_modality["image"] # A list of images

814 if "audio_embeds" in items_by_modality:

774 - audio_embeds_lst = items_by_modality["audio_embeds"]

775 - mm_inputs["audio"] = (

776 - audio_embeds_lst if len(audio_embeds_lst) != 1 else audio_embeds_lst[0]

777 - )

815 + mm_inputs["audio"] = _get_embeds_data(items_by_modality, "audio")

816 if "audio" in items_by_modality:

817 mm_inputs["audio"] = items_by_modality["audio"] # A list of audios

818 if "video" in items_by_modality:

819 mm_inputs["video"] = items_by_modality["video"] # A list of videos

820 +

821 return mm_inputs

822

823 def create_parser(self) -> "BaseMultiModalContentParser":

vllm/v1/engine/input_processor.py

+20 -10

@@ -188,29 +188,39 @@ class InputProcessor:

188 def _validate_single_prompt(single_prompt: dict | str) -> None:

189 if not isinstance(single_prompt, dict):

190 return

191 +

192 mm_data = single_prompt.get("multi_modal_data")

193 mm_uuids = single_prompt.get("multi_modal_uuids")

194 if not mm_data or not mm_uuids:

195 return

196

197 + import torch

198 +

199 + def _get_len(items: object):

200 + if isinstance(items, dict): # Embedding inputs

201 + return _get_len(next(iter(items.values()))) if items else 1

202 +

203 + if isinstance(items, list):

204 + return len(items)

205 + if isinstance(items, torch.Tensor):

206 + # To keep backwards compatibility for single item embedding input

207 + return 1 if getattr(items, "_is_single_item", False) else len(items)

208 +

209 + return 1

210 +

211 for modality, items in mm_data.items():

212 if modality in mm_uuids:

198 - data_len = len(items) if isinstance(items, list) else 1

199 - uuid_len = (

200 - len(mm_uuids[modality])

201 - if isinstance(mm_uuids[modality], list)

202 - else 1

203 - )

213 + data_len = _get_len(items)

214 + uuid_len = _get_len(mm_uuids[modality])

215 if uuid_len != data_len:

216 raise ValueError(

206 - f"multi_modal_uuids for modality '{modality}' "

217 + f"multi_modal_uuids for modality {modality!r} "

218 "must have same length as data: got "

208 - f"{uuid_len} uuids vs "

209 - f"{data_len} items."

219 + f"{uuid_len} uuids vs {data_len} items."

220 )

221 else:

222 raise ValueError(

213 - f"multi_modal_uuids for modality '{modality}' must "

223 + f"multi_modal_uuids for modality {modality!r} must "

224 "be provided if multi_modal_data is provided."

225 )

226