modeling_moss_vl.py

# coding=utf-8
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch MossVL model - Qwen3VL Vision + Text with Cross Attention"""

from dataclasses import dataclass
from typing import Any, Callable, Optional, Union, Tuple, List

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import initialization as init

from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.integrations import use_kernel_forward_from_hub
from transformers.masking_utils import create_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput, CausalLMOutputWithPast
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling, logging
from transformers.utils.deprecation import deprecate_kwarg
from transformers.utils.generic import is_flash_attention_requested
from transformers.utils.output_capturing import OutputRecorder

from .configuration_moss_vl import MossVLConfig, MossVLTextConfig, MossVLVisionConfig



logger = logging.get_logger(__name__)


@dataclass
class MossVLModelOutputWithPast(ModelOutput):
    """
    Output class for MossVL model with additional vision_token_info and rope_deltas fields.
    
    Args:
        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the output of the last layer of the model.
        past_key_values (`Cache`, *optional*):
            Contains pre-computed hidden-states (key and values in the self-attention blocks and 
            cross-attention blocks) that can be used to speed up sequential decoding.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for each layer).
        attentions (`tuple(torch.FloatTensor)`, *optional*):
            Tuple of `torch.FloatTensor` (one for each layer) of attention weights.
        vision_token_info (`List[dict]`, *optional*):
            Information about vision tokens for each sample, used to correctly expand cross-attention masks.
            This is cached during prefill and reused during decode to handle ViT padding correctly.
        rope_deltas (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Position offset due to vision tokens. Used for fast position computation in decode stage.
            rope_deltas = max_position - sequence_length
    """
    
    last_hidden_state: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Cache] = None
    hidden_states: Optional[tuple[torch.FloatTensor]] = None
    attentions: Optional[tuple[torch.FloatTensor]] = None
    vision_token_info: Optional[List[dict]] = None
    rope_deltas: Optional[torch.LongTensor] = None


@dataclass
class MossVLCausalLMOutputWithPast(ModelOutput):
    """
    Output class for MossVL causal language model with additional vision_token_info and rope_deltas fields.
    
    Args:
        loss (`torch.FloatTensor` of shape `(1,)`, *optional*):
            Language modeling loss (for next-token prediction).
        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head.
        past_key_values (`Cache`, *optional*):
            Contains pre-computed hidden-states for speed up sequential decoding.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*):
            Tuple of hidden-states at each layer.
        attentions (`tuple(torch.FloatTensor)`, *optional*):
            Tuple of attention weights.
        vision_token_info (`List[dict]`, *optional*):
            Information about vision tokens for each sample, cached for decode stage.
        rope_deltas (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Position offset due to vision tokens. Used for fast position computation in decode stage.
    """
    
    loss: Optional[torch.FloatTensor] = None
    logits: Optional[torch.FloatTensor] = None
    past_key_values: Optional[Cache] = None
    hidden_states: Optional[tuple[torch.FloatTensor]] = None
    attentions: Optional[tuple[torch.FloatTensor]] = None
    vision_token_info: Optional[List[dict]] = None
    rope_deltas: Optional[torch.LongTensor] = None


# ==================== Vision Components (from Qwen3VL) ====================

class MossVLVisionMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
        self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, hidden_state):
        return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state)))


class MossVLVisionPatchEmbed(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.patch_size = config.patch_size
        self.temporal_patch_size = config.temporal_patch_size
        self.in_channels = config.in_channels
        self.embed_dim = config.hidden_size

        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
        self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype
        hidden_states = hidden_states.view(
            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
        )
        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
        return hidden_states


class MossVLVisionRotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, dim: int, theta: float = 10000.0) -> None:
        super().__init__()
        # Keep dim / theta so that `_init_weights` can rebuild `inv_freq` after
        # from_pretrained materializes the module (it is a non-persistent buffer
        # and therefore never populated by the checkpoint).
        self.dim = dim
        self.theta = theta
        inv_freq = self.compute_inv_freq()
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def compute_inv_freq(self) -> torch.Tensor:
        return 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim))

    def forward(self, seqlen: int) -> torch.Tensor:
        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(seq, self.inv_freq)
        return freqs


class MossVLVisionPatchMerger(nn.Module):
    def __init__(self, config: MossVLVisionConfig, num_deepstack_features=0) -> None:
        super().__init__()
        # spatial_merge，维度变为原始的config.spatial_merge_size**2倍
        base_hidden_size = config.hidden_size * (config.spatial_merge_size**2)
        # 计算输入维度：spatial_merge 后的维度 * (1 + deepstack特征数)
        self.input_hidden_size = base_hidden_size * (1 + num_deepstack_features)
        
        # Use independent LayerNorms for each feature level
        # Total features = 1 (last layer) + num_deepstack_features
        num_features = 1 + num_deepstack_features
        self.norms = nn.ModuleList([
            nn.LayerNorm(config.hidden_size, eps=1e-6)
            for _ in range(num_features)
        ])
        
        self.hidden_size = config.hidden_size
        
        self.linear_fc1 = nn.Linear(self.input_hidden_size, self.input_hidden_size)
        self.act_fn = nn.GELU()
        self.linear_fc2 = nn.Linear(self.input_hidden_size, config.out_hidden_size)

    def forward(
        self,
        last_hidden_state: torch.Tensor,
        deepstack_features: Optional[List[torch.Tensor]] = None,
    ) -> torch.Tensor:
        # 1. Collect all features: [last_hidden_state, deepstack_1, deepstack_2, ...]
        # self.norms[0] corresponds to last_hidden_state
        # self.norms[1:] corresponds to deepstack_features
        if deepstack_features is None:
            deepstack_features = []
        all_inputs = [last_hidden_state] + deepstack_features
        
        # 2. Apply Norm independently
        outs = []
        for i, feat in enumerate(all_inputs):
            outs.append(self.norms[i](feat))
        
        # 3. Concat once
        x = torch.cat(outs, dim=-1)

        # 做merge，维度变为原始的config.spatial_merge_size**2倍，len对应缩小为原来的1/config.spatial_merge_size**2
        x = x.view(-1, self.input_hidden_size)
        x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
        return x


def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb_vision(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    orig_q_dtype = q.dtype
    orig_k_dtype = k.dtype
    q, k = q.float(), k.float()
    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    q_embed = q_embed.to(orig_q_dtype)
    k_embed = k_embed.to(orig_k_dtype)
    return q_embed, k_embed


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: float,
    dropout: float = 0.0,
    **kwargs: Unpack[TransformersKwargs],
):
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights


class MossVLVisionAttention(nn.Module):
    def __init__(self, config: MossVLVisionConfig) -> None:
        super().__init__()
        self.dim = config.hidden_size
        self.num_heads = config.num_heads
        self.head_dim = self.dim // self.num_heads
        self.num_key_value_groups = 1
        self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
        self.proj = nn.Linear(self.dim, self.dim)
        self.scaling = self.head_dim**-0.5
        self.config = config
        self.attention_dropout = 0.0
        self.is_causal = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        seq_length = hidden_states.shape[0]
        query_states, key_states, value_states = (
            self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
        )
        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)

        query_states = query_states.transpose(0, 1).unsqueeze(0)
        key_states = key_states.transpose(0, 1).unsqueeze(0)
        value_states = value_states.transpose(0, 1).unsqueeze(0)

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        if is_flash_attention_requested(self.config):
            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
            attn_output, _ = attention_interface(
                self,
                query_states,
                key_states,
                value_states,
                attention_mask=None,
                scaling=self.scaling,
                dropout=0.0 if not self.training else self.attention_dropout,
                cu_seq_lens_q=cu_seqlens,
                cu_seq_lens_k=cu_seqlens,
                max_length_q=max_seqlen,
                max_length_k=max_seqlen,
                is_causal=False,
                **kwargs,
            )
        else:
            lengths = cu_seqlens[1:] - cu_seqlens[:-1]
            splits = [
                torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
            ]

            attn_outputs = [
                attention_interface(
                    self,
                    q,
                    k,
                    v,
                    attention_mask=None,
                    scaling=self.scaling,
                    dropout=0.0 if not self.training else self.attention_dropout,
                    is_causal=False,
                    **kwargs,
                )[0]
                for q, k, v in zip(*splits)
            ]
            attn_output = torch.cat(attn_outputs, dim=1)

        attn_output = attn_output.reshape(seq_length, -1).contiguous()
        attn_output = self.proj(attn_output)
        return attn_output


class MossVLVisionBlock(GradientCheckpointingLayer):
    def __init__(self, config, attn_implementation: str = "sdpa") -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
        self.attn = MossVLVisionAttention(config=config)
        self.mlp = MossVLVisionMLP(config=config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cu_seqlens: torch.Tensor,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        hidden_states = hidden_states + self.attn(
            self.norm1(hidden_states),
            cu_seqlens=cu_seqlens,
            rotary_pos_emb=rotary_pos_emb,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
        return hidden_states



# ==================== Text Components (from Qwen3 + Cross Attention) ====================

class MossVLTextRotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor

    def __init__(self, config: MossVLTextConfig, device=None):
        super().__init__()
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        rope_parameters = getattr(config, "rope_parameters", None)
        if rope_parameters is None:
            rope_parameters = getattr(config, "rope_scaling", None) or {"rope_type": "default"}

        self.rope_type = rope_parameters.get("rope_type", rope_parameters.get("type", "default"))
        rope_init_fn: Callable = self.compute_default_rope_parameters
        if self.rope_type != "default":
            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)

        self.mrope_section = rope_parameters.get("mrope_section", [24, 20, 20])

    @staticmethod
    def compute_default_rope_parameters(
        config: Optional[MossVLTextConfig] = None,
        device: Optional[torch.device] = None,
        seq_len: Optional[int] = None,
    ) -> tuple[torch.Tensor, float]:
        rope_parameters = getattr(config, "rope_parameters", None) or {}
        base = rope_parameters.get("rope_theta", getattr(config, "rope_theta", 10000.0))
        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
        partial_rotary_factor = rope_parameters.get(
            "partial_rotary_factor", getattr(config, "partial_rotary_factor", 1.0)
        )
        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

    def apply_interleaved_mrope(self, freqs, mrope_section):
        """Apply interleaved MRoPE to 3D rotary embeddings.
        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
        interleaved [THTHWHTHW...TT], preserving frequency continuity.
        args:
            x: (3, bs, seq_len, head_dim // 2)
            mrope_section: (3,)
        returns:
            x_t: (bs, seq_len, head_dim // 2)
        """
        freqs_t = freqs[0]  # just overwrite the first dimension T
        for dim, offset in enumerate((1, 2), start=1):  # H, W
            length = mrope_section[dim] * 3
            idx = slice(offset, length, 3)
            freqs_t[..., idx] = freqs[dim, ..., idx]
        return freqs_t

    @torch.no_grad()
    @dynamic_rope_update
    def forward(self, x, position_ids):
        if position_ids.ndim == 2:
            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
        
        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)

        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().to(x.device) @ position_ids_expanded.float()).transpose(2, 3)
            freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
            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)


@use_kernel_forward_from_hub("RMSNorm")
class MossVLTextRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"

# self attention rotary position embedding
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

# cross attention rotary position embedding
def apply_rotary_pos_emb_cross_attention(states, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    states_embed = (states * cos) + (rotate_half(states) * sin)
    return states_embed


class MossVLTextSelfAttention(nn.Module):
    """Self attention for text decoder"""

    def __init__(self, config: MossVLTextConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True

        self.q_proj = nn.Linear(
            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
        )
        self.k_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.v_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.o_proj = nn.Linear(
            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
        )
        self.q_norm = MossVLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = MossVLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = False,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_values is not None:
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class MossVLTextCrossAttention(nn.Module):
    """Cross attention - for vision-text interaction"""

    def __init__(self, config: MossVLTextConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = False

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)

        self.q_norm = MossVLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = MossVLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
        self,
        hidden_states: torch.Tensor,
        cross_attention_states: torch.Tensor,       
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: bool = None,
        cache_position: Optional[torch.LongTensor] = None,
        query_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        vision_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        **kwargs,
    ) -> torch.Tensor:
        batch_size, seq_length, _ = hidden_states.size()
        
        # Query from text hidden states
        query_states = self.q_proj(hidden_states)
        query_states = query_states.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
        query_states = self.q_norm(query_states)

        if cross_attention_states is not None:
            # Key and Value from vision cross_attention_states
            key_states = self.k_proj(cross_attention_states)
            value_states = self.v_proj(cross_attention_states)
            
            key_states = key_states.view(batch_size, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
            key_states = self.k_norm(key_states)
            value_states = value_states.view(batch_size, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)

            # Apply different RoPE for query (text position) and key (vision position)
            if query_position_embeddings is not None:
                cos, sin = query_position_embeddings
                query_states = apply_rotary_pos_emb_cross_attention(query_states, cos, sin)
            
            if vision_position_embeddings is not None:
                vision_cos, vision_sin = vision_position_embeddings
                key_states = apply_rotary_pos_emb_cross_attention(key_states, vision_cos, vision_sin)


            if past_key_values is not None:
                # if we have a new image + new tokens, we only computed key_states on that new image
                # we still update the cross key states, past_image, new_image. And use it!
                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
            
        elif cache_position[0] != 0:
            key_states, value_states = (
                past_key_values.layers[self.layer_idx].keys,
                past_key_values.layers[self.layer_idx].values,
            ) 
        else:
            raise ValueError(
                "Cross attention layer can't find neither `cross_attn_states` nor cached values for key/values!"
            )

        if is_flash_attention_requested(self.config):
            # Cross attention still relies on an explicit dense mask.
            attention_interface: Callable = ALL_ATTENTION_FUNCTIONS["sdpa"]
        else:
            attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
                self.config._attn_implementation, eager_attention_forward
            )

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            **kwargs,
        )

        attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class MossVLTextMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


class MossVLSelfAttentionDecoderLayer(GradientCheckpointingLayer):
    """Self-attention decoder layer"""

    def __init__(self, config: MossVLTextConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.layer_idx = layer_idx

        self.self_attn = MossVLTextSelfAttention(config=config, layer_idx=layer_idx)
        self.mlp = MossVLTextMLP(config)
        self.input_layernorm = MossVLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = MossVLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        cross_attention_states: Optional[torch.Tensor] = None,
        cross_attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        vision_position_ids: Optional[torch.LongTensor] = None,
        vision_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        output_attentions: bool = False,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, ...]:
        # Self Attention
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
        )
        hidden_states = residual + hidden_states

        # MLP
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (attn_weights,)
        return outputs


class MossVLCrossAttentionDecoderLayer(GradientCheckpointingLayer):
    """Cross-attention decoder layer with tanh-gated attention and MLP"""

    def __init__(self, config: MossVLTextConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.layer_idx = layer_idx

        self.cross_attn = MossVLTextCrossAttention(config=config, layer_idx=layer_idx)
        self.mlp = MossVLTextMLP(config)
        
        self.input_layernorm = MossVLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = MossVLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        # Gates for cross attention (single scalar value).
        # Gate scalar = tanh(gate[0]), initialized to zero so tanh(0)=0 at start.
        self.cross_attn_attn_gate = nn.Parameter(torch.zeros(1))
        self.cross_attn_mlp_gate = nn.Parameter(torch.zeros(1))

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        cross_attention_states: Optional[torch.Tensor] = None,
        cross_attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        full_text_row_masked_out_mask: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        vision_position_ids: Optional[torch.LongTensor] = None,
        vision_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
        output_attentions: bool = False,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, ...]:
        # Cross Attention
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        
        hidden_states, attn_weights = self.cross_attn(
            hidden_states=hidden_states,
            cross_attention_states=cross_attention_states,
            attention_mask=cross_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            query_position_embeddings=position_embeddings,
            vision_position_embeddings=vision_position_embeddings,
        )
        if full_text_row_masked_out_mask is not None:
            hidden_states = full_text_row_masked_out_mask[:, 0] * hidden_states

        hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states

        # MLP
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        if full_text_row_masked_out_mask is not None:
            hidden_states = full_text_row_masked_out_mask[:, 0] * hidden_states

        hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (attn_weights,)
        return outputs




@auto_docstring
class MossVLPreTrainedModel(PreTrainedModel):
    config: MossVLConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["MossVLSelfAttentionDecoderLayer", "MossVLCrossAttentionDecoderLayer", "MossVLVisionBlock"]
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn = True
    _supports_sdpa = True
    _can_compile_fullgraph = True
    _supports_attention_backend = True
    _can_record_outputs = {
        "hidden_states": [MossVLSelfAttentionDecoderLayer, MossVLCrossAttentionDecoderLayer],
        "attentions": [
            OutputRecorder(MossVLTextSelfAttention, index=1, layer_name="self_attn"),  # self-attention layers
            OutputRecorder(MossVLTextCrossAttention, index=1, layer_name="cross_attn"),  # cross-attention layers
        ],
    }

    def _init_weights(self, module):
        """Initialize the weights.
        """
        super()._init_weights(module)
        if isinstance(module, MossVLVisionRotaryEmbedding):
            init.copy_(module.inv_freq, module.compute_inv_freq())





class MossVLVisionModel(MossVLPreTrainedModel):
    config: MossVLVisionConfig
    _no_split_modules = ["MossVLVisionBlock"]

    def __init__(self, config, *inputs, **kwargs) -> None:
        super().__init__(config, *inputs, **kwargs)
        self.spatial_merge_size = config.spatial_merge_size
        self.patch_size = config.patch_size
        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size

        self.patch_embed = MossVLVisionPatchEmbed(config=config)
        self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
        self.num_grid_per_side = int(config.num_position_embeddings**0.5)

        head_dim = config.hidden_size // config.num_heads
        self.rotary_pos_emb = MossVLVisionRotaryEmbedding(head_dim // 2)

        self.blocks = nn.ModuleList([MossVLVisionBlock(config) for _ in range(config.depth)])
        
        # DeepStack: 记录需要提取特征的层索引
        self.deepstack_visual_indexes = config.deepstack_visual_indexes
        num_deepstack_features = len(self.deepstack_visual_indexes)
        
        # Merger: 输入维度 = hidden_size * (1 + num_deepstack_features)
        self.merger = MossVLVisionPatchMerger(
            config=config, 
            num_deepstack_features=num_deepstack_features
        )

        self.gradient_checkpointing = False

    def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor:
        merge_size = self.spatial_merge_size
        max_hw = int(grid_thw[:, 1:].max().item())
        freq_table = self.rotary_pos_emb(max_hw)
        device = freq_table.device

        total_tokens = int(torch.prod(grid_thw, dim=1).sum().item())
        pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device)

        offset = 0
        for num_frames, height, width in grid_thw:
            merged_h, merged_w = height // merge_size, width // merge_size

            block_rows = torch.arange(merged_h, device=device)
            block_cols = torch.arange(merged_w, device=device)
            intra_row = torch.arange(merge_size, device=device)
            intra_col = torch.arange(merge_size, device=device)

            row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
            col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]

            row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)
            col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1)

            coords = torch.stack((row_idx, col_idx), dim=-1)

            if num_frames > 1:
                coords = coords.repeat(num_frames, 1)

            num_tokens = coords.shape[0]
            pos_ids[offset : offset + num_tokens] = coords
            offset += num_tokens

        embeddings = freq_table[pos_ids]
        embeddings = embeddings.flatten(1)
        return embeddings

    def fast_pos_embed_interpolate(self, grid_thw):
        grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]
        device = self.pos_embed.weight.device
        dtype = self.pos_embed.weight.dtype

        idx_parts = [[] for _ in range(4)]
        weight_parts = [[] for _ in range(4)]

        for t, h, w in zip(grid_ts, grid_hs, grid_ws):
            h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h, device=device)
            w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w, device=device)

            h_idxs_floor = h_idxs.int()
            w_idxs_floor = w_idxs.int()
            h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)
            w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1)

            dh = h_idxs - h_idxs_floor
            dw = w_idxs - w_idxs_floor

            base_h = h_idxs_floor * self.num_grid_per_side
            base_h_ceil = h_idxs_ceil * self.num_grid_per_side

            indices = [
                (base_h[None].T + w_idxs_floor[None]).flatten(),
                (base_h[None].T + w_idxs_ceil[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
                (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
            ]

            weights = [
                ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
                ((1 - dh)[None].T * dw[None]).flatten(),
                (dh[None].T * (1 - dw)[None]).flatten(),
                (dh[None].T * dw[None]).flatten(),
            ]

            for i in range(4):
                idx_parts[i].append(indices[i])
                weight_parts[i].append(weights[i])

        idx_tensor = torch.stack([torch.cat(parts) for parts in idx_parts]).to(dtype=torch.long)
        weight_tensor = torch.stack([torch.cat(parts) for parts in weight_parts]).to(dtype=dtype)
        pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None]
        patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]

        patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)])

        patch_pos_embeds_permute = []
        merge_size = self.config.spatial_merge_size
        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
            pos_embed = pos_embed.repeat(t, 1)
            pos_embed = (
                pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
                .permute(0, 1, 3, 2, 4, 5)
                .flatten(0, 4)
            )
            patch_pos_embeds_permute.append(pos_embed)
        patch_pos_embeds = torch.cat(patch_pos_embeds_permute)
        return patch_pos_embeds

    def forward(
        self, 
        hidden_states: torch.Tensor, 
        grid_thw: torch.Tensor, 
        **kwargs
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: input tensor
            grid_thw: [num_images, 3] tensor with (t, h, w) for each image
        Returns:
            hidden_states: [num_tokens, out_hidden_size] - packed hidden states
        """
        hidden_states = self.patch_embed(hidden_states)

        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
        hidden_states = hidden_states + pos_embeds

        rotary_pos_emb = self.rot_pos_emb(grid_thw)

        seq_len, _ = hidden_states.size()
        hidden_states = hidden_states.reshape(seq_len, -1)
        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        position_embeddings = (emb.cos(), emb.sin())

        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
            dim=0,
            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
        )
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

        # DeepStack: 收集不同层的视觉特征
        deepstack_features = []
        for layer_idx, blk in enumerate(self.blocks):
            hidden_states = blk(
                hidden_states,
                cu_seqlens=cu_seqlens,
                position_embeddings=position_embeddings,
                **kwargs,
            )
            # 如果当前层在 deepstack 索引中，保存特征
            if layer_idx in self.deepstack_visual_indexes:
                deepstack_features.append(hidden_states)

        # Merger: 从 hidden_size * (1 + num_deepstack) 映射到 out_hidden_size
        hidden_states = self.merger(hidden_states, deepstack_features)

        return hidden_states






@auto_docstring(
    custom_intro="""
    The MossVL Text Model with self-attention and cross-attention layers for vision-language interaction.
    """
)
class MossVLTextModel(MossVLPreTrainedModel):
    config: MossVLTextConfig
    _no_split_modules = ["MossVLSelfAttentionDecoderLayer", "MossVLCrossAttentionDecoderLayer"]


    def __init__(self, config: MossVLTextConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        
        # Store cross_attention_layers for use in forward pass
        self.cross_attention_layers = config.cross_attention_layers
        
        # Create layers: self-attention or cross-attention at specified indices
        self.layers = nn.ModuleList()
        for layer_idx in range(config.num_hidden_layers):
            if layer_idx in config.cross_attention_layers:
                # Cross attention layer
                self.layers.append(
                    MossVLCrossAttentionDecoderLayer(config, layer_idx)
                )
            else:
                # Self attention layer
                self.layers.append(
                    MossVLSelfAttentionDecoderLayer(config, layer_idx)
                )
        
        self.norm = MossVLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = MossVLTextRotaryEmbedding(config=config)
        self.gradient_checkpointing = False

        self.post_init()


    @auto_docstring
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        cross_attention_states: Optional[torch.Tensor] = None,
        cross_attention_mask: Optional[torch.Tensor] = None,
        vision_position_ids: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Union[tuple, BaseModelOutputWithPast]:
        """
        Args:
            full_text_row_masked_out_mask (`Tuple[torch.Tensor, torch.Tensor]`, *optional*):
                Mask for full text rows that should be masked out in attention computation.
            cross_attention_states (`torch.Tensor`, *optional*):
                Vision features to be used in cross-attention layers. Shape: `(batch_size, vision_seq_len, hidden_size)`.
            cross_attention_mask (`torch.Tensor`, *optional*):
                Attention mask for cross-attention between text and vision. Shape: `(batch_size, 1, text_seq_len, vision_seq_len)`.
            vision_position_ids (`torch.LongTensor`, *optional*):
                Position IDs for vision tokens used in cross-attention. Shape: `(batch_size, vision_seq_len)`.
            cache_position (`torch.LongTensor`, *optional*):
                Absolute cache positions for the current text tokens during incremental decoding.
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if use_cache and past_key_values is None and not torch.jit.is_tracing():
            past_key_values = DynamicCache(config=self.config)

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )


        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        attention_mask = create_causal_mask(
            config=self.config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            cache_position=cache_position,
            past_key_values=past_key_values,
            position_ids=position_ids,
        )

        hidden_states = inputs_embeds

        # Compute text position embeddings (for self-attention and cross-attention query)
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        # Compute vision position embeddings (for cross-attention key/value) if needed
        vision_position_embeddings = None

        if cross_attention_states is not None:
            if vision_position_ids is not None:
                vision_position_embeddings = self.rotary_emb(cross_attention_states, vision_position_ids)

        all_hidden_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        for idx, decoder_layer in enumerate(self.layers):
            # For text-only path we should skip cross attention layers.
            # Let's check if the layer is cross attention layer and if we have cross attention states
            # or cached cross attention states.
            is_cross_attention_layer = idx in self.cross_attention_layers
            is_cross_attention_cache_empty = past_key_values is None or (
                past_key_values is not None and past_key_values.get_seq_length(idx) == 0
            )

            if is_cross_attention_layer and cross_attention_states is None and is_cross_attention_cache_empty:
                continue

            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                full_text_row_masked_out_mask=full_text_row_masked_out_mask,
                past_key_values=past_key_values,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                cross_attention_states=cross_attention_states,
                cross_attention_mask=cross_attention_mask,
                vision_position_ids=vision_position_ids,
                vision_position_embeddings=vision_position_embeddings,
                output_attentions=output_attentions,
                **kwargs,
            )
            hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions += (layer_outputs[1],)

            if output_hidden_states:
                all_hidden_states += (hidden_states,)

        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states = all_hidden_states[:-1] + (hidden_states,)

        if not return_dict:
            outputs = (hidden_states, past_key_values)
            if output_hidden_states:
                outputs += (all_hidden_states,)
            if output_attentions:
                outputs += (all_attentions,)
            return outputs

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values,
            hidden_states=all_hidden_states,
            attentions=all_attentions,
        )


@auto_docstring(
    custom_intro="""
    The MossVL model which consists of a vision encoder (from Qwen3VL) and a language model with cross-attention layers.
    """
)
class MossVLModel(MossVLPreTrainedModel):
    base_model_prefix = ""
    config: MossVLConfig
    _no_split_modules = ["MossVLSelfAttentionDecoderLayer", "MossVLCrossAttentionDecoderLayer", "MossVLVisionBlock"]
    _checkpoint_conversion_mapping = {}
    accepts_loss_kwargs = False
    def __init__(self, config):
        super().__init__(config)
        self.visual = MossVLVisionModel._from_config(config.vision_config)
        self.language_model = MossVLTextModel._from_config(config.text_config)

        # Learnable Separator Token: inserted after each image/frame's vision tokens
        # Initialized from LLM's separator_token_init_id embedding
        self.separator_token = nn.Parameter(
            torch.zeros(config.vision_config.out_hidden_size)
        )

        self.post_init()
    


    def convert_packed_to_batch(
        self,
        hidden_states: torch.Tensor,
        grid_thw: torch.Tensor,
        media_nums_per_sample: Optional[List[int]],
    ) -> Tuple[torch.Tensor, List[dict]]:
        """
        Convert packed vision tokens to batched format with separator tokens.
        
        For each image: inserts 1 separator token after the vision tokens.
        For each video: inserts 1 separator token after EACH frame's vision tokens.
        
        Note: media_nums_per_sample counts each video as 1 media item, 
        but each frame in a video gets its own separator token.
        """
        
        # Calculate number of tokens per media after spatial merge
        tokens_per_media = (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]) // (self.visual.spatial_merge_size ** 2)
        hidden_size = hidden_states.shape[-1]
        
        # If media_nums_per_sample is not provided, assume batch size = 1
        if media_nums_per_sample is None:
            batch_size = 1
            media_nums_per_sample = [grid_thw.shape[0]]
        else:
            batch_size = len(media_nums_per_sample)
        
        # Optimization for batch_size = 1 (common in inference)
        if batch_size == 1:
            # 1. Calculate total length (pure math, fast)
            total_len = 0
            for i in range(grid_thw.shape[0]):
                num_tokens = tokens_per_media[i].item()
                num_frames = grid_thw[i, 0].item()
                total_len += num_tokens + num_frames # + separators
            
            # 2. Handle Padding alignment
            pad_multiple = self.config.vision_seq_pad_multiple
            if total_len % pad_multiple != 0:
                max_seq_len = (total_len + pad_multiple - 1) // pad_multiple * pad_multiple
            else:
                max_seq_len = total_len
            
            # 3. Pre-allocate final tensor
            batched_hidden_states = torch.zeros(
                1, max_seq_len, hidden_size,
                dtype=hidden_states.dtype,
                device=hidden_states.device
            )
            
            # 4. Vectorized fill
            sample_info = {
                'medias': [],
                'total_length': total_len,
                'pad_start': total_len,
                'pad_end': max_seq_len
            }
            
            token_offset = 0
            current_seq_len = 0
            separator_embedding = self.separator_token.to(hidden_states.dtype)
            
            # Iterate through all medias in this single sample
            for media_idx in range(grid_thw.shape[0]):
                num_tokens = tokens_per_media[media_idx].item()
                t, h, w = grid_thw[media_idx].tolist()
                num_frames = t
                tokens_per_frame = num_tokens // num_frames
                
                # --- Vectorized processing start ---
                # Extract vision tokens: (num_tokens, hidden)
                media_vision_tokens = hidden_states[token_offset : token_offset + num_tokens]
                
                # Reshape to (num_frames, tokens_per_frame, hidden)
                media_vision_tokens = media_vision_tokens.view(num_frames, tokens_per_frame, hidden_size)
                              
                # Directly write to destination without creating intermediate large tensors
                chunk_len = num_frames * (tokens_per_frame + 1)
                
                # Get view of the target area: (num_frames, tokens_per_frame + 1, hidden)
                target_view = batched_hidden_states[0, current_seq_len : current_seq_len + chunk_len]
                target_view = target_view.view(num_frames, tokens_per_frame + 1, hidden_size)
                
                # 1. Fill vision tokens
                target_view[:, :tokens_per_frame].copy_(media_vision_tokens)
                
                # 2. Fill separators (Broadcast assignment)
                # separator_embedding is (hidden,), automatically broadcasts to (num_frames, hidden)
                target_view[:, tokens_per_frame] = separator_embedding
                
                # --- Vectorized processing end ---
                
                sample_info['medias'].append({
                    'start': current_seq_len,
                    'end': current_seq_len + chunk_len,
                    'length': chunk_len,
                    'num_frames': num_frames,
                    'grid_h': h,
                    'grid_w': w,
                    'vision_tokens_per_frame': tokens_per_frame,
                    'has_separator': True,
                })
                
                current_seq_len += chunk_len
                token_offset += num_tokens
            
            vision_token_info = [sample_info]
            
            return batched_hidden_states, vision_token_info

        # Calculate tokens per sample including separator tokens
        # For images: +1 separator per image
        # For videos: +num_frames separators per video (one after each frame)
        tokens_per_sample = []
        media_idx = 0
        for num_medias_in_sample in media_nums_per_sample:
            sample_tokens = 0
            for i in range(num_medias_in_sample):
                num_tokens = tokens_per_media[media_idx + i].item()
                num_frames = grid_thw[media_idx + i, 0].item()
                sample_tokens += num_tokens + num_frames  # +num_frames separator tokens
            tokens_per_sample.append(sample_tokens)
            media_idx += num_medias_in_sample
        
        max_seq_len = max(tokens_per_sample)
        pad_multiple = self.config.vision_seq_pad_multiple
        if max_seq_len % pad_multiple != 0:
            max_seq_len = (max_seq_len + pad_multiple - 1) // pad_multiple * pad_multiple
        
        # Initialize batched output with zeros (for padding)
        batched_hidden_states = torch.zeros(
            batch_size, max_seq_len, hidden_size,
            dtype=hidden_states.dtype,
            device=hidden_states.device
        )
        
        # Get separator token (learnable parameter)
        separator_embedding = self.separator_token.to(hidden_states.dtype)
        
        # Track token positions for each sample
        vision_token_info = []
        
        # Split packed tensor and fill batched output
        token_offset = 0
        media_idx = 0
        
        for sample_idx, num_medias_in_sample in enumerate(media_nums_per_sample):
            sample_info = {
                'medias': [],  # List of dicts for each media in this sample
                'total_length': tokens_per_sample[sample_idx],
                'pad_start': tokens_per_sample[sample_idx],
                'pad_end': max_seq_len
            }
            
            seq_offset = 0  # Offset within this sample's sequence
            
            # Process each image/video in this sample
            for _ in range(num_medias_in_sample):
                num_tokens = tokens_per_media[media_idx].item()
                
                t, h, w = grid_thw[media_idx].tolist()
                num_frames = t
                tokens_per_frame = num_tokens // num_frames
                
                # Record start position for this media
                media_start = seq_offset
                
                # Vectorized handling of frames
                # Extract vision tokens for this media: (num_tokens, hidden)
                media_vision_tokens = hidden_states[token_offset : token_offset + num_tokens]
                
                # Reshape to (num_frames, tokens_per_frame, hidden)
                media_vision_tokens = media_vision_tokens.view(num_frames, tokens_per_frame, hidden_size)
                
                # Create separators: (num_frames, 1, hidden)
                separators = separator_embedding.view(1, 1, hidden_size).expand(num_frames, 1, hidden_size)
                
                # Concatenate: (num_frames, tokens_per_frame + 1, hidden)
                media_tokens_with_sep = torch.cat([media_vision_tokens, separators], dim=1)
                
                # Flatten: (num_frames * (tokens_per_frame + 1), hidden)
                media_tokens_with_sep = media_tokens_with_sep.view(-1, hidden_size)
                
                # Assign to batched tensor
                media_length_with_sep = media_tokens_with_sep.shape[0]
                batched_hidden_states[sample_idx, seq_offset : seq_offset + media_length_with_sep] = media_tokens_with_sep
                
                seq_offset += media_length_with_sep
                
                # Total tokens for this media = vision_tokens + num_separators
                media_length = num_tokens + num_frames
                
                # Record this image/video's position within the sample
                # Note: length now includes separator tokens
                sample_info['medias'].append({
                    'start': media_start,
                    'end': media_start + media_length,
                    'length': media_length,
                    'num_frames': num_frames,  # 1 for image, >1 for video
                    'grid_h': h,
                    'grid_w': w,
                    'vision_tokens_per_frame': tokens_per_frame,  # Actual vision tokens per frame (excluding separator)
                    'has_separator': True,  # Flag indicating separator tokens are included
                })
                
                token_offset += num_tokens
                media_idx += 1
            
            vision_token_info.append(sample_info)
        
        return batched_hidden_states, vision_token_info

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def set_decoder(self, decoder):
        self.language_model = decoder

    def get_decoder(self):
        return self.language_model

    def _expand_cross_attention_mask(
        self,
        cross_attention_mask: torch.Tensor,
        vision_token_info: List[dict],
        target_dtype: torch.dtype,
    ) -> torch.Tensor:
        """
        Expand cross_attention_mask from (B, 1, T, N_frames) to (B, 1, T, N_tokens).
        
        Args:
            cross_attention_mask (`torch.Tensor` of shape `(batch_size, 1, text_seq_len, num_frames)`):
                Coarse attention mask where each frame corresponds to one column.
                Can be bool (True=masked) or float (min_value=masked).
            vision_token_info (`List[dict]`):
                Precomputed token info that includes actual token counts after ViT padding.
                Must be provided (either from prefill computation or from cache).
                Each dict contains 'medias' list with 'length', 'num_frames', and 'vision_tokens_per_frame'.
            target_dtype (`torch.dtype`):
                Target dtype for the output mask (typically inputs_embeds.dtype).
        
        Returns:
            `torch.Tensor` of shape `(batch_size, 1, text_seq_len, total_vision_tokens)`:
                Fine-grained attention mask where each vision token has its own column.
                Masked positions have min_value, unmasked positions have 0.0.
        
        Note:
            - vision_token_info contains the actual token counts after ViT padding (pad to multiple of 8)
            - Separator tokens are treated as part of the same frame, sharing the same mask
        """
        if vision_token_info is None:
            raise ValueError(
                "vision_token_info must be provided to _expand_cross_attention_mask. "
                "This should be cached from prefill stage or computed during current forward pass."
            )
        
        batch_size = cross_attention_mask.shape[0]
        
        # Determine target vision length (should be consistent across batch, but take max to be safe)
        max_vision_len = 0
        if vision_token_info:
             max_vision_len = max([info.get('pad_end', 0) for info in vision_token_info])

        if max_vision_len == 0:
             return None

        # Convert bool mask to float mask if needed
        if cross_attention_mask.dtype == torch.bool:
            # True = masked, False = visible
            # Convert to float: True -> min_value, False -> 0.0
            min_value = torch.finfo(target_dtype).min
            float_mask = torch.zeros_like(cross_attention_mask, dtype=target_dtype)
            float_mask.masked_fill_(cross_attention_mask, min_value)
            cross_attention_mask = float_mask
        else:
            # Already float, ensure it's the right dtype
            cross_attention_mask = cross_attention_mask.to(dtype=target_dtype)

        # Pre-allocate final mask with min_dtype (masked)
        # This is memory efficient and handles padding automatically
        min_dtype = torch.finfo(target_dtype).min
        final_mask = torch.full(
            (batch_size, 1, cross_attention_mask.shape[2], max_vision_len),
            min_dtype,
            dtype=target_dtype,
            device=cross_attention_mask.device
        )
        
        for i in range(batch_size):
            medias = vision_token_info[i]['medias']
            if not medias:
                continue
            
            # Collect repetition counts for all frames in this sample
            repeats_parts = []
            for media in medias:
                num_frames = media.get('num_frames', 1)
                length = media['length']
                has_separator = media.get('has_separator', False)
                
                # Determine tokens per frame (including separator)
                if has_separator:
                    vision_tokens_per_frame = media.get('vision_tokens_per_frame', (length // num_frames) - 1)
                    tokens_per_frame_with_sep = vision_tokens_per_frame + 1
                else:
                    tokens_per_frame_with_sep = length // num_frames
                
                # In convert_packed_to_batch we enforce strictly regular frames
                # so we can assume all frames have the same number of tokens
                repeats_parts.append(
                    torch.full(
                        (num_frames,),
                        tokens_per_frame_with_sep,
                        dtype=torch.long,
                        device=cross_attention_mask.device,
                    )
                )
            
            num_valid_frames = sum(part.numel() for part in repeats_parts)
            if num_valid_frames == 0:
                continue
                
            # If cross_attention_mask has more frames (e.g. padded), slice it
            # If it has fewer (shouldn't happen), slice repeats
            valid_mask_frames = min(num_valid_frames, cross_attention_mask.shape[-1])
            repeats_tensor = torch.cat(repeats_parts)
            if valid_mask_frames < num_valid_frames:
                repeats_tensor = repeats_tensor[:valid_mask_frames]
            
            # Extract valid columns for this sample
            # (1, text_len, valid_mask_frames)
            source_mask = cross_attention_mask[i, :, :, :valid_mask_frames]
            
            # Expand using repeat_interleave
            # output shape: (1, text_len, sum(repeats))
            expanded_mask = source_mask.repeat_interleave(repeats_tensor, dim=-1)
            
            # Assign to final_mask
            num_tokens = expanded_mask.shape[-1]
            if num_tokens > max_vision_len:
                num_tokens = max_vision_len
                expanded_mask = expanded_mask[..., :num_tokens]
                
            final_mask[i, :, :, :num_tokens] = expanded_mask
            
        return final_mask

    def compute_position_ids(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        rope_deltas: Optional[torch.LongTensor] = None,
    ) -> torch.Tensor:
        """
        Compute 3D position IDs for text tokens with special handling for image tokens.
        
        Rules:
        - Regular text tokens: increment position (x, x, x) -> (x+1, x+1, x+1)
        - Image token: gets (t, t, t) where t = previous_text_position + 1
        - After processing vision tokens, next text token starts at max(vision_bottom_right) + 1
        
        In decode stage, uses cached rope_deltas to quickly compute new positions.
        
        Args:
            input_ids: (batch_size, seq_len)
            attention_mask: (batch_size, seq_len), optional
            past_key_values: cache object used to infer decode offset from the current text cache length
            
        Returns:
            position_ids: (3, batch_size, seq_len)
        """
        batch_size, seq_len = input_ids.shape
        device = input_ids.device
        image_token_id = self.config.image_token_id
        
        # Decode stage: always advance positions from the current text cache length.
        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
        if past_seen_tokens > 0:
            position_ids = torch.arange(seq_len, device=device, dtype=torch.long)
            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
            position_ids = position_ids + past_seen_tokens

            if rope_deltas is not None:
                position_ids = position_ids + rope_deltas.unsqueeze(1)

            return position_ids.unsqueeze(0).expand(3, -1, -1)
        
        # Prefill stage: compute full position_ids with image token awareness
        # Vectorized implementation
        
        # 1. Identify token types
        is_image_token = (input_ids == image_token_id)
        if attention_mask is not None:
            is_padding = (attention_mask == 0)
        else:
            is_padding = torch.zeros_like(input_ids, dtype=torch.bool)
            
        is_regular_token = ~(is_image_token | is_padding)
        
        # 2. Calculate position increments
        # Regular tokens increment position by 1
        # Image tokens do not increment position (they reuse the "current" position counter)
        # Padding tokens do not increment
        
        # cumulative sum of regular tokens gives the position index
        # We want 0-based index for the first regular token
        # cumsum: [1, 2, 2, 3] -> positions: [0, 1, 2, 2]
        # For image token at index i, we want count of regular tokens before i.
        # This is exactly (cumsum - 1) if the token itself is regular? No.
        
        # Let's use the logic: position[i] = sum(is_regular[:i])
        # We can achieve this by cumsum(is_regular) - is_regular
        
        cumulative_regular = is_regular_token.long().cumsum(dim=1)
        
        # For regular token: position = cumsum - 1 (since it's inclusive) => 0, 1, 2...
        # For image token: position = cumsum (since it's not included in cumsum, cumsum is count of prev regulars)
        # Wait, if is_regular[i] is 0, cumsum[i] == cumsum[i-1].
        # So for image token, position = cumsum[i] is correct.
        # For regular token, position = cumsum[i] - 1 is correct.
        
        # Combine: position = cumsum - is_regular.long()
        base_position_ids = cumulative_regular - is_regular_token.long()
        
        # Apply padding mask (set padding positions to 0)
        base_position_ids = base_position_ids.masked_fill(is_padding, 0)
        
        # Expand to 3D: (3, batch, seq_len)
        position_ids = base_position_ids.unsqueeze(0).expand(3, -1, -1).clone()
        
        return position_ids

    def compute_vision_position_ids(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        vision_token_info: List[dict],
        cross_attention_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Compute 3D position IDs for vision tokens (including separator tokens) and update text position_ids.
        Vectorized implementation for improved efficiency.
        
        Position encoding rules:
        - For text: if not image token, increment position (t-1, t-1, t-1) -> (t, t, t) -> ...
        - For vision: top-left is (t, t, t), increases towards bottom-right to (t, t+h-1, t+w-1)
        - Separator Token after each frame: (x, x, x) where x = max(t+h-1, t+w-1) + 1 = max(t+h, t+w)
        - Image token in text: also gets position (x, x, x) - same as separator
        - Next text token after image: starts at (x+1, x+1, x+1)
        
        Args:
            input_ids: (batch_size, seq_len)
            position_ids: (3, batch_size, seq_len) - will be updated in place
            vision_token_info: metadata about vision tokens (now includes separator positions)
            cross_attention_states: (batch_size, max_vision_seq_len, hidden_size)
            attention_mask: (batch_size, seq_len), optional
            
        Returns:
            vision_pos_ids: (3, batch_size, max_vision_seq_len)
            position_ids: (3, batch_size, seq_len) - updated
            rope_deltas: (batch_size,) - position offset due to vision tokens
        """
        batch_size, max_vision_seq_len, _ = cross_attention_states.shape
        device = cross_attention_states.device
        image_token_id = self.config.image_token_id
        merge_size = self.visual.spatial_merge_size

        # 1. Gather all frame metadata
        # We need to flatten the nested vision_token_info structure to align with image tokens in input_ids
        
        # Find all image tokens in text: (num_occurrences, 2) -> [batch_idx, seq_idx]
        image_token_indices = (input_ids == image_token_id).nonzero()
        
        # Flatten vision_token_info to parallel lists
        # We assume the order of medias in vision_token_info matches the appearance of image tokens in input_ids
        flat_eff_h_parts = []
        flat_eff_w_parts = []
        flat_vis_start_parts = []

        # Processing metadata on CPU (fast enough for typical batch sizes)
        for b_idx, info in enumerate(vision_token_info):
            medias = info.get('medias', [])
            for media in medias:
                num_frames = media['num_frames']
                h, w = media['grid_h'], media['grid_w']
                eh, ew = h // merge_size, w // merge_size
                start = media['start']
                tok_per_frame = media['vision_tokens_per_frame']
                stride = tok_per_frame + 1 # +1 for separator

                frame_offsets = start + torch.arange(num_frames, device=device, dtype=torch.long) * stride
                flat_vis_start_parts.append(frame_offsets)
                flat_eff_h_parts.append(torch.full((num_frames,), eh, device=device, dtype=torch.long))
                flat_eff_w_parts.append(torch.full((num_frames,), ew, device=device, dtype=torch.long))

        # Pre-allocate output
        vision_pos_ids = torch.zeros(
            (3, batch_size, max_vision_seq_len), 
            dtype=torch.long, 
            device=device
        )

        # Handle case where no image tokens or info
        if len(flat_eff_h_parts) == 0 or len(image_token_indices) == 0:
            rope_deltas = position_ids.max(dim=0).values.max(dim=-1).values + 1 - input_ids.shape[1]
            return vision_pos_ids, position_ids, rope_deltas

        flat_eff_h = torch.cat(flat_eff_h_parts)
        flat_eff_w = torch.cat(flat_eff_w_parts)
        flat_vis_starts = torch.cat(flat_vis_start_parts)

        # Align lengths (handle truncation if text has fewer tokens or vice versa)
        num_matches = min(flat_eff_h.shape[0], image_token_indices.shape[0])
        flat_eff_h = flat_eff_h[:num_matches]
        flat_eff_w = flat_eff_w[:num_matches]
        flat_vis_starts = flat_vis_starts[:num_matches]
        
        # Get corresponding text positions
        target_indices = image_token_indices[:num_matches]
        batch_rows = target_indices[:, 0]
        text_cols = target_indices[:, 1]
        
        # 2. Compute Shifts and Update Position IDs
        
        # Calculate max dimensions for each image token: separator_pos = t + max(h, w)
        # Shift amount for subsequent tokens = max(h, w) + 1
        max_hw = torch.maximum(flat_eff_h, flat_eff_w)
        shifts = max_hw + 1
        
        # Create a shift map to apply cumulative shifts
        shift_map = torch.zeros((batch_size, input_ids.shape[1]), dtype=torch.long, device=device)
        shift_map[batch_rows, text_cols] = shifts
        
        # Calculate cumulative shifts along sequence
        cum_shifts = shift_map.cumsum(dim=1)
        
        # Calculate t_vals (start position for each vision grid)
        # t_val = original_pos + shifts_before_this_image
        # cum_shifts at image index includes the image's own shift, so we subtract it
        orig_pos = position_ids[0, batch_rows, text_cols]
        shifts_before = cum_shifts[batch_rows, text_cols] - shifts
        t_vals = orig_pos + shifts_before
        
        # Update text position_ids
        # All tokens get shifted by cum_shifts
        # Image tokens specifically need to be at t_val + max_hw (which is t_val + shift - 1)
        # Our cum_shift update gives: orig_pos + shifts_before + shift = t_val + shift
        # So we subtract 1 from image tokens
        
        # Apply global shift
        # Note: position_ids is (3, B, L), cum_shifts is (B, L). Expand to match.
        new_pos_ids = position_ids + cum_shifts.unsqueeze(0)
        
        # Correct image tokens (subtract 1)
        # We can use boolean mask for efficient update
        img_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        img_token_mask[batch_rows, text_cols] = True
        new_pos_ids[:, img_token_mask] -= 1
        
        # Ensure padding positions remain 0 (if attention_mask provided)
        if attention_mask is not None:
             # Assuming padding is 0 in attention_mask
             padding_mask = (attention_mask == 0).unsqueeze(0)
             new_pos_ids.masked_fill_(padding_mask, 0)
        
        # Update position_ids in-place
        position_ids.copy_(new_pos_ids)
        
        # 3. Populate Vision Pos IDs
        # Group frames by size (eff_h, eff_w) to vectorize grid generation
        # This is efficient because typically there are few distinct aspect ratios
        unique_shapes = torch.unique(torch.stack([flat_eff_h, flat_eff_w], dim=1), dim=0)
        
        for shape in unique_shapes:
            eh, ew = shape[0].item(), shape[1].item()
            
            # Mask for frames of this shape
            mask = (flat_eff_h == eh) & (flat_eff_w == ew)
            
            sub_t_vals = t_vals[mask]
            sub_batch_rows = batch_rows[mask]
            sub_vis_starts = flat_vis_starts[mask]
            
            num_frames_sub = sub_t_vals.shape[0]
            if num_frames_sub == 0: continue
                
            # Generate grids: (num_frames, eh, ew)
            # y ranges 0..eh-1, x ranges 0..ew-1
            # positions: t + y, t + x
            
            y_grid = torch.arange(eh, device=device).view(1, eh, 1).expand(num_frames_sub, -1, ew)
            x_grid = torch.arange(ew, device=device).view(1, 1, ew).expand(num_frames_sub, eh, -1)
            t_grid = sub_t_vals.view(-1, 1, 1).expand(-1, eh, ew)
            
            h_grid = t_grid + y_grid
            w_grid = t_grid + x_grid
            
            # Flatten to assign
            flat_t = t_grid.reshape(-1)
            flat_h = h_grid.reshape(-1)
            flat_w = w_grid.reshape(-1)
            
            # Calculate destination indices in vision_pos_ids
            # (batch, seq_pos)
            tokens_per_frame = eh * ew
            
            # Offsets for each token in the frame 0..N-1
            seq_offsets = torch.arange(tokens_per_frame, device=device).unsqueeze(0) 
            # Add start index: (num_frames, 1) + (1, tokens) -> (num_frames, tokens)
            abs_seq_offsets = seq_offsets + sub_vis_starts.unsqueeze(1)
            
            flat_seq_inds = abs_seq_offsets.reshape(-1)
            flat_batch_inds = sub_batch_rows.unsqueeze(1).expand(-1, tokens_per_frame).reshape(-1)
            
            # Clip to max_vision_seq_len
            valid_mask = flat_seq_inds < max_vision_seq_len
            
            if valid_mask.any():
                final_b = flat_batch_inds[valid_mask]
                final_s = flat_seq_inds[valid_mask]
                
                vision_pos_ids[0, final_b, final_s] = flat_t[valid_mask]
                vision_pos_ids[1, final_b, final_s] = flat_h[valid_mask]
                vision_pos_ids[2, final_b, final_s] = flat_w[valid_mask]
                
        # 4. Handle Separator Tokens
        # Position: t_val + max(eh, ew)
        sep_vals = t_vals + max_hw
        # Index: start + tokens_per_frame = start + eh*ew
        sep_indices = flat_vis_starts + (flat_eff_h * flat_eff_w)
        
        valid_sep_mask = sep_indices < max_vision_seq_len
        
        if valid_sep_mask.any():
            final_b = batch_rows[valid_sep_mask]
            final_s = sep_indices[valid_sep_mask]
            vals = sep_vals[valid_sep_mask]
            
            vision_pos_ids[0, final_b, final_s] = vals
            vision_pos_ids[1, final_b, final_s] = vals
            vision_pos_ids[2, final_b, final_s] = vals
            
        # 5. Compute Rope Deltas
        # rope_deltas[batch_idx] = max_pos + 1 - seq_len
        
        # Use updated position_ids
        # Max pos in each batch - take max across all 3 position dimensions
        # position_ids shape: (3, batch_size, seq_len)
        # We need rope_deltas shape: (batch_size,)
        max_pos = position_ids.max(dim=0).values.max(dim=-1).values  # (batch_size,)
        rope_deltas = max_pos + 1 - input_ids.shape[1]  # (batch_size,)
        
        return vision_pos_ids, position_ids, rope_deltas

    def get_vision_features(
        self, 
        pixel_values: torch.FloatTensor, 
        grid_thw: Optional[torch.LongTensor] = None,
        media_nums_per_sample: Optional[List[int]] = None
    ):
        """
        Args:
            pixel_values: vision pixel values (images and videos merged)
            grid_thw: [num_media, 3] tensor with (t, h, w) for each media item
            media_nums_per_sample: List indicating how many media items each sample has
        Returns:
            vision_embeds: [batch_size, max_seq_len, hidden_size]
            vision_token_info: List[Dict] with media positions and padding info for each sample
        """
        pixel_values = pixel_values.type(self.visual.dtype)
        hidden_states = self.visual(
            pixel_values, 
            grid_thw=grid_thw
        )
        vision_embeds, vision_token_info = self.convert_packed_to_batch(
            hidden_states, 
            grid_thw, 
            media_nums_per_sample
        )
        return vision_embeds, vision_token_info



    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        pixel_values: Optional[torch.Tensor] = None,
        grid_thw: Optional[torch.LongTensor] = None,
        media_nums_per_sample: Optional[List[int]] = None,
        vision_position_ids: Optional[torch.LongTensor] = None,
        cross_attention_mask: Optional[torch.Tensor] = None,
        vision_token_info: Optional[List[dict]] = None,
        rope_deltas: Optional[torch.LongTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, BaseModelOutputWithPast]:
        """
        Args:
            grid_thw (`torch.LongTensor` of shape `(num_media, 3)`, *optional*):
                Grid size for each media item in (temporal, height, width) format. Each row contains `[t, h, w]` 
                representing the number of temporal, height, and width patches for a media item (image or video).
            media_nums_per_sample (`List[int]`, *optional*):
                List indicating how many media items each sample in the batch has. For example, `[2, 1, 3]` means 
                the first sample has 2 media items, the second has 1, and the third has 3.
            vision_position_ids (`torch.LongTensor` of shape `(batch_size, vision_seq_len)`, *optional*):
                Position IDs for vision tokens used in cross-attention. These are computed from text position IDs 
                based on the positions of image/video tokens in the input text.
            cross_attention_mask (`torch.Tensor` of shape `(batch_size, 1, text_seq_len, vision_seq_len)`, *optional*):
                Attention mask for cross-attention between text and vision. Controls which vision tokens each text 
                token can attend to, enforcing causal visibility for video frames.
            vision_token_info (`List[dict]`, *optional*):
                Cached metadata describing how packed vision tokens were regrouped per sample. Reused in decode
                to expand frame-level cross-attention masks to token-level masks without recomputing vision features.
            rope_deltas (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Cached offsets between text sequence length and multimodal RoPE positions. Reused in decode to
                reconstruct text position ids from the current cache length.
        """
        cache_position = kwargs.pop("cache_position", None)
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.get_input_embeddings()(input_ids)

        # Process vision features (images and videos are already merged by processor)
        cross_attention_states = None

        if pixel_values is not None:
            # Determine batch size
            batch_size = inputs_embeds.shape[0]
            
            # Get default media_nums_per_sample if not provided
            if media_nums_per_sample is None:
                # Assume all media belong to first sample if batch_size=1, otherwise raise error
                if batch_size == 1:
                    media_nums_per_sample = [grid_thw.shape[0]]
                else:
                    raise ValueError("media_nums_per_sample must be provided when batch_size > 1")
            
            # Process all vision inputs together through VIT
            # pixel_values and grid_thw are already ordered by appearance in text
            vision_embeds, vision_token_info = self.get_vision_features(
                pixel_values, grid_thw, media_nums_per_sample
            )
            
            # vision_embeds: [batch_size, max_seq_len, hidden_size]
            cross_attention_states = vision_embeds.to(inputs_embeds.device, inputs_embeds.dtype)

        # Generate 3D position IDs for text if not provided
        if position_ids is None:
            # Compute position IDs with image token awareness
            # In decode stage, this uses cached rope_deltas for fast computation
            position_ids = self.compute_position_ids(
                input_ids=input_ids,
                attention_mask=attention_mask,
                past_key_values=past_key_values,
                rope_deltas=rope_deltas,
            )

        # Compute cross_attention_mask, vision_position_ids, and full_text_row_masked_out_mask
        full_text_row_masked_out_mask = None
        
        if cross_attention_mask is not None:
            # Expand mask from frame-level to token-level
            # The processor outputs coarse masks (bool or float) where each frame has one column,
            # we need to expand to fine-grained masks where each vision token has its own column
            # This function also converts bool to float with correct min/max values
            cross_attention_mask = self._expand_cross_attention_mask(
                cross_attention_mask,
                vision_token_info,
                target_dtype=inputs_embeds.dtype
            )
            
            # Handle full_text_row_masked_out_mask logic
            if cross_attention_mask is not None:
                negative_inf_value = torch.finfo(cross_attention_mask.dtype).min
                full_text_row_masked_out_mask = (
                    (cross_attention_mask != negative_inf_value).any(dim=-1).type_as(cross_attention_mask)[..., None]
                )
                cross_attention_mask = cross_attention_mask * full_text_row_masked_out_mask

        if vision_position_ids is None and cross_attention_states is not None and input_ids is not None:
            vision_position_ids, position_ids, rope_deltas = self.compute_vision_position_ids(
                input_ids,
                position_ids,
                vision_token_info,
                cross_attention_states,
                attention_mask
            )

        outputs = self.language_model(
            input_ids=None,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            cross_attention_states=cross_attention_states,
            cross_attention_mask=cross_attention_mask,
            vision_position_ids=vision_position_ids,
            full_text_row_masked_out_mask=full_text_row_masked_out_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )

        if not return_dict:
            last_hidden_state = outputs[0]
            model_outputs = (
                last_hidden_state,
                outputs[1] if len(outputs) > 1 else past_key_values,
            )
            if output_hidden_states:
                model_outputs += (outputs[2],)
            if output_attentions:
                attn_idx = 3 if output_hidden_states else 2
                model_outputs += (outputs[attn_idx],)
            model_outputs += (vision_token_info, rope_deltas)
            return model_outputs

        return MossVLModelOutputWithPast(
            last_hidden_state=outputs.last_hidden_state,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            vision_token_info=vision_token_info,
            rope_deltas=rope_deltas,
        )





@auto_docstring(
    custom_intro="""
    The MossVL model with a language modeling head on top, for conditional generation tasks.
    Combines Qwen3VL vision encoder with LLM via cross-attention layers.
    """
)
class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
    # transformers 5.x expects a dict[target, source]; MossVL does not tie
    # lm_head to the embeddings (config.tie_word_embeddings is False), so the
    # mapping is empty. The legacy list format ["lm_head.weight"] breaks
    # save_pretrained in transformers>=5.
    _tied_weights_keys: dict[str, str] = {}
    config: MossVLConfig
    _checkpoint_conversion_mapping = {}
    accepts_loss_kwargs = False
    def __init__(self, config):
        super().__init__(config)
        self.model = MossVLModel(config)
        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)

        self.post_init()

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    def set_decoder(self, decoder):
        self.model.set_decoder(decoder)

    def get_decoder(self):
        return self.model.get_decoder()


    def get_vision_features(
        self, 
        pixel_values: torch.FloatTensor, 
        grid_thw: Optional[torch.LongTensor] = None,
        media_nums_per_sample: Optional[List[int]] = None
    ):
        """
        Get vision features for images and videos (merged).
        
        Args:
            pixel_values: vision pixel values (images and videos merged)
            grid_thw: [num_media, 3] tensor with (t, h, w) for each media item
            media_nums_per_sample: List indicating how many media items each sample has
        Returns:
            vision_embeds: [batch_size, max_seq_len, hidden_size]
            vision_token_info: List[Dict] with media positions and padding info for each sample
        """
        return self.model.get_vision_features(pixel_values, grid_thw, media_nums_per_sample)

    @property
    def language_model(self):
        return self.model.language_model

    @property
    def visual(self):
        return self.model.visual

    @auto_docstring
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.Tensor] = None,
        grid_thw: Optional[torch.LongTensor] = None,
        media_nums_per_sample: Optional[List[int]] = None,
        vision_position_ids: Optional[torch.LongTensor] = None,
        cross_attention_mask: Optional[torch.Tensor] = None,
        vision_token_info: Optional[List[dict]] = None,
        rope_deltas: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Union[tuple, CausalLMOutputWithPast]:
        """
        Args:
            grid_thw (`torch.LongTensor` of shape `(num_media, 3)`, *optional*):
                Grid size for each media item in (temporal, height, width) format. Each row contains `[t, h, w]` 
                representing the number of temporal, height, and width patches for a media item (image or video).
            media_nums_per_sample (`List[int]`, *optional*):
                List indicating how many media items each sample in the batch has. For example, `[2, 1, 3]` means 
                the first sample has 2 media items, the second has 1, and the third has 3.
            vision_position_ids (`torch.LongTensor` of shape `(batch_size, vision_seq_len)`, *optional*):
                Position IDs for vision tokens used in cross-attention. These are computed from text position IDs 
                based on the positions of image/video tokens in the input text.
            cross_attention_mask (`torch.Tensor` of shape `(batch_size, 1, text_seq_len, vision_seq_len)`, *optional*):
                Attention mask for cross-attention between text and vision. Controls which vision tokens each text 
                token can attend to, enforcing causal visibility for video frames.
            vision_token_info (`List[dict]`, *optional*):
                Cached metadata describing how packed vision tokens were regrouped per sample. Reused across decode
                steps to expand cross-attention masks without re-running the vision encoder.
            rope_deltas (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Cached multimodal RoPE offsets returned by the base model during prefill and reused during decode.
        """
        cache_position = kwargs.pop("cache_position", None)
        outputs = self.model(
            input_ids=input_ids,
            pixel_values=pixel_values,
            grid_thw=grid_thw,
            media_nums_per_sample=media_nums_per_sample,
            position_ids=position_ids,
            attention_mask=attention_mask,
            vision_position_ids=vision_position_ids,
            cross_attention_mask=cross_attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            vision_token_info=vision_token_info,
            rope_deltas=rope_deltas,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
            **kwargs,
        )

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state

        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)

        if not return_dict:
            output = (logits,)
            output += outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return MossVLCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            vision_token_info=outputs.vision_token_info,
            rope_deltas=outputs.rope_deltas,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        position_ids=None,
        use_cache=True,
        pixel_values=None,
        grid_thw=None,
        media_nums_per_sample=None,  # One video is one meida.
        vision_position_ids=None,
        vision_token_info=None,
        rope_deltas=None,
        cross_attention_mask=None,
        **kwargs,
    ):
        """
        Prepare inputs for generation.
        
        Note: Currently only supports offline visual understanding, meaning all multimodal 
        content must be provided before generation starts. We don't support adding new 
        images/videos during generation (streaming mode).
        
        Args:
            media_nums_per_sample: One video counts as one media item (regardless of frame count)
        """
        kwargs.pop("cache_position", None)
        model_inputs = super().prepare_inputs_for_generation(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            position_ids=position_ids,
            pixel_values=pixel_values,
            grid_thw=grid_thw,
            media_nums_per_sample=media_nums_per_sample,
            use_cache=use_cache,
            **kwargs,
        )

        model_input = model_inputs.get("input_ids")
        if model_input is None:
            model_input = model_inputs.get("inputs_embeds")
        current_length = model_input.shape[1]
        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0

        # Let the model recompute multimodal position ids from the current cache length.
        model_inputs["position_ids"] = None
        model_inputs["vision_token_info"] = vision_token_info
        model_inputs["rope_deltas"] = rope_deltas
        
        # Handle cross attention mask
        if cross_attention_mask is not None:
            # Slice to the current text slice on text dimension (dim=2).
            # Shape: [batch, 1, text_len, vision_len] -> [batch, 1, current_len, vision_len]
            cross_attention_mask = cross_attention_mask[:, :, -current_length:, :]
            model_inputs["cross_attention_mask"] = cross_attention_mask

        # Vision inputs are only needed in prefill stage.
        # In decode stage, vision features are retrieved from cross attention cache
        if past_seen_tokens > 0:
            model_inputs["pixel_values"] = None
            model_inputs["grid_thw"] = None 
            model_inputs["media_nums_per_sample"] = None
            model_inputs["vision_position_ids"] = None
            
        else:
            # In prefill stage, include all vision-related inputs
            model_inputs["vision_position_ids"] = vision_position_ids

        return model_inputs

    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder, **kwargs):
        """
        Update model kwargs for generation, extending cross_attention_mask for the newly generated token.
        
        In offline mode (all multimodal content provided before generation):
        - Each newly generated token should have the same cross_attention_mask pattern as the previous token
        - This ensures all generated tokens can attend to all vision tokens that were visible before
        """
        cross_attention_mask_prev = model_kwargs.get("cross_attention_mask", None)
        
        model_kwargs = super()._update_model_kwargs_for_generation(
            outputs=outputs,
            model_kwargs=model_kwargs,
            is_encoder_decoder=is_encoder_decoder,
            **kwargs,
        )
        
        if cross_attention_mask_prev is not None:
            model_kwargs["cross_attention_mask"] = cross_attention_mask_prev

        if getattr(outputs, "vision_token_info", None) is not None:
            model_kwargs["vision_token_info"] = outputs.vision_token_info
        if getattr(outputs, "rope_deltas", None) is not None:
            model_kwargs["rope_deltas"] = outputs.rope_deltas
        
        return model_kwargs


__all__ = [
    "MossVLVisionModel",
    "MossVLForConditionalGeneration",
    "MossVLModel",
    "MossVLPreTrainedModel",
    "MossVLTextModel",
]