Modify merge_multimodal_embeddings to static (#1969)

Author: yeonsily

vllm/model_executor/models/internvl.py

@@ -44,7 +44,8 @@
                          SupportsMultiModal, SupportsPP)
 from .utils import (AutoWeightsLoader, flatten_bn, greedy_plan,
                     init_vllm_registered_model, maybe_prefix,
-                    merge_multimodal_embeddings)
+                    merge_multimodal_embeddings,
+                    merge_multimodal_embeddings_static)
 
 IMG_START = '<img>'
 IMG_END = '</img>'
@@ -1390,6 +1391,21 @@ def get_multimodal_embeddings(
 
         return multimodal_embeddings
 
+    def get_input_embeddings_hpu(
+        self,
+        input_ids: torch.Tensor,
+        image_index_tensor: torch.Tensor,
+        multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
+    ) -> torch.Tensor:
+        inputs_embeds = self.language_model.get_input_embeddings(input_ids)
+        if multimodal_embeddings is not None:
+            inputs_embeds = merge_multimodal_embeddings_static(
+                image_index_tensor,
+                inputs_embeds,
+                multimodal_embeddings,
+            )
+        return inputs_embeds
+
     def get_input_embeddings(
         self,
         input_ids: torch.Tensor,

vllm/model_executor/models/utils.py

@@ -433,6 +433,22 @@ def merge_multimodal_embeddings_from_map(
     return inputs_embeds
 
 
+def merge_multimodal_embeddings_static(
+    is_multimodal_index: torch.Tensor,
+    inputs_embeds: torch.Tensor,
+    multimodal_embeddings: NestedTensors,
+) -> torch.Tensor:
+    assert current_platform.is_hpu(), ("Support HPU only")
+    flattened = _flatten_embeddings(multimodal_embeddings)
+
+    inputs_embeds_s = inputs_embeds.shape
+    inputs_embeds = inputs_embeds.view(inputs_embeds_s[0] * inputs_embeds_s[1],
+                                       inputs_embeds_s[2])
+    inputs_embeds = inputs_embeds.index_copy_(0, is_multimodal_index,
+                                              flattened).view(inputs_embeds_s)
+    return inputs_embeds
+
+
 def _merge_multimodal_embeddings(
     inputs_embeds: torch.Tensor,
     is_multimodal: torch.Tensor,

vllm/worker/hpu_model_runner.py

@@ -637,8 +637,14 @@ def _update_metadata(self,
     def compute_input_embeddings_for_mm_optimized(self, warmup_mode, **kwargs):
         input_ids = kwargs['input_ids']
         vision_embeddings = self.model.get_multimodal_embeddings(**kwargs)
-        inputs_embeds = self.model.get_input_embeddings(
-            input_ids, vision_embeddings)
+        if 'image_index' in kwargs:
+            inputs_embeds = self.model.get_input_embeddings_hpu(
+                input_ids, kwargs['image_index'], vision_embeddings)
+            kwargs.pop("image_index", None)
+        else:
+            inputs_embeds = self.model.get_input_embeddings(
+                input_ids, vision_embeddings)
+
         # TODO: In warmup, we need to warmup the model with dummy image data for
         # multimodal model for prompt, here instead of generating a dummy image,
         # we are just generating attn_mask for the images and pass with
@@ -1772,6 +1778,7 @@ def _prepare_prompt(
                                               pad=0,
                                               dtype=torch.long,
                                               flat=self.use_merged_prefill)
+        image_index_tensor = None
         if self.model_is_mrope:
             input_positions = \
                 make_mrope_positions_tensor_with_pad(input_positions=input_positions,
@@ -1785,6 +1792,11 @@ def _prepare_prompt(
                                               dtype=torch.long,
                                               flat=self.use_merged_prefill)
 
+        if seq_group_metadata.multi_modal_data and self.is_mm_optimized and \
+            'InternVLChatModel' in str(type(self.model.model)):
+            is_image_flatten = (
+                input_tokens_tensor == self.image_token_id).flatten()
+            image_index_tensor = is_image_flatten.nonzero().squeeze(-1)
         slot_mapping = make_cpu_tensor(slot_mapping,
                                        max_len=max_prompt_len,
                                        pad=_PAD_SLOT_ID,
@@ -1872,6 +1884,8 @@ def _prepare_prompt(
             input_positions=input_positions,
         )
         multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
+        if image_index_tensor is not None:
+            multi_modal_kwargs['image_index'] = image_index_tensor
         multi_modal_kwargs = MultiModalKwargs.as_kwargs(multi_modal_kwargs,
                                                         device=self.device)
 
@@ -3872,6 +3886,12 @@ def try_revert_dummy_output_tokens():
                        ('pixel_values')in model_input.multi_modal_kwargs))
                     execute_model_kwargs['attn_metadata'] = attn_metadata
 
+                    if 'image_index' in model_input.multi_modal_kwargs:
+                        execute_model_kwargs[
+                            'image_index'] = model_input.multi_modal_kwargs[
+                                'image_index']
+                        model_input.multi_modal_kwargs.pop('image_index', None)
+
                 if not bypass_model_exec:
                     if self.model_is_mrope or self.is_mm_optimized:
                         if ('pixel_values') in execute_model_kwargs and \