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] Factor out common `_apply_feature_select_strategy` (#26003)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Oct 1, 2025 at 16:31 UTC
1405f0c7bab17b10682d2de337bd544b3e9f2f92
4 files changed
+40
-39
vllm/model_executor/models/llava.py
+3
-16
@@ -41,7 +41,7 @@ from .pixtral import PixtralHFEncoderInfo, PixtralHFVisionModel
41
from .siglip import SiglipVisionModel
42
from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
43
init_vllm_registered_model, maybe_prefix)
44
-from .vision import get_vision_encoder_info
44
+from .vision import get_num_selected_vision_tokens, get_vision_encoder_info
45
46
47
class LlavaImagePixelInputs(TensorSchema):
@@ -147,19 +147,6 @@ class BaseLlavaProcessingInfo(BaseProcessingInfo):
147
def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
148
return {"image": None}
149
150
- def _apply_feature_select_strategy(
151
- self,
152
- strategy: str,
153
- encoder_num_image_tokens: int,
154
- ) -> int:
155
- if strategy == "default":
156
- return encoder_num_image_tokens - 1
157
- if strategy == "full":
158
- return encoder_num_image_tokens
159
-
160
- msg = f"Unexpected feature select strategy: {strategy!r}"
161
- raise NotImplementedError(msg)
162
-
150
def get_num_image_tokens(
151
self,
152
*,
@@ -169,12 +156,12 @@ class BaseLlavaProcessingInfo(BaseProcessingInfo):
156
hf_config = self.get_hf_config()
157
vision_encoder_info = self.get_vision_encoder_info()
158
172
- return self._apply_feature_select_strategy(
173
- hf_config.vision_feature_select_strategy,
159
+ return get_num_selected_vision_tokens(
160
vision_encoder_info.get_num_image_tokens(
161
image_width=image_width,
162
image_height=image_height,
163
),
164
+ hf_config.vision_feature_select_strategy,
165
)
166
167
def get_image_size_with_most_features(self) -> ImageSize:
vllm/model_executor/models/llava_next.py
+3
-2
@@ -27,6 +27,7 @@ from .llava import (BaseLlavaMultiModalProcessor, BaseLlavaProcessingInfo,
27
from .siglip import SiglipVisionModel
28
from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
29
init_vllm_registered_model, maybe_prefix)
30
+from .vision import get_num_selected_vision_tokens
31
32
33
class LlavaNextImagePixelInputs(TensorSchema):
@@ -95,12 +96,12 @@ class LlavaNextProcessingInfo(BaseLlavaProcessingInfo):
96
hf_config = self.get_hf_config()
97
vision_encoder_info = self.get_vision_encoder_info()
98
98
- base_feature_size = self._apply_feature_select_strategy(
99
- hf_config.vision_feature_select_strategy,
99
+ base_feature_size = get_num_selected_vision_tokens(
100
vision_encoder_info.get_num_image_tokens(
101
image_width=image_width,
102
image_height=image_height,
103
),
104
+ hf_config.vision_feature_select_strategy,
105
)
106
107
num_patch_height, num_patch_width = get_anyres_image_grid_shape(
vllm/model_executor/models/tarsier.py
+6
-17
@@ -40,7 +40,8 @@ from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
40
from .siglip import SiglipVisionModel
41
from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
42
maybe_prefix)
43
-from .vision import VisionEncoderInfo, get_vision_encoder_info
43
+from .vision import (VisionEncoderInfo, get_num_selected_vision_tokens,
44
+ get_vision_encoder_info)
45
46
47
class TarsierImagePixelInputs(TensorSchema):
@@ -201,18 +202,6 @@ class TarsierProcessingInfo(BaseProcessingInfo):
202
def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
203
return {"image": None}
204
204
- def _apply_feature_select_strategy(
205
- self,
206
- strategy: str,
207
- encoder_num_image_tokens: int,
208
- ) -> int:
209
- if strategy == "default":
210
- return encoder_num_image_tokens - 1
211
- if strategy == "full":
212
- return encoder_num_image_tokens
213
- msg = f"Unexpected feature select strategy: {strategy!r}"
214
- raise NotImplementedError(msg)
215
-
205
def get_num_image_tokens(
206
self,
207
*,
@@ -221,21 +210,21 @@ class TarsierProcessingInfo(BaseProcessingInfo):
210
) -> int:
211
hf_config = self.get_hf_config()
212
vision_encoder_info = self.get_vision_encoder_info()
224
- num_projected_patches = self._apply_feature_select_strategy(
225
- hf_config.vision_feature_select_strategy,
213
+ num_projected_patches = get_num_selected_vision_tokens(
214
vision_encoder_info.get_num_image_tokens(
215
image_width=image_width,
216
image_height=image_height,
217
),
218
+ hf_config.vision_feature_select_strategy,
219
)
220
if num_projected_patches <= 0:
221
default_size = self.get_image_size_with_most_features()
233
- num_projected_patches_default = self._apply_feature_select_strategy(
234
- hf_config.vision_feature_select_strategy,
222
+ num_projected_patches_default = get_num_selected_vision_tokens(
223
vision_encoder_info.get_num_image_tokens(
224
image_width=default_size.width,
225
image_height=default_size.height,
226
),
227
+ hf_config.vision_feature_select_strategy,
228
)
229
if num_projected_patches_default <= 0:
230
raise ValueError(
vllm/model_executor/models/vision.py
+28
-4
@@ -9,7 +9,6 @@ from typing import (Callable, Final, Generic, Literal, Optional, Protocol,
9
10
import torch
11
from transformers import PretrainedConfig
12
-from typing_extensions import assert_never
12
13
from vllm.distributed import (get_tensor_model_parallel_rank,
14
get_tensor_model_parallel_world_size,
@@ -22,9 +21,13 @@ logger = init_logger(__name__)
21
_C = TypeVar("_C", bound=PretrainedConfig)
22
23
24
+class _RootConfig(Protocol[_C]):
25
+ vision_config: _C
26
+
27
+
28
class VisionEncoderInfo(ABC, Generic[_C]):
29
27
- def __init__(self, hf_config: _C) -> None:
30
+ def __init__(self, hf_config: _RootConfig[_C]) -> None:
31
super().__init__()
32
33
self.hf_config = hf_config
@@ -95,7 +98,7 @@ VisionFeatureSelectStrategy = Union[
98
99
100
def _get_vision_feature_selector(
98
- strategy: VisionFeatureSelectStrategy,
101
+ strategy: Union[VisionFeatureSelectStrategy, str],
102
) -> Callable[[torch.Tensor], torch.Tensor]:
103
if callable(strategy):
104
return strategy
@@ -111,7 +114,28 @@ def _get_vision_feature_selector(
114
if strategy == "full":
115
return lambda feats: feats
116
114
- assert_never(strategy)
117
+ raise ValueError(f"Unexpected feature select strategy: {strategy!r}")
118
+
119
+
120
+def get_num_selected_vision_tokens(
121
+ num_vision_tokens: int,
122
+ strategy: Union[VisionFeatureSelectStrategy, str],
123
+) -> int:
124
+ if callable(strategy):
125
+ dummy_features = torch.empty(1, num_vision_tokens, 64) # [B, L, D]
126
+ dummy_selected_features = strategy(dummy_features)
127
+ return dummy_selected_features.shape[1]
128
+
129
+ if strategy == "class":
130
+ return 1
131
+
132
+ if strategy == "default":
133
+ return num_vision_tokens - 1
134
+
135
+ if strategy == "full":
136
+ return num_vision_tokens
137
+
138
+ raise ValueError(f"Unexpected feature select strategy: {strategy!r}")
139
140
141
def resolve_visual_encoder_outputs(