Merge Libint/intervl_bucket (#1965)

yeonsily · libinta · web-flow · commit bb96123af496 · 2025-09-22T14:00:13.000-07:00
## Essential Elements of an Effective PR Description Checklist
- [ ] The purpose of the PR, such as "Fix some issue (link existing
issues this PR will resolve)".
- [ ] The test plan, such as providing test command.
- [ ] The test results, such as pasting the results comparison before
and after, or e2e results


## Purpose

## Test Plan

## Test Result

&lt;!--- pyml disable-next-line no-emphasis-as-heading --&gt;

---------

Signed-off-by: Libin Tang &lt;libin.tang@intel.com&gt;
Co-authored-by: Libin Tang &lt;libin.tang@intel.com&gt;
Co-authored-by: Libin Tang &lt;litang@habana.ai&gt;
diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py
@@ -577,7 +577,8 @@ def _process_image_input(
 
             for i in batch_breakdown:
                 end_idx = start_idx + i
-                indices = torch.arange(start_idx, end_idx)
+                indices = torch.arange(start_idx,
+                                       end_idx).to(pixel_values.device)
                 batch_sliced_pixel_values = torch.index_select(pixel_values,
                                                                dim=0,
                                                                index=indices)
diff --git a/vllm/model_executor/models/internvl.py b/vllm/model_executor/models/internvl.py
@@ -7,6 +7,7 @@
 # Copyright (c) 2023 OpenGVLab
 # Licensed under The MIT License [see LICENSE for details]
 # --------------------------------------------------------
+import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Mapping, Sequence
 from typing import Any, Literal, Optional, TypedDict, TypeVar, Union
@@ -35,13 +36,15 @@
                                         BaseProcessingInfo, PromptReplacement,
                                         PromptUpdate, PromptUpdateDetails)
 from vllm.multimodal.profiling import BaseDummyInputsBuilder
+from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.tokenizer import AnyTokenizer
 
 from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
                          SupportsMultiModal, SupportsPP)
-from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
-                    maybe_prefix, merge_multimodal_embeddings)
+from .utils import (AutoWeightsLoader, flatten_bn, greedy_plan,
+                    init_vllm_registered_model, maybe_prefix,
+                    merge_multimodal_embeddings)
 
 IMG_START = '<img>'
 IMG_END = '</img>'
@@ -50,6 +53,9 @@
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
 
+is_hpu = current_platform.is_hpu()
+is_lazy = os.environ.get('PT_HPU_LAZY_MODE', '0') == '1' if is_hpu else False
+
 
 class InternVLImagePixelInputs(TypedDict):
     type: Literal["pixel_values"]
@@ -1062,6 +1068,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         self.visual_token_mask = None
         self.make_empty_intermediate_tensors = (
             self.language_model.make_empty_intermediate_tensors)
+        if is_hpu:
+            self.graphed_multimodal_buckets = None
 
     def _patch_quant_config(self, config: PretrainedConfig,
                             quant_config: QuantizationConfig):
@@ -1127,16 +1135,64 @@ def pixel_shuffle(self, x, scale_factor=0.5):
         return x
 
     def extract_feature(self, pixel_values: torch.Tensor) -> torch.Tensor:
-        vit_embeds = self.vision_model(pixel_values=pixel_values)
-        vit_embeds = vit_embeds[:, 1:, :]
-
-        h = w = int(vit_embeds.shape[1]**0.5)
-        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
-        vit_embeds = self.pixel_shuffle(vit_embeds,
-                                        scale_factor=self.downsample_ratio)
-        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1,
-                                        vit_embeds.shape[-1])
-        vit_embeds = self.mlp1(vit_embeds)
+        if is_hpu:
+            if self.vision_buckets.multimodal_buckets:
+                batch_breakdown = greedy_plan(pixel_values.shape[0], \
+                        self.vision_buckets.multimodal_buckets)
+            else:
+                batch_breakdown = [pixel_values.shape[0]]
+
+            start_idx = 0
+            vit_embeds_minibatches = []
+
+            for i in batch_breakdown:
+                end_idx = start_idx + i
+                batch_sliced_pixel_values = \
+                        pixel_values[start_idx:end_idx, ...]
+                if is_lazy:
+                    vit_embeds_minibatch = \
+                        self.vision_model(
+                            pixel_values=batch_sliced_pixel_values,
+                            bypass_hpu_graphs=i
+                            not in self.graphed_multimodal_buckets
+                            and len(self.graphed_multimodal_buckets) > 0)
+                else:
+                    vit_embeds_minibatch = \
+                        self.vision_model(
+                            pixel_values=batch_sliced_pixel_values)
+
+                vit_embeds_minibatch = vit_embeds_minibatch[:, 1:, :]
+
+                h = w = int(vit_embeds_minibatch.shape[1]**0.5)
+                vit_embeds_minibatch = vit_embeds_minibatch.reshape(
+                    vit_embeds_minibatch.shape[0], h, w, -1)
+                vit_embeds_minibatch = self.pixel_shuffle(
+                    vit_embeds_minibatch, scale_factor=self.downsample_ratio)
+                vit_embeds_minibatch = vit_embeds_minibatch.reshape(
+                    vit_embeds_minibatch.shape[0], -1,
+                    vit_embeds_minibatch.shape[-1])
+
+                if is_lazy:
+                    vit_embeds_minibatches += [
+                        self.mlp1(vit_embeds_minibatch,
+                                  bypass_hpu_graphs=i
+                                  not in self.graphed_multimodal_buckets
+                                  and len(self.graphed_multimodal_buckets) > 0)
+                    ]
+                else:
+                    vit_embeds_minibatches += [self.mlp1(vit_embeds_minibatch)]
+                start_idx = end_idx
+            vit_embeds = torch.cat(vit_embeds_minibatches, dim=0)
+        else:
+            vit_embeds = self.vision_model(pixel_values=pixel_values)
+            vit_embeds = vit_embeds[:, 1:, :]
+            h = w = int(vit_embeds.shape[1]**0.5)
+            vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+            vit_embeds = self.pixel_shuffle(vit_embeds,
+                                            scale_factor=self.downsample_ratio)
+            vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1,
+                                            vit_embeds.shape[-1])
+            vit_embeds = self.mlp1(vit_embeds)
         return vit_embeds
 
     def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
@@ -1180,8 +1236,11 @@ def _parse_and_validate_image_input(
 
         image_token_id = kwargs["image_token_id"]
         assert isinstance(image_token_id, torch.Tensor)
-        self.img_context_token_id = image_token_id.flatten().unique().item()
-
+        if is_hpu:
+            self.img_context_token_id = image_token_id.flatten()
+        else:
+            self.img_context_token_id = image_token_id.flatten().unique().item(
+            )
         if pixel_values_flat is not None:
             if not isinstance(pixel_values_flat, (torch.Tensor, list)):
                 raise ValueError("Incorrect type of pixel values. "
@@ -1306,7 +1365,9 @@ def get_language_model(self) -> torch.nn.Module:
 
     def get_multimodal_embeddings(
             self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
-
+        if is_hpu:
+            self.graphed_multimodal_buckets = kwargs.pop(
+                'graphed_multimodal_buckets', [])
         modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
         if not modalities:
             return None
diff --git a/vllm/worker/hpu_model_runner.py b/vllm/worker/hpu_model_runner.py
@@ -112,16 +112,22 @@ class VisionBuckets:
     This class is used to bucket image tokens
     '''
 
-    def __init__(self, is_batch_based):
-        self.is_batch_based = is_batch_based
+    def __init__(self, model):
+        self.is_batch_based = True
         envvar = os.environ.get('VLLM_MULTIMODAL_BUCKETS', "")
         if envvar == 'None':
             self.multimodal_buckets = None
         else:
             if envvar == "":
-                if is_batch_based:
+                if 'InternVLChatModel' in str(type(model)):
+                    multimodal_buckets = list(
+                        range(model.config.min_dynamic_patch,
+                              model.config.max_dynamic_patch +
+                              2))  #As use_thumbnail is true
+                elif 'Gemma3ForConditionalGeneration' in str(type(model)):
                     multimodal_buckets = [1, 2, 4, 8]  # batch sizes for gemma3
                 else:
+                    self.is_batch_based = False
                     multimodal_buckets = [
                         1600, 3136, 4096, 6400, 7744, 9216, 12544
                     ]
@@ -159,9 +165,11 @@ def __call__(cls, *args, **kwargs):
 
 
 def is_mm_optimized(model):
-    return 'Gemma3ForConditionalGeneration' in str(type(model.model)) \
-        if hasattr(model, 'model') else \
-        'Gemma3ForConditionalGeneration' in str(type(model))
+    mm_models = ['Gemma3ForConditionalGeneration', 'InternVLChatModel']
+
+    return any(m in str(type(model.model)) for m in mm_models) \
+        if hasattr(model, 'model') \
+        else any(m in str(type(model)) for m in mm_models)
 
 
 def pad_flat_tensor(tensor, desired_size):
@@ -345,6 +353,7 @@ def __init__(self, model, vllm_config, is_causal, sampler):
         model_config = getattr(self.model, "config", None)
 
         self.model_is_mrope = uses_mrope(model_config)
+
         self.is_mm_optimized = is_mm_optimized(self.model)
         text_config = vllm_config.model_config.hf_config.get_text_config()
         self.interleaved_sliding_window = getattr(
@@ -379,6 +388,12 @@ def __init__(self, model, vllm_config, is_causal, sampler):
                             htorch.hpu.wrap_in_hpu_graph( \
                             self.model.multi_modal_projector, \
                             disable_tensor_cache=True)
+                if hasattr(self.model, 'vision_model'):
+                    self.model.vision_model = htorch.hpu.wrap_in_hpu_graph(
+                        self.model.vision_model, disable_tensor_cache=True)
+                if hasattr(self.model, 'mlp1'):
+                    self.model.mlp1 = htorch.hpu.wrap_in_hpu_graph(
+                        self.model.mlp1, disable_tensor_cache=True)
 
         self._rotary_embed_module = self._get_rotary_embedding_module(
             self.model)
@@ -624,26 +639,30 @@ def compute_input_embeddings_for_mm_optimized(self, warmup_mode, **kwargs):
         vision_embeddings = self.model.get_multimodal_embeddings(**kwargs)
         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
         # attn_metadata, so we can reuse HPU graph without running
         # the whole vision tower.
         if vision_embeddings is not None or (
                 warmup_mode and kwargs['attn_metadata'].is_prompt):
-            input_ids = kwargs['input_ids']
-            positions = kwargs['positions']
-            kwargs = self.model.prepare_attn_masks(
-                mask_dtype=self.dtype,
-                **kwargs,
-            )
-            kwargs['input_ids'] = input_ids
-            kwargs['positions'] = positions
+            if hasattr(self.model, 'prepare_attn_masks'):
+                input_ids = kwargs['input_ids']
+                positions = kwargs['positions']
+                kwargs = self.model.prepare_attn_masks(
+                    mask_dtype=self.dtype,
+                    **kwargs,
+                )
+                kwargs['input_ids'] = input_ids
+                kwargs['positions'] = positions
+                # done compute the visual tokens
+                kwargs.pop('pixel_values', None)
+            else:
+                kwargs.pop('pixel_values_flat', None)
+                kwargs.pop("image_num_patches", None)
+                kwargs.pop("image_token_id", None)
 
         kwargs.update({'inputs_embeds': inputs_embeds})
-        # done compute the visual tokens and others
-        kwargs.pop('pixel_values', None)
         kwargs.pop("num_crops", None)
         kwargs.pop("graphed_multimodal_buckets", None)
         return kwargs
@@ -699,7 +718,6 @@ def forward(self, *args, **kwargs):
         virtual_engine = 0
         if 'virtual_engine' in kwargs:
             virtual_engine = kwargs.pop('virtual_engine')
-
         input_ids = kwargs['input_ids']
         global_attn_masks = kwargs.pop("global_attn_masks") \
                 if kwargs.get("global_attn_masks") else None
@@ -1080,6 +1098,8 @@ def __init__(
                                     and not self.lora_config)
         self.use_delayed_sampling = get_config(
         ).use_delayed_sampling and can_use_delayed_sampling
+        self.mm_tokens_per_image = 1
+        self.image_token_id = 0
 
     def _set_gc_threshold(self) -> None:
         """
@@ -1497,10 +1517,16 @@ def move_to_device(self, tensor):
                                                        non_blocking=True)
 
     def add_vision_buckets_to_mrope_mm_optimized(self):
-        model = self.get_model()
-        self.is_mm_optimized = is_mm_optimized(model)
+        self.is_mm_optimized = is_mm_optimized(self.model)
         if self.model_is_mrope or self.is_mm_optimized:
-            model.vision_buckets = VisionBuckets(self.is_mm_optimized)
+            if hasattr(self.model.model.config, 'mm_tokens_per_image'):
+                self.mm_tokens_per_image = \
+                    self.model.model.config.mm_tokens_per_image
+                self.image_token_id = self.model.model.config.image_token_id
+            elif 'InternVLChatModel' in str(type(self.model.model)):
+                self.image_token_id = 151667
+                self.mm_tokens_per_image = self.model.model.num_image_token
+            self.model.model.vision_buckets = VisionBuckets(self.model.model)
 
     def _prepare_prompt(
         self,
@@ -1631,7 +1657,6 @@ def _prepare_prompt(
                     for idx in range(3):
                         seq_data_mrope_positions[idx] \
                             .extend(mrope_positions[idx])
-
                 multi_modal_kwargs_list.append(mm_kwargs)
 
                 for modality, placeholder_map in placeholder_maps.items():
@@ -2709,17 +2734,28 @@ def create_dummy_multi_modal_seq_group_metadata(self, group_id, img_args,
         else:
             s = self.model.model.config.vision_config.image_size
             pixel_values = torch.randn([img_args, 3, s, s])
-            num_image_tokens = self.model.model.config.mm_tokens_per_image \
-                    * img_args
-            multi_modal_data = {
-                "pixel_values": pixel_values,
-                "num_crops": torch.zeros([img_args], dtype=torch.int32)
-            }
 
-        image_token_id = self.get_model().config.image_token_id
-        prompt_token_ids_image = [image_token_id] * num_image_tokens
+            if 'Gemma3ForConditionalGeneration' in str(type(self.model.model)):
+                multi_modal_data = {
+                    "pixel_values": pixel_values,
+                    "num_crops": torch.zeros([img_args], dtype=torch.int32),
+                }
+            elif 'InternVLChatModel' in str(type(self.model.model)):
+                multi_modal_data = {
+                    "pixel_values_flat":
+                    pixel_values.to(torch.bfloat16),
+                    "image_num_patches":
+                    torch.tensor([pixel_values.shape[0]], dtype=torch.int32),
+                    "image_token_id":
+                    torch.tensor([self.image_token_id], dtype=torch.int64),
+                }
+            else:
+                logger.warning("No support for other models yet")
+            num_image_tokens = self.mm_tokens_per_image * img_args
+        prompt_token_ids_image = [self.image_token_id] * num_image_tokens
         prompt_token_ids = [0] * (
             seq_len - len(prompt_token_ids_image)) + prompt_token_ids_image
+
         prompt_token_ids_array = array('l', prompt_token_ids)  # noqa: F821
         placeholders_by_modality = {
             'image':
@@ -3188,9 +3224,7 @@ def warmup_model(self, kv_caches: List[torch.Tensor]) -> None:
             if graphs:
                 self.graphed_buckets.add(cfg)
             if self.is_mm_run():
-                img_args = (int(seq_len) //
-                            self.model.model.config.mm_tokens_per_image
-                            if self.is_mm_optimized else int(seq_len))
+                img_args = int(seq_len) // self.mm_tokens_per_image
             self.warmup_scenario(
                 int(bs),
                 int(seq_len),
@@ -3539,7 +3573,7 @@ def _get_seq_ids(self, model_input):
     def _get_img_args_from_model_input(self, model_input):
         if (not self.model_is_mrope and not self.is_mm_optimized) or \
             not model_input.multi_modal_kwargs or \
-            'pixel_values' not in model_input.multi_modal_kwargs:
+            ('pixel_values') not in model_input.multi_modal_kwargs:
             return None
         if self.model_is_mrope:
             pixel_values_list = model_input.multi_modal_kwargs['pixel_values']
@@ -3816,18 +3850,17 @@ def try_revert_dummy_output_tokens():
                     'real_seq_len': model_input.seq_lens,
                     'real_batch_size': real_batch_size
                 }
-
                 #Need to set the window_slide mask at this point to decide
                 if is_prompt:
                     attn_metadata = self.model._update_use_window_sdpa(
                         execute_model_kwargs['attn_metadata'], seq_len,
                         bool(model_input.multi_modal_kwargs and \
-                       'pixel_values' in model_input.multi_modal_kwargs))
+                       ('pixel_values')in model_input.multi_modal_kwargs))
                     execute_model_kwargs['attn_metadata'] = attn_metadata
 
                 if not bypass_model_exec:
                     if self.model_is_mrope or self.is_mm_optimized:
-                        if 'pixel_values' in execute_model_kwargs and \
+                        if ('pixel_values') in execute_model_kwargs and \
                                 self.is_mm_optimized:
                             if warmup_mode and not is_pt_profiler_run:
                                 bypass_model_exec = True
