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

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] Allow optional MultiModalKwargsItem in IPC (#23022)

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

Cyrus Leung committed Aug 16, 2025 at 19:30 UTC 4dff91c93da668f4cca3f80aa3a94622d21c34fc

10 files changed +59 -108

tests/v1/core/test_kv_cache_utils.py

+2 -10

@@ -7,9 +7,7 @@ import pytest

7 import torch

9 from vllm.config import ModelConfig, SchedulerConfig, VllmConfig

10 -from vllm.multimodal.inputs import (MultiModalBatchedField,

11 - MultiModalFieldElem, MultiModalKwargsItem,

12 - PlaceholderRange)

10 +from vllm.multimodal.inputs import MultiModalKwargsItem, PlaceholderRange

11 from vllm.sampling_params import SamplingParams

12 from vllm.utils import GiB_bytes, sha256, sha256_cbor_64bit

13 from vllm.v1.core.kv_cache_manager import KVCacheManager

@@ -42,13 +40,7 @@ def make_request(

40 if mm_positions is None:

41 mm_kwargs = None

42 else:

45 - mm_elem = MultiModalFieldElem(

46 - modality="dummy_m",

47 - key="dummy_k",

48 - data=None,

49 - field=MultiModalBatchedField(),

50 - )

51 - mm_item = MultiModalKwargsItem.from_elems([mm_elem])

43 + mm_item = MultiModalKwargsItem.dummy("dummy_m")

44 mm_kwargs = [mm_item] * len(mm_positions)

46 return Request(request_id=request_id,

tests/v1/core/test_prefix_caching.py

+2 -10

@@ -9,9 +9,7 @@ import pytest

9 import torch

11 from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved

12 -from vllm.multimodal.inputs import (MultiModalBatchedField,

13 - MultiModalFieldElem, MultiModalKwargsItem,

14 - PlaceholderRange)

12 +from vllm.multimodal.inputs import MultiModalKwargsItem, PlaceholderRange

13 from vllm.sampling_params import SamplingParams

14 from vllm.utils import sha256, sha256_cbor_64bit

15 from vllm.v1.core.block_pool import BlockPool

@@ -37,13 +35,7 @@ def make_request(

35 if mm_positions is None:

36 mm_kwargs = None

37 else:

40 - mm_elem = MultiModalFieldElem(

41 - modality="dummy_m",

42 - key="dummy_k",

43 - data=None,

44 - field=MultiModalBatchedField(),

45 - )

46 - mm_item = MultiModalKwargsItem.from_elems([mm_elem])

38 + mm_item = MultiModalKwargsItem.dummy("dummy_m")

39 mm_kwargs = [mm_item] * len(mm_positions)

41 return Request(request_id=request_id,

tests/v1/core/test_scheduler.py

+2 -10

@@ -8,9 +8,7 @@ import torch

9 from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,

10 SchedulerConfig, SpeculativeConfig, VllmConfig)

11 -from vllm.multimodal.inputs import (MultiModalBatchedField,

12 - MultiModalFieldElem, MultiModalKwargsItem,

13 - PlaceholderRange)

11 +from vllm.multimodal.inputs import MultiModalKwargsItem, PlaceholderRange

12 from vllm.sampling_params import GuidedDecodingParams, SamplingParams

13 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput

14 from vllm.v1.core.sched.scheduler import Scheduler

@@ -1328,13 +1326,7 @@ def create_requests_with_priority(

1326 for i in range(num_requests):

1327 if mm_positions is not None:

1328 mm_position = mm_positions[i]

1331 - mm_elem = MultiModalFieldElem(

1332 - modality="dummy_m",

1333 - key="dummy_k",

1334 - data=None,

1335 - field=MultiModalBatchedField(),

1336 - )

1337 - mm_item = MultiModalKwargsItem.from_elems([mm_elem])

1329 + mm_item = MultiModalKwargsItem.dummy("dummy_m")

1330 mm_kwargs = [mm_item] * len(mm_position)

1331 else:

1332 mm_position = None

tests/v1/core/utils.py

+2 -10

@@ -6,9 +6,7 @@ import torch

7 from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,

8 SchedulerConfig, SpeculativeConfig, VllmConfig)

9 -from vllm.multimodal.inputs import (MultiModalBatchedField,

10 - MultiModalFieldElem, MultiModalKwargsItem,

11 - PlaceholderRange)

9 +from vllm.multimodal.inputs import MultiModalKwargsItem, PlaceholderRange

10 from vllm.sampling_params import SamplingParams

11 from vllm.v1.core.kv_cache_utils import (get_request_block_hasher,

12 init_none_hash)

@@ -143,13 +141,7 @@ def create_requests(

141 for i in range(num_requests):

142 if mm_positions is not None:

143 mm_position = mm_positions[i]

146 - mm_elem = MultiModalFieldElem(

147 - modality="dummy_m",

148 - key="dummy_k",

149 - data=None,

150 - field=MultiModalBatchedField(),

151 - )

152 - mm_item = MultiModalKwargsItem.from_elems([mm_elem])

144 + mm_item = MultiModalKwargsItem.dummy("dummy_m")

145 mm_kwargs = [mm_item] * len(mm_position)

146 mm_hashes = ["hash"] * len(mm_position)

147 else:

vllm/multimodal/inputs.py

+17 -45

@@ -4,7 +4,7 @@

4 from abc import ABC, abstractmethod

5 from collections import UserDict, defaultdict

6 from collections.abc import Mapping, Sequence

7 -from dataclasses import dataclass, replace

7 +from dataclasses import dataclass

8 from functools import partial

9 from itertools import accumulate

10 from typing import (TYPE_CHECKING, Any, Literal, Optional, TypedDict, TypeVar,

@@ -218,7 +218,7 @@ class MultiModalFieldElem:

218 i.e. the name of the keyword argument to be passed to the model.

219 """

220

221 - data: Optional[NestedTensors]

221 + data: NestedTensors

222 """

223 The tensor data of this field in

224 [`MultiModalKwargs`][vllm.multimodal.inputs.MultiModalKwargs],

@@ -315,13 +315,8 @@ class BaseMultiModalField(ABC):

315 if len(set(field_types)) > 1:

316 raise ValueError(f"Cannot merge different {field_types=}")

317

318 - validated_data = list[NestedTensors]()

319 - for i, elem in enumerate(elems):

320 - assert elem.data is not None, (

321 - f"Cannot merge with empty `elems[{i}]`")

322 - validated_data.append(elem.data)

323 -

324 - return self._reduce_data(validated_data, pin_memory=pin_memory)

318 + batch = [elem.data for elem in elems]

319 + return self._reduce_data(batch, pin_memory=pin_memory)

320

321

322 @dataclass(frozen=True)

@@ -643,6 +638,17 @@ class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]):

638 [`MultiModalDataItems`][vllm.multimodal.parse.MultiModalDataItems].

639 """

640

641 + @staticmethod

642 + def dummy(modality: str):

643 + """Convenience class for testing."""

644 + mm_elem = MultiModalFieldElem(

645 + modality=modality,

646 + key="dummy",

647 + data=torch.empty(1),

648 + field=MultiModalSharedField(1),

649 + )

650 + return MultiModalKwargsItem.from_elems([mm_elem])

651 +

652 @staticmethod

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

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

@@ -654,46 +660,12 @@ class MultiModalKwargsItem(UserDict[str, MultiModalFieldElem]):

660 assert len(modalities) == 1, f"Found different modalities={modalities}"

661 self._modality = next(iter(modalities))

662

657 - self._is_empty = any(elem.data is None for elem in self.values())

658 -

663 @property

664 def modality(self) -> str:

665 return self._modality

666

663 - @property

664 - def is_empty(self) -> bool:

665 - return self._is_empty

666 -

667 - def get_data(self) -> Optional[Mapping[str, NestedTensors]]:

668 - if self._is_empty:

669 - return None

670 -

671 - out_data = dict[str, NestedTensors]()

672 - for key, elem in self.items():

673 - assert elem.data is not None, (

674 - f"Cannot get data of empty `elem[{key!r}]`")

675 - out_data[key] = elem.data

676 -

677 - return out_data

678 -

679 - def require_data(self) -> Mapping[str, NestedTensors]:

680 - if (data := self.get_data()) is None:

681 - raise RuntimeError("Cannot get data of empty item")

682 -

683 - return data

684 -

685 - # These methods create a new item to avoid mutating cached items in place

686 - def with_data(self, data: Mapping[str, NestedTensors]):

687 - return MultiModalKwargsItem({

688 - key: replace(elem, data=data[key])

689 - for key, elem in self.items()

690 - })

691 -

692 - def without_data(self):

693 - return MultiModalKwargsItem({

694 - key: replace(elem, data=None)

695 - for key, elem in self.items()

696 - })

667 + def get_data(self) -> Mapping[str, NestedTensors]:

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

669

670

671 # NOTE: UserDict is for V0 compatibility.

vllm/v1/engine/__init__.py

+2 -1

@@ -3,6 +3,7 @@

4 import enum

5 import time

6 +from collections.abc import Sequence

7 from typing import Any, Optional, Union

9 import msgspec

@@ -47,7 +48,7 @@ class EngineCoreRequest(

49 request_id: str

50 prompt_token_ids: list[int]

50 - mm_kwargs: Optional[list[MultiModalKwargsItem]]

51 + mm_kwargs: Optional[Sequence[Optional[MultiModalKwargsItem]]]

52 mm_hashes: Optional[list[str]]

53 mm_placeholders: Optional[list[PlaceholderRange]]

54 sampling_params: Optional[SamplingParams]

vllm/v1/engine/mm_input_cache.py

+18 -15

@@ -1,11 +1,12 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from collections.abc import Mapping

4 -from typing import TYPE_CHECKING

3 +from collections.abc import Sequence

4 +from typing import TYPE_CHECKING, Optional

6 from vllm.multimodal import MultiModalRegistry

7 from vllm.multimodal.cache import MultiModalCache, MultiModalCacheItemMetadata

8 -from vllm.multimodal.inputs import MultiModalKwargsItem, NestedTensors

8 +from vllm.multimodal.inputs import MultiModalKwargsItem

9 +from vllm.utils import is_list_of

11 if TYPE_CHECKING:

12 from vllm.config import ModelConfig

@@ -58,21 +59,21 @@ class MultiModalInputCacheClient:

60 def get_and_update(

61 self,

61 - mm_kwargs: list[MultiModalKwargsItem],

62 + mm_kwargs: Sequence[MultiModalKwargsItem],

63 mm_hashes: list[str],

63 - ) -> list[MultiModalKwargsItem]:

64 + ) -> list[Optional[MultiModalKwargsItem]]:

65 if not self.enabled:

65 - return mm_kwargs

66 + return list(mm_kwargs)

68 assert len(mm_kwargs) == len(mm_hashes)

69 - out_mm_items = list[MultiModalKwargsItem]()

70 + out_mm_items = list[Optional[MultiModalKwargsItem]]()

71 for mm_item, mm_hash in zip(mm_kwargs, mm_hashes):

72 if self.mm_cache.get(mm_hash) is not None:

72 - out_mm_items.append(mm_item.without_data())

73 + out_mm_items.append(None)

74 else:

75 self.mm_cache[mm_hash] = \

75 - MultiModalCacheItemMetadata.wraps(mm_item.require_data())

76 + MultiModalCacheItemMetadata.wraps(mm_item)

77 out_mm_items.append(mm_item)

79 return out_mm_items

@@ -91,25 +92,27 @@ class MultiModalInputCacheServer:

92 self.enabled = mm_registry.enable_mm_input_cache(model_config)

93 self.mm_cache = MultiModalCache.get_lru_cache(

94 model_config.get_mm_input_cache_gb(),

94 - Mapping[str, NestedTensors],

95 + MultiModalKwargsItem,

96 )

98 def get_and_update(

99 self,

99 - mm_kwargs: list[MultiModalKwargsItem],

100 + mm_kwargs: Sequence[Optional[MultiModalKwargsItem]],

101 mm_hashes: list[str],

102 ) -> list[MultiModalKwargsItem]:

103 if not self.enabled:

103 - return mm_kwargs

104 + mm_kwargs_lst = list(mm_kwargs)

105 + assert is_list_of(mm_kwargs_lst, MultiModalKwargsItem)

106 + return mm_kwargs_lst

107

108 assert len(mm_kwargs) == len(mm_hashes)

109

110 out_mm_items = list[MultiModalKwargsItem]()

111 for mm_item, mm_hash in zip(mm_kwargs, mm_hashes):

109 - if (mm_data := mm_item.get_data()) is None:

110 - out_mm_items.append(mm_item.with_data(self.mm_cache[mm_hash]))

112 + if mm_item is None:

113 + out_mm_items.append(self.mm_cache[mm_hash])

114 else:

112 - self.mm_cache[mm_hash] = mm_data

115 + self.mm_cache[mm_hash] = mm_item

116 out_mm_items.append(mm_item)

117

118 return out_mm_items

vllm/v1/engine/processor.py

+7 -3

@@ -17,6 +17,7 @@ from vllm.multimodal.utils import argsort_mm_positions

17 from vllm.pooling_params import PoolingParams

18 from vllm.sampling_params import SamplingParams

19 from vllm.transformers_utils.tokenizer_group import TokenizerGroup

20 +from vllm.utils import is_list_of

21 from vllm.v1.engine import EngineCoreRequest

22 from vllm.v1.engine.mm_input_cache import MultiModalInputCacheClient

23 from vllm.v1.structured_output.backend_guidance import (

@@ -295,7 +296,7 @@ class Processor:

296 pooling_params = params.clone()

297

298 # Multimodal related.

298 - sorted_mm_inputs: Optional[list[MultiModalKwargsItem]] = None

299 + sorted_mm_inputs: Optional[list[Optional[MultiModalKwargsItem]]] = None

300 sorted_mm_positions: Optional[list[PlaceholderRange]] = None

301 sorted_mm_hashes: Optional[list[str]] = None

302 if decoder_inputs["type"] == "multimodal":

@@ -308,7 +309,7 @@ class Processor:

309 # in the input sequence.

310 sorted_mm_idxs = argsort_mm_positions(decoder_mm_positions)

311

311 - sorted_mm_inputs = [

312 + orig_sorted_mm_inputs = [

313 decoder_mm_inputs.get_item(modality, idx)

314 for modality, idx in sorted_mm_idxs

315 ]

@@ -323,9 +324,12 @@ class Processor:

324

325 if sorted_mm_hashes is not None:

326 sorted_mm_inputs = self.mm_input_cache_client.get_and_update(

326 - sorted_mm_inputs,

327 + orig_sorted_mm_inputs,

328 sorted_mm_hashes,

329 )

330 + else:

331 + assert is_list_of(orig_sorted_mm_inputs, MultiModalKwargsItem)

332 + sorted_mm_inputs = orig_sorted_mm_inputs

333

334 return decoder_inputs.get("prompt"), EngineCoreRequest(

335 request_id=request_id,

vllm/v1/request.py

+5 -2

@@ -125,14 +125,17 @@ class Request:

125 block_hasher: Optional[Callable[["Request"], list["BlockHash"]]]

126 ) -> "Request":

127 if request.mm_kwargs is not None:

128 - assert is_list_of(request.mm_kwargs, MultiModalKwargsItem), (

128 + mm_kwargs_lst = list(request.mm_kwargs)

129 + assert is_list_of(mm_kwargs_lst, MultiModalKwargsItem), (

130 "mm_kwargs was not updated in EngineCore.add_request")

131 + else:

132 + mm_kwargs_lst = None

133

134 return cls(

135 request_id=request.request_id,

136 client_index=request.client_index,

137 prompt_token_ids=request.prompt_token_ids,

135 - multi_modal_kwargs=request.mm_kwargs,

138 + multi_modal_kwargs=mm_kwargs_lst,

139 multi_modal_hashes=request.mm_hashes,

140 multi_modal_placeholders=request.mm_placeholders,

141 sampling_params=request.sampling_params,

vllm/v1/worker/gpu_model_runner.py

+2 -2

@@ -500,8 +500,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

500 second_per_grid_ts = []

501 audio_feature_lengths = []

502 use_audio_in_video = False

503 - for item in self.requests[req_id].mm_kwargs:

504 - mm_input = item.require_data()

503 + for mm_item in self.requests[req_id].mm_kwargs:

504 + mm_input = mm_item.get_data()

505 if mm_input.get("image_grid_thw") is not None:

506 image_grid_thw.append(

507 mm_input["image_grid_thw"].tolist())