From 784fe7e99b80afb564ea29f7e5cb5ee0698390c9 Mon Sep 17 00:00:00 2001 From: Zhonghua Deng Date: Mon, 18 May 2026 10:33:16 +0800 Subject: [PATCH] feat(mimo-v2): add EPD disaggregation support (#24931) --- .../srt/disaggregation/encode_receiver.py | 82 ++- .../srt/disaggregation/encode_server.py | 238 ++++--- python/sglang/srt/managers/mm_utils.py | 72 +- .../sglang/srt/managers/tokenizer_manager.py | 5 + python/sglang/srt/models/mimo_v2.py | 178 ++++- .../srt/multimodal/processors/mimo_v2.py | 669 +++++++++++++++--- python/sglang/srt/server_args.py | 8 +- 7 files changed, 962 insertions(+), 290 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 754d9ad18..63194957e 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -193,7 +193,29 @@ _MODALITY_GRID_ATTRS = { Modality.VIDEO: ("video_grid_thw", False), Modality.AUDIO: ("audio_feature_lens", True), } -_VIDEO_META_ATTRS = ("video_timestamps", "second_per_grid_ts") +# Per-part video metadata for EPD. Tensor attrs cat on dim=0 across parts; +# others chain as lists. video_meta_attrs_for(model_type) resolves the active +# set per instance so non-MiMo runs skip the MiMo audio fields entirely. +_GENERAL_VIDEO_META_ATTRS = ( + "video_timestamps", + "second_per_grid_ts", +) +# MiMo-VL audio-in-video fields; appended only when model_type is MiMo. +_MIMO_VIDEO_AUDIO_META_ATTRS = ( + "video_audio_feature_lens", + "video_audio_segment_lens_flat", + "video_audio_per_video_num_units", + "video_audio_embedding", +) +_VIDEO_META_TENSOR_ATTRS = ("video_audio_feature_lens", "video_audio_embedding") + + +def video_meta_attrs_for(model_type: Optional[str]) -> tuple: + """Video-meta attrs for model_type. MiMo appends its audio-in-video fields.""" + attrs = _GENERAL_VIDEO_META_ATTRS + if model_type and "mimo" in model_type.lower(): + attrs = attrs + _MIMO_VIDEO_AUDIO_META_ATTRS + return attrs def _cat_grid(dims, flatten_items=False): @@ -231,6 +253,7 @@ class MultiModalEmbeddingData(EmbeddingData): modality, embedding, embedding_shape, + model_type: Optional[str] = None, **kwargs, ): super().__init__( @@ -243,6 +266,7 @@ class MultiModalEmbeddingData(EmbeddingData): embedding_shape, **kwargs, ) + self.video_meta_attrs = video_meta_attrs_for(model_type) self.img_grid_thw = [None] * num_parts self.video_grid_thw = [None] * num_parts self.audio_feature_lens = [None] * num_parts @@ -256,8 +280,8 @@ class MultiModalEmbeddingData(EmbeddingData): self.embedding_shape_list = [ embedding_shape if i == part_idx else None for i in range(num_parts) ] - self.video_timestamps = [None] * num_parts - self.second_per_grid_ts = [None] * num_parts + for attr in self.video_meta_attrs: + setattr(self, attr, [None] * num_parts) self._set_part_grid(part_idx, modality, self.get_grid()) if modality == Modality.VIDEO: @@ -274,7 +298,7 @@ class MultiModalEmbeddingData(EmbeddingData): def _set_video_meta_for_part(self, part_idx, source): """Copy video_timestamps and second_per_grid_ts from source (dict or object).""" - for attr_name in _VIDEO_META_ATTRS: + for attr_name in self.video_meta_attrs: val = ( source.get(attr_name) if isinstance(source, dict) @@ -284,11 +308,15 @@ class MultiModalEmbeddingData(EmbeddingData): getattr(self, attr_name)[part_idx] = val @classmethod - def from_embedding_data(cls, embedding_data: EmbeddingData): + def from_embedding_data( + cls, + embedding_data: EmbeddingData, + model_type: Optional[str] = None, + ): """Create MultiModalEmbeddingData from an EmbeddingData instance.""" # Only forward known optional attrs (e.g. video metadata) so they land on the instance extra = {} - for attr in _VIDEO_META_ATTRS: + for attr in video_meta_attrs_for(model_type): val = getattr(embedding_data, attr, None) if val is not None: extra[attr] = val @@ -300,6 +328,7 @@ class MultiModalEmbeddingData(EmbeddingData): modality=embedding_data.modality, embedding=embedding_data.embedding, embedding_shape=embedding_data.shape, + model_type=model_type, **extra, ) mm_data.send_time = embedding_data.send_time @@ -313,11 +342,8 @@ class MultiModalEmbeddingData(EmbeddingData): groups = defaultdict(list) for i, e in enumerate(self.embedding_list): if e is not None: - groups[self.modality_list[i]].append(e.cuda()) - return { - mod: torch.concat(tensors).to("cpu", non_blocking=True) - for mod, tensors in groups.items() - } + groups[self.modality_list[i]].append(e) + return {mod: torch.cat(tensors, dim=0) for mod, tensors in groups.items()} return self.embedding_list @property @@ -333,12 +359,16 @@ class MultiModalEmbeddingData(EmbeddingData): self.audio_feature_lens, flatten_items=True ), } - for attr in _VIDEO_META_ATTRS: + for attr in self.video_meta_attrs: lst = getattr(self, attr, None) if not lst: continue valid = [a for a in lst if a is not None] - if valid: + if not valid: + continue + if attr in _VIDEO_META_TENSOR_ATTRS: + kwargs[attr] = torch.cat(valid, dim=0) + else: kwargs[attr] = list(itertools.chain(*valid)) return kwargs @@ -546,7 +576,7 @@ class WaitingImageRequest: if self.recv_embedding_data is None: self.recv_embedding_data = MultiModalEmbeddingData.from_embedding_data( - recv_obj + recv_obj, model_type=self.model_type ) else: self.recv_embedding_data.add(recv_obj) @@ -634,7 +664,13 @@ class MMReceiverBase(ABC): self.context = zmq.asyncio.Context(20) self.encoder_transfer_backend = server_args.encoder_transfer_backend self.encode_urls = server_args.encoder_urls + self.recv_timeout = envs.SGLANG_ENCODER_RECV_TIMEOUT.get() self.host = get_local_ip_auto(server_args.host) + self.model_type = ( + getattr(hf_config, "model_type", "").lower() + if hf_config is not None + else None + ) if self.encoder_transfer_backend == "mooncake": self.dtype = dtype self.embeddings_engine = get_mooncake_transfer_engine() @@ -724,15 +760,25 @@ class MMReceiverBase(ABC): self.context, zmq.PULL, host=self.host ) mm_data = self._extract_url_data(request_obj) + modalities = [m.get("modality") for m in mm_data] + logger.info( + f"[{req_id}] Sending encode request to E, " + f"modalities={modalities}, num_items={len(mm_data)}" + ) + send_time = time.monotonic() asyncio.create_task( self.encode(req_id, mm_data, embedding_port, "encode", "send") ) - return await asyncio.wait_for( + result = await asyncio.wait_for( self._recv_mm_data(req_id, recv_socket, mm_processor, prompt), - timeout=20, + timeout=self.recv_timeout, ) + elapsed = time.monotonic() - send_time + logger.info(f"[{req_id}] Received embedding from E in {elapsed:.3f}s") + return result except asyncio.TimeoutError: - logger.warning(f"Embedding recv timeout for request {req_id}") + elapsed = time.monotonic() - send_time + logger.warning(f"[{req_id}] Embedding recv timeout after {elapsed:.3f}s") if req_id is not None: self._cleanup_mooncake_buffer(req_id) return None @@ -797,7 +843,7 @@ class MMReceiverBase(ABC): ) if recv_embedding_data is None: recv_embedding_data = MultiModalEmbeddingData.from_embedding_data( - recv_obj + recv_obj, model_type=self.model_type ) else: recv_embedding_data.add(recv_obj) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index c033825b6..ff13cbb1d 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -24,7 +24,10 @@ from sglang.srt.configs.device_config import DeviceConfig from sglang.srt.configs.load_config import LoadConfig from sglang.srt.configs.model_config import ModelConfig from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX -from sglang.srt.disaggregation.encode_receiver import EmbeddingData +from sglang.srt.disaggregation.encode_receiver import ( + EmbeddingData, + video_meta_attrs_for, +) from sglang.srt.distributed.parallel_state import ( get_default_distributed_backend, get_mooncake_transfer_engine, @@ -45,6 +48,7 @@ from sglang.srt.server_args import ( set_global_server_args_for_scheduler, ) from sglang.srt.utils import ( + configure_logger, load_audio, load_image, load_video, @@ -170,15 +174,9 @@ def _get_mm_feature(mm_inputs, modality): ) -def _build_mm_aux_data(mm_inputs): - """ - Build auxiliary data for video modality. - """ - aux_data = { - "video_timestamps": mm_inputs.get("video_timestamps", None), - "second_per_grid_ts": mm_inputs.get("second_per_grid_ts", None), - } - return aux_data +def _build_mm_aux_data(mm_inputs, model_type=None): + # Video aux metadata, scoped to model_type's video-meta attrs. + return {attr: mm_inputs.get(attr) for attr in video_meta_attrs_for(model_type)} class MMEncoder: @@ -832,7 +830,7 @@ class MMEncoder: self.background_tasks.add(task) task.add_done_callback(self.background_tasks.discard) - aux_data = _build_mm_aux_data(mm_inputs) + aux_data = _build_mm_aux_data(mm_inputs, self.model_type) self.embedding_to_send[req_id] = EmbeddingData( req_id, num_parts, @@ -950,104 +948,117 @@ class MMEncoder: return normalized async def _process_mm_items(self, mm_items, modality): - if modality == Modality.IMAGE and self.image_processor: - images = await self._flatten_and_load_images(mm_items) - image_config = self.vision_config.get("image", {}) - if self.model_type in ["kimi_k25", "kimi_vl"]: - images = self._normalize_kimi_encoder_images(images) - processor_input = self.image_processor(images=images, **image_config) - if hasattr(self.model, "thinker"): # for omni models - get_feature_method = self.model.thinker.get_image_feature - else: - get_feature_method = self.model.get_image_feature - elif modality == Modality.VIDEO and self.video_processor: - videos, video_processor_kwargs = await self._flatten_and_load_videos( - mm_items - ) - processor_input = self.video_processor( - videos=videos, **video_processor_kwargs - ) - # Get additional video metadata - if ( - self.model_type - in [ - "qwen3_vl", - "qwen3_vl_moe", - "qwen3_5", - "qwen3_5_moe", - "intern_s2_preview", - ] - and video_processor_kwargs.get("video_metadata", None) is not None - ): - # For qwen3-vl/qwen3.5 models, we need to store the video timestamps - video_metadata = video_processor_kwargs["video_metadata"] - try: - merge_size = ( - self.model_config.hf_config.vision_config.spatial_merge_size - ) - except (AttributeError, KeyError): - merge_size = 2 # Default merge_size + model_preprocessor = getattr(self.model, "preprocess_mm_for_encoder", None) - video_timestamps = [] - for metadata in video_metadata: - video_fps = metadata.get("fps", None) or 24 # original video fps - frames_indices = metadata.get("frames_indices", None) - timestamps = self._calculate_timestamps( - frames_indices, video_fps, merge_size - ) - video_timestamps.append(timestamps) - processor_input["video_timestamps"] = video_timestamps - elif ( - self.model_type in ["qwen2_5_vl", "qwen2_5_omni", "qwen3_omni_moe"] - and processor_input.get("video_grid_thw", None) is not None - ): - # For omni/qwen2_5_vl models, calculate second_per_grid_ts for rotary embedding - video_grid_thw = processor_input["video_grid_thw"] - try: - temporal_patch_size = self.video_processor.temporal_patch_size - except AttributeError: - temporal_patch_size = 2 # Default temporal_patch_size - # get sampled fps, default: 2 - fps_list = [ - self.vision_config.get("video", {}).get("fps", None) or 2 - ] * len(video_grid_thw) - second_per_grid_ts = [(temporal_patch_size / fps) for fps in fps_list] - second_per_grid_ts_tensor = torch.tensor( - second_per_grid_ts, dtype=torch.float32 - ) - processor_input["second_per_grid_ts"] = second_per_grid_ts_tensor - - if hasattr(self.model, "thinker"): # for omni models - get_feature_method = self.model.thinker.get_video_feature - else: - get_feature_method = self.model.get_video_feature - elif modality == Modality.AUDIO and self.audio_processor: - audios = await self._flatten_and_load_audios(mm_items) - audio_config = self.vision_config.get("audio", {}) - processor_input = self.audio_processor.feature_extractor( - audios, **audio_config + if modality == Modality.IMAGE: + processor_input = await self._process_image_items( + mm_items, model_preprocessor ) - processor_input["feature_attention_mask"] = processor_input.pop( - "attention_mask" + elif modality == Modality.VIDEO: + processor_input = await self._process_video_items( + mm_items, model_preprocessor ) - # convert to same format as image/video - input_lengths = torch.tensor( - processor_input["feature_attention_mask"].sum(-1), dtype=torch.long + elif modality == Modality.AUDIO: + processor_input = await self._process_audio_items( + mm_items, model_preprocessor ) - processor_input["audio_feature_lens_raw"] = input_lengths - output_lengths = self._get_feat_extract_output_lengths(input_lengths) - processor_input["audio_feature_lens"] = output_lengths - if hasattr(self.model, "thinker"): # for omni models - get_feature_method = self.model.thinker.get_audio_feature - else: - get_feature_method = self.model.get_audio_feature else: - raise ValueError( - f"Currently only support image, video and audio modalities, {modality} modality has no processor available." - ) + raise ValueError(f"Unsupported modality: {modality}") + target = self.model.thinker if hasattr(self.model, "thinker") else self.model + get_feature_method = getattr(target, f"get_{modality.name.lower()}_feature") return processor_input, get_feature_method + async def _process_image_items(self, mm_items, model_preprocessor): + if not (self.image_processor or model_preprocessor): + raise ValueError("No image processor available") + images = await self._flatten_and_load_images(mm_items) + if model_preprocessor: + return model_preprocessor(images, Modality.IMAGE, self.vision_config) + image_config = self.vision_config.get("image", {}) + if self.model_type in ["kimi_k25", "kimi_vl"]: + images = self._normalize_kimi_encoder_images(images) + return self.image_processor(images=images, **image_config) + + async def _process_video_items(self, mm_items, model_preprocessor): + if model_preprocessor: + return model_preprocessor(mm_items, Modality.VIDEO, self.vision_config) + if not self.video_processor: + raise ValueError("No video processor available") + + videos, video_processor_kwargs = await self._flatten_and_load_videos(mm_items) + processor_input = self.video_processor(videos=videos, **video_processor_kwargs) + + # Get additional video metadata + if ( + self.model_type + in [ + "qwen3_vl", + "qwen3_vl_moe", + "qwen3_5", + "qwen3_5_moe", + "intern_s2_preview", + ] + and video_processor_kwargs.get("video_metadata", None) is not None + ): + video_metadata = video_processor_kwargs["video_metadata"] + try: + merge_size = ( + self.model_config.hf_config.vision_config.spatial_merge_size + ) + except (AttributeError, KeyError): + merge_size = 2 # Default merge_size + + video_timestamps = [] + for metadata in video_metadata: + video_fps = metadata.get("fps", None) or 24 # original video fps + frames_indices = metadata.get("frames_indices", None) + timestamps = self._calculate_timestamps( + frames_indices, video_fps, merge_size + ) + video_timestamps.append(timestamps) + processor_input["video_timestamps"] = video_timestamps + elif ( + self.model_type in ["qwen2_5_vl", "qwen2_5_omni", "qwen3_omni_moe"] + and processor_input.get("video_grid_thw", None) is not None + ): + video_grid_thw = processor_input["video_grid_thw"] + try: + temporal_patch_size = self.video_processor.temporal_patch_size + except AttributeError: + temporal_patch_size = 2 # Default temporal_patch_size + fps_list = [ + self.vision_config.get("video", {}).get("fps", None) or 2 + ] * len(video_grid_thw) + second_per_grid_ts = [(temporal_patch_size / fps) for fps in fps_list] + second_per_grid_ts_tensor = torch.tensor( + second_per_grid_ts, dtype=torch.float32 + ) + processor_input["second_per_grid_ts"] = second_per_grid_ts_tensor + + return processor_input + + async def _process_audio_items(self, mm_items, model_preprocessor): + if model_preprocessor: + return model_preprocessor(mm_items, Modality.AUDIO, self.vision_config) + if not self.audio_processor: + raise ValueError("No audio processor available") + + audios = await self._flatten_and_load_audios(mm_items) + audio_config = self.vision_config.get("audio", {}) + processor_input = self.audio_processor.feature_extractor(audios, **audio_config) + processor_input["feature_attention_mask"] = processor_input.pop( + "attention_mask" + ) + # convert to same format as image/video + input_lengths = torch.tensor( + processor_input["feature_attention_mask"].sum(-1), dtype=torch.long + ) + processor_input["audio_feature_lens_raw"] = input_lengths + output_lengths = self._get_feat_extract_output_lengths(input_lengths) + processor_input["audio_feature_lens"] = output_lengths + return processor_input + async def _encode(self, mm_items, modality: Modality) -> torch.Tensor: try: mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality) @@ -1092,7 +1103,23 @@ class MMEncoder: if self.profiler is not None: self.profiler.step() - aux_data = _build_mm_aux_data(mm_inputs) + aux_data = _build_mm_aux_data(mm_inputs, self.model_type) + + if modality == Modality.VIDEO and mm_inputs.get("video_audio_features"): + target = ( + self.model.thinker if hasattr(self.model, "thinker") else self.model + ) + encode_video_audio_fn = getattr(target, "encode_video_audio", None) + if encode_video_audio_fn is not None: + audio_embedding = encode_video_audio_fn(mm_inputs) + if audio_embedding is not None: + aux_data["video_audio_embedding"] = audio_embedding + else: + logger.warning( + "Videos carry audio tracks but model has no " + "encode_video_audio; dropping audio for EPD encoding." + ) + return ( _get_mm_grid_dim(mm_inputs, modality, self.model_type), mm_embedding, @@ -1417,6 +1444,7 @@ def launch_encoder(server_args, schedule_path, dist_init_method, rank): def launch_server(server_args: ServerArgs): + configure_logger(server_args, prefix=" encode_server") global encoder ctx = mp.get_context("spawn") zmq_ctx = zmq.Context(10) @@ -1453,6 +1481,7 @@ async def get_condition(rid): @app.post("/encode") async def handle_encode_request(request: dict): req_id = request["req_id"] + start_time = time.monotonic() try: def start_background_send(req_id): @@ -1538,6 +1567,11 @@ async def handle_encode_request(request: dict): embedding_port=request["embedding_port"], ) encoder.embedding_to_send.pop(request["req_id"], None) + elapsed = time.monotonic() - start_time + logger.info( + f"[{req_id}] /encode completed in {elapsed:.3f}s, " + f"modality={request['modality']}, tokens={embedding_len}" + ) return ORJSONResponse(content=None) except Exception as e: error_msg = str(e) diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index ca7eca8cf..1c7d3afe1 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -1598,6 +1598,19 @@ def _get_is_default_transport(): return _is_default_tensor_transport +def _wrap_tensor_or_list(value): + """Wrap a CPU tensor (or list of CPU tensors) in ShmPointerMMData.""" + if isinstance(value, torch.Tensor) and value.is_cpu: + return ShmPointerMMData(value) + elif isinstance(value, (list, tuple)): + wrapped = [ + (ShmPointerMMData(t) if isinstance(t, torch.Tensor) and t.is_cpu else t) + for t in value + ] + return type(value)(wrapped) if isinstance(value, tuple) else wrapped + return value + + def wrap_shm_features(obj): """ Scan the object for multimodal tensors and wrap them in SHM pointers. @@ -1607,22 +1620,14 @@ def wrap_shm_features(obj): if hasattr(obj, "mm_inputs") and obj.mm_inputs: for item in obj.mm_inputs.mm_items: - if not hasattr(item, "feature"): - continue - feat = item.feature - if isinstance(feat, torch.Tensor) and feat.is_cpu: - item.feature = ShmPointerMMData(feat) - elif isinstance(feat, (list, tuple)): - wrapped = [ - ( - ShmPointerMMData(t) - if isinstance(t, torch.Tensor) and t.is_cpu - else t - ) - for t in feat - ] - item.feature = ( - type(feat)(wrapped) if isinstance(feat, tuple) else wrapped + if hasattr(item, "feature") and item.feature is not None: + item.feature = _wrap_tensor_or_list(item.feature) + if ( + hasattr(item, "precomputed_embeddings") + and item.precomputed_embeddings is not None + ): + item.precomputed_embeddings = _wrap_tensor_or_list( + item.precomputed_embeddings ) return obj @@ -1646,9 +1651,23 @@ def has_shm_features(recv_reqs): for item in req.mm_inputs.mm_items: if _feature_has_shm(item.feature): return True + if _feature_has_shm(getattr(item, "precomputed_embeddings", None)): + return True return False +def _unwrap_tensor_or_list(value): + """Restore ShmPointerMMData wrappers back into standard torch.Tensors.""" + if isinstance(value, ShmPointerMMData): + return value.materialize() + elif isinstance(value, (list, tuple)): + unwrapped = [ + t.materialize() if isinstance(t, ShmPointerMMData) else t for t in value + ] + return type(value)(unwrapped) if isinstance(value, tuple) else unwrapped + return value + + def unwrap_shm_features(obj): """ Restore ShmPointerMMData wrappers back into standard torch.Tensors. @@ -1663,17 +1682,14 @@ def unwrap_shm_features(obj): return obj # Handle single requests if hasattr(obj, "mm_inputs") and obj.mm_inputs: - mm_items = obj.mm_inputs.mm_items - for item in mm_items: - feat = item.feature - if isinstance(feat, ShmPointerMMData): - item.feature = feat.materialize() - elif isinstance(feat, (list, tuple)): - unwrapped = [ - t.materialize() if isinstance(t, ShmPointerMMData) else t - for t in feat - ] - item.feature = ( - type(feat)(unwrapped) if isinstance(feat, tuple) else unwrapped + for item in obj.mm_inputs.mm_items: + if hasattr(item, "feature") and item.feature is not None: + item.feature = _unwrap_tensor_or_list(item.feature) + if ( + hasattr(item, "precomputed_embeddings") + and item.precomputed_embeddings is not None + ): + item.precomputed_embeddings = _unwrap_tensor_or_list( + item.precomputed_embeddings ) return obj diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 6947055b8..c097aa4b8 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -779,6 +779,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs, ) if mm_inputs is None: + if self.server_args.language_only: + logger.warning( + "Encoder embedding not available, " + "falling back to local mm processing" + ) mm_inputs = await self.mm_processor.process_mm_data_async( image_data=obj.image_data, audio_data=obj.audio_data, diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index 5cc20c384..2b76b5c28 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/python/sglang/srt/models/mimo_v2.py @@ -68,7 +68,11 @@ from sglang.srt.managers.mm_utils import ( MultiModalityDataPaddingPatternMultimodalTokens, general_mm_embed_routine, ) -from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import ( default_weight_loader, @@ -1007,6 +1011,11 @@ class MiMoV2ForCausalLM(nn.Module): "up_proj": ("gate_up_proj", 1), } + # Prefixes for weight routing in encoder_only/language_only modes + _LANGUAGE_WEIGHT_PREFIXES = ("model.", "lm_head.") + _VISION_AUDIO_WEIGHT_PREFIXES = ("visual.", "vision_model.", "audio_") + _VISION_AUDIO_WEIGHT_SUBSTRING = "speech_embeddings" + def __init__( self, config: MiMoV2Config, @@ -1017,27 +1026,36 @@ class MiMoV2ForCausalLM(nn.Module): self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config - self.model = MiMoV2Model( - config, quant_config=quant_config, prefix=add_prefix("model", prefix) - ) + self._encoder_processor = None # lazy-created in preprocess_mm_for_encoder - if self.pp_group.is_last_rank: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=add_prefix("lm_head", prefix), - use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + if not self.config.encoder_only: + self.model = MiMoV2Model( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) ) - else: - # ranks other than the last rank will have a placeholder layer - self.lm_head = PPMissingLayer() - self.logits_processor = LogitsProcessor(config) + if self.pp_group.is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + else: + self.lm_head = PPMissingLayer() + else: + self.model = None + self.lm_head = None + + self.logits_processor = ( + LogitsProcessor(config) if not self.config.encoder_only else None + ) vision_config = getattr(config, "vision_config", None) audio_config = getattr(config, "audio_config", None) self._is_multimodal = vision_config is not None and audio_config is not None + # Always build vision/audio encoders so P can fall back to local + # encoding when the EPD encoder is unreachable. if self._is_multimodal: if hasattr(vision_config, "to_dict"): vision_config = vision_config.to_dict() @@ -1054,11 +1072,15 @@ class MiMoV2ForCausalLM(nn.Module): self.audio_encoder = MiMoAudioEncoder(self.audio_config) self._routed_experts_weights_of_layer = LazyValue( - lambda: { - layer_id: layer.mlp.get_moe_weights() - for layer_id, layer in enumerate(self.model.layers) - if isinstance(layer.mlp, MiMoV2MoE) - } + lambda: ( + { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, MiMoV2MoE) + } + if self.model is not None + else {} + ) ) @property @@ -1066,12 +1088,24 @@ class MiMoV2ForCausalLM(nn.Module): return self._routed_experts_weights_of_layer.value def get_input_embedding(self, input_ids: torch.Tensor) -> torch.Tensor: + assert ( + self.model is not None + ), "get_input_embedding() is not available in encoder_only mode" return self.model.get_input_embedding(input_ids) def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): pattern = MultiModalityDataPaddingPatternMultimodalTokens() return pattern.pad_input_tokens(input_ids, mm_inputs) + def preprocess_mm_for_encoder(self, mm_data, modality, config): + if self._encoder_processor is None: + from sglang.srt.multimodal.processors.mimo_v2 import MiMoProcessor + + self._encoder_processor = MiMoProcessor.from_hf_config( + self.config, mm_config=config + ) + return self._encoder_processor.preprocess_for_encoder(mm_data, modality) + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: pixel_values = torch.cat([item.feature for item in items], dim=0).type( self.visual.dtype @@ -1093,8 +1127,76 @@ class MiMoV2ForCausalLM(nn.Module): def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: return self.audio_encoder.get_audio_feature(items) - def get_input_embeddings(self) -> nn.Embedding: - return self.model.embed_tokens + @torch.inference_mode() + def encode_video_audio(self, mm_inputs: Dict) -> Optional[torch.Tensor]: + # EPD-side hook: encode audio tracks pulled from videos and trim to the + # interleaved per-video segments produced by MiMoProcessor (segment + # starts / lens / per_video_num_units). Returns None if there is no + # audio to encode. The server passes the result through to the receiver + # under aux_data["video_audio_embedding"]. + import numpy as np + + audio_features = mm_inputs.get("video_audio_features") + if not audio_features: + return None + + def _as_tensor(data): + if isinstance(data, torch.Tensor): + return data + if isinstance(data, np.ndarray): + return torch.tensor(data) + if isinstance(data, list) and data and isinstance(data[0], np.ndarray): + return torch.tensor(np.array(data)) + if isinstance(data, list) and data and isinstance(data[0], (int, float)): + return torch.tensor(data) + return data + + audio_feature_lens = mm_inputs["video_audio_feature_lens"] + audio_item = MultimodalDataItem.from_dict( + { + "modality": Modality.AUDIO, + "feature": _as_tensor(audio_features), + } + ) + audio_item.set("audio_feature_lens", _as_tensor(audio_feature_lens)) + + audio_embedding = self.get_audio_feature([audio_item]).cpu() + if audio_embedding.ndim != 2: + audio_embedding = audio_embedding.reshape(-1, audio_embedding.shape[-1]) + + segment_lens_flat = mm_inputs["video_audio_segment_lens_flat"] + segment_starts_flat = mm_inputs["video_audio_segment_starts_flat"] + per_video_num_units = mm_inputs["video_audio_per_video_num_units"] + per_video_audio_token_lens = ( + audio_feature_lens.tolist() + if hasattr(audio_feature_lens, "tolist") + else list(audio_feature_lens) + ) + + trimmed_chunks = [] + emb_offset = 0 + unit_idx = 0 + audio_video_idx = 0 + for num_units in per_video_num_units: + if num_units <= 0: + continue + vid_audio_len = per_video_audio_token_lens[audio_video_idx] + for _ in range(num_units): + start = segment_starts_flat[unit_idx] + seg_len = segment_lens_flat[unit_idx] + trimmed_chunks.append( + audio_embedding[emb_offset + start : emb_offset + start + seg_len] + ) + unit_idx += 1 + emb_offset += vid_audio_len + audio_video_idx += 1 + + return ( + torch.cat(trimmed_chunks, dim=0) if trimmed_chunks else audio_embedding[:0] + ) + + def get_input_embeddings(self) -> Optional[nn.Embedding]: + return self.model.embed_tokens if self.model is not None else None @torch.no_grad() def forward( @@ -1105,6 +1207,10 @@ class MiMoV2ForCausalLM(nn.Module): input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: + assert ( + not self.config.encoder_only + ), "forward() should not be called in encoder_only mode" + if self._is_multimodal: hidden_states, hidden_states_before_norm = general_mm_embed_routine( input_ids=input_ids, @@ -1136,11 +1242,11 @@ class MiMoV2ForCausalLM(nn.Module): @property def start_layer(self): - return self.model.start_layer + return self.model.start_layer if self.model is not None else 0 @property def end_layer(self): - return self.model.end_layer + return self.model.end_layer if self.model is not None else 0 def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ @@ -1168,11 +1274,18 @@ class MiMoV2ForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) skipped_mtp_weights = False + def _is_vision_audio_weight(name): + return ( + name.startswith(self._VISION_AUDIO_WEIGHT_PREFIXES) + or self._VISION_AUDIO_WEIGHT_SUBSTRING in name + ) + for name, loaded_weight in weights: - if not self._is_multimodal and ( - name.startswith(("visual.", "vision_model.", "audio_encoder.")) - or name.startswith("audio_") - or "speech_embeddings" in name + if not self._is_multimodal and _is_vision_audio_weight(name): + continue + + if self.config.encoder_only and name.startswith( + self._LANGUAGE_WEIGHT_PREFIXES ): continue @@ -1374,9 +1487,15 @@ class MiMoV2ForCausalLM(nn.Module): logger.warning(f"Parameter {name} not found in params_dict") def get_embed_and_head(self): + assert ( + self.model is not None and self.lm_head is not None + ), "get_embed_and_head() is not available in encoder_only mode" return self.model.embed_tokens.weight, self.lm_head.weight def set_embed_and_head(self, embed, head): + assert ( + self.model is not None and self.lm_head is not None + ), "set_embed_and_head() is not available in encoder_only mode" del self.model.embed_tokens.weight del self.lm_head.weight self.model.embed_tokens.weight = embed @@ -1385,7 +1504,8 @@ class MiMoV2ForCausalLM(nn.Module): torch.cuda.synchronize() def load_kv_cache_scales(self, quantization_param_path: str) -> None: - self.model.load_kv_cache_scales(quantization_param_path) + if self.model is not None: + self.model.load_kv_cache_scales(quantization_param_path) @classmethod def get_model_config_for_expert_location(cls, config): diff --git a/python/sglang/srt/multimodal/processors/mimo_v2.py b/python/sglang/srt/multimodal/processors/mimo_v2.py index b4bed854b..f8c2c562f 100644 --- a/python/sglang/srt/multimodal/processors/mimo_v2.py +++ b/python/sglang/srt/multimodal/processors/mimo_v2.py @@ -268,6 +268,49 @@ _QWEN2VL_PIXEL_STD = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1) _mean_std_cache = {} +def _decode_frames_and_timestamps(vdw, ele): + # Shared E/D frame-sampling recipe: smart_nframes + linspace + permute. + total_frames, video_fps = len(vdw), vdw.avg_fps + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = list(np.unique(np.linspace(0, total_frames - 1, num=nframes, dtype=np.int64))) + video_tensor = vdw.get_frames_as_tensor(idx).permute(0, 3, 1, 2).float() + timestamps = torch.as_tensor(idx, dtype=torch.float32) / video_fps + return video_tensor, timestamps + + +def _ffprobe_has_audio(src, stdin=None, label=None) -> bool: + # Header-only audio-stream probe for HTTP URLs; avoids full download. + try: + r = subprocess.run( + [ + "ffprobe", + "-v", + "quiet", + "-print_format", + "json", + "-show_streams", + "-select_streams", + "a", + src, + ], + input=stdin, + capture_output=True, + timeout=30, + ) + if r.returncode != 0: + stderr = r.stderr.decode("utf-8", errors="replace") + raise RuntimeError(f"ffprobe failed for {label}: {stderr}") + return bool(json.loads(r.stdout).get("streams")) + except subprocess.TimeoutExpired: + logger.error("ffprobe timed out for %s", label) + raise + except FileNotFoundError as e: + raise RuntimeError("ffprobe not found; install ffmpeg") from e + except json.JSONDecodeError: + logger.error("ffprobe returned invalid JSON for %s", label) + raise + + class MiMoProcessor: def __init__( self, @@ -445,6 +488,207 @@ class MiMoProcessor: for k in kwargs: logger.info(f"[Warning] Ignored unknown parameter {k} for MiMoProcessor") + @classmethod + def from_hf_config(cls, hf_config, mm_config=None, **overrides): + # Params must come from hf_config.processor_config so E and D agree; + # any drift shifts input_ids on the D side. + def _as_dict(obj): + if isinstance(obj, dict): + return obj + return obj.to_dict() if obj and hasattr(obj, "to_dict") else {} + + pc = _as_dict(getattr(hf_config, "processor_config", None)) + ac = _as_dict(getattr(hf_config, "audio_config", None)) + vc = hf_config.vision_config + vget = vc.get if isinstance(vc, dict) else (lambda k, d=None: getattr(vc, k, d)) + patch_size = vget("patch_size", 14) + merge_size = vget("spatial_merge_size", 2) + f = patch_size * merge_size + + kwargs = { + "tokenizer": None, + "patch_size": patch_size, + "merge_size": merge_size, + "temporal_patch_size": vget("temporal_patch_size", 2), + "image_min_pixels": pc.get("image_min_pixels") or 4 * f * f, + "image_max_pixels": pc.get("image_max_pixels") or 4096 * f * f, + "video_min_pixels": pc.get("video_min_pixels") or 4 * f * f, + "video_max_pixels": pc.get("video_max_pixels") or 4096 * f * f, + "video_total_max_pixels": pc.get("video_total_max_pixels") or 16384 * f * f, + "fps": pc.get("fps") or 2, + "num_frames": pc.get("num_frames"), + "max_frames": pc.get("max_frames") or 256, + "min_frames": pc.get("min_frames") or 8, + "video_audio_interleave_length": pc.get("video_audio_interleave_length", 0), + "use_per_grid_t_timestamps": pc.get("use_per_grid_t_timestamps", False), + "use_video_timestamps": pc.get("use_video_timestamps", False), + } + # audio_sampling_rate: processor_config > audio_config > mm_config.audio. + asr = ( + pc.get("audio_sampling_rate") + or ac.get("sampling_rate") + or ac.get("sample_rate") + ) + if asr is not None: + kwargs["audio_sampling_rate"] = asr + + audio_cfg = (mm_config or {}).get("audio", {}) + for k in ( + "audio_sampling_rate", + "audio_hop_length", + "audio_n_mels", + "audio_kernel_size", + "audio_stride_size", + "audio_avg_pooler", + ): + if k in audio_cfg: + kwargs[k] = audio_cfg[k] + if "sampling_rate" in audio_cfg and "audio_sampling_rate" not in kwargs: + kwargs["audio_sampling_rate"] = audio_cfg["sampling_rate"] + + kwargs.update(overrides) + return cls(**kwargs) + + @staticmethod + def has_audio_track(path_or_data) -> bool: + # In-process probe via torchcodec for bytes/path; ffprobe range + # request for HTTP URLs so we do not pre-download the blob here. + if isinstance(path_or_data, str) and path_or_data.startswith( + ("http://", "https://") + ): + return _ffprobe_has_audio(path_or_data, stdin=None, label=path_or_data) + + if isinstance(path_or_data, bytes): + source = io.BytesIO(path_or_data) + elif ( + isinstance(path_or_data, str) + and path_or_data.startswith("data:") + and ";base64," in path_or_data + ): + source = io.BytesIO(base64.b64decode(path_or_data.split(";base64,")[1])) + else: + source = path_or_data # local path or file:// + try: + AudioDecoder(source) + return True + except Exception: + return False + + def _load_video_for_encoder(self, video_data): + # Normalise once to bytes-or-path; reused by frame decode, audio + # detection, and audio preprocessing without re-downloading. + from sglang.srt.utils.common import VideoData, _normalize_video_input + from sglang.srt.utils.video_decoder import VideoDecoderWrapper + + if isinstance(video_data, VideoData): + video_data = video_data.url + if isinstance(video_data, bytes): + video_blob = video_data + else: + video_blob = _normalize_video_input(video_data) + if video_blob is None: + raise ValueError( + f"Unsupported video input type for EPD encoder: {type(video_data)}" + ) + + vdw = VideoDecoderWrapper(video_blob, device="cpu") + try: + video_tuple = _decode_frames_and_timestamps( + vdw, self.default_video_processor_kwargs + ) + finally: + if hasattr(vdw, "close"): + vdw.close() + return video_blob, video_tuple + + def preprocess_for_encoder(self, mm_data, modality): + # EPD encoder-side features. video_audio_* fields appear when any + # video has audio; the D side uses them to rebuild input_ids. + from sglang.srt.managers.schedule_batch import Modality + + if not isinstance(mm_data, (list, tuple)): + mm_data = [mm_data] + + if modality == Modality.IMAGE: + factor = self.patch_size * self.merge_size + min_pixels = self.default_image_processor_kwargs["min_pixels"] + max_pixels = self.default_image_processor_kwargs["max_pixels"] + all_patches, all_grids = [], [] + for img in mm_data: + img_tensor, _, _ = self.get_visual_transform( + img, factor=factor, min_pixels=min_pixels, max_pixels=max_pixels + ) + patches, grid = self._flatten_visual_inputs(img_tensor, "image") + all_patches.append(patches) + all_grids.append(grid) + return { + "pixel_values": torch.cat(all_patches, dim=0), + "image_grid_thw": torch.stack(all_grids), + } + + if modality == Modality.VIDEO: + all_patches, all_grids, all_timestamps = [], [], [] + audio_features, audio_feature_lens = [], [] + seg_lens_flat, seg_starts_flat, per_video_num_units = [], [], [] + + for video_data in mm_data: + video_blob, video_tuple = self._load_video_for_encoder(video_data) + patches, grid, aligned_ts, video_meta = self.process_video( + VideoInput(video=video_tuple) + ) + all_patches.append(patches) + all_grids.append(grid) + step = self.temporal_patch_size * self.temporal_compression_ratio + all_timestamps.extend(aligned_ts[::step].tolist()) + + if self.has_audio_track(video_blob): + audio_spec, audio_token_len = self.preprocess_audio(video_blob) + units = self._build_video_audio_units( + grid, + aligned_ts, + video_meta, + processed_audio=audio_spec, + is_tokenized=False, + audio_token_len=audio_token_len, + ) + audio_features.append(audio_spec) + audio_feature_lens.append(audio_token_len) + seg_lens_flat.extend(u["segment_audio_token_len"] for u in units) + seg_starts_flat.extend(u["audio_start_token_idx"] for u in units) + per_video_num_units.append(len(units)) + else: + per_video_num_units.append(0) + + result = { + "pixel_values_videos": torch.cat(all_patches, dim=0), + "video_grid_thw": torch.stack(all_grids), + "video_timestamps": all_timestamps, + } + if audio_features: + result["video_audio_features"] = audio_features + result["video_audio_feature_lens"] = torch.tensor( + audio_feature_lens, dtype=torch.long + ) + result["video_audio_segment_lens_flat"] = seg_lens_flat + result["video_audio_segment_starts_flat"] = seg_starts_flat + result["video_audio_per_video_num_units"] = per_video_num_units + return result + + if modality == Modality.AUDIO: + all_specs, all_lens = [], [] + for audio in mm_data: + if isinstance(audio, np.ndarray): + audio = (torch.from_numpy(audio).float(), self.audio_sampling_rate) + spec, token_len = self.preprocess_audio(audio) + all_specs.append(spec) + all_lens.append(token_len) + return { + "input_features": all_specs, + "audio_feature_lens_raw": torch.tensor(all_lens, dtype=torch.long), + } + + raise ValueError(f"Unsupported modality for EPD preprocessing: {modality}") + @property def mel_spectrogram(self): self._ensure_audio_dependencies() @@ -928,39 +1172,31 @@ class MiMoProcessor: "verbose": verbose_str, } - def _process_video_audio_content( - self, content_idx, content, video_results, verbose + def _build_video_audio_units( + self, + thw_grid, + timestamps, + video_meta, + processed_audio, + is_tokenized, + audio_token_len, ): - visual_patches, thw_grid, timestamps, video_meta = video_results[content_idx] + # Compute per-grid_t audio-segment boundaries. Tokenizer-free so it + # runs identically on the single-node path and the EPD encoder side. grid_t, grid_h, grid_w = thw_grid - - processed_audio = self.process_audio(content.content) - audio_token_per_second = self.audio_input_id_per_second / self.audio_group_size - + assert ( + len(timestamps) == grid_t * self.temporal_patch_size + ), f"Expected {grid_t} * {self.temporal_patch_size} timestamps, got {len(timestamps)}" if not self.use_video_timestamps: raise NotImplementedError - if isinstance(processed_audio, tuple): - assert ( - content.content.start_time is None and content.content.end_time is None - ), "Audio start_time and end_time must be None when audio is not tokenized" - is_tokenized = False - audio_spec, audio_token_len = processed_audio - audio_input = audio_spec - else: - is_tokenized = True - audio_token_len = processed_audio.shape[0] - audio_input = None - - # Build video-audio units num_media_tokens_per_grid = grid_h * grid_w // (self.merge_size**2) grid_t_timestamps = timestamps[ :: self.temporal_patch_size * self.temporal_compression_ratio ] - text_timestamps = [self.format_timestamp(ts) for ts in grid_t_timestamps] - text_timestamp_ids = [self.tokenizer.encode(ts) for ts in text_timestamps] + audio_token_per_second = self.audio_input_id_per_second / self.audio_group_size - video_audio_units = [] + units = [] for i in range(len(grid_t_timestamps)): audio_start_token_idx = int(grid_t_timestamps[i] * audio_token_per_second) audio_end_token_idx = ( @@ -980,80 +1216,102 @@ class MiMoProcessor: if is_tokenized else None ) - video_audio_units.append( - ( - grid_t_timestamps[i], - text_timestamps[i], - text_timestamp_ids[i], - num_media_tokens_per_grid, - segment_audio_token_len, - segment_audio, - ) + units.append( + { + "timestamp": grid_t_timestamps[i], + "num_video_tokens": num_media_tokens_per_grid, + "segment_audio_token_len": segment_audio_token_len, + "segment_audio": segment_audio, + # Used by encode_server to trim audio_encoder output. + "audio_start_token_idx": audio_start_token_idx, + } ) + return units + + def _build_video_audio_input_ids( + self, + units, + thw_grid, + video_meta, + is_tokenized, + audio_token_len, + verbose=False, + timestamps=None, + ): + # Assemble video+audio input_ids from the unit list produced above. + # Tokenizer-dependent; the language node replays this on EPD. + text_timestamps = [self.format_timestamp(u["timestamp"]) for u in units] + text_timestamp_ids = [self.tokenizer.encode(ts) for ts in text_timestamps] - # Group units by interleave length if self.video_audio_interleave_length == -1: - groups = [list(enumerate(video_audio_units))] + groups = [list(enumerate(units))] elif self.video_audio_interleave_length == 0: - groups = [[(i, u)] for i, u in enumerate(video_audio_units)] + groups = [[(i, u)] for i, u in enumerate(units)] else: assert self.video_audio_interleave_length > 0 groups = [] unit_idx = 0 current_group = [] time_ptr = 0 - while unit_idx < len(video_audio_units): + while unit_idx < len(units): while ( - unit_idx < len(video_audio_units) - and video_audio_units[unit_idx][0] >= time_ptr - and video_audio_units[unit_idx][0] + unit_idx < len(units) + and units[unit_idx]["timestamp"] >= time_ptr + and units[unit_idx]["timestamp"] < time_ptr + self.video_audio_interleave_length ): - current_group.append((unit_idx, video_audio_units[unit_idx])) + current_group.append((unit_idx, units[unit_idx])) unit_idx += 1 if current_group: groups.append(current_group) current_group = [] time_ptr += self.video_audio_interleave_length - # Build input_ids and collect audio segments _input_ids = [self.video_start_token_id] audio_segments = [] verbose_str = "" if verbose: - verbose_str = f"VideoAudio (video_thw_grid={thw_grid}, video_meta={video_meta}, is_audio_tokenized={is_tokenized}, audio_token_len={audio_token_len}): [ " + verbose_str = ( + f"VideoAudio (video_thw_grid={thw_grid}, video_meta={video_meta}, " + f"is_audio_tokenized={is_tokenized}, audio_token_len={audio_token_len}): " + f"[ " + ) for group in groups: + head_idx = group[0][0] if not self.use_per_grid_t_timestamps: - _input_ids += group[0][1][2] + _input_ids += text_timestamp_ids[head_idx] if verbose: - verbose_str += f"{group[0][1][1]} " + verbose_str += f"{text_timestamps[head_idx]} " _video_tokens, _audio_tokens = [], [] video_verbose_str, audio_verbose_str = "", "" for unit_idx, unit in group: - ( - timestamp, - timestamp_text, - timestamp_ids, - video_token_len, - segment_audio_token_len, - segment_audio, - ) = unit if self.use_per_grid_t_timestamps: - _video_tokens += timestamp_ids - _audio_tokens += timestamp_ids - video_verbose_str += timestamp_text + " " - audio_verbose_str += timestamp_text + " " + _video_tokens += text_timestamp_ids[unit_idx] + _audio_tokens += text_timestamp_ids[unit_idx] + video_verbose_str += text_timestamps[unit_idx] + " " + audio_verbose_str += text_timestamps[unit_idx] + " " _video_tokens += ( [self.vision_start_token_id] - + [self.video_token_id] * video_token_len + + [self.video_token_id] * unit["num_video_tokens"] + [self.vision_end_token_id] ) - video_verbose_str += f"[{','.join([f'{ts:.2f}' for ts in timestamps.tolist()[unit_idx*self.temporal_patch_size*self.temporal_compression_ratio : (unit_idx+1)*self.temporal_patch_size*self.temporal_compression_ratio]])}] {video_token_len}*