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

feat(mimo-v2): add EPD disaggregation support (#24931)

Browse Source
This commit is contained in:
Zhonghua Deng
2026-05-18 10:33:16 +08:00
committed by GitHub
parent 43e133208a
commit 784fe7e99b
7 changed files with 962 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/encode_receiver.py
+64 -18
View File
@@ -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)
python/sglang/srt/disaggregation/encode_server.py
+136 -102
View File
@@ -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)
python/sglang/srt/managers/mm_utils.py
+44 -28
View File
@@ -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
python/sglang/srt/managers/tokenizer_manager.py
+5
View File
@@ -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,
python/sglang/srt/models/mimo_v2.py
+149 -29
View File
@@ -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):
python/sglang/srt/multimodal/processors/mimo_v2.py
+559 -110
View File
@@ -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}): [<video_start> "
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"[<video_start> "
)
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]])}] <vision_start> {video_token_len}*<video> <vision_end> "
_audio_tokens += [self.audio_token_id] * segment_audio_token_len
audio_verbose_str += f"{segment_audio_token_len}*<audio> "
if segment_audio is not None:
audio_segments.append(segment_audio)
if verbose and timestamps is not None:
ts_slice = timestamps.tolist()[
unit_idx
* self.temporal_patch_size
* self.temporal_compression_ratio : (unit_idx + 1)
* self.temporal_patch_size
* self.temporal_compression_ratio
]
video_verbose_str += (
f"[{','.join(f'{ts:.2f}' for ts in ts_slice)}] "
f"<vision_start> {unit['num_video_tokens']}*<video> <vision_end> "
)
_audio_tokens += [self.audio_token_id] * unit["segment_audio_token_len"]
audio_verbose_str += f"{unit['segment_audio_token_len']}*<audio> "
if unit["segment_audio"] is not None:
audio_segments.append(unit["segment_audio"])
_input_ids += (
_video_tokens
@@ -1072,13 +1330,55 @@ class MiMoProcessor:
return {
"input_ids": _input_ids,
"audio_segments": audio_segments,
"verbose": verbose_str,
}
def _process_video_audio_content(
self, content_idx, content, video_results, verbose
):
visual_patches, thw_grid, timestamps, video_meta = video_results[content_idx]
processed_audio = self.process_audio(content.content)
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
units = self._build_video_audio_units(
thw_grid,
timestamps,
video_meta,
processed_audio,
is_tokenized,
audio_token_len,
)
built = self._build_video_audio_input_ids(
units,
thw_grid,
video_meta,
is_tokenized,
audio_token_len,
verbose=verbose,
timestamps=timestamps,
)
return {
"input_ids": built["input_ids"],
"pixel_values": visual_patches,
"thw_grid": thw_grid,
"second_per_grid_t": self.temporal_patch_size / video_meta["fps_sampled"],
"audio_input": audio_input,
"audio_segments": audio_segments,
"audio_segments": built["audio_segments"],
"is_tokenized": is_tokenized,
"verbose": verbose_str,
"verbose": built["verbose"],
}
def process(self, contents: list[Content], verbose: bool = False):
@@ -1561,22 +1861,20 @@ class MiMoV2Processor(BaseMultimodalProcessor):
return self.vision_config.spatial_merge_size
def _preprocess_video_sync(self, vdw, preprocess_kwargs=None):
ele = preprocess_kwargs or {}
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))
)
# Seed with processor_config defaults so E/D agree on fps/min/max.
default_kwargs = {
k: v
for k, v in self.mimo_processor.default_video_processor_kwargs.items()
if v is not None and k in ("fps", "min_frames", "max_frames", "num_frames")
}
ele = {**default_kwargs, **(preprocess_kwargs or {})}
try:
video_tensor = vdw.get_frames_as_tensor(idx)
return _decode_frames_and_timestamps(vdw, ele)
except Exception as e:
logger.error(f"Video decode failed in _preprocess_video_sync: {e}")
raise HTTPException(
status_code=432, detail="Video file is corrupted or cannot be decoded"
)
video_tensor = video_tensor.permute(0, 3, 1, 2).float()
timestamps = torch.as_tensor(idx, dtype=torch.float32) / video_fps
return (video_tensor, timestamps)
def process_mm_data(
self, input_text, images=None, videos=None, audios=None, **kwargs
@@ -1613,7 +1911,7 @@ class MiMoV2Processor(BaseMultimodalProcessor):
if "use_audio" in preprocess_kwargs:
use_audio = preprocess_kwargs["use_audio"]
elif isinstance(raw_video_source, str):
use_audio = self.has_audio_track(raw_video_source)
use_audio = self.mimo_processor.has_audio_track(raw_video_source)
else:
use_audio = False
@@ -1853,15 +2151,17 @@ class MiMoV2Processor(BaseMultimodalProcessor):
preprocess_kwargs = (
getattr(raw_video_item, "preprocess_kwargs", {}) or {}
)
use_audio = self.has_audio_track(raw_video_item.url)
use_audio = self.mimo_processor.has_audio_track(
raw_video_item.url
)
raw_video_item_audio = raw_video_item.url
elif isinstance(raw_video_item, dict):
use_audio = self.has_audio_track(
use_audio = self.mimo_processor.has_audio_track(
raw_video_item.get("url", raw_video_item)
)
raw_video_item_audio = raw_video_item
elif isinstance(raw_video_item, str):
use_audio = self.has_audio_track(raw_video_item)
use_audio = self.mimo_processor.has_audio_track(raw_video_item)
raw_video_item_audio = raw_video_item
video_tuple = self._preprocess_video_sync(
@@ -1976,39 +2276,188 @@ class MiMoV2Processor(BaseMultimodalProcessor):
mrope_position_delta=input_sample.rope_deltas,
)
@staticmethod
def has_audio_track(path_or_data: str) -> bool:
try:
is_base64 = path_or_data.startswith("data:") and ";base64," in path_or_data
cmd = [
"ffprobe",
"-v",
"quiet",
"-print_format",
"json",
"-show_streams",
"-select_streams",
"a",
"pipe:0" if is_base64 else path_or_data,
]
inp = (
base64.b64decode(path_or_data.split(";base64,")[1])
if is_base64
else None
def get_mm_data(self, prompt, embeddings, **kwargs):
# EPD: rebuild input_ids from E-side embeddings + segment metadata;
# video+audio reuses _build_video_audio_input_ids for layout parity.
img_grid_thw = kwargs.get("img_grid_thw")
video_grid_thw = kwargs.get("video_grid_thw")
audio_feature_lens = kwargs.get("audio_feature_lens")
video_timestamps = kwargs.get("video_timestamps")
video_audio_feature_lens = kwargs.get("video_audio_feature_lens")
video_audio_segment_lens_flat = kwargs.get("video_audio_segment_lens_flat")
video_audio_per_video_num_units = kwargs.get("video_audio_per_video_num_units")
video_audio_embedding = kwargs.get("video_audio_embedding")
if not isinstance(prompt, str):
prompt = self._tokenizer.decode(prompt)
mp = self.mimo_processor
text_parts = re.split(self.mm_tokens.get_combined_regex(), prompt)
per_video_timestamps = None
if video_timestamps and video_grid_thw is not None:
per_video_timestamps = []
ts_offset = 0
for grid in video_grid_thw:
n_frames = int(grid[0].item()) // mp.temporal_compression_ratio
per_video_timestamps.append(
video_timestamps[ts_offset : ts_offset + n_frames]
)
ts_offset += n_frames
# Un-flatten per-video audio segmentation; None = video has no audio.
num_videos = len(video_grid_thw) if video_grid_thw is not None else 0
per_video_audio_info = [None] * num_videos
if video_audio_per_video_num_units and video_audio_segment_lens_flat:
off, av_idx = 0, 0
for i, nu in enumerate(video_audio_per_video_num_units):
if nu <= 0:
continue
seg_lens = list(video_audio_segment_lens_flat[off : off + nu])
off += nu
per_video_audio_info[i] = {
"segment_lens": seg_lens,
"audio_token_len": (
int(video_audio_feature_lens[av_idx].item())
if video_audio_feature_lens is not None
else sum(seg_lens)
),
}
av_idx += 1
# Merge video-borne audio into AUDIO bucket for uniform slicing.
if video_audio_embedding is not None:
if Modality.AUDIO in embeddings:
raise NotImplementedError(
"Request mixes standalone audio and video-with-audio; "
"EPD merge path for this combination is not yet implemented."
)
embeddings = dict(embeddings)
embeddings[Modality.AUDIO] = video_audio_embedding
merge_size = self.spatial_merge_size
input_ids = []
img_idx = video_idx = audio_idx = 0
for part in text_parts:
mod = self.mm_tokens.get_modality_of_token(part)
if mod == Modality.IMAGE:
grid = img_grid_thw[img_idx]
n = int(grid.prod().item()) // (merge_size**2)
input_ids += (
[mp.vision_start_token_id]
+ [mp.image_token_id] * n
+ [mp.vision_end_token_id]
)
img_idx += 1
elif mod == Modality.VIDEO:
grid = video_grid_thw[video_idx]
ts = per_video_timestamps[video_idx]
n_per_frame = int(grid[1]) * int(grid[2]) // (merge_size**2)
audio_info = per_video_audio_info[video_idx]
if audio_info is not None:
units = [
{
"timestamp": ts[i] if i < len(ts) else 0.0,
"num_video_tokens": n_per_frame,
"segment_audio_token_len": int(seg_len),
"segment_audio": None,
}
for i, seg_len in enumerate(audio_info["segment_lens"])
]
built = mp._build_video_audio_input_ids(
units,
thw_grid=grid,
video_meta=None,
is_tokenized=False,
audio_token_len=audio_info["audio_token_len"],
)
input_ids += built["input_ids"]
else:
ts_ids_per_frame = [
mp.tokenizer.encode(mp.format_timestamp(t)) for t in ts
]
input_ids += (
[mp.video_start_token_id]
+ sum(
[
ts_ids
+ [mp.vision_start_token_id]
+ [mp.video_token_id] * n_per_frame
+ [mp.vision_end_token_id]
for ts_ids in ts_ids_per_frame
],
[],
)
+ [mp.video_end_token_id]
)
video_idx += 1
elif mod == Modality.AUDIO:
n = int(audio_feature_lens[audio_idx].item())
input_ids += (
[mp.audio_start_token_id]
+ [mp.audio_token_id] * n
+ [mp.audio_end_token_id]
)
audio_idx += 1
elif part:
input_ids += mp.tokenizer.encode(part)
input_ids_tensor = torch.tensor(input_ids)
# Slice precomputed embeddings into per-placeholder items
mm_items = []
consumed = {}
for mod, token_id in [
(Modality.IMAGE, mp.image_token_id),
(Modality.VIDEO, mp.video_token_id),
(Modality.AUDIO, mp.audio_token_id),
]:
if mod not in embeddings:
continue
for offset in self.get_mm_items_offset(input_ids_tensor, token_id):
n = offset[1] - offset[0] + 1
s = consumed.get(mod, 0)
mm_items.append(
MultimodalDataItem(
modality=mod,
offsets=[offset],
precomputed_embeddings=embeddings[mod][s : s + n],
)
)
consumed[mod] = s + n
# Position ids
if mp.rope_type == "mrope":
from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index(
spatial_merge_size=self.spatial_merge_size,
image_token_id=mp.image_token_id,
video_token_id=mp.video_token_id,
vision_start_token_id=mp.vision_start_token_id,
model_type="qwen2_5_vl",
input_ids=input_ids_tensor.unsqueeze(0),
image_grid_thw=img_grid_thw,
video_grid_thw=video_grid_thw,
)
r = subprocess.run(cmd, input=inp, capture_output=True, timeout=30)
if r.returncode != 0:
stderr = r.stderr.decode("utf-8", errors="replace")
raise RuntimeError(f"ffprobe failed for {path_or_data}: {stderr}")
return bool(json.loads(r.stdout).get("streams"))
except subprocess.TimeoutExpired:
logger.error("ffprobe timed out for %s", path_or_data)
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", path_or_data)
raise
mrope_positions = mrope_positions.squeeze(1)
else:
mrope_positions = torch.arange(len(input_ids)).expand(3, -1)
mrope_position_delta = torch.zeros((1, 1), dtype=torch.int32)
return MultimodalProcessorOutput(
mm_items=mm_items,
input_ids=input_ids,
im_start_id=self.IM_START_TOKEN_ID,
im_end_id=self.IM_END_TOKEN_ID,
im_token_id=mp.image_token_id,
video_token_id=mp.video_token_id,
audio_token_id=mp.audio_token_id,
audio_start_id=self.AUDIO_START_TOKEN_ID,
audio_end_id=self.AUDIO_END_TOKEN_ID,
mrope_positions=mrope_positions,
mrope_position_delta=mrope_position_delta,
)
@staticmethod
def _make_video_content(
python/sglang/srt/server_args.py
+5 -3
View File
@@ -2095,7 +2095,7 @@ class ServerArgs:
), "Triton kernel MoE is only supported when ep_size == 1"
elif model_arch in MIMO_V2_MODEL_ARCHS:
if model_arch == "MiMoV2ForCausalLM":
if model_arch == "MiMoV2ForCausalLM" and not self.encoder_only:
expected_attn_tp_size = get_mimo_v2_fused_qkv_expected_tp_size(
hf_config
)
@@ -3721,7 +3721,7 @@ class ServerArgs:
self.disaggregation_ib_device
)
# Validate model type: only support Qwen models for now
# Validate model type for encoder disaggregation
hf_config = self.get_model_config().hf_config
model_arch = hf_config.architectures[0]
if (self.encoder_only or self.language_only) and model_arch not in [
@@ -3737,9 +3737,11 @@ class ServerArgs:
"Qwen2_5OmniForConditionalGeneration",
"KimiVLForConditionalGeneration",
"KimiK25ForConditionalGeneration",
"MiMoV2ForCausalLM",
]:
raise ValueError(
f"Model type {model_arch} is not supported for encoder disaggregation, only Qwen models are supported for now."
f"Model type {model_arch} is not supported for encoder disaggregation. "
f"Supported architectures: Qwen2VL, Qwen3VL, Qwen3.5, InternS2, Qwen2Audio, Qwen2.5Omni, Kimi, MiMoV2."
)
def _validate_ib_devices(self, device_str: str) -> Optional[str]: