Support for attention free models in V1

Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Author: christian-pinto

vllm/model_executor/models/prithvi_geospatial_mae.py

@@ -26,13 +26,13 @@
 from vllm.config import VllmConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.interfaces import (IsAttentionFree,
-                                                   SupportsMultiModal,
-                                                   SupportsV0Only)
+                                                   SupportsMultiModalWithRawInput)
 from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.model_executor.pooling_metadata import PoolingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
-                                    MultiModalInputs, MultiModalKwargs)
+from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig, MultiModalFieldElem,
+                                    MultiModalInputs, MultiModalKwargs, MultiModalKwargsItem,
+                                    MultiModalSharedField, PlaceholderRange)
 from vllm.multimodal.parse import MultiModalDataItems
 from vllm.multimodal.processing import (BaseMultiModalProcessor,
                                         BaseProcessingInfo, PromptUpdate)
@@ -75,8 +75,8 @@ def _get_mm_fields_config(
         hf_processor_mm_kwargs: Mapping[str, object],
     ) -> Mapping[str, MultiModalFieldConfig]:
         return dict(
-            pixel_values=MultiModalFieldConfig.batched("image"),
-            location_coords=MultiModalFieldConfig.batched("image"),
+            pixel_values=MultiModalFieldConfig.shared(batch_size=1, modality="image"),
+            location_coords=MultiModalFieldConfig.shared(batch_size=1, modality="image"),
         )
 
     def _get_prompt_updates(
@@ -99,23 +99,32 @@ def apply(
 
         for k, v in mm_data.items():
             mm_kwargs[k] = v
+        mm_place_holders = {
+            "image": [PlaceholderRange(offset=0, length=0)]
+        }
+
+        multimodal_kwargs_items = [
+            MultiModalKwargsItem.from_elems(
+                [MultiModalFieldElem(modality="image", key=key, data=data, field=MultiModalSharedField(1))
+                 for key, data in mm_kwargs.items()]
+            )
+        ]
 
         return MultiModalInputs(
             type="multimodal",
             prompt=prompt,
             prompt_token_ids=[1],
-            mm_kwargs=MultiModalKwargs(mm_kwargs),
+            mm_kwargs=MultiModalKwargs.from_items(multimodal_kwargs_items),
             mm_hashes=None,
-            mm_placeholders={},
+            mm_placeholders=mm_place_holders,
         )
 
 
 @MULTIMODAL_REGISTRY.register_processor(
     PrithviGeoSpatialMAEMultiModalProcessor,
     info=PrithviGeoSpatialMAEProcessingInfo,
     dummy_inputs=PrithviGeoSpatialMAEInputBuilder)
-class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
-                           SupportsV0Only):
+class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModalWithRawInput):
     """ Prithvi Masked Autoencoder"""
 
     @classmethod
@@ -180,7 +189,13 @@ def _parse_and_validate_multimodal_data(
             location_coords = None
 
         return pixel_values, location_coords
-
+    
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        # We do not really use any input tokens and therefore no embeddings to be calculated
+        # However, due to the mandatory token ids in the input prompt we pass one token and the
+        # size of the dummy embedding tensors must reflect that.
+        return torch.empty(input_ids.shape)
+    
     def forward(
         self,
         input_ids: Optional[torch.Tensor],
@@ -202,7 +217,7 @@ def pooler(
         hidden_states: torch.Tensor,
         pooling_metadata: PoolingMetadata,
     ) -> Optional[PoolerOutput]:
-        return PoolerOutput([PoolingSequenceGroupOutput(hidden_states)])
+        return PoolerOutput([PoolingSequenceGroupOutput(hidden_states[0])])
 
     def load_weights(self, weights: Iterable[tuple[str,
                                                    torch.Tensor]]) -> set[str]: