feat: inputs_embeds for new models

Author: Nan2018

vllm/model_executor/models/exaone.py

@@ -54,7 +54,8 @@
 from vllm.utils import is_hip
 
 from .interfaces import SupportsLoRA
-from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
+from .utils import (PPMissingLayer, get_inputs_embeds, is_pp_missing_parameter,
+                    make_layers)
 
 
 class ExaoneGatedMLP(nn.Module):
@@ -365,12 +366,13 @@ def forward(
         attn_metadata: AttentionMetadata,
         intermediate_tensors: Optional[IntermediateTensors],
         inputs_embeds: Optional[torch.Tensor] = None,
+        inputs_embeds_masks: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
         if get_pp_group().is_first_rank:
-            if inputs_embeds is not None:
-                hidden_states = inputs_embeds
-            else:
-                hidden_states = self.get_input_embeddings(input_ids)
+            hidden_states = get_inputs_embeds(input_ids,
+                                              self.get_input_embeddings,
+                                              inputs_embeds,
+                                              inputs_embeds_masks)
             residual = None
         else:
             assert intermediate_tensors is not None
@@ -484,9 +486,12 @@ def forward(
         kv_caches: List[torch.Tensor],
         attn_metadata: AttentionMetadata,
         intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        inputs_embeds_masks: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
         model_output = self.transformer(input_ids, positions, kv_caches,
-                                        attn_metadata, intermediate_tensors)
+                                        attn_metadata, intermediate_tensors,
+                                        inputs_embeds, inputs_embeds_masks)
         return model_output
 
     def compute_logits(

vllm/model_executor/models/granite.py

@@ -52,7 +52,8 @@
 from vllm.utils import is_hip
 
 from .interfaces import SupportsLoRA
-from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
+from .utils import (PPMissingLayer, get_inputs_embeds, is_pp_missing_parameter,
+                    make_layers)
 
 
 class GraniteMLP(nn.Module):
@@ -304,12 +305,13 @@ def forward(
         attn_metadata: AttentionMetadata,
         intermediate_tensors: Optional[IntermediateTensors],
         inputs_embeds: Optional[torch.Tensor] = None,
+        inputs_embeds_masks: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
         if get_pp_group().is_first_rank:
-            if inputs_embeds is not None:
-                hidden_states = inputs_embeds
-            else:
-                hidden_states = self.get_input_embeddings(input_ids)
+            hidden_states = get_inputs_embeds(input_ids,
+                                              self.get_input_embeddings,
+                                              inputs_embeds,
+                                              inputs_embeds_masks)
             residual = None
         else:
             assert intermediate_tensors is not None
@@ -418,9 +420,12 @@ def forward(
         kv_caches: List[torch.Tensor],
         attn_metadata: AttentionMetadata,
         intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        inputs_embeds_masks: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, IntermediateTensors]:
         model_output = self.model(input_ids, positions, kv_caches,
-                                  attn_metadata, intermediate_tensors)
+                                  attn_metadata, intermediate_tensors,
+                                  inputs_embeds, inputs_embeds_masks)
         return model_output
 
     def compute_logits(

vllm/model_executor/models/phimoe.py

@@ -47,6 +47,7 @@
 from vllm.sequence import IntermediateTensors
 
 from .interfaces import SupportsLoRA
+from .utils import get_inputs_embeds
 
 
 class PhiMoEConfig(PretrainedConfig):
@@ -462,8 +463,12 @@ def forward(
         positions: torch.Tensor,
         kv_caches: List[torch.Tensor],
         attn_metadata: AttentionMetadata,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        inputs_embeds_masks: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        hidden_states = self.embed_tokens(input_ids)
+        hidden_states = get_inputs_embeds(input_ids, self.embed_tokens,
+                                          inputs_embeds, inputs_embeds_masks)
+
         residual = None
         for i in range(len(self.layers)):
             layer = self.layers[i]
@@ -540,9 +545,12 @@ def forward(
         kv_caches: List[torch.Tensor],
         attn_metadata: AttentionMetadata,
         intermediate_tensors: Optional[IntermediateTensors] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        inputs_embeds_masks: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         hidden_states = self.model(input_ids, positions, kv_caches,
-                                   attn_metadata)
+                                   attn_metadata, inputs_embeds,
+                                   inputs_embeds_masks)
         return hidden_states
 
     def compute_logits(self, hidden_states: torch.Tensor,