[VLM] feat: early-return in mm processor if the input is preprocessed (#26117)

python/sglang/srt/managers/mm_utils.py

@@ -1271,6 +1271,24 @@ def _slice_value(value, start, end):
         return value
 
 
+def _grid_rows_to_cpu_list(value):
+    if isinstance(value, torch.Tensor):
+        value = value.detach()
+        if value.device.type != "cpu":
+            value = value.cpu()
+        return value.tolist()
+    if isinstance(value, np.ndarray):
+        return value.tolist()
+    return value
+
+
+def _prod_grid_values(grid):
+    result = 1
+    for value in grid:
+        result *= int(value)
+    return result
+
+
 def _slice_model_data(
     data: dict,
     index: int,
@@ -1356,10 +1374,10 @@ def get_new_expanded_mm_items(original_mm_items):
                         expanded_mm_items.append(item)
                     continue
 
+                image_grid_rows = _grid_rows_to_cpu_list(image_grid_thw)
                 patches_per_item = []
-                for grid in image_grid_thw:
+                for grid in image_grid_rows:
-                    grid_tensor = torch.as_tensor(grid, dtype=torch.long)
+                    patches_per_item.append(_prod_grid_values(grid))
-                    patches_per_item.append(int(torch.prod(grid_tensor).item()))
 
                 cumulative = torch.cumsum(
                     torch.tensor(patches_per_item, dtype=torch.long), dim=0
@@ -1407,17 +1425,14 @@ def get_new_expanded_mm_items(original_mm_items):
                 # grid_len = num_videos, num_items = sum(T for each video) = total frames
                 grid_len = _get_length(video_grid_thw)
                 num_videos = grid_len
+                video_grid_rows = _grid_rows_to_cpu_list(video_grid_thw)
 
                 # Calculate total frames and frames per video
                 frames_per_video = []
                 total_frames = 0
                 for i in range(num_videos):
-                    grid = video_grid_thw[i]
+                    grid = video_grid_rows[i]
-                    if isinstance(grid, torch.Tensor):
+                    T = int(grid[0])  # T is the first element [T, H, W]
-                        T = int(grid[0].item())  # T is the first element [T, H, W]
-                    else:
-                        grid_tensor = torch.as_tensor(grid, dtype=torch.long)
-                        T = int(grid_tensor[0].item())
                     frames_per_video.append(T)
                     total_frames += T
 
@@ -1429,12 +1444,8 @@ def get_new_expanded_mm_items(original_mm_items):
                 # Calculate patches per video: T * H * W for each video
                 patches_per_video = []
                 for i in range(num_videos):
-                    grid = video_grid_thw[i]
+                    grid = video_grid_rows[i]
-                    if isinstance(grid, torch.Tensor):
+                    patches_per_video.append(_prod_grid_values(grid))
-                        patches_per_video.append(int(torch.prod(grid).item()))
-                    else:
-                        grid_tensor = torch.as_tensor(grid, dtype=torch.long)
-                        patches_per_video.append(int(torch.prod(grid_tensor).item()))
 
                 # Calculate cumulative patches to get slice indices for each video
                 cumulative = torch.cumsum(

python/sglang/srt/multimodal/processors/base_processor.py

@@ -50,6 +50,10 @@ class BaseMultiModalProcessorOutput:
     # input_text with all multimodality placeholder token expanded
     input_text: str
 
+    # original pre-tokenized ids, useful for processor_output/precomputed inputs,
+    # when they already carry the input ids
+    input_ids: Optional[Union[List[int], torch.Tensor]] = None
+
     # frames loaded from image, in given order
     images: Optional[list[Union[Image.Image, dict]]] = dataclasses.field(
         default_factory=list
@@ -518,15 +522,8 @@ class BaseMultimodalProcessor(ABC):
 
         Class method that can be pickled for multiprocessing
         """
-        if isinstance(data, dict):
+        if cls._is_preprocessed_input(data):
-            data_format = data.get("format")
+            return data
-            if data_format in (
-                MultimodalInputFormat.PROCESSOR_OUTPUT.name,
-                MultimodalInputFormat.PRECOMPUTED_EMBEDDING.name,
-                "processor_output",
-                "precomputed_embedding",
-            ):
-                return data
         try:
             if modality == Modality.IMAGE:
                 img, _ = load_image(data, cls.gpu_image_decode)
@@ -546,6 +543,49 @@ class BaseMultimodalProcessor(ABC):
         except Exception as e:
             raise RuntimeError(f"Error while loading data {data}: {e}")
 
+    @staticmethod
+    def _get_preprocessed_input_format(data):
+        """returns the detailed format if the provided data is already preprocessed.
+        returns none if the provided data is not preprocessed
+        """
+        if not isinstance(data, dict):
+            return None
+        data_format = data.get("format")
+        if isinstance(data_format, MultimodalInputFormat):
+            return data_format
+        if data_format in (
+            MultimodalInputFormat.PROCESSOR_OUTPUT.name,
+            "processor_output",
+        ):
+            return MultimodalInputFormat.PROCESSOR_OUTPUT
+        if data_format in (
+            MultimodalInputFormat.PRECOMPUTED_EMBEDDING.name,
+            "precomputed_embedding",
+        ):
+            return MultimodalInputFormat.PRECOMPUTED_EMBEDDING
+        return None
+
+    @classmethod
+    def _is_preprocessed_input(cls, data):
+        """returns if the data is already preprocessed (by the vlm processor)"""
+        return cls._get_preprocessed_input_format(data) is not None
+
+    @classmethod
+    def _all_mm_data_is_preprocessed(cls, *data_lists):
+        has_mm_data = False
+        for data_list in data_lists:
+            if not data_list:
+                continue
+            if not isinstance(data_list, list):
+                data_list = [data_list]
+            for item in data_list:
+                if item is None:
+                    continue
+                has_mm_data = True
+                if not cls._is_preprocessed_input(item):
+                    return False
+        return has_mm_data
+
     def _submit_mm_data_loading_tasks_simple(
         self,
         data_list: Optional[list],
@@ -670,10 +710,8 @@ class BaseMultimodalProcessor(ABC):
 
         formatted_indices = []
         for idx, item in enumerate(data_list):
-            if isinstance(item, dict):
+            if BaseMultimodalProcessor._is_preprocessed_input(item):
-                fmt = item.get("format")
+                formatted_indices.append(idx)
-                if fmt in {"processor_output", "precomputed_embedding"}:
-                    formatted_indices.append(idx)
 
         if formatted_indices:
             if len(data_list) != 1:
@@ -708,12 +746,7 @@ class BaseMultimodalProcessor(ABC):
     def _process_loaded_mm_data(self, modality, raw_data, result):
         images, videos, audios = [], [], []
 
-        is_precomputed = isinstance(raw_data, dict) and raw_data.get("format") in [
+        is_precomputed = self._is_preprocessed_input(raw_data)
-            MultimodalInputFormat.PROCESSOR_OUTPUT.name,
-            MultimodalInputFormat.PRECOMPUTED_EMBEDDING.name,
-            "processor_output",
-            "precomputed_embedding",
-        ]
 
         if modality == Modality.IMAGE:
             if is_precomputed:
@@ -744,6 +777,19 @@ class BaseMultimodalProcessor(ABC):
 
         BaseMultimodalProcessor.validate_mm_data(image_data, video_data, audio_data)
 
+        input_ids = prompt if isinstance(prompt, list) else None
+        if input_ids is not None and self._all_mm_data_is_preprocessed(
+            image_data, video_data, audio_data
+        ):
+            # fast path for preprocessed data: early return
+            return BaseMultiModalProcessorOutput(
+                input_text="",
+                input_ids=input_ids,
+                images=list(image_data or []),
+                videos=list(video_data or []),
+                audios=list(audio_data or []),
+            )
+
         multimodal_tokens_pattern = multimodal_tokens.get_combined_regex()
         if isinstance(prompt, list) and return_text:
             assert len(prompt) and isinstance(prompt[0], int)
@@ -782,6 +828,7 @@ class BaseMultimodalProcessor(ABC):
                 return_text=return_text,
                 discard_alpha_channel=discard_alpha_channel,
                 audio_sample_rate=audio_sample_rate,
+                input_ids=input_ids,
             )
         # For models other than MiniCPMO and MiniCPMV,
         # totally align multimodal_tokens, fast path
@@ -794,6 +841,7 @@ class BaseMultimodalProcessor(ABC):
             return_text=return_text,
             discard_alpha_channel=discard_alpha_channel,
             audio_sample_rate=audio_sample_rate,
+            input_ids=input_ids,
         )
 
     async def fast_load_mm_data(
@@ -806,6 +854,7 @@ class BaseMultimodalProcessor(ABC):
         return_text: Optional[bool] = True,
         discard_alpha_channel: bool = True,
         audio_sample_rate: Optional[int] = None,
+        input_ids: Optional[Union[List[int], torch.Tensor]] = None,
     ) -> BaseMultiModalProcessorOutput:
         """
         A fast version of `load_mm_data` that loads multimodal data directly.
@@ -878,6 +927,7 @@ class BaseMultimodalProcessor(ABC):
             audios=audios,
             videos=videos,
             input_text=prompt_str,
+            input_ids=input_ids,
         )
 
     async def legacy_load_mm_data(
@@ -890,6 +940,7 @@ class BaseMultimodalProcessor(ABC):
         return_text: Optional[bool] = True,
         discard_alpha_channel: bool = True,
         audio_sample_rate: Optional[int] = None,
+        input_ids: Optional[Union[List[int], torch.Tensor]] = None,
     ) -> BaseMultiModalProcessorOutput:
         """
         Each frame of video/image will be replaced by a single image token
@@ -997,6 +1048,7 @@ class BaseMultimodalProcessor(ABC):
             audios=audios,
             videos=videos,
             input_text="".join(new_text_parts),
+            input_ids=input_ids,
         )
 
     @staticmethod
@@ -1083,6 +1135,15 @@ class BaseMultimodalProcessor(ABC):
 
         return collected_items, input_ids, ret
 
+    @staticmethod
+    def _ensure_input_ids_is_tensor(input_ids) -> Optional[torch.Tensor]:
+        """make sure the input_ids is a flattened tensor"""
+        if input_ids is None:
+            return None
+        if isinstance(input_ids, torch.Tensor):
+            return input_ids.flatten().to(dtype=torch.long)
+        return torch.tensor(input_ids, dtype=torch.long).flatten()
+
     def process_and_combine_mm_data(
         self,
         base_output: BaseMultiModalProcessorOutput,
@@ -1136,16 +1197,19 @@ class BaseMultimodalProcessor(ABC):
             ret = None
 
         # Handle dict items (processed or precomputed)
+        dict_ret = None
         for modality, dict_item in dict_items:
-            input_format = dict_item.get("format", None)
+            input_format = self._get_preprocessed_input_format(dict_item)
-            if input_format == "processor_output":
+            if input_format is not None and dict_ret is None:
+                dict_ret = dict_item
+            if input_format == MultimodalInputFormat.PROCESSOR_OUTPUT:
                 items = self.collect_mm_items_from_processor_output(dict_item)
                 for item in items:
                     item.format = MultimodalInputFormat.PROCESSOR_OUTPUT
                 all_collected_items.extend(items)
-            elif input_format == "precomputed_embedding":
+            elif input_format == MultimodalInputFormat.PRECOMPUTED_EMBEDDING:
-                feature = dict_item["feature"]
+                dict_item = dict(dict_item)
-                del dict_item["feature"]
+                feature = dict_item.pop("feature")
                 all_collected_items.append(
                     MultimodalDataItem(
                         modality=modality,
@@ -1155,6 +1219,18 @@ class BaseMultimodalProcessor(ABC):
                     )
                 )
         # Fallback tokenization if no raw items were processed
+        if ret is None and dict_ret is not None:
+            ret = dict_ret
+
+        if input_ids is None:
+            input_ids = self._ensure_input_ids_is_tensor(base_output.input_ids)
+
+        if input_ids is None:
+            for _, dict_item in dict_items:
+                input_ids = self._ensure_input_ids_is_tensor(dict_item.get("input_ids"))
+                if input_ids is not None:
+                    break
+
         if input_ids is None:
             input_ids = self._tokenizer(
                 base_output.input_text,