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
[Misc] Improve type annotations for jsontree (#25577)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Sep 24, 2025 at 22:49 UTC
9313be501721c42078e06490ff0b8a7eee0a4fb6
5 files changed
+88
-39
vllm/model_executor/models/aya_vision.py
+12
-11
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# Adapted from https://github.com/huggingface/transformers/tree/main/src/transformers/models/aya_vision
4
from collections.abc import Iterable, Mapping, Sequence
5
-from typing import Annotated, Literal, Optional, Union, cast
5
+from typing import Annotated, Literal, Optional, Union
6
7
import torch
8
from torch import nn
@@ -347,12 +347,16 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,
347
loader = AutoWeightsLoader(self)
348
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
349
350
- def _image_pixels_to_features(self, vision_tower: SiglipVisionModel,
351
- pixel_values: torch.Tensor,
352
- **kwargs) -> torch.Tensor:
353
- target_dtype = vision_tower.get_input_embeddings().weight.dtype
354
- image_features = vision_tower(pixel_values.to(dtype=target_dtype),
355
- **kwargs)
350
+ def _image_pixels_to_features(
351
+ self,
352
+ vision_tower: SiglipVisionModel,
353
+ pixel_values: torch.Tensor,
354
+ **kwargs,
355
+ ) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
356
+ target_dtype: torch.dtype = \
357
+ vision_tower.get_input_embeddings().weight.dtype
358
+ image_features: Union[torch.Tensor, tuple[torch.Tensor, ...]] = \
359
+ vision_tower(pixel_values.to(dtype=target_dtype), **kwargs)
360
361
def select_features(leaf: torch.Tensor):
362
return self._select_image_features(
@@ -360,10 +364,7 @@ class AyaVisionForConditionalGeneration(nn.Module, SupportsMultiModal,
364
strategy=self.config.vision_feature_select_strategy,
365
)
366
363
- return cast(
364
- Union[torch.Tensor, tuple[torch.Tensor, ...]],
365
- json_map_leaves(select_features, image_features),
366
- )
367
+ return json_map_leaves(select_features, image_features)
368
369
def _select_image_features(self, image_features: torch.Tensor, *,
370
strategy: str) -> torch.Tensor:
vllm/model_executor/models/llava.py
+4
-6
@@ -4,7 +4,7 @@
4
from abc import abstractmethod
5
from collections.abc import Iterable, Mapping, Sequence
6
from typing import (Annotated, Final, Literal, Optional, Protocol, TypeVar,
7
- Union, cast)
7
+ Union)
8
9
import torch
10
import torch.nn as nn
@@ -623,7 +623,8 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
623
) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
624
# NOTE: we skip the step to select the vision feature layer since
625
# this is already done inside the vision tower
626
- image_features = vision_tower(pixel_values)
626
+ image_features: Union[torch.Tensor, tuple[torch.Tensor, ...]] = \
627
+ vision_tower(pixel_values)
628
629
def select_features(leaf: torch.Tensor):
630
return self._select_image_features(
@@ -631,10 +632,7 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
632
strategy=self.config.vision_feature_select_strategy,
633
)
634
634
- return cast(
635
- Union[torch.Tensor, tuple[torch.Tensor, ...]],
636
- json_map_leaves(select_features, image_features),
637
- )
635
+ return json_map_leaves(select_features, image_features)
636
637
def _process_image_pixels(
638
self,
vllm/model_executor/models/minimax_vl_01.py
+4
-6
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterable, Mapping
4
-from typing import Annotated, Literal, Optional, Union, cast
4
+from typing import Annotated, Literal, Optional, Union
5
6
import torch
7
import torch.nn as nn
@@ -254,7 +254,8 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal,
254
) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
255
# NOTE: we skip the step to select the vision feature layer since
256
# this is already done inside the vision tower
257
- image_features = tuple(vision_tower(p) for p in pixel_values)
257
+ image_features: tuple[torch.Tensor, ...] = \
258
+ tuple(vision_tower(p) for p in pixel_values)
259
260
def select_features(leaf: torch.Tensor):
261
return self._select_image_features(
@@ -262,10 +263,7 @@ class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal,
263
strategy=self.config.vision_feature_select_strategy,
264
)
265
265
- return cast(
266
- Union[torch.Tensor, tuple[torch.Tensor, ...]],
267
- json_map_leaves(select_features, image_features),
268
- )
266
+ return json_map_leaves(select_features, image_features)
267
268
# adapted from https://huggingface.co/MiniMaxAI/MiniMax-VL-01/blob/main/modeling_minimax_vl_01.py#L616-L631
269
def pack_image_features(self, image_features: list[torch.Tensor],
vllm/model_executor/models/tarsier.py
+4
-11
@@ -4,7 +4,7 @@
4
import math
5
from collections.abc import Iterable, Mapping, Sequence
6
from typing import (Annotated, Final, Literal, Optional, Protocol, TypeVar,
7
- Union, cast)
7
+ Union)
8
9
import torch
10
import torch.nn as nn
@@ -490,11 +490,8 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,
490
pixel_values: Union[torch.Tensor, list[torch.Tensor]],
491
) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]:
492
# From vLLM LLaVA, vision tower output handling
493
- image_hidden_states = vision_tower(pixel_values)
494
- if not isinstance(image_hidden_states, torch.Tensor):
495
- raise TypeError(
496
- f"image_hidden_states type: {type(image_hidden_states)}"
497
- " is not supported")
493
+ image_hidden_states: Union[torch.Tensor, tuple[torch.Tensor, ...]] = \
494
+ vision_tower(pixel_values)
495
496
def select_features_fn(leaf: torch.Tensor):
497
return self._select_image_features(
@@ -502,11 +499,7 @@ class TarsierForConditionalGeneration(nn.Module, SupportsMultiModal,
499
strategy=self.config.vision_feature_select_strategy,
500
)
501
505
- selected_features = cast(
506
- Union[torch.Tensor, tuple[torch.Tensor, ...]],
507
- json_map_leaves(select_features_fn, image_hidden_states),
508
- )
509
- return selected_features
502
+ return json_map_leaves(select_features_fn, image_hidden_states)
503
504
def _add_tarsier_split_tokens(
505
self, projected_image_features: torch.Tensor) -> torch.Tensor:
vllm/utils/jsontree.py
+64
-5
@@ -4,7 +4,7 @@
4
5
from collections.abc import Iterable
6
from functools import reduce
7
-from typing import Callable, TypeVar, Union, overload
7
+from typing import Callable, TypeVar, Union, cast, overload
8
9
_T = TypeVar("_T")
10
_U = TypeVar("_U")
@@ -30,10 +30,42 @@ def json_iter_leaves(value: JSONTree[_T]) -> Iterable[_T]:
30
yield value
31
32
33
+@overload
34
+def json_map_leaves(
35
+ func: Callable[[_T], _U],
36
+ value: Union[_T, dict[str, _T]],
37
+) -> Union[_U, dict[str, _U]]:
38
+ ...
39
+
40
+
41
+@overload
42
+def json_map_leaves(
43
+ func: Callable[[_T], _U],
44
+ value: Union[_T, list[_T]],
45
+) -> Union[_U, list[_U]]:
46
+ ...
47
+
48
+
49
+@overload
50
+def json_map_leaves(
51
+ func: Callable[[_T], _U],
52
+ value: Union[_T, tuple[_T, ...]],
53
+) -> Union[_U, tuple[_U, ...]]:
54
+ ...
55
+
56
+
57
+@overload
58
def json_map_leaves(
59
func: Callable[[_T], _U],
60
value: JSONTree[_T],
61
) -> JSONTree[_U]:
62
+ ...
63
+
64
+
65
+def json_map_leaves(
66
+ 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]]:
69
"""Apply a function to each leaf in a nested JSON structure."""
70
if isinstance(value, dict):
71
return {k: json_map_leaves(func, v) for k, v in value.items()}
@@ -45,6 +77,33 @@ def json_map_leaves(
77
return func(value)
78
79
80
+@overload
81
+def json_reduce_leaves(
82
+ func: Callable[[_T, _T], _T],
83
+ value: Union[_T, dict[str, _T]],
84
+ /,
85
+) -> _T:
86
+ ...
87
+
88
+
89
+@overload
90
+def json_reduce_leaves(
91
+ func: Callable[[_T, _T], _T],
92
+ value: Union[_T, list[_T]],
93
+ /,
94
+) -> _T:
95
+ ...
96
+
97
+
98
+@overload
99
+def json_reduce_leaves(
100
+ func: Callable[[_T, _T], _T],
101
+ value: Union[_T, tuple[_T, ...]],
102
+ /,
103
+) -> _T:
104
+ ...
105
+
106
+
107
@overload
108
def json_reduce_leaves(
109
func: Callable[[_T, _T], _T],
@@ -65,10 +124,10 @@ def json_reduce_leaves(
124
125
126
def json_reduce_leaves(
68
- func: Callable[..., Union[_T, _U]],
69
- value: JSONTree[_T],
70
- initial: _U = ..., # type: ignore[assignment]
71
- /,
127
+ func: Callable[..., Union[_T, _U]],
128
+ value: Union[dict[str, _T], list[_T], tuple[_T, ...], JSONTree[_T]],
129
+ initial: _U = cast(_U, ...), # noqa: B008
130
+ /,
131
) -> Union[_T, _U]:
132
"""
133
Apply a function of two arguments cumulatively to each leaf in a