chore: import upstream snapshot with attribution

python/sglang/srt/models/mimo_vl.py

@@ -0,0 +1,515 @@
+"""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.runtime_context import get_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_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,
+            # MiMo-VL's merger MLP is square (intermediate == context_dim * merge**2),
+            # so no dim padding is needed. The Qwen2.5-VL formula num_heads * head_dim
+            # over-sizes it here because MiMo uses qk_channels (64) for head_dim rather
+            # than hidden_size // num_heads, which would mismatch the checkpoint.
+            padded_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 = int(grid_thw[:, 1:].max())
+        # transformers 5.12's rotary forward takes 1-D position_ids on the input device (grid_thw is CPU).
+        rotary_pos_emb_full = self.rotary_pos_emb(
+            torch.arange(max_grid_size, device=self.device)
+        )
+        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,
+        )