[EPD][VLM] support video/audio input (#17824)

Co-authored-by: siyu <liusy58@linux.alibaba.com>
This commit is contained in:
Zheng Wengang
2026-03-16 14:18:21 +08:00
committed by GitHub
parent 738cbde902
commit 135af6dc92
12 changed files with 1924 additions and 343 deletions

View File

@@ -256,15 +256,20 @@ class BaseMultimodalProcessor(ABC):
def spatial_merge_size(self):
return self.hf_config.vision_config.spatial_merge_size
def build_input_ids(self, prompt, img_grid_thw):
def build_input_ids(
self, prompt, img_grid_thw=None, video_grid_thw=None, audio_seq_lens=None
):
"""
Use prompt and img_grid_thw to build input_ids
Use prompt, img_grid_thw, video_grid_thw, and audio_seq_lens to build input_ids.
Supports image, video, and audio tokens.
"""
if not isinstance(prompt, list):
prompt = self._tokenizer.encode(prompt)
img_token_id = self.IM_TOKEN_ID
spatial_merge_size = self.spatial_merge_size
img_token_id = getattr(self, "IM_TOKEN_ID", None)
video_token_id = getattr(self, "VIDEO_TOKEN_ID", None)
audio_token_id = getattr(self, "audio_token_id", None)
spatial_merge_size = getattr(self, "spatial_merge_size", 1)
input_ids = []
offsets = []
@@ -273,34 +278,81 @@ class BaseMultimodalProcessor(ABC):
# Use img_token_id instead of im_start_id, because a dummy im_start_id
# may be generated by the tokenizer.
img_start_indices = list(
filter(lambda i: prompt[i + 1] == img_token_id, range(len(prompt) - 1))
)
vision_start_indices = []
for i in range(len(prompt) - 1):
if img_token_id is not None and prompt[i + 1] == img_token_id:
vision_start_indices.append((i, Modality.IMAGE))
elif video_token_id is not None and prompt[i + 1] == video_token_id:
vision_start_indices.append((i, Modality.VIDEO))
elif audio_token_id is not None and prompt[i + 1] == audio_token_id:
vision_start_indices.append((i, Modality.AUDIO))
# get modality list with order preserved
modality_list = [modality for _, modality in vision_start_indices]
for cur_img_idx, img_start_idx in enumerate(img_start_indices):
assert cur_idx <= img_start_idx
# include img_start_id
input_ids.extend(prompt[cur_idx : img_start_idx + 1])
img_offset_start = len(input_ids)
img_token_num = img_grid_thw[cur_img_idx].prod() // (spatial_merge_size**2)
input_ids.extend([img_token_id] * img_token_num)
# jump to img_end_id
cur_idx = img_start_idx + 2
offsets.append((img_offset_start, len(input_ids) - 1))
img_idx = 0
video_idx = 0
audio_idx = 0
for mm_start_idx, modality in vision_start_indices:
if modality == Modality.IMAGE:
mm_token_num = img_grid_thw[img_idx].prod() // (spatial_merge_size**2)
mm_token_id = img_token_id
img_idx += 1
elif modality == Modality.VIDEO:
mm_token_num = video_grid_thw[video_idx].prod() // (
spatial_merge_size**2
)
mm_token_id = video_token_id
video_idx += 1
elif modality == Modality.AUDIO:
mm_token_num = int(audio_seq_lens[audio_idx].item())
mm_token_id = audio_token_id
audio_idx += 1
else:
raise ValueError(f"Invalid modality: {modality}")
assert cur_idx <= mm_start_idx
input_ids.extend(prompt[cur_idx : mm_start_idx + 1])
mm_offset_start = len(input_ids)
input_ids.extend([mm_token_id] * mm_token_num)
cur_idx = (
mm_start_idx + 2
) # jump to img_end_id, video_end_id, or audio_end_id
offsets.append((mm_offset_start, len(input_ids) - 1))
else:
input_ids.extend(prompt[cur_idx:])
return input_ids, offsets
return input_ids, offsets, modality_list
def get_mm_data(self, prompt, embeddings, img_grid_thw):
input_ids, offsets = self.build_input_ids(prompt, img_grid_thw)
mm_items = [
MultimodalDataItem(
modality=Modality.IMAGE,
offsets=offsets,
precomputed_embeddings=embeddings,
def get_mm_data(self, prompt, embeddings, **kwargs):
img_grid_thw = kwargs.get("img_grid_thw", None)
video_grid_thw = kwargs.get("video_grid_thw", None)
audio_feature_lens = kwargs.get("audio_feature_lens", None)
input_ids, offsets, modality_list = self.build_input_ids(
prompt,
img_grid_thw=img_grid_thw,
video_grid_thw=video_grid_thw,
audio_seq_lens=audio_feature_lens,
)
assert all(isinstance(modality, Modality) for modality in modality_list)
mm_items = []
consumed_per_modality = {}
for modality, offset in zip(modality_list, offsets):
num_tokens = offset[1] - offset[0] + 1
embedding_start = consumed_per_modality.get(modality, 0)
embedding_slice = embeddings[modality][
embedding_start : embedding_start + num_tokens
]
consumed_per_modality[modality] = embedding_start + num_tokens
mm_items.append(
MultimodalDataItem(
modality=modality,
offsets=offset,
precomputed_embeddings=embedding_slice,
)
)
]
return {
"input_ids": input_ids,
@@ -308,6 +360,7 @@ class BaseMultimodalProcessor(ABC):
"im_start_id": self.IM_START_TOKEN_ID,
"im_end_id": self.IM_END_TOKEN_ID,
"im_token_id": self.IM_TOKEN_ID,
"video_token_id": getattr(self, "VIDEO_TOKEN_ID", None),
}
def process_mm_data(

View File

@@ -1,6 +1,6 @@
import re
from sglang.srt.managers.schedule_batch import Modality
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
from sglang.srt.models.qwen2_audio import Qwen2AudioForConditionalGeneration
from sglang.srt.multimodal.processors.base_processor import (
BaseMultimodalProcessor,
@@ -31,6 +31,52 @@ class Qwen2AudioMultimodalProcessor(BaseMultimodalProcessor):
self.ATTR_NAME_TO_MODALITY.update({"feature_attention_mask": Modality.AUDIO})
def get_mm_data(self, prompt, embeddings, **kwargs):
audio_feature_lens = kwargs.get("audio_feature_lens", None)
# Convert audio_feature_lens to token counts for build_input_ids
output_lengths = None
input_lengths = None
if audio_feature_lens is not None:
if audio_feature_lens.dim() > 1:
audio_feature_lens = audio_feature_lens.flatten()
input_lengths = (audio_feature_lens - 1) // 2 + 1
output_lengths = (input_lengths - 2) // 2 + 1
input_ids, offsets, modality_list = self.build_input_ids(
prompt,
audio_seq_lens=output_lengths,
)
mm_items = []
consumed_per_modality = {}
for modality, offset in zip(modality_list, offsets):
num_tokens = offset[1] - offset[0] + 1
embedding_start = consumed_per_modality.get(modality, 0)
embedding_slice = embeddings[modality][
embedding_start : embedding_start + num_tokens
]
consumed_per_modality[modality] = embedding_start + num_tokens
mm_items.append(
MultimodalDataItem(
modality=modality,
offsets=offset,
precomputed_embeddings=embedding_slice,
)
)
if mm_items:
mm_items[0].audio_feature_lens = output_lengths
return {
"mm_items": mm_items,
"input_ids": input_ids,
"audio_start_id": self.audio_start_id,
"audio_token_id": self.audio_token_id,
"audio_end_id": self.audio_end_id,
}
async def process_mm_data_async(
self,
audio_data,

View File

@@ -254,6 +254,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
self.IM_START_TOKEN_ID = hf_config.vision_start_token_id
self.IM_END_TOKEN_ID = hf_config.vision_end_token_id
self.IM_TOKEN_ID = hf_config.image_token_id
self.VIDEO_TOKEN_ID = hf_config.video_token_id
self.vision_start_token_id = hf_config.vision_start_token_id
self.vision_end_token_id = getattr(hf_config, "vision_end_token_id", None)
@@ -271,12 +272,138 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
image_token_regex=re.compile(
r"<\|vision_start\|>(?:<\|image_pad\|>)+<\|vision_end\|>"
),
video_token_id=hf_config.video_token_id,
video_token_id=self.VIDEO_TOKEN_ID,
audio_token_id=self.audio_token_id,
).build(_processor)
def get_mm_data(self, prompt, embeddings, img_grid_thw):
input_ids, offsets = self.build_input_ids(prompt, img_grid_thw)
def build_input_ids_with_timestamps(
self, prompt, embeddings, img_grid_thw, video_grid_thw, video_timestamps
):
"""
Build input_ids with timestamps for qwen3_vl models.
"""
if not isinstance(prompt, list):
prompt = self._processor.tokenizer.encode(prompt)
img_token_id = getattr(self, "IM_TOKEN_ID", None)
video_token_id = getattr(self, "VIDEO_TOKEN_ID", None)
audio_token_id = getattr(self, "audio_token_id", None)
spatial_merge_size = getattr(self, "spatial_merge_size", 1)
vision_start_token_id = getattr(self, "vision_start_token_id", None)
vision_end_token_id = getattr(self, "vision_end_token_id", None)
input_ids = []
offsets = []
modality_list = []
cur_idx = 0
vision_start_indices = []
for i in range(len(prompt) - 1):
if img_token_id is not None and prompt[i + 1] == img_token_id:
vision_start_indices.append((i, Modality.IMAGE))
elif video_token_id is not None and prompt[i + 1] == video_token_id:
vision_start_indices.append((i, Modality.VIDEO))
img_idx = 0
video_idx = 0
model_type = getattr(self, "model_type", None)
for mm_start_idx, modality in vision_start_indices:
modality_list.append(modality)
video_tokens = None
if modality == Modality.IMAGE:
mm_token_num = img_grid_thw[img_idx].prod() // (spatial_merge_size**2)
mm_token_id = img_token_id
img_idx += 1
elif modality == Modality.VIDEO:
curr_timestamps = video_timestamps[video_idx]
num_frames = video_grid_thw[video_idx][0]
frame_seqlen = video_grid_thw[video_idx][1:].prod().item() // (
spatial_merge_size**2
)
video_tokens = []
_current_offset = len(input_ids) + mm_start_idx + 1 - cur_idx
# take single frame as one mm_item
for frame_idx in range(num_frames):
if frame_idx > 0:
modality_list.append(Modality.VIDEO)
curr_time = curr_timestamps[frame_idx]
timestamp_text = f"<{curr_time:.1f} seconds>"
timestamp_tokens = self._processor.tokenizer.encode(
timestamp_text, add_special_tokens=False
)
video_tokens.extend(timestamp_tokens)
_current_offset += len(timestamp_tokens)
if vision_start_token_id is not None:
video_tokens.append(vision_start_token_id)
_current_offset += 1
video_tokens.extend([video_token_id] * frame_seqlen)
if vision_end_token_id is not None:
video_tokens.append(vision_end_token_id)
offsets.append(
(_current_offset, _current_offset + frame_seqlen - 1)
)
_current_offset += (
frame_seqlen + 1
if vision_end_token_id is not None
else frame_seqlen
) # for vision_end_token_id
mm_token_num = len(video_tokens)
mm_token_id = None
video_idx += 1
else:
logger.warning(
f"{modality} modality is not supported for qwen3_vl models with timestamps."
)
continue
assert cur_idx <= mm_start_idx
input_ids.extend(prompt[cur_idx : mm_start_idx + 1])
if modality == Modality.VIDEO:
input_ids.extend(video_tokens)
else:
mm_offset_start = len(input_ids)
input_ids.extend([mm_token_id] * mm_token_num)
offsets.append((mm_offset_start, len(input_ids) - 1))
cur_idx = mm_start_idx + 2 # jump to vision_end_id
else:
input_ids.extend(prompt[cur_idx:])
return input_ids, offsets, modality_list
def get_mm_data(self, prompt, embeddings, **kwargs):
img_grid_thw = kwargs.get("img_grid_thw", None)
video_grid_thw = kwargs.get("video_grid_thw", None)
audio_feature_lens = kwargs.get("audio_feature_lens", None)
video_timestamps = kwargs.get("video_timestamps", None)
second_per_grid_ts = kwargs.get("second_per_grid_ts", None)
audio_seq_lens = None
if audio_feature_lens is not None:
if self.model_type == "qwen3_omni_moe":
# apply _get_feat_extract_lengths to get seq_lens
input_lengths_leave = audio_feature_lens % 100
feat_lengths = (input_lengths_leave - 1) // 2 + 1
audio_seq_lens = (
((feat_lengths - 1) // 2 + 1 - 1) // 2
+ 1
+ (audio_feature_lens // 100) * 13
)
elif self.model_type == "qwen2_5_omni":
audio_seq_lens = (audio_feature_lens - 1) // 2 + 1
audio_seq_lens = (audio_seq_lens - 2) // 2 + 1
if (
self.model_type in ["qwen3_vl", "qwen3_vl_moe"]
and video_timestamps is not None
):
input_ids, offsets, modality_list = self.build_input_ids_with_timestamps(
prompt, embeddings, img_grid_thw, video_grid_thw, video_timestamps
)
else:
input_ids, offsets, modality_list = self.build_input_ids(
prompt, img_grid_thw, video_grid_thw, audio_seq_lens=audio_seq_lens
)
assert all(isinstance(modality, Modality) for modality in modality_list)
mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index(
spatial_merge_size=self.hf_config.vision_config.spatial_merge_size,
image_token_id=self.mm_tokens.image_token_id,
@@ -285,19 +412,41 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
model_type=self.model_type,
input_ids=torch.tensor(input_ids, dtype=torch.long).unsqueeze(0),
image_grid_thw=img_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
use_audio_in_video=False,
audio_seqlens=(
audio_feature_lens if self.model_type == "qwen3_omni_moe" else None
),
audio_token_id=getattr(self.hf_config, "audio_token_id", None),
audio_start_token_id=self.audio_start_token_id,
position_id_per_seconds=getattr(
self.hf_config, "position_id_per_seconds", None
),
tokens_per_second=getattr(
self.hf_config.vision_config, "tokens_per_second", None
),
)
mrope_positions = mrope_positions.squeeze(1)
mm_items = [
MultimodalDataItem(
modality=Modality.IMAGE,
offsets=offsets,
precomputed_embeddings=embeddings,
mm_items = []
consumed_per_modality = {}
for modality, offset in zip(modality_list, offsets):
num_tokens = offset[1] - offset[0] + 1
embedding_start = consumed_per_modality.get(modality, 0)
embedding_slice = embeddings[modality][
embedding_start : embedding_start + num_tokens
]
consumed_per_modality[modality] = embedding_start + num_tokens
logger.info(f"Get embedding slice for {modality}, num_tokens={num_tokens}")
mm_items.append(
MultimodalDataItem(
modality=modality,
offsets=offset,
precomputed_embeddings=embedding_slice,
)
)
]
return {
"input_ids": input_ids,