From 651af06a0b5e0abae1c92d4c87dad22e6600cefb Mon Sep 17 00:00:00 2001 From: Zhonghua Deng Date: Fri, 1 May 2026 00:02:26 +0800 Subject: [PATCH] [Feature] Xiaomi MiMo-V2.5 day0 support (#23811) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 张袁 Co-authored-by: 刘安岐 Co-authored-by: Xinyuan Tong Co-authored-by: Shangming Cai Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com> --- python/sglang/srt/configs/model_config.py | 59 +- python/sglang/srt/layers/attention/vision.py | 78 +- python/sglang/srt/layers/linear.py | 2 +- python/sglang/srt/managers/io_struct.py | 4 +- python/sglang/srt/managers/mm_utils.py | 47 +- python/sglang/srt/models/mimo_audio.py | 1350 +++++++++++ python/sglang/srt/models/mimo_v2.py | 235 +- python/sglang/srt/models/mimo_v2_nextn.py | 19 +- python/sglang/srt/models/mimo_vl.py | 507 ++++ .../srt/multimodal/processors/mimo_v2.py | 2039 +++++++++++++++++ python/sglang/srt/parser/conversation.py | 13 +- python/sglang/srt/server_args.py | 43 +- python/sglang/srt/speculative/eagle_info.py | 4 +- python/sglang/srt/utils/common.py | 13 +- .../test/server_fixtures/mmmu_fixture.py | 5 +- .../8-gpu-models/test_mimo_models.py | 38 + 16 files changed, 4369 insertions(+), 87 deletions(-) create mode 100644 python/sglang/srt/models/mimo_audio.py create mode 100644 python/sglang/srt/models/mimo_vl.py create mode 100644 python/sglang/srt/multimodal/processors/mimo_v2.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index a4492bd17..cf8bbff0b 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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 ( - "MiMoV2ForCausalLM", - "MiMoV2FlashForCausalLM", - ): + if is_draft_model and self.hf_config.architectures[0] in MIMO_V2_MODEL_ARCHS: 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", - "MiMoV2FlashForCausalLM", + *MIMO_V2_MODEL_ARCHS, "MiMoV2MTP", "Gemma4ForCausalLM", "Gemma4ForConditionalGeneration", @@ -417,14 +443,7 @@ class ModelConfig: return True # MiMoV2 creates sinks only when the config flags are set. - if any( - a in archs - for a in ( - "MiMoV2FlashForCausalLM", - "MiMoV2ForCausalLM", - "MiMoV2MTP", - ) - ): + if any(a in archs for a in (*MIMO_V2_MODEL_ARCHS, "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", - "MiMoV2FlashForCausalLM", + *MIMO_V2_MODEL_ARCHS, "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( - x in model_architectures - for x in ( - "MiMoV2ForCausalLM", - "MiMoV2FlashForCausalLM", - ) - ): + elif any(arch in MIMO_V2_MODEL_ARCHS for arch in model_architectures): 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 diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 0eb6de6f8..29ccd7606 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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( - q, - k, - v, + fa_kwargs = dict( 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( - q, - k, - v, + fa_kwargs = dict( 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) - k = k.view(original_shape) + q = q.view(original_q_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 diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 477a99994..ddd5b6a4f 100644 --- a/python/sglang/srt/layers/linear.py +++ b/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), diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index aafc1d5e5..e5cdc3b25 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/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 diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 918b564fb..c71b0e1c4 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/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 ( - hasattr(item, "feature") - and isinstance(item.feature, torch.Tensor) - and item.feature.is_cpu - ): - item.feature = ShmPointerMMData(item.feature) + 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 + ) 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): - item.feature = item.feature.materialize() + 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 + ) return obj diff --git a/python/sglang/srt/models/mimo_audio.py b/python/sglang/srt/models/mimo_audio.py new file mode 100644 index 000000000..a90547920 --- /dev/null +++ b/python/sglang/srt/models/mimo_audio.py @@ -0,0 +1,1350 @@ +"""MiMo audio: tokenizer, encoding utilities, and audio encoder.""" + +# Audio tokenizer adapted from https://github.com/XiaomiMiMo/MiMo-Audio-Tokenizer.git + +import logging +import math +import os +import typing as tp +from dataclasses import dataclass +from functools import wraps +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers.activations import ACT2FN +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_utils import PreTrainedModel +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config +from transformers.models.qwen2.modeling_qwen2 import Qwen2Model + +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import is_cuda + +if is_cuda(): + from sgl_kernel.flash_attn import flash_attn_varlen_func +else: + + def flash_attn_varlen_func(*args, **kwargs): + raise RuntimeError("MiMoAudioTokenizer requires CUDA to run.") + + +logger = logging.getLogger(__name__) + + +def _compute_default_rope_parameters( + config=None, device=None, seq_len=None, **rope_kwargs +): + if config is not None and len(rope_kwargs) > 0: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive" + ) + if len(rope_kwargs) > 0: + base = rope_kwargs["base"] + dim = rope_kwargs["dim"] + elif config is not None: + base = config.rope_theta + partial_rotary_factor = ( + config.partial_rotary_factor + if hasattr(config, "partial_rotary_factor") + else 1.0 + ) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + head_dim = config.hidden_size // config.num_attention_heads + logger.info( + "audio.head_dim not set; defaulting to hidden_size/num_heads = %d", + head_dim, + ) + dim = int(head_dim * partial_rotary_factor) + attention_factor = 1.0 + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, dim, 2, dtype=torch.int64).to( + device=device, dtype=torch.float + ) + / dim + ) + ) + return inv_freq, attention_factor + + +_ROPE_INIT_FUNCTIONS = { + "default": _compute_default_rope_parameters, +} + + +def _dynamic_rope_update(rope_forward): + def longrope_frequency_update(self, position_ids, device): + seq_len = torch.max(position_ids) + 1 + if hasattr(self.config, "original_max_position_embeddings"): + original_max_position_embeddings = ( + self.config.original_max_position_embeddings + ) + else: + original_max_position_embeddings = self.config.max_position_embeddings + if seq_len > original_max_position_embeddings: + if not hasattr(self, "long_inv_freq"): + self.long_inv_freq, _ = self.rope_init_fn( + self.config, device, seq_len=original_max_position_embeddings + 1 + ) + self.register_buffer("inv_freq", self.long_inv_freq, persistent=False) + else: + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + + def dynamic_frequency_update(self, position_ids, device): + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn( + self.config, device, seq_len=seq_len + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.max_seq_len_cached = seq_len + + if ( + seq_len < self.original_max_seq_len + and self.max_seq_len_cached > self.original_max_seq_len + ): + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @wraps(rope_forward) + def wrapper(self, x, position_ids): + if "dynamic" in self.rope_type: + dynamic_frequency_update(self, position_ids, device=x.device) + elif self.rope_type == "longrope": + longrope_frequency_update(self, position_ids, device=x.device) + return rope_forward(self, x, position_ids) + + return wrapper + + +class AudioRotaryEmbedding(nn.Module): + def __init__(self, base, dim, max_seq_len, rope_type="default", device=None): + super().__init__() + self.max_seq_len = max_seq_len + self.rope_type = rope_type + self.rope_init_fn = _ROPE_INIT_FUNCTIONS[self.rope_type] + inv_freq, self.attention_scaling = self.rope_init_fn( + device=device, base=base, dim=dim + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @_dynamic_rope_update + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[:, None].float().expand(-1, 1).to(x.device) + position_ids_expanded = position_ids[None, :].float() + device_type = ( + x.device.type + if isinstance(x.device.type, str) and x.device.type != "mps" + else "cpu" + ) + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(0, 1) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class EuclideanCodebook(nn.Module): + """Codebook with Euclidean distance (inference-only).""" + + def __init__( + self, dim: int, codebook_size: int, kmeans_init: bool = False, **kwargs + ): + super().__init__() + init_fn = self._uniform_init if not kmeans_init else torch.zeros + embed = init_fn(codebook_size, dim) + + self.codebook_size = codebook_size + + self.register_buffer("inited", torch.Tensor([not kmeans_init])) + self.register_buffer("cluster_size", torch.zeros(codebook_size)) + self.register_buffer("embed", embed) + self.register_buffer("embed_avg", embed.clone()) + + def preprocess(self, x): + x = rearrange(x, "... d -> (...) d") + return x + + def quantize(self, x): + embed = self.embed.t() + dist_val = -( + x.pow(2).sum(1, keepdim=True) + - 2 * x @ embed + + embed.pow(2).sum(0, keepdim=True) + ) + embed_ind = dist_val.max(dim=-1).indices + return embed_ind + + def postprocess_emb(self, embed_ind, shape): + return embed_ind.view(*shape[:-1]) + + def dequantize(self, embed_ind): + quantize = F.embedding(embed_ind, self.embed) + return quantize + + def encode(self, x): + shape = x.shape + x = self.preprocess(x) + embed_ind = self.quantize(x) + embed_ind = self.postprocess_emb(embed_ind, shape) + return embed_ind + + def decode(self, embed_ind): + quantize = self.dequantize(embed_ind) + return quantize + + @staticmethod + def _uniform_init(*shape: int): + t = torch.empty(shape) + nn.init.kaiming_uniform_(t) + return t + + +class VectorQuantization(nn.Module): + """Vector quantization with euclidean distance (inference-only).""" + + def __init__( + self, + dim: int, + codebook_size: int, + codebook_dim: tp.Optional[int] = None, + kmeans_init: bool = True, + **kwargs, + ): + super().__init__() + _codebook_dim: int = codebook_dim if codebook_dim is not None else dim + + requires_projection = _codebook_dim != dim + self.project_in = ( + nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity() + ) + self.project_out = ( + nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity() + ) + + self._codebook = EuclideanCodebook( + dim=_codebook_dim, + codebook_size=codebook_size, + kmeans_init=kmeans_init, + ) + self.codebook_size = codebook_size + + @property + def codebook(self): + return self._codebook.embed + + def encode(self, x): + x = self.project_in(x) + embed_in = self._codebook.encode(x) + return embed_in + + def decode(self, embed_ind): + quantize = self._codebook.decode(embed_ind) + quantize = self.project_out(quantize) + return quantize + + +class ResidualVectorQuantization(nn.Module): + """Residual vector quantization implementation. + Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf + """ + + def __init__(self, *, num_quantizers, codebook_size, **kwargs): + super().__init__() + if isinstance(codebook_size, int): + codebook_size = [codebook_size] * num_quantizers + elif len(codebook_size) < num_quantizers: + codebook_size += [codebook_size[-1]] * (num_quantizers - len(codebook_size)) + self.layers = nn.ModuleList( + [ + VectorQuantization(codebook_size=codebook_size[i], **kwargs) + for i in range(num_quantizers) + ] + ) + + def encode( + self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None + ) -> torch.Tensor: + residual = x + all_indices = [] + n_q = len(self.layers) if n_q is None else n_q + st = 0 if st is None else st + for layer in self.layers[st:n_q]: + indices = layer.encode(residual) + quantized = layer.decode(indices) + residual = residual - quantized + all_indices.append(indices) + out_indices = torch.stack(all_indices) + return out_indices + + def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor: + quantized_out = self.layers[st].decode(q_indices[0]) + for i in range(1, len(q_indices)): + layer = self.layers[st + i] + quantized = layer.decode(q_indices[i]) + quantized_out = quantized_out + quantized + return quantized_out + + +class ResidualVectorQuantizer(nn.Module): + """Residual Vector Quantizer (inference-only).""" + + def __init__( + self, + dimension: int = 256, + n_q: int = 8, + bins: int | list = 1024, + kmeans_init: bool = True, + **kwargs, + ): + super().__init__() + self.n_q = n_q + self.vq = ResidualVectorQuantization( + dim=dimension, + codebook_size=bins, + num_quantizers=n_q, + kmeans_init=kmeans_init, + ) + + def encode( + self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None + ) -> torch.Tensor: + n_q = n_q if n_q else self.n_q + st = st or 0 + codes = self.vq.encode(x, n_q=n_q, st=st) + return codes + + def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor: + quantized = self.vq.decode(codes, st=st) + return quantized + + +class MiMoAudioTokenizerConfig(PretrainedConfig): + model_type = "mimo_audio_tokenizer" + + def __init__( + self, + max_audio_seconds: int = 1800, + stride_size: int = 2, + avg_pooler: int = 1, + d_model: int = 768, + scale_embedding: bool = True, + kernel_size: int = 3, + activation_function: str = "gelu", + encoder_layers: int = 8, + encoder_skip_layer_id: int = None, + encoder_attention_heads: int = 12, + encoder_ffn_dim: int = 3072, + encoder_causal: bool = False, + encoder_attn_window_size: list = None, + decoder_layers: int = 8, + decoder_attention_heads: int = 12, + decoder_ffn_dim: int = 3072, + decoder_kernel_size: int = 3, + decoder_stride_size: int = 2, + decoder_causal: bool = True, + decoder_attn_window_size: list = None, + nfft: int = 1024, + vocoder_dim: int = 512, + vocoder_intermediate_dim: int = 4096, + vocoder_num_layers: int = 30, + n_mels: int = 80, + sampling_rate: int = 24000, + hop_length: int = 240, + window_size: int = 1024, + vocoder_padding: str = "same", + fmin: int = 0, + fmax: int = None, + num_quantizers: int = 12, + codebook_size: list = None, + threshold_ema_dead_code: int = 10, + position_embedding_type: str = "rope", + rope_theta: int = 10000, + rope_type: str = "default", + ln_type: str = "LayerNorm", + vocoder_attention_heads: int = 4, + vocoder_attn_window_size: list = None, + use_istft_only: bool = False, + hybrid_attention: bool = False, + hybrid_block_size: int = 8, + swa_per_block: int = 2, + **kwargs, + ): + super().__init__(**kwargs) + self.max_audio_seconds = max_audio_seconds + self.stride_size = stride_size + self.avg_pooler = avg_pooler + self.d_model = d_model + self.scale_embedding = scale_embedding + self.kernel_size = kernel_size + self.activation_function = activation_function + self.encoder_layers = encoder_layers + self.encoder_skip_layer_id = encoder_skip_layer_id + self.encoder_attention_heads = encoder_attention_heads + self.encoder_ffn_dim = encoder_ffn_dim + self.encoder_causal = encoder_causal + self.encoder_attn_window_size = ( + encoder_attn_window_size + if encoder_attn_window_size is not None + else [-1, -1] + ) + self.decoder_layers = decoder_layers + self.decoder_attention_heads = decoder_attention_heads + self.decoder_ffn_dim = decoder_ffn_dim + self.decoder_kernel_size = decoder_kernel_size + self.decoder_stride_size = decoder_stride_size + self.decoder_causal = decoder_causal + self.decoder_attn_window_size = ( + decoder_attn_window_size + if decoder_attn_window_size is not None + else [-1, -1] + ) + self.nfft = nfft + self.vocoder_dim = vocoder_dim + self.vocoder_intermediate_dim = vocoder_intermediate_dim + self.vocoder_num_layers = vocoder_num_layers + self.n_mels = n_mels + self.sampling_rate = sampling_rate + self.hop_length = hop_length + self.window_size = window_size + self.vocoder_padding = vocoder_padding + self.fmin = fmin + self.fmax = fmax + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size if codebook_size is not None else [1024] + self.threshold_ema_dead_code = threshold_ema_dead_code + self.position_embedding_type = position_embedding_type + self.rope_theta = rope_theta + self.rope_type = rope_type + self.ln_type = ln_type + self.vocoder_attention_heads = vocoder_attention_heads + self.vocoder_attn_window_size = ( + vocoder_attn_window_size + if vocoder_attn_window_size is not None + else [40, 10] + ) + self.use_istft_only = use_istft_only + self.hybrid_attention = hybrid_attention + self.hybrid_block_size = hybrid_block_size + self.swa_per_block = swa_per_block + + +def get_sequence_mask(inputs, inputs_length): + if inputs.dim() == 3: + bsz, tgt_len, _ = inputs.size() + else: + bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length) + sequence_mask = torch.arange(0, tgt_len).to(inputs.device) + sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view( + bsz, tgt_len, 1 + ) + unpacking_index = torch.cumsum(sequence_mask.to(torch.int64).view(-1), dim=0) - 1 + return sequence_mask, unpacking_index + + +def unpack_hidden_states( + hidden_states, lengths, sequence_mask=None, unpacking_index=None +): + bsz = lengths.shape[0] + if sequence_mask is None or unpacking_index is None: + sequence_mask, unpacking_index = get_sequence_mask(hidden_states, lengths) + hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( + bsz, torch.max(lengths), hidden_states.shape[-1] + ) + return torch.where(sequence_mask, hidden_states, 0) + + +def get_position_ids(lengths): + total_len = lengths.sum() + offset = torch.cat([torch.zeros(1).to(lengths), lengths[:-1].cumsum(dim=0)]) + offset = torch.repeat_interleave(offset, lengths) + return torch.arange(0, total_len).to(offset) - offset + + +LAYER_NORM = {"LayerNorm": nn.LayerNorm} + + +class AudioEncoderAttention(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + window_size: Tuple[int, int] = (-1, -1), + causal: bool = False, + ): + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.window_size = window_size + self.causal = causal + + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True) + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True) + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_position_embeddings=None, + ): + bsz, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states).view( + bsz, self.num_heads, self.head_dim + ) + key_states = self.k_proj(hidden_states).view(bsz, self.num_heads, self.head_dim) + value_states = self.v_proj(hidden_states).view( + bsz, self.num_heads, self.head_dim + ) + + if rope_position_embeddings is not None: + cos, sin = rope_position_embeddings + query_states, key_states = self.apply_rotary_pos_emb( + query_states, key_states, cos, sin + ) + + attn_output = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens, + cu_seqlens, + max_seqlen, + max_seqlen, + causal=self.causal, + window_size=self.window_size, + ) + + attn_output = attn_output.reshape(bsz, self.embed_dim) + attn_output = self.out_proj(attn_output) + return attn_output + + @staticmethod + def _rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + @classmethod + def apply_rotary_pos_emb(cls, q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (cls._rotate_half(q) * sin) + k_embed = (k * cos) + (cls._rotate_half(k) * sin) + return q_embed, k_embed + + +class AudioEncoderTransformerLayer(nn.Module): + def __init__( + self, + config: MiMoAudioTokenizerConfig, + causal: bool, + attn_window_size: Tuple[int, int] = (-1, -1), + ): + super().__init__() + self.embed_dim = config.d_model + + self.self_attn = AudioEncoderAttention( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + window_size=attn_window_size, + causal=causal, + ) + self.self_attn_layer_norm = LAYER_NORM[config.ln_type](self.embed_dim) + + self.activation_fn = ACT2FN[config.activation_function] + self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) + self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) + self.final_layer_norm = LAYER_NORM[config.ln_type](self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states = self.self_attn( + hidden_states, + cu_seqlens, + max_seqlen, + rope_position_embeddings=rope_position_embeddings, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = self.fc2(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class AudioEncoder(nn.Module): + def __init__( + self, + config: MiMoAudioTokenizerConfig, + ): + super().__init__() + self.config = config + self.max_source_positions = ( + config.max_audio_seconds * config.sampling_rate // config.hop_length + ) // config.stride_size + self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 + self.skip_layer_idx = config.encoder_skip_layer_id + + self.conv1 = nn.Conv1d( + config.n_mels, + config.d_model, + kernel_size=config.kernel_size, + padding=1, + ) + self.conv2 = nn.Conv1d( + config.d_model, + config.d_model, + kernel_size=config.kernel_size, + stride=config.stride_size, + padding=1, + ) + + self.position_embedding = AudioRotaryEmbedding( + config.rope_theta, + config.d_model // config.encoder_attention_heads, + self.max_source_positions, + config.rope_type, + ) + + attn_window_sizes = [] + if config.hybrid_attention: + for i in range(config.encoder_layers): + if i % config.swa_per_block < config.swa_per_block - 1: + attn_window_sizes.append(tuple(config.encoder_attn_window_size)) + else: + attn_window_sizes.append((-1, -1)) + else: + attn_window_sizes = [ + tuple(config.encoder_attn_window_size) + ] * config.encoder_layers + + self.layers = nn.ModuleList( + [ + AudioEncoderTransformerLayer( + config=config, + causal=config.encoder_causal, + attn_window_size=attn_window_sizes[i], + ) + for i in range(config.encoder_layers) + ] + ) + + self.layer_norm = LAYER_NORM[config.ln_type](config.d_model) + + if config.avg_pooler != 1: + self.down_sample_layer = nn.Sequential( + nn.Conv1d( + config.d_model, + config.d_model, + config.avg_pooler, + config.avg_pooler, + bias=False, + ), + nn.GELU(), + ) + self.down_sample_norm = LAYER_NORM[config.ln_type](config.d_model) + else: + self.down_sample_layer = None + + if config.num_quantizers != 0: + self.quantizer = ResidualVectorQuantizer( + dimension=config.d_model, + n_q=config.num_quantizers, + bins=config.codebook_size, + threshold_ema_dead_code=config.threshold_ema_dead_code, + ) + else: + self.quantizer = None + + def get_features(self, input_features, output_length): + input_features = input_features.to(self.conv1.weight) + inputs_embeds = nn.functional.gelu(self.conv1(input_features)) + inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) + inputs_embeds = inputs_embeds.permute(0, 2, 1) + bsz, tgt_len, _ = inputs_embeds.size() + hidden_states = inputs_embeds + + position_ids = get_position_ids(output_length).long().to(input_features.device) + rope_position_embeddings = self.position_embedding(input_features, position_ids) + + attention_mask, unpacking_index = get_sequence_mask( + hidden_states, output_length + ) + hidden_states = torch.masked_select(hidden_states, attention_mask).view( + torch.sum(output_length), self.config.d_model + ) + + cu_seqlens = F.pad( + torch.cumsum(output_length, dim=0), (1, 0), "constant", 0 + ).to(device=hidden_states.device, dtype=torch.int32) + max_seqlen = torch.max(output_length).to(torch.int32).item() + + skip_connect_hidden_states = 0.0 + for idx, encoder_layer in enumerate(self.layers): + hidden_states = encoder_layer( + hidden_states, + cu_seqlens, + max_seqlen, + rope_position_embeddings=rope_position_embeddings, + ) + if (self.skip_layer_idx is not None) and idx == self.skip_layer_idx - 1: + skip_connect_hidden_states = hidden_states.clone() + + hidden_states += skip_connect_hidden_states + hidden_states = self.layer_norm(hidden_states) + + if self.down_sample_layer is not None: + hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( + bsz, tgt_len, self.config.d_model + ) + if hidden_states.size(1) % self.config.avg_pooler: + pad_len = ( + self.config.avg_pooler + - hidden_states.size(1) % self.config.avg_pooler + ) + hidden_states = torch.nn.functional.pad( + hidden_states, (0, 0, 0, pad_len), mode="constant", value=0.0 + ) + tgt_len += pad_len + tgt_len = tgt_len // self.config.avg_pooler + hidden_states = self.down_sample_layer(hidden_states.transpose(1, 2)) + output_length = ( + output_length // self.config.avg_pooler + + (output_length % self.config.avg_pooler != 0).int() + ) + hidden_states = hidden_states.transpose(1, 2) + attention_mask, unpacking_index = get_sequence_mask( + hidden_states, output_length + ) + hidden_states = torch.masked_select(hidden_states, attention_mask).view( + torch.sum(output_length), self.config.d_model + ) + hidden_states = self.down_sample_norm(hidden_states) + + return ( + hidden_states, + output_length, + attention_mask, + unpacking_index, + tgt_len, + bsz, + ) + + def get_output_length(self, mel_len): + tgt_len = mel_len + 3 - self.config.kernel_size + return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1 + + @torch.no_grad() + def encode( + self, + input_features, + input_lens=None, + output_length=None, + return_codes_only=False, + n_q=None, + use_quantizer=True, + ): + if output_length is None: + output_length = self.get_output_length(input_lens) + input_features = unpack_hidden_states(input_features, input_lens) + hidden_states, output_length, attention_mask, unpacking_index, tgt_len, bsz = ( + self.get_features( + input_features=input_features.transpose(1, 2), + output_length=output_length, + ) + ) + + dtype = hidden_states.dtype + if use_quantizer and self.quantizer is not None: + self.quantizer.float() + codes = self.quantizer.encode(hidden_states.float(), n_q=n_q) + if return_codes_only: + return codes, output_length + hidden_states = self.quantizer.decode(codes) + hidden_states = hidden_states.to(dtype) + else: + codes = None + + hidden_states_packed = hidden_states.clone() + hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( + bsz, tgt_len, self.config.d_model + ) + hidden_states = torch.where(attention_mask, hidden_states, 0) + return hidden_states, hidden_states_packed, output_length, codes + + @torch.no_grad() + def decode_vq(self, codes): + self.quantizer.float() + return self.quantizer.decode(codes) + + +class MiMoAudioTokenizer(PreTrainedModel): + config_class = MiMoAudioTokenizerConfig + + def __init__(self, config: MiMoAudioTokenizerConfig): + super().__init__(config) + self.config = config + self.sampling_rate = config.sampling_rate + self.encoder = AudioEncoder(config=config) + self.downsample_rate = int(config.hop_length * 2 * config.avg_pooler) + + def get_output_length(self, mel_len): + tgt_len = mel_len + 3 - self.config.kernel_size + return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1 + + @torch.no_grad() + def encode(self, mels, input_lens, use_quantizer=True): + input_features = mels + encoder_output_length = self.get_output_length(input_lens) + hidden_states, hidden_states_packed, encoder_output_length, codes = ( + self.encoder.encode( + input_features, input_lens=input_lens, use_quantizer=use_quantizer + ) + ) + return hidden_states, hidden_states_packed, encoder_output_length, codes + + +def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int): + if features.size(0) != lengths.sum().item(): + raise ValueError( + f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}" + ) + + split_points = [] + current_sum = 0 + + for i, seq_len in enumerate(lengths): + if current_sum + seq_len > max_length and current_sum > 0: + split_points.append(i) + current_sum = seq_len.item() + else: + current_sum += seq_len.item() + + # Convert split points to group sizes + group_sizes = [] + prev = 0 + for point in split_points: + group_sizes.append(point - prev) + prev = point + if prev < len(lengths): + group_sizes.append(len(lengths) - prev) + + len_groups = torch.split(lengths, group_sizes) + feature_sizes = [group.sum().item() for group in len_groups] + feature_groups = torch.split(features, feature_sizes) + + return feature_groups, len_groups + + +@torch.no_grad() +def encode_batch( + audio_tokenizer_encoder, + input_features: torch.Tensor, + input_lens: torch.Tensor, + max_length: int = 256000, +): + feature_groups, len_groups = group_by_length(input_features, input_lens, max_length) + + encoded_parts = [] + for features, lengths in zip(feature_groups, len_groups): + codes, _ = audio_tokenizer_encoder.encode( # codes are also packed + input_features=features, input_lens=lengths, return_codes_only=True + ) + encoded_parts.append(codes) + + return torch.cat(encoded_parts, dim=-1) + + +def _segment_lengths_for_mel(mel: torch.Tensor, segment_size: int): + """Split mel into segments of segment_size with a possible shorter remainder.""" + input_len = mel.size(0) + segs = [segment_size] * (input_len // segment_size) + if input_len % segment_size > 0: + segs.append(input_len % segment_size) + return segs + + +@torch.no_grad() +def tokenize_audio_batch(mels, audio_tokenizer_encoder, segment_size=6000, device=None): + """ + Tokenize multiple mels in one encode_batch call. + Returns list of code tensors, each [T_i, C] for that mel. + """ + if not mels: + return [] + if device is None: + device = next(audio_tokenizer_encoder.parameters()).device + # Build segment lengths per mel + input_len_seg_per_mel = [_segment_lengths_for_mel(m, segment_size) for m in mels] + input_lens_flat = [s for segs in input_len_seg_per_mel for s in segs] + input_features = torch.cat([m.to(device) for m in mels], dim=0) + input_lens_t = torch.tensor(input_lens_flat, dtype=torch.long, device=device) + codes_packed = encode_batch( + audio_tokenizer_encoder, + input_features=input_features, + input_lens=input_lens_t, + ) + codes = codes_packed.transpose(0, 1).detach() # [total_code_T, C] + # Code length per mel: must match encoder's actual output (get_output_length + optional avg_pooler downsampling) + code_lengths = [] + for segs in input_len_seg_per_mel: + out_len = audio_tokenizer_encoder.get_output_length( + torch.tensor(segs, dtype=torch.long, device=device) + ) + if getattr(audio_tokenizer_encoder, "down_sample_layer", None) is not None: + avg = audio_tokenizer_encoder.config.avg_pooler + out_len = out_len // avg + (out_len % avg != 0).long() + code_lengths.append(out_len.sum().item()) + code_list = torch.split(codes, code_lengths) + return list(code_list) + + +@dataclass +class MiMoAudioEncoderConfig: + tokenizer_version: str = "v1" + speech_vocab_size: str = "1025-1025-129-129-129-129-129-129" + speech_zeroemb_idx: str = "1024-1024-128-128-128-128-128-128" + group_size: int = 4 + audio_channels: int = 8 + input_local_layers: int = 6 + input_local_dim: int = 1024 + input_full_attention: bool = True + input_local_attn_heads: int = 64 + input_local_head_dim: int = 16 + input_local_intermediate_size: int = 4096 + input_local_hidden_dropout: float = 0.0 + out_hidden_size: int = 4096 # mimo vl hidden dim + rope_theta: float = 640000.0 + partial_rotary_factor: float = 0.334 + projection_layers: int = 1 + add_post_norm: bool = False + audio_segment_size: int = 6000 + + +class AudioProjection(nn.Module): + def __init__( + self, + input_size, + hidden_size, + output_size, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.input_size = input_size + self.hidden_size = hidden_size + self.output_size = output_size + self.mlp = nn.Sequential( + nn.Linear(self.input_size, self.hidden_size, bias=False), + nn.GELU(), + nn.Linear(self.hidden_size, self.output_size, bias=False), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.mlp(x) + + +class MiMoV2AudioConfig: + def __init__( + self, + speech_vocab_size: str | int = "1280", + speech_lm_head_sizes: str | int | None = None, + speech_zeroemb_idx: str | int = "1280", + delay_pattern: str = "0-1-2-3-4-5-6-7-7-7-7-7-7-7-7-7-7-7-7-7", + group_size: int = 4, + audio_channels: int = 20, + input_local_dim: int = 1024, + input_local_layers: int = 6, + input_local_attn_heads: int = 16, + input_local_intermediate_size: int = 4096, + input_local_rope_theta: float = 640000.0, + input_local_partial_rotary_factor: float = 1.0, + output_local_dim: int = 1024, + output_local_layers: int = 6, + output_local_attn_heads: int = 16, + output_local_intermediate_size: int = 4096, + output_local_rope_theta: float = 640000.0, + output_local_partial_rotary_factor: float = 1.0, + input_projection_layers: int = 2, + output_projection_layers: int = 2, + add_encoder_post_norm: bool = True, + audio_config: dict = None, + **kwargs, + ): + for key, value in kwargs.items(): + setattr(self, key, value) + + if audio_config is not None: + self._load_from_audio_config(audio_config) + else: + self.speech_vocab_size = speech_vocab_size + self.speech_lm_head_sizes = ( + speech_lm_head_sizes + if speech_lm_head_sizes is not None + else speech_vocab_size + ) + self.speech_zeroemb_idx = speech_zeroemb_idx + self.delay_pattern = delay_pattern + self.group_size = group_size + self.audio_channels = audio_channels + self.input_local_dim = input_local_dim + self.input_local_layers = input_local_layers + self.input_local_attn_heads = input_local_attn_heads + self.input_local_intermediate_size = input_local_intermediate_size + self.input_local_rope_theta = input_local_rope_theta + self.input_local_partial_rotary_factor = input_local_partial_rotary_factor + self.output_local_dim = output_local_dim + self.output_local_layers = output_local_layers + self.output_local_attn_heads = output_local_attn_heads + self.output_local_intermediate_size = output_local_intermediate_size + self.output_local_rope_theta = output_local_rope_theta + self.output_local_partial_rotary_factor = output_local_partial_rotary_factor + self.input_projection_layers = input_projection_layers + self.output_projection_layers = output_projection_layers + self.add_encoder_post_norm = add_encoder_post_norm + + self._attn_implementation_internal = "sdpa" + + def _load_from_audio_config(self, audio_config: dict): + """Load audio parameters from audio_config dict in checkpoint. + + Uses naming that matches megatron2hf conversion output to minimize manual mapping. + """ + self.group_size = audio_config.get("group_size", 4) + self.audio_channels = audio_config.get("audio_channels", 20) + self.speech_vocab_size = audio_config.get("speech_vocab_size", "1280") + self.speech_lm_head_sizes = audio_config.get( + "speech_lm_head_sizes", self.speech_vocab_size + ) + self.speech_zeroemb_idx = audio_config.get("speech_zeroemb_idx", "1280") + # Per-channel decode delays; len must equal audio_channels. + self.delay_pattern = audio_config.get( + "audio_output_delay_pattern", "0-1-2-3-4-5-6-7-7-7-7-7-7-7-7-7-7-7-7-7" + ) + + self.input_local_dim = audio_config.get("input_local_dim", 1024) + self.input_local_layers = audio_config.get("input_local_layers", 6) + self.input_local_attn_heads = audio_config.get("input_local_attn_heads", 16) + self.input_local_intermediate_size = audio_config.get( + "input_local_intermediate_size", 4096 + ) + self.input_local_rope_theta = audio_config.get( + "input_local_rope_theta", 640000.0 + ) + self.input_local_partial_rotary_factor = audio_config.get( + "input_local_partial_rotary_factor", 1.0 + ) + + self.output_local_dim = audio_config.get("output_local_dim", 1024) + self.output_local_layers = audio_config.get("output_local_layers", 6) + self.output_local_attn_heads = audio_config.get("output_local_attn_heads", 16) + self.output_local_intermediate_size = audio_config.get( + "output_local_intermediate_size", 4096 + ) + self.output_local_rope_theta = audio_config.get( + "output_local_rope_theta", 640000.0 + ) + self.output_local_partial_rotary_factor = audio_config.get( + "output_local_partial_rotary_factor", 1.0 + ) + + self.input_projection_layers = audio_config.get("input_projection_layers", 2) + self.output_projection_layers = audio_config.get("output_projection_layers", 2) + + self.add_encoder_post_norm = audio_config.get("add_encoder_post_norm", True) + + def _parse_maybe_list(self, value: str | int, length: int) -> list[int]: + if isinstance(value, str) and "-" in value: + return [int(s) for s in value.split("-")] + return [int(value)] * length + + def parsed_speech_empty_ids(self): + return self._parse_maybe_list(self.speech_zeroemb_idx, self.audio_channels) + + def parsed_speech_vocab_sizes(self): + return self._parse_maybe_list(self.speech_vocab_size, self.audio_channels) + + def parsed_speech_lm_head_sizes(self): + return self._parse_maybe_list(self.speech_lm_head_sizes, self.audio_channels) + + def parsed_delay_pattern(self): + return self._parse_maybe_list(self.delay_pattern, self.audio_channels) + + def input_local_config(self): + """Create config for input local transformer.""" + config = Qwen2Config() + for attr in dir(self): + if not attr.startswith("_") and hasattr(config, attr): + setattr(config, attr, getattr(self, attr)) + + config.hidden_size = self.input_local_dim + config.num_hidden_layers = self.input_local_layers + config.num_attention_heads = self.input_local_attn_heads + config.num_key_value_heads = self.input_local_attn_heads + config.head_dim = getattr( + self, + "input_local_head_dim", + self.input_local_dim // self.input_local_attn_heads, + ) + config.intermediate_size = self.input_local_intermediate_size + config.rope_theta = self.input_local_rope_theta + config.partial_rotary_factor = self.input_local_partial_rotary_factor + config._attn_implementation_internal = "sdpa" + + return config + + def output_local_config(self): + """Create config for output local transformer.""" + config = Qwen2Config() + for attr in dir(self): + if not attr.startswith("_") and hasattr(config, attr): + setattr(config, attr, getattr(self, attr)) + + config.hidden_size = self.output_local_dim + config.num_hidden_layers = self.output_local_layers + config.num_attention_heads = self.output_local_attn_heads + config.num_key_value_heads = self.output_local_attn_heads + config.head_dim = self.output_local_dim // self.output_local_attn_heads + config.intermediate_size = self.output_local_intermediate_size + config.rope_theta = self.output_local_rope_theta + config.partial_rotary_factor = self.output_local_partial_rotary_factor + config._attn_implementation_internal = "sdpa" + + return config + + +class MiMoAudioEncoder(nn.Module): + config: MiMoAudioEncoderConfig + + def __init__(self, config): + super().__init__() + if not isinstance(config, MiMoV2AudioConfig): + config_dict = ( + vars(config) if hasattr(config, "__dict__") else config.__dict__ + ) + config = MiMoV2AudioConfig(**config_dict) + self.config = config + self.server_args = get_global_server_args() + self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder + self.speech_empty_ids = self.parsed_speech_empty_ids() + self.audio_channels = config.audio_channels + self.audio_group_size = config.group_size + self.audio_segment_size = config.audio_segment_size + speech_vocab_size = self._parse_maybe_list( + self.config.speech_vocab_size, self.config.audio_channels + ) + input_local_config = Qwen2Config( + hidden_size=self.config.input_local_dim, + num_hidden_layers=self.config.input_local_layers, + num_attention_heads=self.config.input_local_attn_heads, + num_key_value_heads=self.config.input_local_attn_heads, + intermediate_size=self.config.input_local_intermediate_size, + attention_dropout=self.config.input_local_hidden_dropout, + rope_theta=self.config.rope_theta, + partial_rotary_factor=self.config.partial_rotary_factor, + ) + input_local_config.head_dim = self.config.input_local_head_dim + + self.input_local_transformer = Qwen2Model(input_local_config) + + if not self.config.add_post_norm: + self.input_local_transformer.norm = nn.Identity() + + self.speech_embeddings = nn.ModuleList( + [ + nn.Embedding( + speech_vocab_size[i], + self.config.input_local_dim, + padding_idx=self.speech_empty_ids[i], + ) + for i in range(self.config.audio_channels) + ] + ) + + if self.config.projection_layers == 1: + self.projection = nn.Linear( + self.config.input_local_dim * self.config.group_size, + self.config.out_hidden_size, + bias=False, + ) + elif self.config.projection_layers == 2: + self.projection = AudioProjection( + self.config.input_local_dim * self.config.group_size, + self.config.input_local_dim * self.config.group_size * 4, + self.config.out_hidden_size, + ) + else: + raise ValueError( + f"Invalid projection layers: {self.config.projection_layers}" + ) + + model_path = self.server_args.model_path + if not os.path.isdir(model_path): + from huggingface_hub import snapshot_download + + model_path = snapshot_download( + model_path, + allow_patterns=["audio_tokenizer/*"], + ) + audio_tokenizer_path = os.path.join(model_path, "audio_tokenizer") + dev = torch.device(f"cuda:{torch.cuda.current_device()}") + self.audio_tokenizer = self._load_audio_tokenizer(audio_tokenizer_path, dev) + + @staticmethod + def _load_audio_tokenizer(path: str, device: torch.device) -> MiMoAudioTokenizer: + """Load MiMoAudioTokenizer manually to avoid new-transformers compat issues.""" + import json + import os + + from safetensors.torch import load_file + + config_path = os.path.join(path, "config.json") + with open(config_path) as f: + config_dict = json.load(f) + config = MiMoAudioTokenizer.config_class(**config_dict) + model = MiMoAudioTokenizer(config) + # Load weights from safetensors or pytorch bin + safetensors_path = os.path.join(path, "model.safetensors") + bin_path = os.path.join(path, "pytorch_model.bin") + if os.path.exists(safetensors_path): + state_dict = load_file(safetensors_path, device="cpu") + elif os.path.exists(bin_path): + state_dict = torch.load(bin_path, map_location="cpu", weights_only=True) + else: + raise FileNotFoundError( + f"No model weights found in {path} " + "(expected model.safetensors or pytorch_model.bin)" + ) + model.load_state_dict(state_dict, strict=False) + model = model.to(device=device, dtype=torch.bfloat16) + model.eval() + model.requires_grad_(False) + return model + + def parsed_speech_empty_ids(self): + return self._parse_maybe_list( + self.config.speech_zeroemb_idx, self.config.audio_channels + ) + + def _parse_maybe_list(self, value: str | int, length: int) -> List[int]: + if isinstance(value, str) and "-" in value: + return [int(s) for s in value.split("-")] + return [int(value)] * length + + # adapted from mimo-audio + def apply_input_local_transformer(self, speech_embeddings: torch.Tensor): + output = self.input_local_transformer( + inputs_embeds=speech_embeddings, + return_dict=True, + is_causal=not self.config.input_full_attention, # for SDPA + ) + return output.last_hidden_state # [T//group_size, group_size, input_local_dim] + + def apply_speech_embeddings(self, audio_codes: torch.Tensor) -> torch.Tensor: + num_segments = audio_codes.shape[0] + _audio_embeddings = torch.zeros( + (num_segments, self.config.group_size, self.config.input_local_dim), + dtype=next(self.speech_embeddings[0].parameters()).dtype, + device=audio_codes.device, + ) + for i in range(self.config.audio_channels): + _audio_embeddings.add_(self.speech_embeddings[i](audio_codes[:, :, i])) + return _audio_embeddings + + def process_audio(self, audio): + T = audio.shape[0] + audio = audio[:, : self.audio_channels] + padded_T = ( + (T + self.audio_group_size - 1) + // self.audio_group_size + * self.audio_group_size + ) + padded_audio = torch.cat( + [ + audio, + torch.zeros( + padded_T - T, + self.audio_channels, + dtype=torch.int32, + device=audio.device, + ) + + audio[-1, :], + ], + dim=0, + ) # pad using the last embedding + padded_audio = padded_audio.reshape( + padded_T // self.audio_group_size, + self.audio_group_size, + self.audio_channels, + ) + return padded_audio + + def get_audio_feature(self, items) -> torch.Tensor: + # items: already audio-only MultimodalDataItem list from caller. + # Each item.feature is either one mel tensor or a list of mel tensors (e.g. long audio split into chunks). + all_mels = [] + for item in items: + f = item.feature + if isinstance(f, (list, tuple)): + all_mels.extend(f) + else: + all_mels.append(f) + if not all_mels: + device = next(self.projection.parameters()).device + dtype = next(self.projection.parameters()).dtype + return torch.empty( + 0, self.config.out_hidden_size, device=device, dtype=dtype + ) + # Batch tokenize: one encode_batch call for all mels + device = next(self.audio_tokenizer.encoder.parameters()).device + code_list = tokenize_audio_batch( + all_mels, + self.audio_tokenizer.encoder, + segment_size=self.audio_segment_size, + device=device, + ) + codecs_to_concat = [] + for codecs in code_list: + padded_codes = self.process_audio( + codecs + ) # [T//group_size, group_size, audio_channels] + codecs_to_concat.append(padded_codes) + audio_codes = torch.cat( + codecs_to_concat, dim=0 + ) # [T//group_size, group_size, audio_channels] + + _audio_embeddings = self.apply_speech_embeddings(audio_codes) + audio_embeds = self.apply_input_local_transformer( + _audio_embeddings + ) # [T//group_size, group_size, input_local_dim] + B = audio_embeds.shape[0] + audio_embeds = self.projection(audio_embeds.reshape(B, -1)) + return audio_embeds diff --git a/python/sglang/srt/models/mimo_v2.py b/python/sglang/srt/models/mimo_v2.py index 7a842aa59..5cc20c384 100644 --- a/python/sglang/srt/models/mimo_v2.py +++ b/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( - input_ids, - positions, - forward_batch, - input_embeds, - pp_proxy_tensors=pp_proxy_tensors, - ) + if self._is_multimodal: + hidden_states, hidden_states_before_norm = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + 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] - default_weight_loader(param, loaded_weight) + expected_fused_tp_size = get_mimo_v2_fused_qkv_expected_tp_size( + 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: diff --git a/python/sglang/srt/models/mimo_v2_nextn.py b/python/sglang/srt/models/mimo_v2_nextn.py index ad81b69e4..737d8ec7d 100644 --- a/python/sglang/srt/models/mimo_v2_nextn.py +++ b/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] - default_weight_loader(param, loaded_weight) + load_mimo_v2_qkv_proj_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 diff --git a/python/sglang/srt/models/mimo_vl.py b/python/sglang/srt/models/mimo_vl.py new file mode 100644 index 000000000..a7b1fb7c3 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/multimodal/processors/mimo_v2.py b/python/sglang/srt/multimodal/processors/mimo_v2.py new file mode 100644 index 000000000..b4bed854b --- /dev/null +++ b/python/sglang/srt/multimodal/processors/mimo_v2.py @@ -0,0 +1,2039 @@ +"""MiMoV2 multimodal processor -- protocol, utilities, and processor.""" + +import asyncio +import base64 +import copy +import io +import json +import math +import os +import re +import subprocess +import time +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from io import BytesIO +from typing import List, Literal, Optional, Union + +import numpy as np +import pybase64 +import requests +import torch +import torch.nn.functional as F +from fastapi import HTTPException +from PIL import Image +from torchcodec.decoders import AudioDecoder +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( + Qwen2_5_VLVisionConfig, +) + +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalProcessorOutput, +) +from sglang.srt.models.mimo_v2 import MiMoV2ForCausalLM +from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor, + MultimodalSpecialTokens, +) +from sglang.srt.multimodal.processors.qwen_vl import smart_nframes +from sglang.srt.utils import ImageData, VideoData +from sglang.utils import logger + +try: + import torchaudio + from torchaudio.transforms import MelSpectrogram +except ImportError: + logger.warning( + "torchaudio is not installed; audio inputs will fail at request time" + ) + torchaudio = None + MelSpectrogram = None + + +@dataclass +class ImageInput: + image: Image.Image | str | bytes | torch.Tensor + max_pixels: Optional[int] = None + min_pixels: Optional[int] = None + + def __post_init__(self): + if not isinstance(self.image, (Image.Image, str, bytes, torch.Tensor)): + raise ValueError( + f"image must be a PIL.Image.Image, str, bytes, or torch.Tensor, but got {type(self.image)}" + ) + + +@dataclass +class VideoInput: + video: str | bytes | tuple[torch.Tensor, torch.Tensor] + min_pixels: Optional[int] = None + max_pixels: Optional[int] = None + total_max_pixels: Optional[int] = None + fps: Optional[float] = None + num_frames: Optional[int] = None + max_frames: Optional[int] = None + min_frames: Optional[int] = None + do_include_last_frame: Optional[bool] = False + start_time: Optional[float] = None + end_time: Optional[float] = None + segment_type: Literal["individual", "partial"] = "individual" + + def __post_init__(self): + if not isinstance(self.video, (str, bytes, tuple)): + raise ValueError( + f"video must be a str, bytes, or tuple, but got {type(self.video)}" + ) + if isinstance(self.video, tuple): + if len(self.video) != 2: + raise ValueError( + f"video must be a tuple of 2 elements (pixels, timestamps), but got {len(self.video)} elements" + ) + if not isinstance(self.video[0], torch.Tensor) or not isinstance( + self.video[1], torch.Tensor + ): + raise ValueError( + f"video must be a tuple of Tensors (pixels, timestamps), but got {type(self.video[0])} and {type(self.video[1])}" + ) + if ( + self.video[0].ndim != 4 + or self.video[1].ndim != 1 + or self.video[0].shape[0] != self.video[1].shape[0] + ): + raise ValueError( + f"video must be a tuple of (pixels-TCHW, timestamps-T), but got {self.video[0].shape} and {self.video[1].shape}" + ) + assert self.segment_type in ["individual", "partial"] + assert self.segment_type == "partial" or ( + self.start_time is None and self.end_time is None + ) + + +@dataclass +class AudioInput: + """ + if audio is str or bytes, only load it as mel spectrogram. + if audio is tuple, it is (waveform, original_sr) + if audio is torch.Tensor, it is tokenized input ids with shape (T, n_vq+). + if audio is np.ndarray, it is a pre-loaded waveform (1D, already resampled). + """ + + audio: str | bytes | tuple | torch.Tensor | np.ndarray + + def __post_init__(self): + if not isinstance(self.audio, (str, bytes, tuple, torch.Tensor, np.ndarray)): + raise ValueError( + f"audio must be a str, bytes, tuple, torch.Tensor, or np.ndarray, but got {type(self.audio)}" + ) + if isinstance(self.audio, tuple): + if ( + len(self.audio) != 2 + or not isinstance(self.audio[0], torch.Tensor) + or not isinstance(self.audio[1], (int, float)) + ): + raise ValueError( + f"audio must be a tuple of (waveform-T, original_sr-int/float), but got {len(self.audio)} elements and {type(self.audio[0])} and {type(self.audio[1])}" + ) + if self.audio[0].ndim != 1: + raise ValueError( + f"waveform must be a 1D tensor, but got {self.audio[0].ndim}D tensor" + ) + if self.audio[1] <= 0: + raise ValueError( + f"original_sr must be a positive number, but got {self.audio[1]}" + ) + if isinstance(self.audio, torch.Tensor) and self.audio.ndim != 2: + raise ValueError( + f"audio must be a 2D tensor, but got {self.audio.ndim}D tensor" + ) + + +@dataclass +class VideoAudioInput: + video: str | bytes | tuple[torch.Tensor, torch.Tensor] + audio: str | bytes | torch.Tensor + min_pixels: Optional[int] = None + max_pixels: Optional[int] = None + total_max_pixels: Optional[int] = None + fps: Optional[float] = None + num_frames: Optional[int] = None + max_frames: Optional[int] = None + min_frames: Optional[int] = None + do_include_last_frame: Optional[bool] = False + start_time: Optional[float] = None + end_time: Optional[float] = None + segment_type: Literal["individual", "partial"] = "individual" + + def __post_init__(self): + if not isinstance(self.video, (str, bytes, tuple)): + raise ValueError( + f"video must be a str, bytes, or tuple, but got {type(self.video)}" + ) + if isinstance(self.video, tuple): + if len(self.video) != 2: + raise ValueError( + f"video must be a tuple of 2 elements (pixels, timestamps), but got {len(self.video)} elements" + ) + if not isinstance(self.video[0], torch.Tensor) or not isinstance( + self.video[1], torch.Tensor + ): + raise ValueError( + f"video must be a tuple of Tensors (pixels, timestamps), but got {type(self.video[0])} and {type(self.video[1])}" + ) + if ( + self.video[0].ndim != 4 + or self.video[1].ndim != 1 + or self.video[0].shape[0] != self.video[1].shape[0] + ): + raise ValueError( + f"video must be a tuple of (pixels-TCHW, timestamps-T), but got {self.video[0].shape} and {self.video[1].shape}" + ) + assert self.segment_type in ["individual", "partial"] + assert self.segment_type == "partial" or ( + self.start_time is None and self.end_time is None + ) + + if not isinstance(self.audio, (str, bytes, torch.Tensor)): + raise ValueError( + f"audio must be a str, bytes, or torch.Tensor, but got {type(self.audio)}" + ) + if isinstance(self.audio, torch.Tensor) and self.audio.ndim != 2: + raise ValueError( + f"audio must be a 2D tensor, but got {self.audio.ndim}D tensor" + ) + + +TextInput = str | list[int] + + +@dataclass +class MiMoInputSample: + input_ids: torch.Tensor + labels: Optional[torch.Tensor] + pixel_values: list[torch.Tensor] + pixel_values_videos: list[torch.Tensor] + image_thw_grids: list[torch.Tensor] + video_thw_grids: list[torch.Tensor] + audio_inputs: list[torch.Tensor] + position_ids: Optional[torch.Tensor] = None + rope_deltas: Optional[torch.Tensor] = None + extra: dict = field(default_factory=dict) + + +@dataclass +class Content: + type: Literal["text", "image", "video", "audio", "video_audio"] + content: TextInput | ImageInput | VideoInput | AudioInput | VideoAudioInput + is_target: Optional[bool] = None + + def __post_init__(self): + if self.type not in ["text", "image", "video", "audio", "video_audio"]: + raise ValueError( + f"type must be one of text, image, video, audio, video_audio, but got {self.type}" + ) + if self.type == "text": + if not isinstance(self.content, (str, list)) or ( + isinstance(self.content, list) + and not all(isinstance(item, int) for item in self.content) + ): + raise ValueError( + f"content must be a str or a list of ints, but got {type(self.content)}" + ) + elif self.type == "image": + if not isinstance(self.content, ImageInput): + raise ValueError( + f"content must be a ImageInput, but got {type(self.content)}" + ) + elif self.type == "video": + if not isinstance(self.content, VideoInput): + raise ValueError( + f"content must be a VideoInput, but got {type(self.content)}" + ) + elif self.type == "audio": + if not isinstance(self.content, AudioInput): + raise ValueError( + f"content must be a AudioInput, but got {type(self.content)}" + ) + elif self.type == "video_audio": + if not isinstance(self.content, VideoAudioInput): + raise ValueError( + f"content must be a VideoAudioInput, but got {type(self.content)}" + ) + + +_QWEN2VL_PIXEL_MEAN = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1) +_QWEN2VL_PIXEL_STD = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1) +_mean_std_cache = {} + + +class MiMoProcessor: + def __init__( + self, + tokenizer, + patch_size=14, + merge_size=2, + temporal_patch_size=2, + temporal_compression_ratio=1, + video_tokens_per_second=2, + use_video_timestamps=False, + video_audio_interleave_length=0, + use_per_grid_t_timestamps=True, + audio_kernel_size=3, + audio_stride_size=2, + audio_avg_pooler=2, + audio_sampling_rate=24000, + audio_nfft=960, + audio_hop_length=240, + audio_window_size=960, + audio_fmin=0, + audio_fmax=None, + audio_n_mels=128, + audio_segment_size=6000, + audio_channels=8, + audio_group_size=4, + audio_input_id_per_second=25, + audio_zeroemb_idx=4096, + image_min_pixels=None, + image_max_pixels=None, + video_min_pixels=None, + video_max_pixels=None, + video_total_max_pixels=None, + fps=None, + num_frames=None, + max_frames=None, + min_frames=None, + image_token_id=None, + video_token_id=None, + audio_token_id=None, + vision_start_token_id=None, + vision_end_token_id=None, + audio_start_token_id=None, + audio_end_token_id=None, + video_start_token_id=None, + video_end_token_id=None, + pad_token_id=None, + rope_type="rope", + video_process_num_threads=16, + device=None, + **kwargs, + ): + self.tokenizer = tokenizer + self.video_process_num_threads = video_process_num_threads + + if device is None: + self.device = None + else: + self.device = torch.device(device) if isinstance(device, str) else device + + self.rope_type = rope_type + if self.rope_type == "1d": + self.rope_type = "rope" + assert self.rope_type in ["rope", "mrope"] + + self.use_video_timestamps = use_video_timestamps + assert self.use_video_timestamps + assert ( + not self.use_video_timestamps or self.rope_type == "rope" + ), "use_video_timestamps only supports 1d rope" + self.video_audio_interleave_length = video_audio_interleave_length + self.use_per_grid_t_timestamps = False + assert ( + self.video_audio_interleave_length == -1 or self.rope_type == "rope" + ), "video_audio_interleave_length != -1 only supports 1d rope" + assert ( + self.video_audio_interleave_length == -1 + or self.video_audio_interleave_length >= 0 + ) + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.audio_token_id = audio_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + self.audio_start_token_id = audio_start_token_id + self.audio_end_token_id = audio_end_token_id + self.video_start_token_id = video_start_token_id + self.video_end_token_id = video_end_token_id + self.pad_token_id = pad_token_id + + self.patch_size = patch_size + self.merge_size = merge_size + self.temporal_patch_size = temporal_patch_size + self.temporal_compression_ratio = temporal_compression_ratio + + self.video_tokens_per_second = video_tokens_per_second + + self.audio_sampling_rate = audio_sampling_rate + self.audio_nfft = audio_nfft + self.audio_hop_length = audio_hop_length + self.audio_window_size = audio_window_size + self.audio_fmin = audio_fmin + self.audio_fmax = audio_fmax + self.audio_n_mels = audio_n_mels + + self.audio_segment_size = audio_segment_size + + self.audio_kernel_size = audio_kernel_size + self.audio_stride_size = audio_stride_size + self.audio_avg_pooler = audio_avg_pooler + + self.mel_spectrogram_kwargs = dict( + sample_rate=audio_sampling_rate, + n_fft=audio_nfft, + hop_length=audio_hop_length, + win_length=audio_window_size, + f_min=audio_fmin, + f_max=audio_fmax, + n_mels=audio_n_mels, + power=1.0, + center=True, + ) + self._mel_spectrogram = None + self._resamplers = OrderedDict() + self._resamplers_max = 16 + + self.audio_channels = audio_channels + self.audio_group_size = audio_group_size + self.audio_input_id_per_second = audio_input_id_per_second + if isinstance(audio_zeroemb_idx, int): + self.audio_zeroemb_idxs = torch.tensor( + [audio_zeroemb_idx] * self.audio_channels, dtype=torch.int32 + ) + elif isinstance(audio_zeroemb_idx, list): + if len(audio_zeroemb_idx) == 1: + self.audio_zeroemb_idxs = torch.tensor( + audio_zeroemb_idx * self.audio_channels, dtype=torch.int32 + ) + elif len(audio_zeroemb_idx) == self.audio_channels: + self.audio_zeroemb_idxs = torch.tensor( + audio_zeroemb_idx, dtype=torch.int32 + ) + else: + raise ValueError( + f"audio_zeroemb_idx must be a list of 1 or {self.audio_channels} integers, but got {len(audio_zeroemb_idx)}" + ) + else: + raise ValueError( + f"audio_zeroemb_idx must be an integer or a list of {self.audio_channels} integers, but got {type(audio_zeroemb_idx)}" + ) + + assert image_min_pixels is not None + assert image_max_pixels is not None + assert video_min_pixels is not None + assert video_max_pixels is not None + assert video_total_max_pixels is not None + assert fps is not None or num_frames is not None + + self.default_image_processor_kwargs = { + "min_pixels": image_min_pixels, + "max_pixels": image_max_pixels, + } + + self.default_video_processor_kwargs = { + "min_pixels": video_min_pixels, + "max_pixels": video_max_pixels, + "total_max_pixels": video_total_max_pixels, + "fps": fps, + "num_frames": num_frames, + "max_frames": max_frames, + "min_frames": min_frames, + } + + self.http_session = requests.Session() + for k in kwargs: + logger.info(f"[Warning] Ignored unknown parameter {k} for MiMoProcessor") + + @property + def mel_spectrogram(self): + self._ensure_audio_dependencies() + if self._mel_spectrogram is None: + self._mel_spectrogram = MelSpectrogram(**self.mel_spectrogram_kwargs) + return self._mel_spectrogram + + @staticmethod + def _ensure_audio_dependencies(): + if torchaudio is None or MelSpectrogram is None: + raise RuntimeError( + "torchaudio is required for audio inputs; install torchaudio" + ) + + def prepare_image_kwargs(self, image: ImageInput): + kwargs = {} + for k in ["min_pixels", "max_pixels"]: + if getattr(image, k) is not None: + kwargs[k] = getattr(image, k) + else: + kwargs[k] = self.default_image_processor_kwargs[k] + return kwargs + + def prepare_video_kwargs(self, video: VideoInput | VideoAudioInput): + kwargs = {} + for k in ["min_pixels", "max_pixels", "total_max_pixels"]: + if getattr(video, k) is not None: + kwargs[k] = getattr(video, k) + else: + kwargs[k] = self.default_video_processor_kwargs[k] + if video.num_frames is not None: + kwargs["num_frames"] = video.num_frames + elif video.fps is not None: + kwargs["fps"] = video.fps + if video.max_frames is not None: + kwargs["max_frames"] = video.max_frames + if video.min_frames is not None: + kwargs["min_frames"] = video.min_frames + elif self.default_video_processor_kwargs["num_frames"] is not None: + kwargs["num_frames"] = self.default_video_processor_kwargs["num_frames"] + elif self.default_video_processor_kwargs["fps"] is not None: + kwargs["fps"] = self.default_video_processor_kwargs["fps"] + if self.default_video_processor_kwargs["max_frames"] is not None: + kwargs["max_frames"] = self.default_video_processor_kwargs["max_frames"] + if self.default_video_processor_kwargs["min_frames"] is not None: + kwargs["min_frames"] = self.default_video_processor_kwargs["min_frames"] + else: + raise ValueError("Video sampling strategy not specified") + return kwargs + + def preprocess_audio(self, audio: str | bytes): + self._ensure_audio_dependencies() + """ + - Input: audio filename string, bytes, or tuple of (waveform, original_sr) + - Output: + - mel spectrogram: torch.Tensor (T, n_mels) + - number of tokens: int + """ + assert isinstance( + audio, (str, bytes, tuple) + ), f"audio must be a str, bytes or tuple, but got {type(audio)}" + if isinstance(audio, tuple): + waveform, original_sr = audio + else: + if isinstance(audio, bytes): + file = io.BytesIO(audio) + elif isinstance(audio, str): + if audio.startswith("data:"): + file = io.BytesIO( + pybase64.b64decode(audio.split(",")[1], validate=True) + ) + elif audio.startswith("http://") or audio.startswith("https://"): + dl_start = time.perf_counter() + timeout = int(os.getenv("REQUEST_TIMEOUT", "5")) + try: + response = self.http_session.get( + audio, stream=True, timeout=timeout + ) + dl_elapsed_ms = (time.perf_counter() - dl_start) * 1000 + if dl_elapsed_ms > 1000.0: + content_len = len(response.content) + logger.warning( + f"Slow audio download: {dl_elapsed_ms:.2f}ms, " + f"size={content_len / 1024:.1f}KB, url={audio}" + ) + file = io.BytesIO(response.content) + response.close() + except Exception as e: + dl_elapsed_ms = (time.perf_counter() - dl_start) * 1000 + logger.error( + f"Failed to download audio: {dl_elapsed_ms:.2f}ms, " + f"error={type(e).__name__}: {e}, url={audio}" + ) + raise + else: + file = audio + try: + samples = AudioDecoder(file).get_all_samples() + except RuntimeError as e: + audio_source = ( + audio + if isinstance(audio, str) + and (audio.startswith("http://") or audio.startswith("https://")) + else "" + ) + logger.error(f"Failed to decode audio: {e}, source={audio_source}") + raise ValueError( + f"Invalid audio format: source={audio_source}, detail={e}" + ) from e + waveform = samples.data + original_sr = samples.sample_rate + + if original_sr != self.audio_sampling_rate: + if original_sr in self._resamplers: + self._resamplers.move_to_end(original_sr) + else: + if len(self._resamplers) >= self._resamplers_max: + self._resamplers.popitem(last=False) + self._resamplers[original_sr] = torchaudio.transforms.Resample( + orig_freq=original_sr, new_freq=self.audio_sampling_rate + ) + waveform = self._resamplers[original_sr](waveform) + if waveform.ndim == 2: + waveform = waveform.mean(dim=0) + spec = self.mel_spectrogram(waveform[None, :]) + spec = torch.log(torch.clip(spec, min=1e-7)).squeeze() + spec = spec.transpose(0, 1) + + audio_token_len = spec.shape[0] + 3 - self.audio_kernel_size + audio_token_len = ( + audio_token_len + 2 - self.audio_kernel_size + ) // self.audio_stride_size + 1 + audio_token_len = audio_token_len // self.audio_avg_pooler + int( + audio_token_len % self.audio_avg_pooler != 0 + ) + audio_token_len = math.ceil(audio_token_len / self.audio_group_size) + + return spec, audio_token_len + + def process_image(self, image: ImageInput): + kwargs = self.prepare_image_kwargs(image) + image = image.image + if isinstance(image, (str, bytes)): + image = self.fetch_image(image) + image_transformed_tensor, _, _ = self.get_visual_transform( + image, + factor=self.patch_size * self.merge_size, + min_pixels=kwargs["min_pixels"], + max_pixels=kwargs["max_pixels"], + device=self.device, + ) + return image_transformed_tensor + + def process_video( + self, video_input: VideoInput | VideoAudioInput, temporal_padding_factor=None + ): + + def smart_resize_video( + num_total_frames, min_pixels, max_pixels, total_max_pixels, **kwargs + ): + max_pixels_per_frame = ( + total_max_pixels + * self.temporal_patch_size + * self.temporal_compression_ratio + // num_total_frames + ) + max_pixels = max(min_pixels, min(max_pixels_per_frame, max_pixels)) + return min_pixels, max_pixels + + def segment_frame_selector(all_timestamps, start_time, end_time): + """Select frame indices in [start_time, end_time). If none found, pick the nearest frame to the left.""" + if not isinstance(all_timestamps, torch.Tensor): + all_timestamps = torch.tensor(all_timestamps) + + mask = (all_timestamps >= start_time) & (all_timestamps < end_time) + candidate_indices = torch.where(mask)[0] + + if len(candidate_indices) == 0: + left_mask = all_timestamps <= start_time + left_indices = torch.where(left_mask)[0] + if len(left_indices) > 0: + selected_frame_indices = left_indices[-1:].clone() + else: + raise ValueError( + f"No frames before start_time {start_time} in all_timestamps {all_timestamps.tolist()}" + ) + else: + selected_frame_indices = candidate_indices + + assert ( + len(selected_frame_indices) > 0 + ), f"No frames selected for segment {start_time} - {end_time} in all_timestamps {all_timestamps.tolist()}" + return selected_frame_indices + + kwargs = self.prepare_video_kwargs(video_input) + video = video_input.video + + if not isinstance(video, tuple): + raise ValueError( + f"video must be a tuple of (video_tensor, timestamps), but got {type(video)}. " + "Video download and decoding should be done by sglang load_video before calling process_video." + ) + + video_tensor, timestamps_sampled = video + if len(timestamps_sampled) < 2: + logger.info( + "[Warning] Less than two frames are sampled, using default fps (1 fps)" + ) + fps_sampled = 1 + else: + fps_sampled = 1 / (timestamps_sampled[1] - timestamps_sampled[0]) + num_frames_sampled = video_tensor.shape[0] + + start_time = ( + video_input.start_time + if video_input.start_time is not None + else timestamps_sampled[0] + ) + end_time = ( + video_input.end_time + if video_input.end_time is not None + else timestamps_sampled[-1] + (1 / fps_sampled) + ) + + if video_input.segment_type == "individual": + start_time_seg = start_time + end_time_seg = end_time + timestamps_seg = timestamps_sampled + frames = video_tensor + num_frames_seg = num_frames_sampled + else: + selected_indices = segment_frame_selector( + timestamps_sampled, start_time, end_time + ) + + timestamps_seg = timestamps_sampled[selected_indices] + frames = video_tensor[selected_indices] + num_frames_seg = len(timestamps_seg) + start_time_seg = ( + timestamps_seg[0].item() + if isinstance(timestamps_seg[0], torch.Tensor) + else timestamps_seg[0] + ) + end_time_seg = ( + timestamps_seg[-1].item() + if isinstance(timestamps_seg[-1], torch.Tensor) + else timestamps_seg[-1] + ) + (1 / fps_sampled).item() + + video_meta = { + "fps_sampled": fps_sampled, + "segment_start_time": start_time_seg, + "segment_end_time": end_time_seg, + } + + min_pixels, max_pixels = smart_resize_video(num_frames_sampled, **kwargs) + + assert ( + num_frames_seg > 0 + ), f"Sampled frame number must be >0. start_time {video_input.start_time}, end_time {video_input.end_time}, start_time_seg {start_time_seg}, end_time_seg {end_time_seg}. Full timestamps {timestamps_sampled.tolist()}. " + + temporal_padding_factor = ( + self.temporal_patch_size * self.temporal_compression_ratio + if temporal_padding_factor is None + else temporal_padding_factor + ) + + if num_frames_seg % temporal_padding_factor == 0: + aligned_frames = frames + aligned_timestamps = timestamps_seg + else: + aligned_num_frames = ( + (num_frames_seg + temporal_padding_factor - 1) + // temporal_padding_factor + ) * temporal_padding_factor + num_frames_needed = aligned_num_frames - num_frames_seg + aligned_frames = torch.cat( + [ + frames, + frames[-1:].repeat(num_frames_needed, *[1] * (frames.ndim - 1)), + ], + dim=0, + ) + aligned_timestamps = torch.cat( + [timestamps_seg, timestamps_seg[-1:].repeat(num_frames_needed)], dim=0 + ) + + video_transformed_tensor, _, _ = self.get_visual_transform_batch( + aligned_frames, + factor=self.patch_size * self.merge_size, + min_pixels=min_pixels, + max_pixels=max_pixels, + device=self.device, + ) + + visual_patches, thw_grid = self._flatten_visual_inputs( + video_transformed_tensor, "video" + ) + return visual_patches, thw_grid, aligned_timestamps, video_meta + + def process_audio(self, audio: AudioInput): + audio = audio.audio + if isinstance(audio, np.ndarray): + waveform = torch.from_numpy(audio).float() + audio = (waveform, self.audio_sampling_rate) + if isinstance(audio, (str, bytes, tuple)): + audio_spec, audio_token_len = self.preprocess_audio(audio) + return audio_spec, audio_token_len + + assert ( + audio.shape[1] >= self.audio_channels + ), f"audio must have at least {self.audio_channels} channels, but got {audio.shape[1]}" + T = audio.shape[0] + audio = audio[:, : self.audio_channels].to(torch.long) + padded_T = ( + (T + self.audio_group_size - 1) + // self.audio_group_size + * self.audio_group_size + ) + padded_audio = torch.cat( + [ + audio, + torch.zeros(padded_T - T, self.audio_channels, dtype=torch.long) + + audio[-1, :], + ], + dim=0, + ) + padded_audio = padded_audio.reshape( + padded_T // self.audio_group_size, + self.audio_group_size, + self.audio_channels, + ) + return padded_audio + + def _process_videos_parallel(self, contents): + video_contents_info = [] + for idx, content in enumerate(contents): + if content.type in ("video", "video_audio"): + video_contents_info.append((idx, content.content)) + + video_results = {} + if not video_contents_info: + return video_results + + num_threads = min(self.video_process_num_threads, len(video_contents_info)) + if num_threads > 1 and len(video_contents_info) > 1: + with ThreadPoolExecutor(max_workers=num_threads) as executor: + future_to_idx = { + executor.submit(self.process_video, video_input): idx + for idx, video_input in video_contents_info + } + for future in as_completed(future_to_idx): + idx = future_to_idx[future] + try: + video_results[idx] = future.result() + except Exception as e: + raise RuntimeError( + f"Error processing video at index {idx}: {e}" + ) from e + else: + for idx, video_input in video_contents_info: + video_results[idx] = self.process_video(video_input) + return video_results + + def _process_text_content(self, content, verbose): + if isinstance(content.content, str): + _input_ids = self.tokenizer.encode(content.content) + else: + _input_ids = content.content + _labels = _input_ids if content.is_target else None + + verbose_str = "" + if verbose: + if isinstance(content.content, str): + verbose_str = f"Text: [{content.content}]\n" + else: + verbose_str = f"Text: [{self.tokenizer.decode(content.content)}]\n" + + return {"input_ids": _input_ids, "labels": _labels, "verbose": verbose_str} + + def _process_image_content(self, content, verbose): + image_tensor = self.process_image(content.content) + visual_patches, thw_grid = self._flatten_visual_inputs(image_tensor, "image") + grid_t, grid_h, grid_w = thw_grid + num_media_tokens = (grid_t * grid_h * grid_w) // (self.merge_size**2) + _input_ids = ( + [self.vision_start_token_id] + + [self.image_token_id] * num_media_tokens + + [self.vision_end_token_id] + ) + + verbose_str = "" + if verbose: + verbose_str = f"Image (shape={image_tensor.shape}, image_thw_grid={thw_grid}): [ {num_media_tokens}* ]\n" + + return { + "input_ids": _input_ids, + "pixel_values": visual_patches, + "thw_grid": thw_grid, + "verbose": verbose_str, + } + + def _process_video_content(self, content_idx, video_results, verbose): + visual_patches, thw_grid, timestamps, video_meta = video_results[content_idx] + grid_t, grid_h, grid_w = thw_grid + num_media_tokens = ( + (grid_t * grid_h * grid_w) + // (self.merge_size**2) + // self.temporal_compression_ratio + ) + + assert ( + len(timestamps) == grid_t * self.temporal_patch_size + ), f"Expected {grid_t} * {self.temporal_patch_size} = {grid_t * self.temporal_patch_size} timestamps, but got {len(timestamps)}" + + if not self.use_video_timestamps: + raise NotImplementedError + + num_media_tokens_per_grid = grid_h * grid_w // (self.merge_size**2) + text_timestamps = [ + self.format_timestamp(ts) + for ts in timestamps[ + :: self.temporal_patch_size * self.temporal_compression_ratio + ] + ] + text_timestamp_ids = [self.tokenizer.encode(ts) for ts in text_timestamps] + _input_ids = ( + [self.video_start_token_id] + + sum( + [ + ts_ids + + [self.vision_start_token_id] + + [self.video_token_id] * num_media_tokens_per_grid + + [self.vision_end_token_id] + for ts_ids in text_timestamp_ids + ], + [], + ) + + [self.video_end_token_id] + ) + + verbose_str = "" + if verbose: + verbose_str = f"Video (video_thw_grid={thw_grid}, video_meta={video_meta}): [ " + for i, ts in enumerate(text_timestamps): + verbose_str += f"{ts} {timestamps.tolist()[i*self.temporal_patch_size*self.temporal_compression_ratio : (i+1)*self.temporal_patch_size*self.temporal_compression_ratio]} {num_media_tokens_per_grid}* " + verbose_str += "]\n" + + return { + "input_ids": _input_ids, + "pixel_values": visual_patches, + "thw_grid": thw_grid, + "second_per_grid_t": self.temporal_patch_size / video_meta["fps_sampled"], + "verbose": verbose_str, + } + + def _process_audio_content(self, content, verbose): + processed_audio = self.process_audio(content.content) + if isinstance(processed_audio, tuple): + 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 = processed_audio + _input_ids = ( + [self.audio_start_token_id] + + [self.audio_token_id] * audio_token_len + + [self.audio_end_token_id] + ) + + verbose_str = "" + if verbose: + verbose_str = f"Audio (is_tokenized={is_tokenized}): [ {audio_token_len}*