[gaudi] Vlm rebase and issue fix in benchmark test (#3263)

Signed-off-by: Wang, Yi A <[email protected]>

Author: sywangyi

backends/gaudi/server/text_generation_server/models/__init__.py

@@ -83,9 +83,6 @@
     from text_generation_server.models.custom_modeling.flash_neox_modeling import (
         FlashGPTNeoXForCausalLM,
     )
-    from text_generation_server.models.pali_gemma import (
-        PaliGemmaBatch,
-    )
     from text_generation_server.models.custom_modeling.flash_pali_gemma_modeling import (
         PaliGemmaForConditionalGeneration,
     )
@@ -153,7 +150,6 @@
     )
 
     VLM_BATCH_TYPES = {
-        PaliGemmaBatch,
         FlashVlmCausalLMBatch,
         FlashMllamaCausalLMBatch,
     }
@@ -635,6 +631,7 @@ def get_model(
                 default_dtype=torch.bfloat16,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
+                support_chunking=False,
             )
         elif model_type == BAICHUAN:
             return FlashCausalLM(
@@ -784,6 +781,8 @@ def get_model(
                 kv_cache_dtype=kv_cache_dtype,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
+                # TODO: Fix bug in rust image_text_replacement implementation
+                support_chunking=False,
             )
         elif model_type == QWEN2_5_VL:
             return FlashVlmCausalLM(
@@ -799,6 +798,8 @@ def get_model(
                 lora_adapter_ids=lora_adapter_ids,
                 config_class=Qwen2_5_VLConfig,
                 processor_class=Qwen2_5_VLProcessor,
+                # TODO: Fix bug in rust image_text_replacement implementation
+                support_chunking=False,
             )
         elif model_type == QWEN3:
             return FlashCausalLM(
@@ -824,6 +825,7 @@ def get_model(
                 default_dtype=torch.bfloat16,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
+                support_chunking=False,
             )
         elif model_type == IDEFICS2:
             return FlashVlmCausalLM(
@@ -868,7 +870,6 @@ def get_model(
                 default_dtype=torch.bfloat16,
                 trust_remote_code=trust_remote_code,
                 lora_adapter_ids=lora_adapter_ids,
-                batch_class=PaliGemmaBatch,
             )
         elif model_type == LLAVA_NEXT:
             return FlashVlmCausalLM(

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_llama4_modeling.py

@@ -1356,55 +1356,36 @@ def get_image_features(
         hidden_state = self.vision_model(pixel_values)
         return hidden_state
 
-    def forward(
+    def get_vision_embeds(
         self,
-        input_ids: torch.LongTensor = None,
-        pixel_values: torch.FloatTensor = None,
-        pixel_attention_mask=None,
-        position_ids: Optional[torch.LongTensor] = None,
-        cu_seqlen_prefill: Optional[torch.Tensor] = None,
-        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]] = None,
-        slots: torch.Tensor = None,
-        seqlen: Seqlen = None,
-        hpu_attention_meta: Optional[HPUPagedAttentionMetadata] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        vision_feature_layer: Optional[Union[int, List[int]]] = None,
-        vision_feature_select_strategy: Optional[str] = None,
-        image_sizes: torch.Tensor = None,
-        lm_head_indices: Optional[torch.Tensor] = None,
-        adapter_data: Optional[torch.Tensor] = None,
-        **lm_kwargs,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-
-        def _get_padding_mask(input_ids, pad_token_id=0):
-            return (input_ids != pad_token_id).long()
+        pixel_values: torch.FloatTensor,
+        pixel_attention_mask: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[torch.Tensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+    ):
+        image_features = self.get_image_features(
+            pixel_values=pixel_values,
+            vision_feature_layer=self.config.vision_config.vision_feature_layer,
+            vision_feature_select_strategy=self.config.vision_config.vision_feature_select_strategy,
+            image_sizes=image_sizes,
+        )
+        vision_flat = image_features.view(-1, image_features.size(-1))
+        image_features = self.multi_modal_projector(vision_flat)
+        return image_features
 
-        attention_mask = _get_padding_mask(input_ids)
-        attention_mask = attention_mask.view(seqlen.input_lengths.shape[0], -1)
+    def get_inputs_embeds(
+        self,
+        input_ids: torch.Tensor,
+        vision_embeds: torch.Tensor = None,
+        pixel_values: torch.FloatTensor = None,
+        image_sizes: Optional[torch.LongTensor] = None,
+    ):
         inputs_embeds = self.text_model.model.embed_tokens(input_ids)
-        vision_feature_layer = (
-            vision_feature_layer
-            if vision_feature_layer is not None
-            else self.config.vision_config.vision_feature_layer
-        )
-        vision_feature_select_strategy = (
-            vision_feature_select_strategy
-            if vision_feature_select_strategy is not None
-            else self.config.vision_config.vision_feature_select_strategy
-        )
-
-        if pixel_values is not None:
-            image_features = self.get_image_features(
-                pixel_values=pixel_values,
-                vision_feature_layer=vision_feature_layer,
-                vision_feature_select_strategy=vision_feature_select_strategy,
-                image_sizes=image_sizes,
-            )
-            original_inputs_embeds_shape = inputs_embeds.shape
-
-            vision_flat = image_features.view(-1, image_features.size(-1))
-            projected_vision_flat = self.multi_modal_projector(vision_flat)
 
+        if vision_embeds is not None:
+            # When we generate, we don't want to replace the potential image_token_id that we generated by images
+            # that simply don't exist
+            original_inputs_embeds_shape = inputs_embeds.shape
             special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(
                 -1
             )
@@ -1414,19 +1395,33 @@ def _get_padding_mask(input_ids, pad_token_id=0):
             final_mask_1d = final_mask[..., 0].reshape(-1)
             num_tokens_to_fill = final_mask_1d.sum()
 
-            if num_tokens_to_fill != projected_vision_flat.size(0):
+            if num_tokens_to_fill != vision_embeds.size(0):
                 raise ValueError(
                     f"Mismatch: final_mask wants {num_tokens_to_fill} embeddings, "
-                    f"but multi_modal_projector returned {projected_vision_flat.size(0)}"
+                    f"but multi_modal_projector returned {vision_embeds.size(0)}"
                 )
 
             expanded_mask = final_mask_1d.unsqueeze(-1).expand(
                 -1, inputs_embeds.size(-1)
             )
-            inputs_embeds = inputs_embeds.masked_scatter(
-                expanded_mask, projected_vision_flat
-            )
+            inputs_embeds = inputs_embeds.masked_scatter(expanded_mask, vision_embeds)
             inputs_embeds = inputs_embeds.view(original_inputs_embeds_shape)
+        return inputs_embeds
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: Optional[torch.LongTensor] = None,
+        cu_seqlen_prefill: Optional[torch.Tensor] = None,
+        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]] = None,
+        slots: torch.Tensor = None,
+        seqlen: Seqlen = None,
+        hpu_attention_meta: Optional[HPUPagedAttentionMetadata] = None,
+        lm_head_indices: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        adapter_data: Optional[torch.Tensor] = None,
+        **lm_kwargs,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
 
         logits, speculative_logits = self.text_model(
             inputs_embeds,

backends/gaudi/server/text_generation_server/models/custom_modeling/flash_llava_next.py

@@ -163,111 +163,124 @@ def _merge_input_ids_with_image_features(
             )
         return inputs_embeds
 
-    def forward(
+    def get_vision_embeds(
         self,
-        input_ids: torch.Tensor,
-        position_ids: torch.Tensor,
-        cu_seqlen_prefill: Optional[torch.Tensor],
-        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
-        slots: torch.Tensor,
-        seqlen: Seqlen,
-        hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
-        lm_head_indices: Optional[torch.Tensor] = None,
-        pixel_values: torch.FloatTensor = None,
-        # Unused for this model
-        pixel_attention_mask=None,
-        image_sizes: Optional[torch.LongTensor] = None,
-        adapter_data: Optional[torch.Tensor] = None,
+        pixel_values: torch.FloatTensor,
+        pixel_attention_mask: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[torch.Tensor] = None,
         image_grid_thw: Optional[torch.LongTensor] = None,
     ):
-        inputs_embeds = self.text_model.embed_tokens(input_ids)
-        if pixel_values is not None and len(pixel_values) > 0:
-            # num_special_image_tokens = (input_ids == self.config.image_token_index).sum()
-            # assert num_special_image_tokens == len(pixel_values), f"Received {num_special_image_tokens} for {len(pixel_values)} images, this is invalid"
-            # 1. Extract the input embeddings
-
-            # 2. Merge text and images
-            num_images, num_patches, channels, height, width = pixel_values.shape
-            pixel_values = pixel_values.view(
-                num_images * num_patches, channels, height, width
-            )
-            image_features = self.vision_tower(pixel_values)
+        # num_special_image_tokens = (input_ids == self.config.image_token_index).sum()
+        # assert num_special_image_tokens == len(pixel_values), f"Received {num_special_image_tokens} for {len(pixel_values)} images, this is invalid"
+        # 1. Extract the input embeddings
+
+        # 2. Merge text and images
+        num_images, num_patches, channels, height, width = pixel_values.shape
+        pixel_values = pixel_values.view(
+            num_images * num_patches, channels, height, width
+        )
+        image_features = self.vision_tower(pixel_values)
 
-            # selected_image_feature = image_features.hidden_states[self.config.vision_feature_layer]
-            # Already done within the clip model
-            selected_image_feature = image_features.last_hidden_state
+        # selected_image_feature = image_features.hidden_states[self.config.vision_feature_layer]
+        # Already done within the clip model
+        selected_image_feature = image_features.last_hidden_state
 
-            if self.config.vision_feature_select_strategy == "default":
-                selected_image_feature = selected_image_feature[:, 1:]
-            elif self.config.vision_feature_select_strategy == "full":
-                selected_image_feature = selected_image_feature
-            else:
-                raise RuntimeError(
-                    f"Strategy `{self.config.vision_feature_select_strategy}` is not supported/valid."
-                )
+        if self.config.vision_feature_select_strategy == "default":
+            selected_image_feature = selected_image_feature[:, 1:]
+        elif self.config.vision_feature_select_strategy == "full":
+            selected_image_feature = selected_image_feature
+        else:
+            raise RuntimeError(
+                f"Strategy `{self.config.vision_feature_select_strategy}` is not supported/valid."
+            )
 
-            image_features = self.multi_modal_projector(selected_image_feature)
+        image_features = self.multi_modal_projector(selected_image_feature)
 
-            # split up image_features for each of the individual images
-            # hence we get a list of image_features, each of shape (5, num_patches, hidden_size)
-            # if we assume each image has 5 image features (base image + 4 patches)
-            split_sizes = [num_patches] * num_images
-            image_features = torch.split(image_features, split_sizes, dim=0)
+        # split up image_features for each of the individual images
+        # hence we get a list of image_features, each of shape (5, num_patches, hidden_size)
+        # if we assume each image has 5 image features (base image + 4 patches)
+        split_sizes = [num_patches] * num_images
+        image_features = torch.split(image_features, split_sizes, dim=0)
 
-            # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
-            height = width = (
-                self.config.vision_config.image_size
-                // self.config.vision_config.patch_size
-            )
+        # NOTE we only support multimodal_patch_merge_type == "spatial_unpad"
+        height = width = (
+            self.config.vision_config.image_size // self.config.vision_config.patch_size
+        )
 
-            new_image_features = []
-            for image_idx, image_feature in enumerate(image_features):
-                if image_feature.shape[0] > 1:
-                    base_image_feature = image_feature[0]
-                    image_feature = image_feature[1:]
-
-                    if height * width != base_image_feature.shape[0]:
-                        raise ValueError(
-                            "The number of patches is not consistent with the image size."
-                        )
-
-                    # Dimensions are intentionally swapped to be bug-compatible with
-                    # upstream: https://github.com/LLaVA-VL/LLaVA-NeXT/issues/59
-                    num_patch_width, num_patch_height = get_anyres_image_grid_shape(
-                        image_sizes[image_idx],
-                        self.config.image_grid_pinpoints,
-                        self.config.vision_config.image_size,
-                    )
-                    image_feature = image_feature.view(
-                        num_patch_height, num_patch_width, height, width, -1
+        new_image_features = []
+        for image_idx, image_feature in enumerate(image_features):
+            if image_feature.shape[0] > 1:
+                base_image_feature = image_feature[0]
+                image_feature = image_feature[1:]
+
+                if height * width != base_image_feature.shape[0]:
+                    raise ValueError(
+                        "The number of patches is not consistent with the image size."
                     )
-                    image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
-                    image_feature = image_feature.flatten(1, 2).flatten(2, 3)
-                    image_feature = unpad_image(image_feature, image_sizes[image_idx])
-                    image_feature = torch.cat(
-                        (
-                            image_feature,
-                            self.image_newline[:, None, None].expand(
-                                *image_feature.shape[:-1], 1
-                            ),
+
+                # Dimensions are intentionally swapped to be bug-compatible with
+                # upstream: https://github.com/LLaVA-VL/LLaVA-NeXT/issues/59
+                num_patch_width, num_patch_height = get_anyres_image_grid_shape(
+                    image_sizes[image_idx],
+                    self.config.image_grid_pinpoints,
+                    self.config.vision_config.image_size,
+                )
+                image_feature = image_feature.view(
+                    num_patch_height, num_patch_width, height, width, -1
+                )
+                image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
+                image_feature = image_feature.flatten(1, 2).flatten(2, 3)
+                image_feature = unpad_image(image_feature, image_sizes[image_idx])
+                image_feature = torch.cat(
+                    (
+                        image_feature,
+                        self.image_newline[:, None, None].expand(
+                            *image_feature.shape[:-1], 1
                         ),
-                        dim=-1,
-                    )
-                    image_feature = image_feature.flatten(1, 2).transpose(0, 1)
-                    image_feature = torch.cat(
-                        (base_image_feature, image_feature), dim=0
-                    )
-                else:
-                    image_feature = image_feature[0]
-                    image_feature = torch.cat(
-                        (image_feature, self.image_newline[None]), dim=0
-                    )
-                new_image_features.append(image_feature)
-            image_features = torch.stack(new_image_features, dim=0)
+                    ),
+                    dim=-1,
+                )
+                image_feature = image_feature.flatten(1, 2).transpose(0, 1)
+                image_feature = torch.cat((base_image_feature, image_feature), dim=0)
+            else:
+                image_feature = image_feature[0]
+                image_feature = torch.cat(
+                    (image_feature, self.image_newline[None]), dim=0
+                )
+            new_image_features.append(image_feature)
+        image_features = torch.stack(new_image_features, dim=0)
+        return image_features.view(-1, image_features.shape[-1])
+
+    def get_inputs_embeds(
+        self,
+        input_ids: torch.Tensor,
+        vision_embeds: torch.Tensor = None,
+        pixel_values: torch.FloatTensor = None,
+        image_sizes: Optional[torch.LongTensor] = None,
+    ):
+        inputs_embeds = self.text_model.embed_tokens(input_ids)
 
+        if vision_embeds is not None:
+            # When we generate, we don't want to replace the potential image_token_id that we generated by images
+            # that simply don't exist
             inputs_embeds = self._merge_input_ids_with_image_features(
-                input_ids, inputs_embeds, image_features
+                input_ids, inputs_embeds, vision_embeds
             )
+        return inputs_embeds
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        cu_seqlen_prefill: Optional[torch.Tensor],
+        kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
+        slots: torch.Tensor,
+        seqlen: Seqlen,
+        hpu_attention_meta: Optional[HPUPagedAttentionMetadata],
+        lm_head_indices: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
+        adapter_data: Optional[torch.Tensor] = None,
+    ):
 
         hidden_states = self.text_model.model(
             inputs_embeds=inputs_embeds,