sgl-project--sglang/python/sglang/srt/models/minimax_vl_common.py

# SPDX-License-Identifier: Apache-2.0

import logging
from dataclasses import dataclass, fields
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn

from sglang.srt.layers.activation import get_act_fn
from sglang.srt.layers.attention.vision import (
    BATCH_BUCKETS,
    FLASHINFER_MAX_SEQLEN_BUCKETS,
    FLASHINFER_WORKSPACE_SIZE_BYTES,
    VisionAttention,
)
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.linear import (
    ColumnParallelLinear,
    RowParallelLinear,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.rotary_embedding.utils import rotate_half
from sglang.srt.managers.schedule_batch import MultimodalDataItem
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.utils import add_prefix, get_compiler_backend, round_up

logger = logging.getLogger(__name__)


@dataclass
class CLIPVisionConfig:
    hidden_size: int
    intermediate_size: int
    num_hidden_layers: int
    num_attention_heads: int
    num_channels: int
    image_size: int
    patch_size: int
    hidden_act: str
    layer_norm_eps: float
    img_token_compression_config: dict
    position_embedding_type: str
    rope_mode: str
    rope_theta: float
    vision_segment_max_frames: Optional[int]

    @classmethod
    def from_dict(cls, d: dict) -> "CLIPVisionConfig":
        valid_keys = {f.name for f in fields(cls)}
        filtered = {k: v for k, v in d.items() if k in valid_keys}
        if "rope_theta" not in filtered and isinstance(d.get("rope_parameters"), dict):
            rope_theta = d["rope_parameters"].get("rope_theta")
            if rope_theta is not None:
                filtered["rope_theta"] = rope_theta
        return cls(**filtered)


class MiniMaxVLMultiModalProjector(nn.Module):
    def __init__(
        self,
        vision_hidden_size: int,
        text_hidden_size: int,
        projector_hidden_act: str,
        multimodal_projector_bias: bool,
        projector_hidden_size: Optional[int] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        use_data_parallel: bool = False,
    ):
        super().__init__()

        mid_size = (
            projector_hidden_size
            if projector_hidden_size is not None
            else text_hidden_size
        )

        tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size
        tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank

        self.linear_1 = ColumnParallelLinear(
            vision_hidden_size,
            mid_size,
            bias=multimodal_projector_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.linear_1",
            tp_size=tp_size,
            tp_rank=tp_rank,
        )
        assert (
            projector_hidden_act == "gelu"
        ), f"Only gelu activation is supported, got {projector_hidden_act}"
        self.act = get_act_fn(projector_hidden_act)
        self.linear_2 = RowParallelLinear(
            mid_size,
            text_hidden_size,
            bias=multimodal_projector_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.linear_2",
            tp_size=tp_size,
            tp_rank=tp_rank,
            use_dp_attention_reduce=is_dp_attention_enabled(),
        )

    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
        hidden_states, _ = self.linear_1(image_features)
        hidden_states = self.act(hidden_states)
        hidden_states, _ = self.linear_2(hidden_states)
        return hidden_states


class MiniMaxVLPatchMerger(nn.Module):
    def __init__(
        self,
        spatial_merge_size: int,
        text_hidden_size: int,
        projector_hidden_act: str,
        patch_merge_bias: bool,
        projector_hidden_size: Optional[int] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        use_data_parallel: bool = False,
    ):
        super().__init__()
        self.spatial_merge_size = spatial_merge_size

        mid_size = (
            projector_hidden_size
            if projector_hidden_size is not None
            else text_hidden_size
        )

        tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size
        tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank

        self.linear_1 = ColumnParallelLinear(
            text_hidden_size * spatial_merge_size**2,
            mid_size,
            bias=patch_merge_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.linear_1",
            tp_size=tp_size,
            tp_rank=tp_rank,
        )
        assert (
            projector_hidden_act == "gelu"
        ), f"Only gelu activation is supported, got {projector_hidden_act}"
        self.act = get_act_fn(projector_hidden_act)
        self.linear_2 = RowParallelLinear(
            mid_size,
            text_hidden_size,
            bias=patch_merge_bias,
            quant_config=quant_config,
            prefix=f"{prefix}.linear_2",
            tp_size=tp_size,
            tp_rank=tp_rank,
            use_dp_attention_reduce=is_dp_attention_enabled(),
        )

    def forward(self, image_features: torch.Tensor) -> torch.Tensor:
        image_features = image_features.reshape(
            image_features.shape[0] // (self.spatial_merge_size**2), -1
        )
        hidden_states, _ = self.linear_1(image_features)
        hidden_states = self.act(hidden_states)
        hidden_states, _ = self.linear_2(hidden_states)
        return hidden_states


def _prepare_rotary_cos_sin(
    freqs: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    cos = freqs.cos().repeat(1, 2).unsqueeze(-2).float()
    sin = freqs.sin().repeat(1, 2).unsqueeze(-2).float()
    return cos, sin


@torch.compile(dynamic=True, backend=get_compiler_backend())
def _minimax_rope_applier(
    q: torch.Tensor,
    k: torch.Tensor,
    position_embeddings: Tuple[torch.Tensor, torch.Tensor],
    x_shape=None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """3D RoPE uses rope_dim=60 < head_dim=64; trailing dims pass through unrotated."""
    cos, sin = position_embeddings
    rot_dim = cos.shape[-1]

    q_rot = q[..., :rot_dim].float()
    q_pass = q[..., rot_dim:]
    k_rot = k[..., :rot_dim].float()
    k_pass = k[..., rot_dim:]

    q_rot = (q_rot * cos) + (rotate_half(q_rot) * sin)
    k_rot = (k_rot * cos) + (rotate_half(k_rot) * sin)

    q = torch.cat((q_rot.to(q_pass.dtype), q_pass), dim=-1)
    k = torch.cat((k_rot.to(k_pass.dtype), k_pass), dim=-1)
    return q, k


class CLIPVisionEmbeddings(nn.Module):
    def __init__(self, config: CLIPVisionConfig):
        super().__init__()

        self.config = config
        self.embed_dim = config.hidden_size
        self.patch_size = config.patch_size
        self.input_num_channels = config.num_channels

        self.temporal_patch_size = config.img_token_compression_config.get(
            "temporal_patch_size", 2
        )

        self.patch_embedding = nn.Conv3d(
            in_channels=self.input_num_channels,
            out_channels=self.embed_dim,
            kernel_size=(self.temporal_patch_size, self.patch_size, self.patch_size),
            stride=(self.temporal_patch_size, self.patch_size, self.patch_size),
            bias=False,
        )

    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
        if self.patch_embedding.weight.dtype != pixel_values.dtype:
            self.patch_embedding = self.patch_embedding.to(pixel_values.dtype)

        assert (
            pixel_values.dim() == 2
        ), f"pixel_values must be 2D, got {pixel_values.dim()}D"
        pixel_values = pixel_values.reshape(
            pixel_values.shape[0],
            self.input_num_channels,
            self.temporal_patch_size,
            self.patch_size,
            self.patch_size,
        )
        patch_embeds = self.patch_embedding(pixel_values)
        patch_embeds = patch_embeds.reshape(patch_embeds.shape[0], -1)
        return patch_embeds


class CLIPEncoderLayer(nn.Module):
    def __init__(
        self,
        config: CLIPVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        use_data_parallel: bool = False,
        workspace_buffer: Optional[torch.Tensor] = None,
    ) -> None:
        super().__init__()

        self.embed_dim = config.hidden_size
        self.use_data_parallel = use_data_parallel
        tp_size = 1 if use_data_parallel else get_parallel().attn_tp_size
        tp_rank = 0 if use_data_parallel else get_parallel().attn_tp_rank

        self.self_attn = VisionAttention(
            embed_dim=config.hidden_size,
            num_heads=config.num_attention_heads,
            projection_size=config.hidden_size,
            use_qkv_parallel=True,
            flatten_batch=True,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
            use_data_parallel=use_data_parallel,
            use_dp_attention_reduce=is_dp_attention_enabled(),
            customized_position_embedding_applier=_minimax_rope_applier,
            workspace_buffer=workspace_buffer,
        )

        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.fc1 = ColumnParallelLinear(
            config.hidden_size,
            config.intermediate_size,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp.fc1",
            tp_size=tp_size,
            tp_rank=tp_rank,
        )
        hidden_act = getattr(config, "hidden_act", "gelu")
        assert (
            hidden_act == "gelu"
        ), f"Only gelu activation is supported, got {hidden_act}"
        self.act = get_act_fn(hidden_act)
        self.fc2 = RowParallelLinear(
            config.intermediate_size,
            config.hidden_size,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp.fc2",
            tp_size=tp_size,
            tp_rank=tp_rank,
            use_dp_attention_reduce=is_dp_attention_enabled(),
        )
        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seq_len: torch.Tensor,
        rotary_pos_emb: torch.Tensor,
        max_seqlen: Optional[int] = None,
        sequence_lengths: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.layer_norm1(hidden_states)
        hidden_states = self.self_attn(
            x=hidden_states,
            cu_seqlens=cu_seq_len,
            position_embeddings=rotary_pos_emb,
            max_seqlen=max_seqlen,
            sequence_lengths=sequence_lengths,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states, _ = self.fc1(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states, _ = self.fc2(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states


class CLIPEncoder(nn.Module):
    def __init__(
        self,
        config: CLIPVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        use_data_parallel: bool = False,
        workspace_buffer: Optional[torch.Tensor] = None,
    ) -> None:
        super().__init__()

        self.config = config
        self.use_data_parallel = use_data_parallel

        self.layers = nn.ModuleList(
            [
                CLIPEncoderLayer(
                    config=config,
                    quant_config=quant_config,
                    prefix=f"{prefix}.layers.{layer_idx}",
                    use_data_parallel=use_data_parallel,
                    workspace_buffer=workspace_buffer,
                )
                for layer_idx in range(config.num_hidden_layers)
            ]
        )

    def forward(
        self,
        inputs_embeds,
        cu_seq_len: torch.Tensor,
        rotary_pos_emb: torch.Tensor,
        max_seqlen: Optional[int] = None,
        sequence_lengths: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        hidden_states = inputs_embeds
        cos_sin = _prepare_rotary_cos_sin(rotary_pos_emb)

        for encoder_layer in self.layers:
            hidden_states = encoder_layer(
                hidden_states,
                cu_seq_len,
                cos_sin,
                max_seqlen=max_seqlen,
                sequence_lengths=sequence_lengths,
            )

        return hidden_states


class MiniMaxVLVisionTransformer(nn.Module):
    def __init__(
        self,
        config: CLIPVisionConfig,
        quant_config: Optional[QuantizationConfig] = None,
        *,
        require_post_norm: Optional[bool] = None,
        prefix: str = "",
        use_data_parallel: bool = False,
    ) -> None:
        super().__init__()

        self.config = config
        self.use_data_parallel = use_data_parallel
        embed_dim = config.hidden_size

        self.temporal_patch_size = config.img_token_compression_config.get(
            "temporal_patch_size", 2
        )
        self.spatial_merge_size = config.img_token_compression_config.get(
            "spatial_merge_size", 2
        )

        self.embeddings = CLIPVisionEmbeddings(config)
        # NOTE: Typo "layrnorm" matches the original transformers code and the
        # weight names used in the published checkpoints; do not "fix" it.
        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

        workspace_buffer: Optional[torch.Tensor] = None
        if (
            get_server_args().mm_attention_backend == "flashinfer_cudnn"
            and torch.cuda.is_available()
        ):
            workspace_buffer = torch.empty(
                FLASHINFER_WORKSPACE_SIZE_BYTES,
                dtype=torch.uint8,
                device=torch.device("cuda", torch.cuda.current_device()),
            )

        self.encoder = CLIPEncoder(
            config=config,
            quant_config=quant_config,
            prefix=f"{prefix}.encoder",
            use_data_parallel=use_data_parallel,
            workspace_buffer=workspace_buffer,
        )

        assert (
            self.config.position_embedding_type == "rope"
        ), "Only rope position embedding is supported"
        assert self.config.rope_mode == "3d", "Only 3D RoPE is supported"
        rope_theta = getattr(config, "rope_theta")
        assert rope_theta is not None, "rope_theta must be set"
        self.vision_segment_max_frames = getattr(config, "vision_segment_max_frames")

        head_dim = embed_dim // config.num_attention_heads
        rope_dims = 2 * (head_dim // 2)

        self.t_dim = int(2 * ((rope_dims // 3) // 2))
        self.h_dim = int(2 * ((rope_dims // 3) // 2))
        self.w_dim = int(2 * ((rope_dims // 3) // 2))

        inv_freq_t = 1.0 / (
            rope_theta
            ** (torch.arange(0, self.t_dim, 2, dtype=torch.float32) / self.t_dim)
        )
        inv_freq_h = 1.0 / (
            rope_theta
            ** (torch.arange(0, self.h_dim, 2, dtype=torch.float32) / self.h_dim)
        )
        inv_freq_w = 1.0 / (
            rope_theta
            ** (torch.arange(0, self.w_dim, 2, dtype=torch.float32) / self.w_dim)
        )

        self.register_buffer("inv_freq_t", inv_freq_t, persistent=False)
        self.register_buffer("inv_freq_h", inv_freq_h, persistent=False)
        self.register_buffer("inv_freq_w", inv_freq_w, persistent=False)

        num_hidden_layers = config.num_hidden_layers
        if len(self.encoder.layers) > config.num_hidden_layers:
            raise ValueError(
                f"The original encoder only has {num_hidden_layers} "
                f"layers, but you requested {len(self.encoder.layers)} layers."
            )

        if require_post_norm is None:
            require_post_norm = len(self.encoder.layers) == num_hidden_layers
        self.post_layernorm = (
            nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
            if require_post_norm
            else None
        )

    def _get_3d_rope_embed(
        self, grid_t: int, grid_h: int, grid_w: int, spatial_merge_size: int
    ) -> torch.Tensor:
        tokens_per_frame = grid_h * grid_w

        tpos_ids = (
            torch.arange(grid_t, device=self.inv_freq_t.device)
            .unsqueeze(1)
            .expand(-1, tokens_per_frame)
            .flatten()
        )

        hpos_ids = (
            torch.arange(grid_h, device=self.inv_freq_h.device)
            .unsqueeze(1)
            .expand(-1, grid_w)
        )
        hpos_ids = hpos_ids.reshape(
            grid_h // spatial_merge_size,
            spatial_merge_size,
            grid_w // spatial_merge_size,
            spatial_merge_size,
        )
        hpos_ids = hpos_ids.permute(0, 2, 1, 3)
        hpos_ids = hpos_ids.unsqueeze(0).expand(grid_t, -1, -1, -1, -1).flatten()

        wpos_ids = (
            torch.arange(grid_w, device=self.inv_freq_w.device)
            .unsqueeze(0)
            .expand(grid_h, -1)
        )
        wpos_ids = wpos_ids.reshape(
            grid_h // spatial_merge_size,
            spatial_merge_size,
            grid_w // spatial_merge_size,
            spatial_merge_size,
        )
        wpos_ids = wpos_ids.permute(0, 2, 1, 3)
        wpos_ids = wpos_ids.unsqueeze(0).expand(grid_t, -1, -1, -1, -1).flatten()

        max_t = max(grid_t, 1)
        max_hw = max(grid_h, grid_w)

        seq_t = torch.arange(
            max_t, device=self.inv_freq_t.device, dtype=self.inv_freq_t.dtype
        )
        seq_hw = torch.arange(
            max_hw, device=self.inv_freq_h.device, dtype=self.inv_freq_h.dtype
        )

        freqs_t = torch.outer(seq_t, self.inv_freq_t)
        freqs_h = torch.outer(seq_hw, self.inv_freq_h)
        freqs_w = torch.outer(seq_hw, self.inv_freq_w)

        emb_t = freqs_t[tpos_ids]
        emb_h = freqs_h[hpos_ids]
        emb_w = freqs_w[wpos_ids]

        return torch.cat([emb_t, emb_h, emb_w], dim=-1)

    def _get_rope_embed_3d(self, grid_thw, spatial_merge_size: int) -> torch.Tensor:
        all_rope_embeds = [
            self._get_3d_rope_embed(grid_t, grid_h, grid_w, spatial_merge_size)
            for grid_t, grid_h, grid_w in grid_thw
        ]
        return torch.cat(all_rope_embeds, dim=0)

    def _apply_max_frames_limit(
        self, origin_grid_thw: list[list[int]]
    ) -> List[List[int]]:
        if self.vision_segment_max_frames is None:
            return origin_grid_thw
        max_frames = self.vision_segment_max_frames
        ret_grid_thw = []
        for grid_t, grid_h, grid_w in origin_grid_thw:
            if grid_t <= max_frames:
                ret_grid_thw.append([grid_t, grid_h, grid_w])
            else:
                for i in range(0, grid_t, max_frames):
                    sub_grid_t = min(max_frames, grid_t - i)
                    ret_grid_thw.append([sub_grid_t, grid_h, grid_w])
        return ret_grid_thw

    def _compute_cu_seq_len(
        self,
        grid_thw: list[list[int]],
        device: torch.device,
    ) -> torch.Tensor:
        grid_thw = self._apply_max_frames_limit(grid_thw)
        cu_seq_len = [0]
        for grid_t, grid_h, grid_w in grid_thw:
            cu_seq_len.append(grid_t * grid_h * grid_w)
        cu_seq_len = torch.tensor(cu_seq_len, device=device).to(torch.int32)
        cu_seq_len = torch.cumsum(cu_seq_len, dim=0).to(torch.int32)
        return cu_seq_len

    @staticmethod
    def _bucket_flashinfer_batch_size(batch_size: int) -> int:
        return next(
            (b for b in BATCH_BUCKETS if b >= batch_size),
            round_up(batch_size, BATCH_BUCKETS[0]),
        )

    @staticmethod
    def _bucket_flashinfer_max_seqlen(real_max_seqlen: int) -> int:
        if real_max_seqlen <= 0:
            return FLASHINFER_MAX_SEQLEN_BUCKETS[0]
        return next(
            (s for s in FLASHINFER_MAX_SEQLEN_BUCKETS if s >= real_max_seqlen),
            round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]),
        )

    @classmethod
    def _compute_flashinfer_sequence_lengths_padded(
        cls,
        token_cu_seqlens: np.ndarray,
    ) -> np.ndarray:
        assert token_cu_seqlens.ndim == 1 and token_cu_seqlens.size >= 2
        B = int(token_cu_seqlens.size - 1)
        seq_lens = (token_cu_seqlens[1:] - token_cu_seqlens[:-1]).astype(np.int32)
        B_padded = cls._bucket_flashinfer_batch_size(B)
        if B_padded != B:
            pad = np.zeros((B_padded - B,), dtype=np.int32)
            seq_lens = np.concatenate([seq_lens, pad], axis=0)
        return seq_lens

    @classmethod
    def _compute_flashinfer_batch_offsets_packed(
        cls,
        token_cu_seqlens: np.ndarray,
        *,
        elem_per_token: int,
    ) -> np.ndarray:
        assert token_cu_seqlens.ndim == 1 and token_cu_seqlens.size >= 2
        B = int(token_cu_seqlens.size - 1)
        B_padded = cls._bucket_flashinfer_batch_size(B)
        token_indptr = token_cu_seqlens.astype(np.int64, copy=False)
        if B_padded != B:
            pad = np.full((B_padded - B,), token_indptr[-1], dtype=token_indptr.dtype)
            token_indptr = np.concatenate([token_indptr, pad], axis=0)
        elem_indptr = (token_indptr * int(elem_per_token)).astype(np.int32)
        return np.concatenate([elem_indptr, elem_indptr, elem_indptr], axis=0)

    def _build_flashinfer_cudnn_inputs(
        self,
        cu_seq_len: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
        device = cu_seq_len.device
        token_cu_seqlens_np = cu_seq_len.detach().cpu().numpy().astype(np.int32)

        real_seq_lens = token_cu_seqlens_np[1:] - token_cu_seqlens_np[:-1]
        max_seqlen = self._bucket_flashinfer_max_seqlen(
            int(real_seq_lens.max()) if real_seq_lens.size > 0 else 0
        )

        seq_lens_padded = self._compute_flashinfer_sequence_lengths_padded(
            token_cu_seqlens_np
        )

        attn_tp_size = 1 if self.use_data_parallel else get_parallel().attn_tp_size
        elem_per_token = self.config.hidden_size // attn_tp_size

        offsets_packed = self._compute_flashinfer_batch_offsets_packed(
            token_cu_seqlens_np,
            elem_per_token=elem_per_token,
        )

        sequence_lengths = (
            torch.from_numpy(seq_lens_padded)
            .to(device=device, dtype=torch.int32, non_blocking=True)
            .view(-1, 1, 1, 1)
        )
        cu_seqlens_packed = torch.from_numpy(offsets_packed).to(
            device=device, dtype=torch.int32, non_blocking=True
        )
        return cu_seqlens_packed, sequence_lengths, int(max_seqlen)

    def forward(
        self,
        pixel_values: torch.Tensor,
        grid_thw: list[list[int]],
    ) -> torch.Tensor:
        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        assert pixel_values.dtype == torch.bfloat16, "pixel_values must be bfloat16"

        hidden_states = self.embeddings(pixel_values)
        hidden_states = self.pre_layrnorm(hidden_states)

        grid_thw = self._apply_max_frames_limit(grid_thw)

        cu_seq_len = self._compute_cu_seq_len(grid_thw, hidden_states.device)
        rotary_pos_emb = self._get_rope_embed_3d(grid_thw, self.spatial_merge_size)

        assert (
            rotary_pos_emb.device == hidden_states.device
        ), "rotary_pos_emb and hidden_states must be on the same device"

        max_seqlen: Optional[int] = None
        sequence_lengths: Optional[torch.Tensor] = None
        encoder_cu_seq_len = cu_seq_len
        if get_server_args().mm_attention_backend == "flashinfer_cudnn":
            (
                encoder_cu_seq_len,
                sequence_lengths,
                max_seqlen,
            ) = self._build_flashinfer_cudnn_inputs(cu_seq_len)

        return self.encoder(
            inputs_embeds=hidden_states,
            cu_seq_len=encoder_cu_seq_len,
            rotary_pos_emb=rotary_pos_emb,
            max_seqlen=max_seqlen,
            sequence_lengths=sequence_lengths,
        )


class MiniMaxVLVisionModel(nn.Module):
    def __init__(
        self,
        config: CLIPVisionConfig,
        text_hidden_size: int,
        projector_hidden_size: Optional[int] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        multimodal_projector_bias: bool = True,
        patch_merge_bias: bool = True,
    ) -> None:
        super().__init__()
        self.config = config
        self.quant_config = quant_config

        self.use_data_parallel = get_server_args().mm_enable_dp_encoder
        self.vision_config = config

        self.vision_model = MiniMaxVLVisionTransformer(
            config=config,
            quant_config=quant_config,
            prefix=add_prefix("vision_model", prefix),
            use_data_parallel=self.use_data_parallel,
        )

        self.multi_modal_projector = MiniMaxVLMultiModalProjector(
            vision_hidden_size=config.hidden_size,
            text_hidden_size=text_hidden_size,
            projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
            multimodal_projector_bias=multimodal_projector_bias,
            projector_hidden_size=projector_hidden_size,
            quant_config=quant_config,
            prefix=add_prefix("multi_modal_projector", prefix),
            use_data_parallel=self.use_data_parallel,
        )

        spatial_merge_size = config.img_token_compression_config.get(
            "spatial_merge_size", 2
        )
        self.spatial_merge_size = spatial_merge_size
        self.patch_merge_mlp = MiniMaxVLPatchMerger(
            spatial_merge_size=spatial_merge_size,
            text_hidden_size=text_hidden_size,
            projector_hidden_act=getattr(config, "projector_hidden_act", "gelu"),
            patch_merge_bias=patch_merge_bias,
            projector_hidden_size=projector_hidden_size,
            quant_config=quant_config,
            prefix=add_prefix("patch_merge_mlp", prefix),
            use_data_parallel=self.use_data_parallel,
        )
        self.dtype = self.vision_model.embeddings.patch_embedding.weight.dtype
        # Required by run_dp_sharded_mrope_vision_model when input is empty.
        self.out_hidden_size = text_hidden_size

    def forward(
        self,
        pixel_values: torch.Tensor,
        grid_thw: list[list[int]],
    ) -> torch.Tensor:
        hidden_states = self.vision_model(pixel_values=pixel_values, grid_thw=grid_thw)
        if hidden_states.dim() == 3:
            hidden_states = hidden_states.squeeze(0)
        hidden_states = self.multi_modal_projector(hidden_states)
        hidden_states = self.patch_merge_mlp(hidden_states)
        return hidden_states


def _run_vision_tower(
    vision_tower: MiniMaxVLVisionModel,
    pixel_values: torch.Tensor,
    grid_thw: list[list[int]],
    use_data_parallel: bool,
) -> torch.Tensor:
    if use_data_parallel:
        return run_dp_sharded_mrope_vision_model(
            vision_tower,
            pixel_values,
            grid_thw,
            rope_type="rope_3d",
        )
    return vision_tower(pixel_values, grid_thw=grid_thw)


def get_image_feature(
    vision_tower: MiniMaxVLVisionModel,
    items: List[MultimodalDataItem],
    use_data_parallel: bool,
) -> torch.Tensor:
    pixel_values = torch.cat([item.feature for item in items], dim=0).type(
        vision_tower.dtype
    )
    image_grid_thw: list[list[int]] = []
    for item in items:
        image_grid_thw.extend(item.image_grid_thw.tolist())
    return _run_vision_tower(
        vision_tower, pixel_values, image_grid_thw, use_data_parallel
    )


def get_video_feature(
    vision_tower: MiniMaxVLVisionModel,
    items: List[MultimodalDataItem],
    use_data_parallel: bool,
) -> torch.Tensor:
    pixel_values = torch.cat([item.feature for item in items], dim=0).type(
        vision_tower.dtype
    )
    video_grid_thw: list[list[int]] = []
    for item in items:
        video_grid_thw.extend(item.video_grid_thw.tolist())
    assert pixel_values.dim() == 2, pixel_values.dim()
    return _run_vision_tower(
        vision_tower, pixel_values, video_grid_thw, use_data_parallel
    )


def _parse_vit_layer_idx(name: str) -> Optional[int]:
    parts = name.split(".")
    for i, p in enumerate(parts):
        if p == "layers" and i + 1 < len(parts):
            return int(parts[i + 1])
    return None


def load_vision_weight(
    name: str,
    loaded_weight: torch.Tensor,
    params_dict: dict,
    vit_qkv_weights: dict,
    vit_qkv_biases: dict,
) -> None:
    if (
        "self_attn.q_proj" in name
        or "self_attn.k_proj" in name
        or "self_attn.v_proj" in name
    ):
        if name.endswith(".weight"):
            target = vit_qkv_weights
        elif name.endswith(".bias"):
            target = vit_qkv_biases
        else:
            return
        layer_idx = _parse_vit_layer_idx(name)
        if layer_idx is None:
            return
        qkv_type = "q" if "q_proj" in name else ("k" if "k_proj" in name else "v")
        target.setdefault(layer_idx, {})[qkv_type] = loaded_weight
        return

    param_name = name
    if "vision_tower.vision_model." in param_name:
        param_name = param_name.replace(".mlp.fc1.", ".fc1.")
        param_name = param_name.replace(".mlp.fc2.", ".fc2.")
        param_name = param_name.replace(".self_attn.out_proj.", ".self_attn.proj.")
    if name.startswith("patch_merge_mlp.") or name.startswith("multi_modal_projector."):
        param_name = "vision_tower." + param_name

    if param_name in params_dict:
        param = params_dict[param_name]
        weight_loader = getattr(param, "weight_loader", default_weight_loader)
        weight_loader(param, loaded_weight)


def merge_vit_qkv_weights(
    vit_qkv_weights: dict,
    vit_qkv_biases: dict,
    params_dict: dict,
) -> None:
    for layer_idx, qkv_dict in vit_qkv_weights.items():
        if {"q", "k", "v"} <= qkv_dict.keys():
            merged = torch.cat([qkv_dict["q"], qkv_dict["k"], qkv_dict["v"]], dim=0)
            param_name = (
                f"vision_tower.vision_model.encoder.layers.{layer_idx}"
                ".self_attn.qkv_proj.weight"
            )
            if param_name in params_dict:
                param = params_dict[param_name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, merged)

    for layer_idx, qkv_dict in vit_qkv_biases.items():
        if {"q", "k", "v"} <= qkv_dict.keys():
            merged = torch.cat([qkv_dict["q"], qkv_dict["k"], qkv_dict["v"]], dim=0)
            param_name = (
                f"vision_tower.vision_model.encoder.layers.{layer_idx}"
                ".self_attn.qkv_proj.bias"
            )
            if param_name in params_dict:
                param = params_dict[param_name]
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
                weight_loader(param, merged)