diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 50a641ac6..01d4cd78b 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/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: - grid_tensor = torch.as_tensor(grid, dtype=torch.long) - patches_per_item.append(int(torch.prod(grid_tensor).item())) + for grid in image_grid_rows: + patches_per_item.append(_prod_grid_values(grid)) 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] - if isinstance(grid, torch.Tensor): - 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()) + grid = video_grid_rows[i] + T = int(grid[0]) # T is the first element [T, H, W] 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] - if isinstance(grid, torch.Tensor): - 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())) + grid = video_grid_rows[i] + patches_per_video.append(_prod_grid_values(grid)) # Calculate cumulative patches to get slice indices for each video cumulative = torch.cumsum( diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index a647e9e0d..842cd903c 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/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): - data_format = data.get("format") - if data_format in ( - MultimodalInputFormat.PROCESSOR_OUTPUT.name, - MultimodalInputFormat.PRECOMPUTED_EMBEDDING.name, - "processor_output", - "precomputed_embedding", - ): - return data + if cls._is_preprocessed_input(data): + 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): - fmt = item.get("format") - if fmt in {"processor_output", "precomputed_embedding"}: - formatted_indices.append(idx) + if BaseMultimodalProcessor._is_preprocessed_input(item): + 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 [ - MultimodalInputFormat.PROCESSOR_OUTPUT.name, - MultimodalInputFormat.PRECOMPUTED_EMBEDDING.name, - "processor_output", - "precomputed_embedding", - ] + is_precomputed = self._is_preprocessed_input(raw_data) 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) - if input_format == "processor_output": + input_format = self._get_preprocessed_input_format(dict_item) + 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": - feature = dict_item["feature"] - del dict_item["feature"] + elif input_format == MultimodalInputFormat.PRECOMPUTED_EMBEDDING: + dict_item = dict(dict_item) + 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,