Optimized merge_multimodal_embeddings on Gaudi

Signed-off-by: gyou2021 <ganmei.you@intel.com>

Author: gyou2021

vllm/model_executor/models/utils.py

@@ -476,6 +476,14 @@ def merge_multimodal_embeddings(
     Note:
         This updates ``inputs_embeds`` in place.
     """
+    if current_platform.is_hpu():
+        return _hpu_merge_multimodal_embeddings(
+            input_ids,
+            inputs_embeds,
+            multimodal_embeddings,
+            placeholder_token_id,
+        )
+        
     if isinstance(placeholder_token_id, list):
         placeholder_token_id = torch.tensor(placeholder_token_id,
                                             device=input_ids.device)
@@ -737,3 +745,27 @@ def fast_topk(values, topk, dim):
     else:
         # Use topk for efficiency with larger k values
         return torch.topk(values, topk, dim=dim)
+
+def _hpu_merge_multimodal_embeddings(
+    input_ids: torch.Tensor,
+    inputs_embeds: torch.Tensor,
+    multimodal_embeddings: NestedTensors,
+    placeholder_token_id: torch.tensor,
+) -> torch.Tensor:
+    """
+    Merge ``multimodal_embeddings`` into ``inputs_embeds`` by overwriting the
+    positions in ``inputs_embeds`` corresponding to placeholder tokens in
+    ``input_ids``.
+    merge_multimodal_embeddings on HPU to avoid dynamicity.    
+    Note:
+        This updates ``inputs_embeds`` in place.
+    """
+    batch_size, seq_length, hidden_size = inputs_embeds.shape
+    inputs_embeds = inputs_embeds.reshape(-1, hidden_size)
+    multimodal_embeddings = multimodal_embeddings.reshape(-1, hidden_size)
+    placeholder_token_id = torch.tensor(placeholder_token_id,
+                                        device=input_ids.device)
+    mask = torch.isin(input_ids.reshape(-1), placeholder_token_id)
+    inputs_embeds.index_put_((mask, ), multimodal_embeddings)
+    inputs_embeds = inputs_embeds.reshape(batch_size, seq_length, hidden_size)
+    return inputs_embeds