[Feature] Xiaomi MiMo-V2.5 day0 support (#23811)

Co-authored-by: 张袁 <zhangyuan36@xiaomi.com>
Co-authored-by: 刘安岐 <liuanqi6@xiaomi.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>

python/sglang/srt/configs/model_config.py

@@ -39,6 +39,36 @@ from sglang.utils import is_in_ci
 
 logger = logging.getLogger(__name__)
 
+MIMO_V2_MODEL_ARCHS = (
+    "MiMoV2ForCausalLM",
+    "MiMoV2FlashForCausalLM",
+)
+MIMO_V2_MULTIMODAL_ARCHS = ("MiMoV2ForCausalLM",)
+
+
+def get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
+    layout = getattr(hf_config, "attention_projection_layout", None)
+    if layout is None:
+        return None
+    if layout != "fused_qkv":
+        raise ValueError(
+            "MiMoV2 hf_config has unsupported "
+            f"attention_projection_layout={layout!r}; expected 'fused_qkv' "
+            "or unset."
+        )
+
+    num_key_value_heads = getattr(hf_config, "num_key_value_heads", None)
+    text_config = getattr(hf_config, "text_config", None)
+    if num_key_value_heads is None and text_config is not None:
+        num_key_value_heads = getattr(text_config, "num_key_value_heads", None)
+    if num_key_value_heads is None:
+        raise ValueError(
+            "MiMoV2 hf_config has attention_projection_layout='fused_qkv' "
+            "but num_key_value_heads is missing; this value is required to "
+            "derive the fused qkv_proj TP size."
+        )
+    return num_key_value_heads
+
 
 class AttentionArch(IntEnum):
     MLA = auto()
@@ -337,10 +367,7 @@ class ModelConfig:
 
         if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM":
             self.hf_config.architectures[0] = "MiMoMTP"
-        if is_draft_model and self.hf_config.architectures[0] in (
+        if is_draft_model and self.hf_config.architectures[0] in MIMO_V2_MODEL_ARCHS:
-            "MiMoV2ForCausalLM",
-            "MiMoV2FlashForCausalLM",
-        ):
             self.hf_config.architectures[0] = "MiMoV2MTP"
         if is_draft_model and self.hf_config.architectures[0] == "Step3p5ForCausalLM":
             self.hf_config.architectures[0] = "Step3p5MTP"
@@ -397,8 +424,7 @@ class ModelConfig:
         self.has_attention_sinks = self._detect_attention_sinks()
 
         self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [
-            "MiMoV2ForCausalLM",
+            *MIMO_V2_MODEL_ARCHS,
-            "MiMoV2FlashForCausalLM",
             "MiMoV2MTP",
             "Gemma4ForCausalLM",
             "Gemma4ForConditionalGeneration",
@@ -417,14 +443,7 @@ class ModelConfig:
             return True
 
         # MiMoV2 creates sinks only when the config flags are set.
-        if any(
+        if any(a in archs for a in (*MIMO_V2_MODEL_ARCHS, "MiMoV2MTP")):
-            a in archs
-            for a in (
-                "MiMoV2FlashForCausalLM",
-                "MiMoV2ForCausalLM",
-                "MiMoV2MTP",
-            )
-        ):
             return getattr(
                 self.hf_text_config, "add_swa_attention_sink_bias", False
             ) or getattr(self.hf_text_config, "add_full_attention_sink_bias", False)
@@ -1383,6 +1402,7 @@ multimodal_model_archs = [
     "LlavaVidForCausalLM",
     "Lfm2VlForConditionalGeneration",
     "LightOnOCRForConditionalGeneration",
+    *MIMO_V2_MULTIMODAL_ARCHS,
     "MiniCPMO",
     "MiniCPMV",
     "Mistral3ForConditionalGeneration",
@@ -1528,8 +1548,7 @@ def is_hybrid_swa_model(model_architectures: List[str]):
     hybrid_swa_archs = {
         "Llama4ForConditionalGeneration",
         "GptOssForCausalLM",
-        "MiMoV2ForCausalLM",
+        *MIMO_V2_MODEL_ARCHS,
-        "MiMoV2FlashForCausalLM",
         "MiMoV2MTP",
         "Step3p5ForCausalLM",
         "Step3p5MTP",
@@ -1559,13 +1578,7 @@ def get_hybrid_layer_ids(
         full_attention_layer_ids = [
             i for i, x in enumerate(layer_types) if x == "full_attention"
         ]
-    elif any(
+    elif any(arch in MIMO_V2_MODEL_ARCHS for arch in model_architectures):
-        x in model_architectures
-        for x in (
-            "MiMoV2ForCausalLM",
-            "MiMoV2FlashForCausalLM",
-        )
-    ):
         hybrid_layer_pattern = getattr(hf_text_config, "hybrid_layer_pattern", None)
         swa_attention_layer_ids = [
             i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1

python/sglang/srt/layers/attention/vision.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import dataclasses
 import functools
 import math
+import warnings
 from functools import lru_cache, partial
 from typing import Any, Callable, Optional, Tuple
 
@@ -406,34 +407,39 @@ class VisionFlash3Attention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
+        window_size = kwargs.get("window_size", (-1, -1))
+        s_aux = kwargs.get("s_aux", None)
+
         if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
             max_seqlen = cu_seqlens[1]
-            output = flash_attn_varlen_func(
+            fa_kwargs = dict(
-                q,
-                k,
-                v,
                 cu_seqlens_q=cu_seqlens[0],
                 cu_seqlens_k=cu_seqlens[0],
                 max_seqlen_q=max_seqlen,
                 max_seqlen_k=max_seqlen,
                 softmax_scale=softmax_scale,
+                window_size=window_size,
             )
+            if s_aux is not None:
+                fa_kwargs["sinks"] = s_aux
+            output = flash_attn_varlen_func(q, k, v, **fa_kwargs)
         else:
             cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
             cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
             seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
             max_seqlen = seq_lens.max().item()
 
-            output = flash_attn_varlen_func(
+            fa_kwargs = dict(
-                q,
-                k,
-                v,
                 cu_seqlens_q=cu_seqlens,
                 cu_seqlens_k=cu_seqlens,
                 max_seqlen_q=max_seqlen,
                 max_seqlen_k=max_seqlen,
                 softmax_scale=softmax_scale,
+                window_size=window_size,
             )
+            if s_aux is not None:
+                fa_kwargs["sinks"] = s_aux
+            output = flash_attn_varlen_func(q, k, v, **fa_kwargs)
 
         return output
 
@@ -749,7 +755,8 @@ class VisionAttention(nn.Module):
         num_heads: int,
         projection_size: int,
         use_qkv_parallel: bool,
-        head_size: Optional[int] = None,
+        num_kv_heads: Optional[int] = None,
+        head_dim: Optional[int] = None,
         qkv_backend: Optional[str] = None,
         quant_config: Optional[QuantizationConfig] = None,
         dropout: float = 0.0,
@@ -770,13 +777,23 @@ class VisionAttention(nn.Module):
         use_dp_attention_reduce: bool = False,
         aux_stream: Optional[torch.cuda.Stream] = None,
         workspace_buffer: Optional[torch.Tensor] = None,
+        use_sink: bool = False,
+        window_size: Tuple[int, int] = (-1, -1),
         **kwargs,
     ):
         super().__init__()
+        if head_dim is None and "head_size" in kwargs:
+            head_dim = kwargs.pop("head_size")
+            warnings.warn(
+                "VisionAttention(head_size=...) is deprecated; use head_dim=...",
+                DeprecationWarning,
+                stacklevel=2,
+            )
         self.tp_size = 1 if use_data_parallel else get_attention_tp_size()
         self.tp_rank = 0 if use_data_parallel else get_attention_tp_rank()
         self.dropout = dropout
-        self.head_size = head_size if head_size is not None else embed_dim // num_heads
+        num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.head_size = head_dim if head_dim is not None else embed_dim // num_heads
         self.hidden_size_per_attention_head = dist_utils.divide(
             projection_size, num_heads
         )
@@ -784,7 +801,7 @@ class VisionAttention(nn.Module):
             num_dummy_heads + num_heads, self.tp_size
         )
         self.num_attention_kv_heads_per_partition = dist_utils.divide(
-            num_dummy_heads + num_heads, self.tp_size
+            num_dummy_heads + num_kv_heads, self.tp_size
         )
 
         self.q_size = self.num_attention_heads_per_partition * self.head_size
@@ -838,7 +855,7 @@ class VisionAttention(nn.Module):
                 hidden_size=embed_dim,
                 head_size=self.head_size,
                 total_num_heads=num_dummy_heads + num_heads,
-                total_num_kv_heads=num_dummy_heads + num_heads,
+                total_num_kv_heads=num_dummy_heads + num_kv_heads,
                 bias=qkv_bias,
                 quant_config=quant_config,
                 tp_rank=self.tp_rank,
@@ -870,6 +887,20 @@ class VisionAttention(nn.Module):
         self.aux_stream = aux_stream
         self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] if aux_stream else []
 
+        self.window_size = window_size
+        if use_sink:
+            # Allocate the full (unsharded) sink tensor for weight loading;
+            # only the local TP slice is used in forward.
+            self.sinks = nn.Parameter(
+                torch.empty(
+                    self.num_attention_heads_per_partition * self.tp_size,
+                    dtype=torch.bfloat16,
+                ),
+                requires_grad=False,
+            )
+        else:
+            self.sinks = None
+
     def _init_qk_norm(
         self, norm_dim: int, eps: float, var_hidden_size: Optional[int] = None
     ):
@@ -989,6 +1020,7 @@ class VisionAttention(nn.Module):
         rotary_pos_emb_cos: Optional[torch.Tensor] = None,
         rotary_pos_emb_sin: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
+        full_attn: bool = True,
         **kwargs,
     ) -> torch.Tensor:
         r"""
@@ -1070,19 +1102,20 @@ class VisionAttention(nn.Module):
             sin = rotary_pos_emb_sin
 
         if cos is not None and sin is not None:
-            original_shape = q.shape
+            original_q_shape = q.shape
+            original_k_shape = k.shape
 
-            # [total_tokens, head, head_size]
+            # [total_tokens, head, head_size] for q / [total_tokens, kv_head, head_size] for k
             q = q.view(-1, head, self.head_size)
-            k = k.view(-1, head, self.head_size)
+            k = k.view(-1, kv_head, self.head_size)
 
             if cos.size(-1) * 2 == self.head_size:
                 cos = torch.cat([cos, cos], dim=-1)
                 sin = torch.cat([sin, sin], dim=-1)
 
             q, k = apply_rotary_pos_emb(q, k, cos, sin)
-            q = q.view(original_shape)
+            q = q.view(original_q_shape)
-            k = k.view(original_shape)
+            k = k.view(original_k_shape)
 
         if q.dim() == 4:
             # [b, s, head, head_size] --> [b * s, head, head_size]
@@ -1118,6 +1151,15 @@ class VisionAttention(nn.Module):
             else:
                 q, k = self._apply_qk_norm(q, k)
 
+        if full_attn or self.sinks is None:
+            effective_window_size = (-1, -1)
+            s_aux = None
+        else:
+            effective_window_size = self.window_size
+            q_head_start = self.tp_rank * self.num_attention_heads_per_partition
+            q_head_end = (self.tp_rank + 1) * self.num_attention_heads_per_partition
+            s_aux = self.sinks[q_head_start:q_head_end]
+
         output = self.qkv_backend.forward(
             q=q,
             k=k,
@@ -1130,6 +1172,8 @@ class VisionAttention(nn.Module):
             max_seqlen=max_seqlen,
             output_ws=attn_output_ws,
             softmax_scale=self.softmax_scale,
+            window_size=effective_window_size,
+            s_aux=s_aux,
         )
 
         assert output.dim() == 3, output.shape

python/sglang/srt/layers/linear.py

@@ -1047,7 +1047,7 @@ class QKVParallelLinear(ColumnParallelLinear):
         block_n, _ = self.quant_method.quant_config.weight_block_size
         q_size = self.total_num_heads * self.head_size // block_n
         k_size = self.total_num_kv_heads * self.head_size // block_n
-        v_size = self.total_num_kv_heads * self.head_size // block_n
+        v_size = self.total_num_kv_heads * self.v_head_size // block_n
         shard_offsets = [
             # (shard_id, shard_offset, shard_size)
             ("q", 0, q_size),

python/sglang/srt/managers/io_struct.py

@@ -38,7 +38,7 @@ from sglang.srt.observability.req_time_stats import (
     SchedulerReqTimeStats,
 )
 from sglang.srt.sampling.sampling_params import SamplingParams
-from sglang.srt.utils import ImageData
+from sglang.srt.utils import ImageData, VideoData
 
 # Handle serialization of Image for pydantic
 if TYPE_CHECKING:
@@ -118,7 +118,7 @@ class SessionParams:
 # Individual data item types for each modality
 ImageDataInputItem = Union[Image, str, ImageData, Dict]
 AudioDataInputItem = Union[str, Dict]
-VideoDataInputItem = Union[str, Dict]
+VideoDataInputItem = Union[str, VideoData, Dict]
 # Union type for any multimodal data item
 MultimodalDataInputItem = Union[
     ImageDataInputItem, VideoDataInputItem, AudioDataInputItem

python/sglang/srt/managers/mm_utils.py

@@ -1608,15 +1608,35 @@ def wrap_shm_features(obj):
 
     if hasattr(obj, "mm_inputs") and obj.mm_inputs:
         for item in obj.mm_inputs.mm_items:
-            if (
+            if not hasattr(item, "feature"):
-                hasattr(item, "feature")
+                continue
-                and isinstance(item.feature, torch.Tensor)
+            feat = item.feature
-                and item.feature.is_cpu
+            if isinstance(feat, torch.Tensor) and feat.is_cpu:
-            ):
+                item.feature = ShmPointerMMData(feat)
-                item.feature = ShmPointerMMData(item.feature)
+            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
+                )
     return obj
 
 
+def _feature_has_shm(feat) -> bool:
+    """Check whether a single feature (tensor, ShmPointer, or list) contains ShmPointerMMData."""
+    if isinstance(feat, ShmPointerMMData):
+        return True
+    if isinstance(feat, (list, tuple)):
+        return any(isinstance(t, ShmPointerMMData) for t in feat)
+    return False
+
+
 def has_shm_features(recv_reqs):
     """Return True if any request in the list contains ShmPointerMMData."""
     for req in recv_reqs:
@@ -1625,7 +1645,7 @@ def has_shm_features(recv_reqs):
                 return True
         elif hasattr(req, "mm_inputs") and req.mm_inputs:
             for item in req.mm_inputs.mm_items:
-                if isinstance(item.feature, ShmPointerMMData):
+                if _feature_has_shm(item.feature):
                     return True
     return False
 
@@ -1646,6 +1666,15 @@ def unwrap_shm_features(obj):
     if hasattr(obj, "mm_inputs") and obj.mm_inputs:
         mm_items = obj.mm_inputs.mm_items
         for item in mm_items:
-            if isinstance(item.feature, ShmPointerMMData):
+            feat = item.feature
-                item.feature = item.feature.materialize()
+            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
+                )
     return obj

python/sglang/srt/models/mimo_v2.py

@@ -13,13 +13,14 @@
 # ==============================================================================
 
 import logging
-from typing import Any, Dict, Iterable, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
 from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
+from sglang.srt.configs.model_config import get_mimo_v2_fused_qkv_expected_tp_size
 from sglang.srt.distributed import (
     get_moe_expert_parallel_world_size,
     get_pp_group,
@@ -63,11 +64,18 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from sglang.srt.managers.mm_utils import (
+    MultiModalityDataPaddingPatternMultimodalTokens,
+    general_mm_embed_routine,
+)
+from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     kv_cache_scales_loader,
 )
+from sglang.srt.models.mimo_audio import MiMoAudioEncoder, MiMoAudioEncoderConfig
+from sglang.srt.models.mimo_vl import MiMoVisionTransformer, MiMoVLVisionConfig
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     LazyValue,
@@ -81,6 +89,38 @@ MiMoV2Config = None
 logger = logging.getLogger(__name__)
 
 
+def load_mimo_v2_qkv_proj_weight(
+    name, param, loaded_weight, expected_fused_tp_size: Optional[int] = None
+):
+    if loaded_weight.shape == param.shape:
+        # The checkpoint already stores this rank's qkv_proj shard.
+        default_weight_loader(param, loaded_weight)
+        return
+
+    if loaded_weight.ndim != param.ndim or loaded_weight.shape[1:] != param.shape[1:]:
+        raise ValueError(
+            f"qkv_proj weight {name}: unexpected shape {tuple(loaded_weight.shape)}; "
+            f"expected sharded {tuple(param.shape)}"
+        )
+
+    tp_size = get_attention_tp_size()
+    tp_rank = get_attention_tp_rank()
+    if expected_fused_tp_size is not None and tp_size != expected_fused_tp_size:
+        raise ValueError(
+            f"MiMoV2 fused qkv_proj checkpoint is TP={expected_fused_tp_size}-"
+            f"interleaved; got attention tp_size={tp_size} while loading {name}."
+        )
+
+    fused_shape = (param.shape[0] * tp_size, *param.shape[1:])
+    if tuple(loaded_weight.shape) != fused_shape:
+        raise ValueError(
+            f"qkv_proj weight {name}: unexpected shape {tuple(loaded_weight.shape)}; "
+            f"expected fused {fused_shape} or sharded {tuple(param.shape)}"
+        )
+
+    default_weight_loader(param, loaded_weight.chunk(tp_size, dim=0)[tp_rank])
+
+
 class MiMoV2MLP(nn.Module):
     def __init__(
         self,
@@ -995,6 +1035,24 @@ class MiMoV2ForCausalLM(nn.Module):
 
         self.logits_processor = LogitsProcessor(config)
 
+        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
+        if self._is_multimodal:
+            if hasattr(vision_config, "to_dict"):
+                vision_config = vision_config.to_dict()
+            if hasattr(audio_config, "to_dict"):
+                audio_config = audio_config.to_dict()
+
+            self.visual = MiMoVisionTransformer(
+                MiMoVLVisionConfig.from_dict(vision_config),
+                norm_eps=getattr(config, "rms_norm_eps", 1e-6),
+                quant_config=None,
+                prefix=add_prefix("visual", prefix),
+            )
+            self.audio_config = MiMoAudioEncoderConfig(**audio_config)
+            self.audio_encoder = MiMoAudioEncoder(self.audio_config)
+
         self._routed_experts_weights_of_layer = LazyValue(
             lambda: {
                 layer_id: layer.mlp.get_moe_weights()
@@ -1010,6 +1068,31 @@ class MiMoV2ForCausalLM(nn.Module):
     def get_input_embedding(self, input_ids: torch.Tensor) -> torch.Tensor:
         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 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
+        )
+        image_grid_thw = torch.cat([item.image_grid_thw for item in items], dim=0)
+        assert pixel_values.dim() == 2, pixel_values.dim()
+        assert image_grid_thw.dim() == 2, image_grid_thw.dim()
+        return self.visual(pixel_values, grid_thw=image_grid_thw)
+
+    def get_video_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
+        pixel_values = torch.cat([item.feature for item in items], dim=0).type(
+            self.visual.dtype
+        )
+        video_grid_thw = torch.cat([item.video_grid_thw for item in items], dim=0)
+        assert pixel_values.dim() == 2, pixel_values.dim()
+        assert video_grid_thw.dim() == 2, video_grid_thw.dim()
+        return self.visual(pixel_values, grid_thw=video_grid_thw)
+
+    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
 
@@ -1022,13 +1105,23 @@ class MiMoV2ForCausalLM(nn.Module):
         input_embeds: torch.Tensor = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> torch.Tensor:
-        hidden_states, hidden_states_before_norm = self.model(
+        if self._is_multimodal:
-            input_ids,
+            hidden_states, hidden_states_before_norm = general_mm_embed_routine(
-            positions,
+                input_ids=input_ids,
-            forward_batch,
+                forward_batch=forward_batch,
-            input_embeds,
+                language_model=self.model,
-            pp_proxy_tensors=pp_proxy_tensors,
+                multimodal_model=self,
-        )
+                positions=positions,
+                pp_proxy_tensors=pp_proxy_tensors,
+            )
+        else:
+            hidden_states, hidden_states_before_norm = self.model(
+                input_ids,
+                positions,
+                forward_batch,
+                input_embeds,
+                pp_proxy_tensors=pp_proxy_tensors,
+            )
 
         if self.pp_group.is_last_rank:
             return self.logits_processor(
@@ -1058,6 +1151,11 @@ class MiMoV2ForCausalLM(nn.Module):
             ("gate_up_proj", "gate_proj", 0),
             ("gate_up_proj", "up_proj", 1),
         ]
+        stacked_params_mapping_vit = [
+            # (param_name, shard_name, shard_id)
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
 
         # (param_name, weight_name, expert_id, shard_id)
         expert_params_mapping = DeepEPMoE.make_expert_params_mapping(
@@ -1068,8 +1166,105 @@ class MiMoV2ForCausalLM(nn.Module):
         )
 
         params_dict = dict(self.named_parameters())
+        skipped_mtp_weights = False
 
         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
+            ):
+                continue
+
+            if self._is_multimodal and "audio" in name:
+                if "projection" in name:
+                    if (
+                        "audio_encoder.audio_projection" in name
+                        and "audio_encoder.projection" not in name
+                    ):
+                        name = name.replace(
+                            "audio_encoder.audio_projection", "audio_encoder.projection"
+                        )
+                    elif (
+                        "audio_projection" in name
+                        and "audio_encoder.projection" not in name
+                    ):
+                        name = name.replace(
+                            "audio_projection", "audio_encoder.projection"
+                        )
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    continue
+
+                if "input_local_transformer" in name:
+                    if (
+                        "audio_input_local_transformer" in name
+                        and "audio_encoder.input_local_transformer" not in name
+                    ):
+                        name = name.replace(
+                            "audio_input_local_transformer",
+                            "audio_encoder.input_local_transformer",
+                        )
+                    if name not in params_dict:
+                        logger.warning(
+                            f"Parameter {name} not found in params_dict, skipping"
+                        )
+                        continue
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+                    continue
+
+            if self._is_multimodal and "speech_embeddings" in name:
+                if (
+                    "speech_embeddings" in name
+                    and "audio_encoder.speech_embeddings" not in name
+                ):
+                    name = name.replace(
+                        "speech_embeddings", "audio_encoder.speech_embeddings"
+                    )
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight[: param.shape[0], :])
+                continue
+
+            if self._is_multimodal and "visual" in name:
+                name = name.replace("vision_model.", "")
+                name = name.replace(r"attn.qkv.", r"attn.qkv_proj.")
+                match_stacked_vit = False
+                for param_name, weight_name, shard_id in stacked_params_mapping_vit:
+                    if weight_name not in name:
+                        continue
+                    name = name.replace(weight_name, param_name)
+                    # Skip loading extra bias for GPTQ models.
+                    if name.endswith(".bias") and name not in params_dict:
+                        match_stacked_vit = True
+                        continue
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(param, loaded_weight, shard_id)
+                    match_stacked_vit = True
+                    break
+                if match_stacked_vit:
+                    continue
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+                if name.endswith("patch_embed.proj.weight"):
+                    patch_embed = self.get_submodule(name.rsplit(".", 2)[0])
+                    if hasattr(patch_embed, "sync_proj_weight_linear_format"):
+                        patch_embed.sync_proj_weight_linear_format()
+                continue
+
             layer_id = get_layer_id(name)
             if (
                 layer_id is not None
@@ -1099,21 +1294,35 @@ class MiMoV2ForCausalLM(nn.Module):
                 else:
                     continue
 
-            # TODO: skip mtp weights for now, need to implement mtp
             if "mtp" in name:
+                if not skipped_mtp_weights:
+                    logger.info(
+                        "Skipping draft-only MiMo-V2 MTP weights while loading the "
+                        "target model; MiMoV2MTP loads these weights in the draft "
+                        "model runner."
+                    )
+                    skipped_mtp_weights = True
                 continue
 
             # Support fused qkv_proj checkpoint (Pro format)
             if "qkv_proj" in name:
                 if name in params_dict:
-                    tp_size = get_attention_tp_size()
-                    tp_rank = get_attention_tp_rank()
                     param = params_dict[name]
-                    loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
+                    expected_fused_tp_size = get_mimo_v2_fused_qkv_expected_tp_size(
-                    default_weight_loader(param, loaded_weight)
+                        self.config
+                    )
+                    load_mimo_v2_qkv_proj_weight(
+                        name, param, loaded_weight, expected_fused_tp_size
+                    )
                 continue
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
+                if (
+                    "compression_attention" in name
+                    or "hybrid_softmax_attention" in name
+                    or "compressed_softmax_attn" in name
+                ):
+                    continue
                 if weight_name not in name:
                     continue
                 if ("mlp.experts." in name) and name not in params_dict:

python/sglang/srt/models/mimo_v2_nextn.py

@@ -19,6 +19,7 @@ import torch
 from torch import nn
 from transformers import PretrainedConfig
 
+from sglang.srt.configs.model_config import get_mimo_v2_fused_qkv_expected_tp_size
 from sglang.srt.distributed import get_tensor_model_parallel_world_size
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.communicator import (
@@ -28,7 +29,6 @@ from sglang.srt.layers.communicator import (
 )
 from sglang.srt.layers.dp_attention import (
     get_attention_tp_rank,
-    get_attention_tp_size,
     is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
@@ -44,6 +44,7 @@ from sglang.srt.models.mimo_v2 import (
     MiMoV2Attention,
     MiMoV2ForCausalLM,
     MiMoV2MLP,
+    load_mimo_v2_qkv_proj_weight,
 )
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix
@@ -304,20 +305,24 @@ class MiMoV2MTP(MiMoV2ForCausalLM):
             # Support fused qkv_proj checkpoint (Pro format)
             if "qkv_proj" in name:
                 if name in params_dict:
-                    tp_size = get_attention_tp_size()
-                    tp_rank = get_attention_tp_rank()
                     param = params_dict[name]
-                    loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
+                    load_mimo_v2_qkv_proj_weight(
-                    default_weight_loader(param, loaded_weight)
+                        name,
+                        param,
+                        loaded_weight,
+                        expected_fused_tp_size=get_mimo_v2_fused_qkv_expected_tp_size(
+                            self.config
+                        ),
+                    )
                 continue
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
 
-                if weight_name not in name:
+                if f".{weight_name}." not in name:
                     continue
                 if "mtp_block" not in name:
                     break
-                name = name.replace(weight_name, param_name)
+                name = name.replace(f".{weight_name}.", f".{param_name}.")
                 # Skip loading extra bias for GPTQ models.
                 if name.endswith(".bias") and name not in params_dict:
                     continue

python/sglang/srt/models/mimo_vl.py

@@ -0,0 +1,507 @@
+"""Inference-only MiMo vision model: attention + ViT."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Optional, Tuple, Type
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VisionRotaryEmbedding,
+)
+
+from sglang.srt.layers.attention.vision import VisionAttention
+from sglang.srt.layers.layernorm import RMSNorm
+from sglang.srt.layers.quantization import QuantizationConfig
+from sglang.srt.models.qwen2_5_vl import Qwen2_5_VisionPatchMerger, Qwen2_5_VLMLP
+from sglang.srt.server_args import get_global_server_args
+from sglang.srt.utils import add_prefix
+
+
+class MiMoVLVisionConfig(PretrainedConfig):
+    model_type = "mimovl"
+    base_config_key = "vision_config"
+
+    def __init__(
+        self,
+        depth=28,
+        hidden_size=1280,
+        hidden_act="silu",
+        intermediate_size=4608,
+        num_heads=32,
+        in_channels=3,
+        patch_size=16,
+        spatial_merge_size=2,
+        temporal_patch_size=2,
+        tokens_per_second=2,
+        window_size=128,
+        out_hidden_size=2048,
+        fullatt_block_indexes=[7, 15, 23, 31],
+        initializer_range=0.02,
+        kv_channels=64,
+        qk_channels=64,
+        num_query_groups=4,
+        num_key_value_heads=8,
+        vit_window_attn_types=None,
+        visual_token_window_size=64,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.depth = depth
+        self.hidden_size = hidden_size
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        if num_key_value_heads is None:
+            num_key_value_heads = num_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spatial_merge_size = spatial_merge_size
+        self.temporal_patch_size = temporal_patch_size
+        self.tokens_per_second = tokens_per_second
+        self.window_size = window_size
+        self.fullatt_block_indexes = fullatt_block_indexes
+        self.out_hidden_size = out_hidden_size
+        self.initializer_range = initializer_range
+        self.kv_channels = kv_channels
+        self.qk_channels = qk_channels
+        self.num_query_groups = num_query_groups
+        self.vit_window_attn_types = vit_window_attn_types or [-1] * depth
+        self.visual_token_window_size = visual_token_window_size
+
+
+class MiMoVisionPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 16,
+        temporal_patch_size: int = 2,
+        in_channels: int = 3,
+        embed_dim: int = 1536,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.in_channels = in_channels
+        self.embed_dim = embed_dim
+
+        kernel_size = [temporal_patch_size, patch_size, patch_size]
+        self.proj = nn.Conv3d(
+            in_channels,
+            embed_dim,
+            kernel_size=kernel_size,
+            stride=kernel_size,
+            bias=False,
+        )
+        self.proj_weight_linear_format = None
+
+    @torch.no_grad()
+    def sync_proj_weight_linear_format(self):
+        self.proj_weight_linear_format = self.proj.weight.view(self.embed_dim, -1)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        target_dtype = self.proj.weight.dtype
+        hidden_states = F.linear(
+            hidden_states.to(dtype=target_dtype), self.proj_weight_linear_format
+        )
+        return hidden_states
+
+
+class MiMoVisionBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        num_heads: int,
+        hidden_act="silu",
+        norm_layer: Type[nn.Module] = None,
+        attn_implementation: Optional[str] = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+        num_dummy_heads: int = 0,
+        rms_norm_eps: float = 1e-6,
+        use_sink: bool = False,
+        window_size: Tuple[int, int] = (-1, -1),
+        num_kv_heads: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        use_data_parallel: bool = False,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        self.norm1 = RMSNorm(dim, eps=rms_norm_eps)
+        self.norm2 = RMSNorm(dim, eps=rms_norm_eps)
+        self.use_data_parallel = use_data_parallel
+
+        if attn_implementation is None:
+            softmax_in_single_precision = False
+            qkv_backend = None
+            flatten_batch = True
+        elif attn_implementation == "sdpa":
+            softmax_in_single_precision = False
+            qkv_backend = "sdpa"
+            flatten_batch = True
+        elif attn_implementation == "flash_attention_2":
+            softmax_in_single_precision = False
+            qkv_backend = "triton_attn"
+            flatten_batch = True
+        elif attn_implementation == "eager":
+            softmax_in_single_precision = True
+            qkv_backend = "sdpa"
+            flatten_batch = True
+        elif attn_implementation == "flash_attention_3":
+            softmax_in_single_precision = False
+            qkv_backend = "fa3"
+            flatten_batch = True
+
+        self.attn = VisionAttention(
+            embed_dim=dim,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            projection_size=dim,
+            use_qkv_parallel=True,
+            proj_bias=True,
+            qkv_bias=True,
+            qkv_backend=qkv_backend,
+            softmax_in_single_precision=softmax_in_single_precision,
+            flatten_batch=flatten_batch,
+            quant_config=quant_config,
+            prefix=add_prefix("attn", prefix),
+            num_dummy_heads=num_dummy_heads,
+            use_sink=use_sink,
+            window_size=window_size,
+            use_data_parallel=use_data_parallel,
+        )
+        self.mlp = Qwen2_5_VLMLP(
+            dim,
+            intermediate_dim,
+            hidden_act=hidden_act,
+            quant_config=quant_config,
+            prefix=add_prefix("mlp", prefix),
+            use_data_parallel=use_data_parallel,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+        position_embeddings: torch.Tensor,
+        full_attn: bool = True,
+    ) -> torch.Tensor:
+        S, B, H = x.shape
+        # norm1: flatten to 2D -> [S*B, H], then reshape back
+        x2d = x.reshape(-1, H)
+        hidden_states = self.norm1(x2d).reshape(S, B, H)
+
+        # Attention expects [B, S, H]
+        hidden_states = rearrange(hidden_states, "s b h -> b s h")
+        attn = self.attn(
+            hidden_states,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            position_embeddings=position_embeddings,
+            full_attn=full_attn,
+        )
+        attn = rearrange(attn, "b s h -> s b h")
+
+        # norm2 with fused residual-add: also 2D
+        attn2d = attn.reshape(-1, H)
+        x_norm_2d, x_after_add_2d = self.norm2(x2d, residual=attn2d)
+        x_norm = x_norm_2d.reshape(S, B, H)
+        x_after_add = x_after_add_2d.reshape(S, B, H)
+
+        # MLP and final residual
+        mlp_out = self.mlp(x_norm)
+        x = x_after_add + mlp_out
+        return x
+
+
+class MiMoVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        vision_config: MiMoVLVisionConfig,
+        norm_eps: float = 1e-6,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.server_args = get_global_server_args()
+        self.vit_window_attn_types = vision_config.vit_window_attn_types
+        patch_size: int = vision_config.patch_size
+        temporal_patch_size: int = vision_config.temporal_patch_size
+        spatial_merge_size: int = vision_config.spatial_merge_size
+        self.spatial_merge_size = spatial_merge_size
+        self.spatial_merge_unit: int = spatial_merge_size * spatial_merge_size
+        in_channels: int = vision_config.in_channels
+        hidden_size: int = vision_config.hidden_size
+        depth: int = vision_config.depth
+        num_heads: int = vision_config.num_heads
+        num_kv_heads = getattr(vision_config, "num_key_value_heads", None)
+        if num_kv_heads is None:
+            num_kv_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.qk_channels = getattr(vision_config, "qk_channels", None)
+        self.kv_channels = getattr(vision_config, "kv_channels", None)
+        self.fullatt_block_indexes = vision_config.fullatt_block_indexes
+        self.window_size = vision_config.window_size
+        self.patch_size = vision_config.patch_size
+        self.use_data_parallel = self.server_args.mm_enable_dp_encoder
+        mlp_hidden_size: int = vision_config.intermediate_size
+        self.patch_embed = MiMoVisionPatchEmbed(
+            patch_size=patch_size,
+            temporal_patch_size=temporal_patch_size,
+            in_channels=in_channels,
+            embed_dim=hidden_size,
+        )
+        self.use_sink = getattr(vision_config, "use_sink", False)
+        norm_layer = partial(nn.LayerNorm, eps=norm_eps)
+        head_dim = (
+            self.qk_channels
+            if self.qk_channels is not None
+            else hidden_size // num_heads
+        )
+        self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)
+        self.visual_token_window_size = getattr(
+            vision_config, "visual_token_window_size", -1
+        )
+        self.blocks = nn.ModuleList(
+            [
+                MiMoVisionBlock(
+                    dim=hidden_size,
+                    intermediate_dim=mlp_hidden_size,
+                    num_heads=num_heads,
+                    hidden_act=vision_config.hidden_act,
+                    norm_layer=norm_layer,
+                    attn_implementation="flash_attention_3",
+                    quant_config=quant_config,
+                    prefix=add_prefix(f"blocks.{i}", prefix),
+                    use_sink=(
+                        self.use_sink if i not in self.fullatt_block_indexes else False
+                    ),
+                    window_size=(
+                        self.visual_token_window_size,
+                        self.visual_token_window_size,
+                    ),
+                    num_kv_heads=num_kv_heads,
+                    head_dim=self.qk_channels,
+                    use_data_parallel=self.use_data_parallel,
+                )
+                for i in range(depth)
+            ]
+        )
+
+        self.vision_config = vision_config
+        self.merger = Qwen2_5_VisionPatchMerger(
+            dim=vision_config.out_hidden_size,
+            context_dim=hidden_size,
+            spatial_merge_size=spatial_merge_size,
+            quant_config=quant_config,
+            prefix=add_prefix("merger", prefix),
+            use_data_parallel=self.use_data_parallel,
+        )
+        self._post_init()
+
+    def apply_index(self, tensor: torch.Tensor, index: torch.Tensor):
+        tensor = tensor.unflatten(0, (-1, self.spatial_merge_unit))
+        tensor = tensor[index]
+        tensor = tensor.flatten(0, 1)
+        return tensor
+
+    def _post_init(self):
+        for name, param in self.named_parameters():
+            if "bias" in name:
+                param.data.zero_()
+
+    def get_window_index_1d(self, grid_thw, col=True):
+        window_index: list = []
+        window_index_id = 0
+        for grid_t, grid_h, grid_w in grid_thw:
+            llm_grid_h, llm_grid_w = (
+                grid_h // self.spatial_merge_size,
+                grid_w // self.spatial_merge_size,
+            )
+            index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(
+                grid_t, llm_grid_h, llm_grid_w
+            )
+            if col:
+                index_new = index.transpose(1, 2).reshape(-1)
+            else:
+                index_new = index.reshape(-1)
+            window_index.append(index_new + window_index_id)
+            window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()
+        window_index = torch.cat(
+            window_index,
+            dim=0,
+        )
+        return window_index
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.patch_embed.proj.weight.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.blocks[0].mlp.gate_up_proj.weight.device
+
+    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
+        pos_ids = []
+        for i in range(grid_thw.size(0)):
+            t, h, w = grid_thw[i].tolist()
+            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+
+            hpos_ids = hpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+            hpos_ids = hpos_ids.flatten()
+
+            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+            wpos_ids = wpos_ids.reshape(
+                h // self.spatial_merge_size,
+                self.spatial_merge_size,
+                w // self.spatial_merge_size,
+                self.spatial_merge_size,
+            )
+            wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+            wpos_ids = wpos_ids.flatten()
+
+            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+        pos_ids = torch.cat(pos_ids, dim=0)
+        max_grid_size = grid_thw[:, 1:].max()
+        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+        rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+        return rotary_pos_emb
+
+    def _prepare_forward(
+        self,
+        x: torch.Tensor,
+        grid_thw: torch.Tensor,
+    ):
+        # patchify
+        x = x.to(device=self.device, dtype=self.dtype)
+        x = self.patch_embed(x)
+        # compute position embedding
+        rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+        window_index_1d_col = self.get_window_index_1d(grid_thw, col=True).to(
+            device=x.device
+        )
+        reverse_window_index_1d_col = torch.argsort(window_index_1d_col).to(
+            device=x.device
+        )
+
+        rotary_pos_emb = rotary_pos_emb.to(device=x.device)
+        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+
+        def get_position_embeddings(emb, x):
+            position_embeddings = (emb.cos(), emb.sin())
+            position_embeddings = (
+                position_embeddings[0].to(x.device),
+                position_embeddings[1].to(x.device),
+            )
+            return position_embeddings
+
+        seqlens = torch.repeat_interleave(
+            grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
+        )
+        cu_seqlens = torch.cat(
+            [
+                torch.tensor([0], device=x.device, dtype=torch.int32),
+                seqlens.cumsum(dim=0).to(device=x.device, dtype=torch.int32),
+            ]
+        )
+        max_seqlen = seqlens.max().item()
+
+        row_based_embeddings = get_position_embeddings(emb, x)
+        col_based_embeddings = get_position_embeddings(
+            self.apply_index(emb, window_index_1d_col), x
+        )
+
+        # transformers
+        x = x.unsqueeze(1)  # [S, 1, H]
+
+        return (
+            x,
+            row_based_embeddings,
+            col_based_embeddings,
+            window_index_1d_col,
+            reverse_window_index_1d_col,
+            cu_seqlens,
+            max_seqlen,
+        )
+
+    def run_blocks(
+        self,
+        x: torch.Tensor,
+        row_based_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        col_based_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        window_index_1d_col: torch.Tensor,
+        reverse_window_index_1d_col: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+    ) -> torch.Tensor:
+        for layer_num, blk in enumerate(self.blocks):
+            window_attn_type = self.vit_window_attn_types[layer_num]
+
+            # window_attn_type = 1: col-based SWA
+            if window_attn_type == 1 and (
+                layer_num == 0 or self.vit_window_attn_types[layer_num - 1] != 1
+            ):
+                x = self.apply_index(x, window_index_1d_col)
+
+            if (
+                layer_num > 0
+                and window_attn_type != 1
+                and self.vit_window_attn_types[layer_num - 1] == 1
+            ):
+                x = self.apply_index(x, reverse_window_index_1d_col)
+
+            position_embeddings = (
+                col_based_embeddings if window_attn_type == 1 else row_based_embeddings
+            )
+            full_attn = layer_num in self.fullatt_block_indexes
+
+            x = blk(
+                x,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
+                position_embeddings=position_embeddings,
+                full_attn=full_attn,
+            )
+        x = self.merger(x)
+        return x
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        grid_thw: torch.Tensor,
+    ) -> torch.Tensor:
+        (
+            x,
+            row_based_embeddings,
+            col_based_embeddings,
+            window_index_1d_col,
+            reverse_window_index_1d_col,
+            cu_seqlens,
+            max_seqlen,
+        ) = self._prepare_forward(x, grid_thw)
+
+        return self.run_blocks(
+            x,
+            row_based_embeddings,
+            col_based_embeddings,
+            window_index_1d_col,
+            reverse_window_index_1d_col,
+            cu_seqlens,
+            max_seqlen,
+        )

python/sglang/srt/parser/conversation.py

@@ -35,7 +35,7 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 from typing_extensions import Literal
 
 from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
-from sglang.srt.utils import ImageData, read_system_prompt_from_file
+from sglang.srt.utils import ImageData, VideoData, read_system_prompt_from_file
 
 
 class SeparatorStyle(IntEnum):
@@ -97,7 +97,7 @@ class Conversation:
     audio_token: str = "<audio>"
 
     image_data: Optional[List[ImageData]] = None
-    video_data: Optional[List[str]] = None
+    video_data: Optional[List[Union[str, VideoData]]] = None
     modalities: Optional[List[str]] = None
     stop_token_ids: Optional[int] = None
 
@@ -413,9 +413,14 @@ class Conversation:
         """Append a new image."""
         self.image_data.append(ImageData(url=image, detail=detail))
 
-    def append_video(self, video: str):
+    def append_video(self, video: str, preprocess_kwargs: Optional[Dict] = None):
         """Append a new video."""
-        self.video_data.append(video)
+        if preprocess_kwargs:
+            self.video_data.append(
+                VideoData(video, preprocess_kwargs=preprocess_kwargs)
+            )
+        else:
+            self.video_data.append(video)
 
     def append_audio(self, audio: str):
         """Append a new audio."""

python/sglang/srt/server_args.py

@@ -77,6 +77,10 @@ logger = logging.getLogger(__name__)
 
 # Define constants
 DEFAULT_UVICORN_ACCESS_LOG_EXCLUDE_PREFIXES = ()
+MIMO_V2_MODEL_ARCHS = (
+    "MiMoV2ForCausalLM",
+    "MiMoV2FlashForCausalLM",
+)
 
 SAMPLING_BACKEND_CHOICES = {"flashinfer", "pytorch", "ascend"}
 
@@ -1631,7 +1635,10 @@ class ServerArgs:
         )
 
     def _handle_model_specific_adjustments(self):
-        from sglang.srt.configs.model_config import is_deepseek_nsa
+        from sglang.srt.configs.model_config import (
+            get_mimo_v2_fused_qkv_expected_tp_size,
+            is_deepseek_nsa,
+        )
 
         if parse_connector_type(self.model_path) == ConnectorType.INSTANCE:
             return
@@ -1967,13 +1974,33 @@ class ServerArgs:
                     self.ep_size == 1
                 ), "Triton kernel MoE is only supported when ep_size == 1"
 
-        elif any(
+        elif model_arch in MIMO_V2_MODEL_ARCHS:
-            x in model_arch
+            if model_arch == "MiMoV2ForCausalLM":
-            for x in (
+                expected_attn_tp_size = get_mimo_v2_fused_qkv_expected_tp_size(
-                "MiMoV2ForCausalLM",
+                    hf_config
-                "MiMoV2FlashForCausalLM",
+                )
-            )
+                attn_dp_size = self.dp_size if self.enable_dp_attention else 1
-        ):
+                effective_attn_tp_size = (
+                    self.tp_size // attn_dp_size // self.attn_cp_size
+                )
+                if (
+                    expected_attn_tp_size is not None
+                    and effective_attn_tp_size != expected_attn_tp_size
+                ):
+                    raise ValueError(
+                        "MiMoV2ForCausalLM requires effective attention TP "
+                        f"size {expected_attn_tp_size} because its fused "
+                        "qkv_proj weights are "
+                        f"TP={expected_attn_tp_size}-interleaved; got "
+                        f"{effective_attn_tp_size} "
+                        f"(tp_size={self.tp_size}, dp_size={self.dp_size}, "
+                        f"enable_dp_attention={self.enable_dp_attention}, "
+                        f"attn_cp_size={self.attn_cp_size}). "
+                        "Set --tp, --dp, --enable-dp-attention, and "
+                        "--attention-context-parallel-size so the effective "
+                        f"attention TP size is {expected_attn_tp_size}."
+                    )
+
             if self.speculative_algorithm == "EAGLE":
                 self.enable_multi_layer_eagle = True
                 logger.info(

python/sglang/srt/speculative/eagle_info.py

@@ -74,10 +74,12 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
     grammar: BaseGrammarObject = None
 
     # Shape info for padding
-    num_tokens_per_req: int = -1
+    num_tokens_per_req: int = -1  # -1 auto-fills from draft_token_num.
 
     def __post_init__(self):
         super().__init__(SpecInputType.EAGLE_VERIFY)
+        if self.num_tokens_per_req < 0:
+            self.num_tokens_per_req = self.draft_token_num
 
     def get_spec_adjust_token_coefficient(self) -> Tuple[int, int]:
         return self.draft_token_num, self.draft_token_num

python/sglang/srt/utils/common.py

@@ -787,6 +787,13 @@ class ImageData:
     url: str
     detail: Optional[Literal["auto", "low", "high"]] = "auto"
     max_dynamic_patch: Optional[int] = None
+    preprocess_kwargs: Optional[Dict] = None
+
+
+@dataclass
+class VideoData:
+    url: str
+    preprocess_kwargs: Optional[Dict] = None
 
 
 image_extension_names = (".png", ".jpg", ".jpeg", ".webp", ".gif")
@@ -924,7 +931,11 @@ def _normalize_video_input(
         return None
 
 
-def load_video(video_file: Union[str, bytes], use_gpu: bool = True):
+def load_video(video_file: Union[str, bytes, VideoData], use_gpu: bool = True):
+    if isinstance(video_file, VideoData):
+        # preprocess_kwargs is consumed by the multimodal processor, not here.
+        video_file = video_file.url
+
     if isinstance(video_file, (list, tuple, torch.Tensor, np.ndarray)):
         return video_file
 

python/sglang/test/server_fixtures/mmmu_fixture.py

@@ -22,6 +22,8 @@ class MMMUServerBase(CustomTestCase):
     This fixture handles server lifecycle for single-model MMMU tests.
     For multi-model tests that need to start/stop servers within test methods,
     use MMMUMultiModelTestBase instead.
+    Set server_api_key = None to launch without auth when sharing the server
+    with mixins whose clients do not send API keys.
     """
 
     model = None
@@ -29,6 +31,7 @@ class MMMUServerBase(CustomTestCase):
     timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
     other_args: list[str] = []
     mem_fraction_static: float = DEFAULT_MEM_FRACTION_STATIC
+    server_api_key = "sk-123456"
 
     @classmethod
     def setUpClass(cls):
@@ -53,7 +56,7 @@ class MMMUServerBase(CustomTestCase):
             cls.model,
             cls.base_url,
             timeout=cls.timeout,
-            api_key=cls.api_key,
+            api_key=cls.server_api_key,
             other_args=server_args,
             env=process_env,
         )

test/registered/8-gpu-models/test_mimo_models.py

@@ -4,6 +4,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
 from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin
 from sglang.test.server_fixtures.default_fixture import DefaultServerBase
+from sglang.test.server_fixtures.mmmu_fixture import MMMUServerBase
 
 register_cuda_ci(est_time=330, suite="stage-c-test-8-gpu-h200")
 
@@ -48,5 +49,42 @@ class TestMiMoV2Flash(GSM8KMixin, SpecDecodingMixin, DefaultServerBase):
     num_accepted_drafts_thres = 3.2
 
 
+MIMO_V2_MODEL = "XiaomiMiMo/MiMo-V2.5"
+MIMO_V2_OTHER_ARGS = [
+    "--tp",
+    "8",
+    "--dp",
+    "2",
+    "--enable-dp-attention",
+    "--mm-enable-dp-encoder",
+    "--attention-backend",
+    "fa3",
+    "--mm-attention-backend",
+    "fa3",
+    "--reasoning-parser",
+    "mimo",
+]
+MIMO_V2_MTP_OTHER_ARGS = MIMO_V2_OTHER_ARGS + [
+    "--speculative-algorithm",
+    "EAGLE",
+    "--speculative-num-steps",
+    "3",
+    "--speculative-eagle-topk",
+    "1",
+    "--speculative-num-draft-tokens",
+    "4",
+    "--enable-multi-layer-eagle",
+]
+
+
+class TestMiMoV2(GSM8KMixin, MMMUServerBase):
+    gsm8k_accuracy_thres = 0.75
+    gsm8k_accept_length_thres = 2.5
+    model = MIMO_V2_MODEL
+    mem_fraction_static = 0.65
+    server_api_key = None
+    other_args = MIMO_V2_MTP_OTHER_ARGS
+
+
 if __name__ == "__main__":
     unittest.main()