minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-05-23 17:02:44 +08:00
committed by GitHub
parent 19b60a4f9e
commit 774b29dade
2 changed files with 126 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/mm_utils.py
+26 -15
View File
@@ -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(
python/sglang/srt/multimodal/processors/base_processor.py
+100 -24
View File
@@ -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,