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

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 14, 05:30 UTC

[Refactor] Defer tensor data construction in MultiModalKwargs (#23030)

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

Cyrus Leung committed Aug 17, 2025 at 12:05 UTC 5c32143b9db19ae728087019678843fa238afa82

12 files changed +73 -104

tests/multimodal/test_cache.py

+1 -1

@@ -25,7 +25,7 @@ def _dummy_item(modality: str, size_by_key: dict[str, int]):

27 def _dummy_kw(size_by_key_modality: dict[str, dict[str, int]]):

28 - return MultiModalKwargs.from_items([

28 + return MultiModalKwargs([

29 _dummy_item(modality, size_by_key)

30 for modality, size_by_key in size_by_key_modality.items()

31 ])

tests/v1/test_serial_utils.py

+1 -33

@@ -100,38 +100,6 @@ class MyRequest(msgspec.Struct):

100

101

102 def test_multimodal_kwargs():

103 - d = {

104 - "foo":

105 - torch.zeros(20000, dtype=torch.float16),

106 - "bar": [torch.zeros(i * 1000, dtype=torch.int8) for i in range(3)],

107 - "baz": [

108 - torch.rand((256), dtype=torch.float16),

109 - [

110 - torch.rand((1, 12), dtype=torch.float32),

111 - torch.rand((3, 5, 7), dtype=torch.float64),

112 - ], [torch.rand((4, 4), dtype=torch.float16)]

113 - ],

114 - }

115 -

116 - # pack mm kwargs into a mock request so that it can be decoded properly

117 - req = MyRequest(mm=[MultiModalKwargs(d)])

118 -

119 - encoder = MsgpackEncoder()

120 - decoder = MsgpackDecoder(MyRequest)

121 -

122 - encoded = encoder.encode(req)

123 -

124 - assert len(encoded) == 6

125 -

126 - total_len = sum(memoryview(x).cast("B").nbytes for x in encoded)

127 -

128 - # expected total encoding length, should be 44559, +-20 for minor changes

129 - assert 44539 <= total_len <= 44579

130 - decoded: MultiModalKwargs = decoder.decode(encoded).mm[0]

131 - assert all(nested_equal(d[k], decoded[k]) for k in d)

132 -

133 -

134 -def test_multimodal_items_by_modality():

103 e1 = MultiModalFieldElem("audio", "a0",

104 torch.zeros(1000, dtype=torch.bfloat16),

105 MultiModalBatchedField())

@@ -151,7 +119,7 @@ def test_multimodal_items_by_modality():

119 audio = MultiModalKwargsItem.from_elems([e1])

120 video = MultiModalKwargsItem.from_elems([e2])

121 image = MultiModalKwargsItem.from_elems([e3, e4])

154 - mm = MultiModalKwargs.from_items([audio, video, image])

122 + mm = MultiModalKwargs([audio, video, image])

123

124 # pack mm kwargs into a mock request so that it can be decoded properly

125 req = MyRequest([mm])

vllm/inputs/registry.py

+1 -1

@@ -240,6 +240,6 @@ class InputRegistry:

240

241 return DummyData(

242 seq_data=SequenceData.from_seqs(dec_data.prompt_token_ids),

243 - multi_modal_data=dec_data.multi_modal_data,

243 + multi_modal_data=dec_data.multi_modal_data.get_data(),

244 multi_modal_placeholders=dec_data.multi_modal_placeholders,

245 )

vllm/model_executor/models/prithvi_geospatial_mae.py

+1 -1

@@ -136,7 +136,7 @@ class PrithviGeoSpatialMAEMultiModalProcessor(BaseMultiModalProcessor):

136 type="multimodal",

137 prompt=prompt,

138 prompt_token_ids=[1],

139 - mm_kwargs=MultiModalKwargs.from_items(multimodal_kwargs_items),

139 + mm_kwargs=MultiModalKwargs(multimodal_kwargs_items),

140 mm_hashes=None,

141 mm_placeholders=mm_placeholders,

142 )

vllm/multimodal/base.py

+1 -1

@@ -99,7 +99,7 @@ class MultiModalPlaceholderMap:

99 seq_mm_placeholders = seq_group.multi_modal_placeholders

100

101 if not seq_mm_data or not seq_mm_placeholders:

102 - return MultiModalKwargs({}), {}

102 + return MultiModalKwargs(), {}

103

104 placeholder_maps = dict[str, MultiModalPlaceholderMap]()

105

vllm/multimodal/cache.py

+1 -1

@@ -46,7 +46,7 @@ class MultiModalCache:

46 ) -> int:

47 # MultiModalKwargs is not a subclass of dict

48 if isinstance(leaf, MultiModalKwargs):

49 - return cls.get_item_size(leaf.data, debug=debug)

49 + return cls.get_item_size(leaf.get_data(), debug=debug)

51 # MultiModalKwargsItem is not a subclass of dict

52 if isinstance(leaf, MultiModalKwargsItem):

vllm/multimodal/inputs.py

+54 -42

@@ -653,7 +653,7 @@ class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]):

653 def from_elems(elems: Sequence[MultiModalFieldElem]):

654 return MultiModalKwargsItem({elem.key: elem for elem in elems})

655

656 - def __init__(self, data: Mapping[str, MultiModalFieldElem]) -> None:

656 + def __init__(self, data: Mapping[str, MultiModalFieldElem] = {}) -> None:

657 super().__init__(data)

658

659 modalities = {elem.modality for elem in self.data.values()}

@@ -668,9 +668,7 @@ class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]):

668 return {key: elem.data for key, elem in self.items()}

669

670

671 -# NOTE: UserDict is for V0 compatibility.

672 -# V1 should access individual items via `get_item`.

673 -class MultiModalKwargs(UserDict[str, NestedTensors]):

671 +class MultiModalKwargs:

672 """

673 A dictionary that represents the keyword arguments to

674 [`torch.nn.Module.forward`][].

@@ -714,40 +712,16 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):

712 elems = [v[item_idx] for v in elems_in_modality.values()]

713 items.append(MultiModalKwargsItem.from_elems(elems))

714

717 - return MultiModalKwargs.from_items(items)

715 + return MultiModalKwargs(items)

716

719 - @staticmethod

720 - def from_items(

721 - items: Sequence[MultiModalKwargsItem],

722 - *,

723 - pin_memory: bool = False,

724 - ):

725 - """Construct a new

726 - [`MultiModalKwargs`][vllm.multimodal.inputs.MultiModalKwargs]

727 - from multiple items."""

728 - elems_by_key = defaultdict[str, list[MultiModalFieldElem]](list)

729 - for item in items:

730 - for key, elem in item.items():

731 - elems_by_key[key].append(elem)

732 -

733 - data = {

734 - key: elems[0].field.reduce_data(elems, pin_memory=pin_memory)

735 - for key, elems in elems_by_key.items() if len(elems) > 0

736 - }

737 -

738 - return MultiModalKwargs(data, items=items)

739 -

740 - def __init__(

741 - self,

742 - data: Mapping[str, NestedTensors],

743 - *,

744 - items: Optional[Sequence[MultiModalKwargsItem]] = None,

745 - ) -> None:

746 - super().__init__(data)

717 + def __init__(self, items: Sequence[MultiModalKwargsItem] = ()) -> None:

718 + super().__init__()

719

748 - items_by_modality = full_groupby(items or [], key=lambda x: x.modality)

720 + items_by_modality = full_groupby(items, key=lambda x: x.modality)

721 self._items_by_modality = dict(items_by_modality)

722

723 + self._data: Optional[Mapping[str, NestedTensors]] = None

724 +

725 @property

726 def modalities(self):

727 return self._items_by_modality.keys()

@@ -839,22 +813,41 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):

813

814 return cast(BatchedTensorInputs, json_mapped)

815

842 - def __delitem__(self, key: str) -> None:

843 - super().__delitem__(key)

816 + def keys(self):

817 + return self.get_data().keys()

818 +

819 + def values(self):

820 + return self.get_data().values()

821 +

822 + def items(self):

823 + return self.get_data().items()

824 +

825 + def get(self, key: str, /, default=None):

826 + return self.get_data().get(key, default)

827 +

828 + def pop(self, key: str, *args, **kwargs):

829 + data = dict(self.get_data())

830 + res = data.pop(key, *args, **kwargs)

831

832 for items in self._items_by_modality.values():

833 for item in items:

847 - item.pop(key, None)

834 + item.pop(key, *args, **kwargs)

835 +

836 + self._data = None

837 +

838 + return res

839 +

840 + def __iter__(self):

841 + return iter(self.get_data())

842 +

843 + def __getitem__(self, key: str):

844 + return self.get_data()[key]

845

846 def __eq__(self, other: object) -> bool:

847 if not isinstance(other, self.__class__):

848 return False

852 - if self._items_by_modality != other._items_by_modality:

853 - return False

849

855 - ks = self.keys()

856 - return (ks == other.keys()

857 - and all(nested_tensors_equal(self[k], other[k]) for k in ks))

850 + return self._items_by_modality == other._items_by_modality

851

852 def _validate_modality(self, method_name: str, modality: str) -> None:

853 if not self._items_by_modality:

@@ -888,6 +881,25 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):

881 self._validate_modality("get_items", modality)

882 return self._items_by_modality[modality]

883

884 + def get_data(self,

885 + *,

886 + pin_memory: bool = False) -> Mapping[str, NestedTensors]:

887 + if self._data is not None:

888 + return self._data

889 +

890 + elems_by_key = defaultdict[str, list[MultiModalFieldElem]](list)

891 + for items in self._items_by_modality.values():

892 + for item in items:

893 + for key, elem in item.items():

894 + elems_by_key[key].append(elem)

895 +

896 + data = {

897 + key: elems[0].field.reduce_data(elems, pin_memory=pin_memory)

898 + for key, elems in elems_by_key.items() if len(elems) > 0

899 + }

900 + self._data = data

901 + return data

902 +

903

904 MultiModalPlaceholderDict: TypeAlias = Mapping[str, Sequence[PlaceholderRange]]

905 """

vllm/multimodal/processing.py

+1 -1

@@ -1480,7 +1480,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):

1480 mm_missing_kwargs=mm_missing_kwargs,

1481 )

1482

1483 - mm_kwargs = MultiModalKwargs.from_items([

1483 + mm_kwargs = MultiModalKwargs([

1484 item for cache_items in mm_cache_items_merged.values()

1485 for item in cache_items

1486 ])

vllm/multimodal/utils.py

+7 -5

@@ -402,12 +402,14 @@ def group_mm_kwargs_by_modality(

402 for modality, items in groupby(mm_kwargs, key=lambda item: item.modality):

403 items_lst = list(items)

404

405 - # mm_kwargs_group = MultiModalKwargs.from_items(items_lst,

406 - # pin_memory=pin_memory)

405 + # mm_kwargs_group = MultiModalKwargs(items_lst) \

406 + # .get_data(pin_memory=pin_memory)

407

408 # if device is not None:

409 - # mm_kwargs_group = json_map_leaves(lambda x: x.to(device=device),

410 - # mm_kwargs_group.data)

409 + # mm_kwargs_group = json_map_leaves(

410 + # lambda x: x.to(device=device),

411 + # mm_kwargs_group,

412 + # )

413

414 # TODO: Once V0 is removed, we can use the merging logic above

415 # to avoid creating an extra batch dimension (except for fields

@@ -415,7 +417,7 @@ def group_mm_kwargs_by_modality(

417 # We will also need to update each model to remove `flatten_bn`.

418 mm_kwargs_group = MultiModalKwargs.as_kwargs(

419 MultiModalKwargs.batch(

418 - [MultiModalKwargs.from_items([item]) for item in items_lst],

420 + [MultiModalKwargs([item]) for item in items_lst],

421 pin_memory=pin_memory,

422 ),

423 device=device,

vllm/sequence.py

+2 -2

@@ -524,7 +524,7 @@ class Sequence:

524 if self.inputs["type"] == "multimodal":

525 return self.inputs["mm_kwargs"]

526

527 - return MultiModalKwargs({})

527 + return MultiModalKwargs()

528

529 @property

530 def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:

@@ -780,7 +780,7 @@ class SequenceGroup:

780 return self.first_seq.multi_modal_data

781 elif self.encoder_seq is not None:

782 return self.encoder_seq.multi_modal_data

783 - return MultiModalKwargs({})

783 + return MultiModalKwargs()

784

785 @property

786 def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:

vllm/v1/serial_utils.py

+2 -15

@@ -117,16 +117,9 @@ class MsgpackEncoder:

117 return self._encode_mm_item(obj)

118

119 if isinstance(obj, MultiModalKwargs):

120 - mm: MultiModalKwargs = obj

121 - if not mm.modalities:

122 - # just return the main dict if there are no modalities.

123 - return dict(mm)

124 -

125 - # ignore the main dict, it will be re-indexed.

126 - # Any tensors *not* indexed by modality will be ignored.

120 return [

121 self._encode_mm_item(item)

129 - for itemlist in mm._items_by_modality.values()

122 + for itemlist in obj._items_by_modality.values()

123 for item in itemlist

124 ]

125

@@ -268,13 +261,7 @@ class MsgpackDecoder:

261 if issubclass(t, MultiModalKwargsItem):

262 return self._decode_mm_item(obj)

263 if issubclass(t, MultiModalKwargs):

271 - if isinstance(obj, list):

272 - return MultiModalKwargs.from_items(

273 - self._decode_mm_items(obj))

274 - return MultiModalKwargs({

275 - k: self._decode_nested_tensors(v)

276 - for k, v in obj.items()

277 - })

264 + return MultiModalKwargs(self._decode_mm_items(obj))

265 if t is UtilityResult:

266 return self._decode_utility_result(obj)

267 return obj

vllm/v1/worker/gpu_input_batch.py

+1 -1

@@ -58,7 +58,7 @@ class CachedRequestState:

58 @deprecated("`mm_inputs` is superseded by `mm_kwargs` and will be "

59 "removed in v0.13. Please use `mm_kwargs` instead.")

60 def mm_inputs(self) -> list[MultiModalKwargs]:

61 - return [MultiModalKwargs.from_items([item]) for item in self.mm_kwargs]

61 + return [MultiModalKwargs([item]) for item in self.mm_kwargs]

63 def get_token_id(self, idx: int) -> int:

64 if idx < self.num_prompt_tokens: