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
[CI Failure] fix tests/entrypoints/openai/test_skip_tokenizer.py (#22708)
Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Aug 12, 2025 at 20:42 UTC
f7ad6a1eb3deb9ca70a6bce3705dbd16cf9d8b28
1 file changed
+9
-4
vllm/model_executor/models/prithvi_geospatial_mae.py
+9
-4
@@ -25,11 +25,11 @@ import torch.nn as nn
25
from transformers import BatchFeature
26
27
from vllm.config import VllmConfig
28
-from vllm.model_executor.layers.pooler import (AllPool, PoolerHead,
29
- PoolerIdentity, SimplePooler)
28
+from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
29
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
30
from vllm.model_executor.models.interfaces import (
32
- IsAttentionFree, MultiModalEmbeddings, SupportsMultiModalWithRawInput)
31
+ IsAttentionFree, MultiModalEmbeddings, SupportsMultiModalWithRawInput,
32
+ default_pooling_type)
33
from vllm.model_executor.models.utils import AutoWeightsLoader
34
from vllm.multimodal import MULTIMODAL_REGISTRY
35
from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
@@ -142,6 +142,7 @@ class PrithviGeoSpatialMAEMultiModalProcessor(BaseMultiModalProcessor):
142
)
143
144
145
+@default_pooling_type("All")
146
@MULTIMODAL_REGISTRY.register_processor(
147
PrithviGeoSpatialMAEMultiModalProcessor,
148
info=PrithviGeoSpatialMAEProcessingInfo,
@@ -198,7 +199,11 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree,
199
"Only SemanticSegmentationTask is supported for now "
200
"by PrithviGeospatialMAE.")
201
201
- self.pooler = SimplePooler(AllPool(), PoolerHead(PoolerIdentity()))
202
+ pooler_config = vllm_config.model_config.pooler_config
203
+ assert pooler_config is not None
204
+
205
+ self.pooler = DispatchPooler(
206
+ {"encode": Pooler.for_encode(pooler_config)}, )
207
208
def _parse_and_validate_multimodal_data(
209
self, **kwargs) -> tuple[torch.Tensor, Optional[torch.Tensor]]: