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

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

[mypy] Further improve MM type annotations (#25654)

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

Cyrus Leung committed Sep 25, 2025 at 18:57 UTC 12c1287d6418d9a97155bb036b9e79cd3aff803b

6 files changed +90 -48

vllm/model_executor/models/transformers.py

+5 -2

@@ -415,9 +415,12 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):

415 self._get_mm_fields_config(processed_data, hf_processor_mm_kwargs,

416 num_image_patches),

417 )

418 +

419 # Use overrides if provided; fallback to data-dependent hashing.

419 - mm_hashes = (mm_uuids if mm_uuids is not None else self._hash_mm_items(

420 - mm_items, hf_processor_mm_kwargs, tokenization_kwargs))

420 + mm_hashes = self._hash_mm_items(mm_items,

421 + hf_processor_mm_kwargs,

422 + tokenization_kwargs,

423 + mm_uuids=mm_uuids)

424

425 return MultiModalInputs(

426 type="multimodal",

vllm/multimodal/inputs.py

+15 -8

@@ -14,7 +14,7 @@ import numpy as np

14 from typing_extensions import NotRequired, TypeAlias, TypeVar, deprecated

16 from vllm.utils import LazyLoader, full_groupby, is_list_of

17 -from vllm.utils.jsontree import JSONTree, json_map_leaves

17 +from vllm.utils.jsontree import json_map_leaves

19 if TYPE_CHECKING:

20 import torch

@@ -203,7 +203,7 @@ def nested_tensors_equal(a: NestedTensors, b: NestedTensors) -> bool:

203 return a == b

204

205

206 -BatchedTensorInputs: TypeAlias = Mapping[str, NestedTensors]

206 +BatchedTensorInputs: TypeAlias = dict[str, NestedTensors]

207 """

208 A dictionary containing nested tensors which have been batched via

209 [`MultiModalKwargs.batch`][vllm.multimodal.inputs.MultiModalKwargs.batch].

@@ -377,6 +377,7 @@ class MultiModalBatchedField(BaseMultiModalField):

377 pin_memory: bool,

378 ) -> NestedTensors:

379 if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"):

380 + batch = cast(list[torch.Tensor], batch)

381 if len(batch) == 1:

382 # An optimization when `batch` contains only one tensor:

383 # - produce exactly same result as `torch.stack(batch)`

@@ -422,6 +423,7 @@ class MultiModalFlatField(BaseMultiModalField):

423 pin_memory: bool,

424 ) -> NestedTensors:

425 if len(batch) > 0 and is_list_of(batch, torch.Tensor, check="all"):

426 + batch = cast(list[torch.Tensor], batch)

427 if len(batch) == 1:

428 # An optimization when `batch` contains only one tensor:

429 # - produce exactly same result as `torch.concat(batch)`

@@ -764,6 +766,15 @@ class MultiModalKwargsItems(UserDict[str, Sequence[_I]]):

766

767 return super().__getitem__(modality) # type: ignore[return-value]

768

769 + def require_data(self) -> "MultiModalKwargsItems[MultiModalKwargsItem]":

770 + for modality, items in self.items():

771 + for i, item in enumerate(items):

772 + if item is None:

773 + raise RuntimeError(

774 + f"Found empty mm_items[{modality}][{i}]")

775 +

776 + return self # type: ignore[return-value]

777 +

778 def get_data(self, *, pin_memory: bool = False) -> "MultiModalKwargs":

779 elems_by_key = defaultdict[str, list[MultiModalFieldElem]](list)

780 for modality, items in self.items():

@@ -897,15 +908,11 @@ class MultiModalKwargs(UserDict[str, NestedTensors]):

908 *,

909 device: torch.types.Device,

910 ) -> BatchedTensorInputs:

900 - json_inputs = cast(JSONTree[torch.Tensor], batched_inputs)

901 -

902 - json_mapped = json_map_leaves(

911 + return json_map_leaves(

912 lambda x: x.to(device=device, non_blocking=True),

904 - json_inputs,

913 + batched_inputs,

914 )

915

907 - return cast(BatchedTensorInputs, json_mapped)

908 -

916 def __getitem__(self, key: str):

917 if key not in self:

918 raise KeyError(f"Keyword argument {key!r} not found. "

vllm/multimodal/processing.py

+9 -10

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

1585 *,

1586 mm_uuids: Optional[MultiModalUUIDDict] = None,

1587 ) -> MultiModalHashes:

1588 - """Create MM hashes to be returned (only used in V1).

1588 + """Create MM hashes to be returned.

1589

1590

1591 Note: When overrides are provided via callers of `apply`,

@@ -2098,23 +2098,22 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):

2098 encoder_inputs: MultiModalInputs,

2099 ):

2100 tokenizer = self.info.get_tokenizer()

2101 - decoder_prompt = self.create_decoder_prompt(prompt, mm_data)

2102 - if isinstance(decoder_prompt, str):

2101 + decoder_prompt_raw = self.create_decoder_prompt(prompt, mm_data)

2102 + if isinstance(decoder_prompt_raw, str):

2103 + decoder_prompt = decoder_prompt_raw

2104 decoder_prompt_ids = encode_tokens(tokenizer,

2104 - decoder_prompt,

2105 + decoder_prompt_raw,

2106 add_special_tokens=False)

2107 else:

2107 - decoder_prompt_ids = decoder_prompt

2108 - decoder_prompt = decode_tokens(tokenizer, decoder_prompt)

2108 + decoder_prompt = decode_tokens(tokenizer, decoder_prompt_raw)

2109 + decoder_prompt_ids = decoder_prompt_raw

2110

2111 mm_inputs = MultiModalEncDecInputs(

2112 encoder_prompt=encoder_inputs["prompt"],

2113 encoder_prompt_token_ids=encoder_inputs["prompt_token_ids"],

2114 **encoder_inputs)

2114 - mm_inputs.update({

2115 - "prompt": decoder_prompt,

2116 - "prompt_token_ids": decoder_prompt_ids

2117 - })

2115 + mm_inputs["prompt"] = decoder_prompt

2116 + mm_inputs["prompt_token_ids"] = decoder_prompt_ids

2117 return mm_inputs

2118

2119 def apply(

vllm/multimodal/profiling.py

+3 -3

@@ -13,7 +13,7 @@ import vllm.envs as envs

13 from vllm.logger import init_logger

15 from .inputs import (MultiModalDataDict, MultiModalEncDecInputs,

16 - MultiModalInputs, MultiModalKwargsOptionalItems,

16 + MultiModalInputs, MultiModalKwargsItems,

17 MultiModalPlaceholderDict)

18 from .processing import (BaseMultiModalProcessor, BaseProcessingInfo,

19 EncDecMultiModalProcessor)

@@ -43,7 +43,7 @@ class DummyDecoderData(NamedTuple):

43 """Dummy data used for profiling."""

45 prompt_token_ids: list[int]

46 - multi_modal_data: MultiModalKwargsOptionalItems

46 + multi_modal_data: MultiModalKwargsItems

47 multi_modal_placeholders: MultiModalPlaceholderDict

@@ -239,7 +239,7 @@ class MultiModalProfiler(Generic[_I]):

239

240 return DummyDecoderData(

241 prompt_token_ids=prompt_token_ids,

242 - multi_modal_data=mm_inputs["mm_kwargs"],

242 + multi_modal_data=mm_inputs["mm_kwargs"].require_data(),

243 multi_modal_placeholders=mm_inputs["mm_placeholders"],

244 )

245

vllm/multimodal/utils.py

+24 -20

@@ -19,6 +19,7 @@ from typing_extensions import deprecated

20 import vllm.envs as envs

21 from vllm.connections import HTTPConnection, global_http_connection

22 +from vllm.utils.jsontree import json_map_leaves

24 from .audio import AudioMediaIO

25 from .base import MediaIO

@@ -383,6 +384,7 @@ def group_mm_kwargs_by_modality(

384 *,

385 device: torch.types.Device = None,

386 pin_memory: bool = False,

387 + merge_by_field_config: bool = False,

388 ) -> Iterable[tuple[str, int, BatchedTensorInputs]]:

389 """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same

390 modality together into the same `MultiModalKwargs` instance.

@@ -400,29 +402,31 @@ 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

403 - # mm_kwargs_group = MultiModalKwargsItems.from_items(items_lst) \

404 - # .get_data(pin_memory=pin_memory)

405 -

406 - # if device is not None:

407 - # mm_kwargs_group = json_map_leaves(

408 - # lambda x: x.to(device=device),

409 - # mm_kwargs_group,

410 - # )

411 -

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

405 + # TODO: Enable `merge_by_field_config` for all models

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

407 # that are meant to be stacked anyway).

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

416 - mm_kwargs_group = MultiModalKwargs.as_kwargs(

417 - MultiModalKwargs.batch(

418 - [

419 - MultiModalKwargsItems.from_seq([item]).get_data()

420 - for item in items_lst

421 - ],

422 - pin_memory=pin_memory,

423 - ),

424 - device=device,

425 - )

409 + if merge_by_field_config:

410 + mm_kwargs_group: BatchedTensorInputs = dict(

411 + MultiModalKwargsItems.from_seq(items_lst).get_data(

412 + pin_memory=pin_memory))

413 +

414 + if device is not None:

415 + mm_kwargs_group = json_map_leaves(

416 + lambda x: x.to(device=device),

417 + mm_kwargs_group,

418 + )

419 + else:

420 + mm_kwargs_group = MultiModalKwargs.as_kwargs(

421 + MultiModalKwargs.batch(

422 + [

423 + MultiModalKwargsItems.from_seq([item]).get_data()

424 + for item in items_lst

425 + ],

426 + pin_memory=pin_memory,

427 + ),

428 + device=device,

429 + )

430

431 yield modality, len(items_lst), mm_kwargs_group

432

vllm/utils/jsontree.py

+34 -5

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

5 from collections.abc import Iterable

6 from functools import reduce

7 -from typing import Callable, TypeVar, Union, cast, overload

7 +from typing import TYPE_CHECKING, Callable, TypeVar, Union, cast, overload

8 +

9 +if TYPE_CHECKING:

10 + import torch

11 +

12 + from vllm.multimodal.inputs import BatchedTensorInputs

14 _T = TypeVar("_T")

15 _U = TypeVar("_U")

@@ -17,6 +22,19 @@ JSONTree = Union[

22 ]

23 """A nested JSON structure where the leaves need not be JSON-serializable."""

25 +_JSONTree = Union[

26 + dict[str, "JSONTree[_T]"],

27 + list["JSONTree[_T]"],

28 + tuple["JSONTree[_T]", ...],

29 + dict[str, _T],

30 + list[_T],

31 + tuple[_T, ...],

32 + _T,

33 +]

34 +"""

35 +Same as `JSONTree` but with additional `Union` members to satisfy overloads.

36 +"""

37 +

39 def json_iter_leaves(value: JSONTree[_T]) -> Iterable[_T]:

40 """Iterate through each leaf in a nested JSON structure."""

@@ -30,6 +48,14 @@ def json_iter_leaves(value: JSONTree[_T]) -> Iterable[_T]:

48 yield value

51 +@overload

52 +def json_map_leaves(

53 + func: Callable[["torch.Tensor"], "torch.Tensor"],

54 + value: "BatchedTensorInputs",

55 +) -> "BatchedTensorInputs":

56 + ...

57 +

58 +

59 @overload

60 def json_map_leaves(

61 func: Callable[[_T], _U],

@@ -64,11 +90,14 @@ def json_map_leaves(

91 def json_map_leaves(

92 func: Callable[[_T], _U],

67 - value: Union[dict[str, _T], list[_T], tuple[_T, ...], JSONTree[_T]],

68 -) -> Union[dict[str, _U], list[_U], tuple[_U, ...], JSONTree[_U]]:

93 + value: Union["BatchedTensorInputs", _JSONTree[_T]],

94 +) -> Union["BatchedTensorInputs", _JSONTree[_U]]:

95 """Apply a function to each leaf in a nested JSON structure."""

96 if isinstance(value, dict):

71 - return {k: json_map_leaves(func, v) for k, v in value.items()}

97 + return {

98 + k: json_map_leaves(func, v) # type: ignore[arg-type]

99 + for k, v in value.items()

100 + }

101 elif isinstance(value, list):

102 return [json_map_leaves(func, v) for v in value]

103 elif isinstance(value, tuple):

@@ -125,7 +154,7 @@ def json_reduce_leaves(

154

155 def json_reduce_leaves(

156 func: Callable[..., Union[_T, _U]],

128 - value: Union[dict[str, _T], list[_T], tuple[_T, ...], JSONTree[_T]],

157 + value: _JSONTree[_T],

158 initial: _U = cast(_U, ...), # noqa: B008

159 /,

160 ) -> Union[_T, _U]: