diff --git "a/modeling_moss_vl.py" "b/modeling_moss_vl.py"
new file mode 100644--- /dev/null
+++ "b/modeling_moss_vl.py"
@@ -0,0 +1,3110 @@
+# 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
+import queue
+import threading
+from typing import Any, Callable, Dict, Optional, Union, Tuple, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.generation.stopping_criteria import StoppingCriteria, StoppingCriteriaList
+from transformers.generation.streamers import TextIteratorStreamer
+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 OutputRecorder
+
+from .configuration_moss_vl import MossVLConfig, MossVLTextConfig, MossVLVisionConfig
+
+
+
+logger = logging.get_logger(__name__)
+
+_OFFLINE_SYSTEM_PROMPTS = {
+ "no_thinking": {
+ "text_image": "You are a helpful AI assistant. Respond to the user's request based on the provided text and/or images.",
+ "video": "You are a helpful AI assistant specializing in video analysis. Respond to the user's request based on the provided video content.",
+ },
+ "deep_thinking": {
+ "text_image": "A conversation between User and Assistant. The user makes a request, and the assistant responds to it based on the provided text and/or images. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process hereanswer here.",
+ "video": "A conversation between User and Assistant specializing in video analysis. The user makes a request, and the assistant responds to it based on the provided video content. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process hereanswer here.",
+ },
+}
+
+
+class _OfflineCancelStoppingCriteria(StoppingCriteria):
+ def __init__(self, cancel_event: threading.Event):
+ self.cancel_event = cancel_event
+
+ def __call__(self, input_ids, scores, **kwargs) -> bool:
+ return self.cancel_event.is_set()
+
+
+class _OfflineQueueStreamer(TextIteratorStreamer):
+ def __init__(self, tokenizer, output_text_queue: "queue.Queue[str]"):
+ super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
+ self.output_text_queue = output_text_queue
+ self.collected_chunks: List[str] = []
+
+ def on_finalized_text(self, text: str, stream_end: bool = False):
+ if text:
+ self.collected_chunks.append(text)
+ self.output_text_queue.put(text)
+ super().on_finalized_text(text, stream_end=stream_end)
+
+
+_OFFLINE_THINKING_MODE_ALIASES = {
+ "no_thinking": "no_thinking",
+ "default": "no_thinking",
+ "standard": "no_thinking",
+ "deep_thinking": "deep_thinking",
+ "thinking": "deep_thinking",
+ "reasoning": "deep_thinking",
+}
+
+_OFFLINE_SYSTEM_PROMPT_TYPE_ALIASES = {
+ "text_image": "text_image",
+ "text-image": "text_image",
+ "image_text": "text_image",
+ "image-text": "text_image",
+ "text": "text_image",
+ "image": "text_image",
+ "video": "video",
+}
+
+
+@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
+
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ 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: List[torch.Tensor] = []) -> 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
+
+ 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 = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ if self.config._attn_implementation == "flash_attention_2":
+ 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__()
+ # BC: "rope_type" was originally "type"
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.rope_type = config.rope_scaling.get("rope_type", "default")
+ else:
+ self.rope_type = "default"
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.original_inv_freq = self.inv_freq
+
+
+ if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+ self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])
+ else:
+ self.mrope_section = [24, 20, 20]
+
+ 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:
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ 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, # vision_cache_position
+ 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, {"cache_position": cache_position}
+ )
+
+ 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!"
+ )
+
+ attention_interface: Callable = eager_attention_forward
+ if self.config._attn_implementation != "eager":
+ # 如果是flash attention，走sdpa_attention_forward
+ if self.config._attn_implementation == "flash_attention_3" or self.config._attn_implementation == "flash_attention_2":
+ attention_interface = ALL_ATTENTION_FUNCTIONS["sdpa"]
+ else:
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+ 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_cache_position: Optional[torch.LongTensor] = None,
+ vision_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ # Self Attention
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ hidden_states, _ = 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
+
+ return hidden_states
+
+
+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_cache_position: Optional[torch.LongTensor] = None,
+ vision_position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ # Cross Attention
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, _ = 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=vision_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
+
+ return hidden_states
+
+
+
+
+@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.
+
+ Note: For loading pretrained weights:
+ - Cross attention: can be initialized from the previous layer's self attention weights
+ """
+ std = getattr(self.config, "initializer_range", 0.02)
+ if hasattr(self.config, "text_config") and hasattr(self.config.text_config, "initializer_range"):
+ std = self.config.text_config.initializer_range
+
+ if isinstance(module, MossVLVisionPatchMerger):
+ # Initialize merger weights
+ # Input: hidden_size * (1 + num_deepstack_features) -> Output: out_hidden_size
+ # This projection handles concatenated features, so we might want specific initialization
+ module.linear_fc1.weight.data.normal_(mean=0.0, std=std)
+ module.linear_fc2.weight.data.normal_(mean=0.0, std=std)
+ if module.linear_fc1.bias is not None:
+ module.linear_fc1.bias.data.zero_()
+ if module.linear_fc2.bias is not None:
+ module.linear_fc2.bias.data.zero_()
+
+ # Initialize separate LayerNorms
+ if hasattr(module, "norms"):
+ for norm in module.norms:
+ if hasattr(norm, "weight") and norm.weight is not None:
+ norm.weight.data.fill_(1.0)
+ if hasattr(norm, "bias") and norm.bias is not None:
+ norm.bias.data.zero_()
+
+
+
+
+
+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]
+
+ idx_list = [[] for _ in range(4)]
+ weight_list = [[] 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)
+ w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w)
+
+ 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_list[i].extend(indices[i].tolist())
+ weight_list[i].extend(weights[i].tolist())
+
+ idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=self.pos_embed.weight.device)
+ weight_tensor = torch.tensor(
+ weight_list, dtype=self.pos_embed.weight.dtype, device=self.pos_embed.weight.device
+ )
+ 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,
+ vision_cache_position: Optional[torch.LongTensor] = 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)`.
+ vision_cache_position (`torch.LongTensor`, *optional*):
+ Cache position for vision tokens. Shape: `(vision_seq_len,)`.
+ """
+ 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,
+ input_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 vision_cache_position is None:
+ # TODO：use cache_position now
+ vision_cache_position = cache_position
+
+ 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)
+
+
+ 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_cache_position=vision_cache_position,
+ vision_position_embeddings=vision_position_embeddings,
+ **kwargs,
+ )
+ hidden_states = layer_outputs
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@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)
+ self.vision_token_info = None # cache vision_token_info here for decode stage
+ self.rope_deltas = None # cache position deltas for decode stage
+
+ # 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 = []
+ 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.extend([tokens_per_frame_with_sep] * num_frames)
+
+ num_valid_frames = len(repeats)
+ 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])
+ if valid_mask_frames < num_valid_frames:
+ repeats = repeats[:valid_mask_frames]
+
+ # Extract valid columns for this sample
+ # (1, text_len, valid_mask_frames)
+ source_mask = cross_attention_mask[i, :, :, :valid_mask_frames]
+
+ # Convert repeats to tensor
+ repeats_tensor = torch.tensor(repeats, device=cross_attention_mask.device)
+
+ # 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,
+ cache_position: 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
+ cache_position: (seq_len,), position in cache
+
+ 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: use cached rope_deltas for fast computation
+ if cache_position is not None and cache_position[0] != 0 and self.rope_deltas is not None:
+ # In decode, position = cache_position + rope_deltas
+ # rope_deltas is per-sample: (batch_size,)
+ position_ids = torch.arange(seq_len, device=device, dtype=torch.long)
+ position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) # (batch, seq_len)
+
+ # Add cache_position offset
+ if cache_position is not None:
+ position_ids = position_ids + cache_position[0]
+
+ # Add rope_deltas (position offset due to vision tokens)
+ # self.rope_deltas shape: (batch_size,), need to unsqueeze for broadcasting
+ position_ids = position_ids + self.rope_deltas.unsqueeze(1) # (batch, seq_len)
+
+ # Expand to 3D: (3, batch, seq_len)
+ position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+ return position_ids
+
+ # 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 = position_ids.device if position_ids is not None else input_ids.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().to(device)
+
+ # 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 = []
+ flat_eff_w = []
+ flat_vis_starts = []
+ flat_batch_indices = []
+
+ # 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
+
+ # Generate entries for all frames in this media
+ for f in range(num_frames):
+ flat_eff_h.append(eh)
+ flat_eff_w.append(ew)
+ flat_vis_starts.append(start + f * stride)
+ flat_batch_indices.append(b_idx)
+
+ # 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) == 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
+
+ # Align lengths (handle truncation if text has fewer tokens or vice versa)
+ num_matches = min(len(flat_eff_h), len(image_token_indices))
+
+ # Convert to tensors
+ flat_eff_h = torch.tensor(flat_eff_h[:num_matches], device=device, dtype=torch.long)
+ flat_eff_w = torch.tensor(flat_eff_w[:num_matches], device=device, dtype=torch.long)
+ flat_vis_starts = torch.tensor(flat_vis_starts[:num_matches], device=device, dtype=torch.long)
+
+ # 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
+
+ def get_vision_features_chunked(
+ self,
+ pixel_values: torch.FloatTensor,
+ grid_thw: Optional[torch.LongTensor] = None,
+ media_nums_per_sample: Optional[List[int]] = None,
+ vision_chunked_length: Optional[int] = None,
+ ):
+ """
+ Chunk the visual encoder forward by media items, then reuse the same
+ packed-to-batch conversion logic. This keeps output semantics identical
+ to `get_vision_features(...)` while reducing prefill memory pressure.
+ """
+ if (
+ vision_chunked_length is None
+ or vision_chunked_length <= 0
+ or grid_thw is None
+ or grid_thw.shape[0] <= vision_chunked_length
+ ):
+ return self.get_vision_features(pixel_values, grid_thw, media_nums_per_sample)
+
+ pixel_values = pixel_values.type(self.visual.dtype)
+ token_counts = (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).tolist()
+
+ hidden_state_chunks = []
+ token_offset = 0
+ for media_start in range(0, grid_thw.shape[0], vision_chunked_length):
+ media_end = min(media_start + vision_chunked_length, grid_thw.shape[0])
+ chunk_grid_thw = grid_thw[media_start:media_end]
+ chunk_token_count = sum(token_counts[media_start:media_end])
+ chunk_pixel_values = pixel_values[token_offset:token_offset + chunk_token_count]
+ token_offset += chunk_token_count
+
+ hidden_state_chunks.append(
+ self.visual(
+ chunk_pixel_values,
+ grid_thw=chunk_grid_thw,
+ )
+ )
+
+ hidden_states = torch.cat(hidden_state_chunks, dim=0)
+ 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,
+ cache_position: Optional[torch.LongTensor] = 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_chunked_length (`int`, *optional*):
+ Number of media items to process per visual-encoder chunk during prefill. This only changes
+ how the vision tower is executed, not the final prompt or decoding logic.
+ """
+ vision_chunked_length = kwargs.pop("vision_chunked_length", None)
+ 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
+ num_vision_tokens = 0
+
+ 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_chunked(
+ pixel_values,
+ grid_thw,
+ media_nums_per_sample,
+ vision_chunked_length=vision_chunked_length,
+ )
+
+ # vision_embeds: [batch_size, max_seq_len, hidden_size]
+ cross_attention_states = vision_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ num_vision_tokens = cross_attention_states.shape[1]
+
+ # Cache vision_token_info for decode stage (prefill only)
+
+ self.vision_token_info = vision_token_info
+ else:
+ # In decode stage, use cached vision_token_info
+ vision_token_info = self.vision_token_info
+
+ # 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,
+ cache_position=cache_position,
+ )
+
+ # 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
+ )
+
+ # Cache rope_deltas for decode stage (only in prefill)
+ # rope_deltas = max_position - sequence_length
+ # This allows fast position computation in decode: position = cache_position + rope_deltas
+ if cache_position is not None and cache_position[0] == 0:
+ self.rope_deltas = rope_deltas
+
+
+
+ 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,
+ **kwargs,
+ )
+
+ 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=self.vision_token_info,
+ rope_deltas=self.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):
+ _tied_weights_keys = ["lm_head.weight"]
+ 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,
+ cache_position: Optional[torch.LongTensor] = None,
+ vision_chunked_length: Optional[int] = None,
+ logits_to_keep: Union[int, torch.Tensor] = 0,
+ **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_chunked_length (`int`, *optional*):
+ Number of media items to process per visual-encoder chunk during prefill. This only changes
+ how the vision tower is executed, not the final prompt or decoding logic.
+ """
+ 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,
+ cache_position=cache_position,
+ vision_chunked_length=vision_chunked_length,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+
+ 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)
+
+ 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,
+ cache_position=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,
+ cross_attention_mask=None,
+ vision_chunked_length=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)
+ """
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ cache_position=cache_position,
+ 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,
+ )
+
+
+ # For decoding stage, if position_ids are generated by GenerationMixin (2D),
+ # we can set them to None to let forward recompute them from cache_position.
+ model_inputs["position_ids"] = None
+
+ # Handle cross attention mask
+ if cross_attention_mask is not None:
+ # Slice to current sequence length on text dimension (dim=2)
+ # Shape: [batch, 1, text_len, vision_len] -> [batch, 1, cache_len, vision_len]
+ cross_attention_mask = cross_attention_mask[:, :, -cache_position.shape[0]:, :]
+ model_inputs["cross_attention_mask"] = cross_attention_mask
+
+ # Vision inputs are only needed in prefill stage (cache_position[0] == 0)
+ # In decode stage, vision features are retrieved from cross attention cache
+ if cache_position[0] != 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
+ model_inputs["vision_chunked_length"] = vision_chunked_length
+
+ 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,
+ )
+
+ # Extend cross_attention_mask for the new token
+ # Copy the last token's mask pattern for the newly generated token
+ if cross_attention_mask_prev is not None:
+ model_kwargs["cross_attention_mask"] = torch.cat(
+ [cross_attention_mask_prev, cross_attention_mask_prev[:, :, -1:, :]],
+ dim=2 # Concatenate along text sequence dimension
+ )
+
+ return model_kwargs
+
+ @staticmethod
+ def _offline_flatten_content_with_vision_tokens(content) -> str:
+ if isinstance(content, str):
+ return content
+ if not isinstance(content, list):
+ return str(content) if content else ""
+
+ parts = []
+ for item in content:
+ if isinstance(item, dict):
+ if item.get("type") == "image" or "image" in item:
+ parts.append("<|image|>")
+ elif item.get("type") == "video" or "video" in item:
+ parts.append("<|video|>")
+ if "text" in item:
+ parts.append(str(item["text"]))
+ elif isinstance(item, str):
+ parts.append(item)
+ return "".join(parts)
+
+ @staticmethod
+ def _offline_sanitize_prompt_text(processor, text: Any) -> str:
+ if text is None:
+ return ""
+
+ sanitized = str(text)
+ replacements = [
+ (getattr(processor, "image_placeholder", None), ""),
+ (getattr(processor, "video_placeholder", None), ""),
+ (getattr(processor, "image_token", None), ""),
+ (getattr(processor, "video_token", None), ""),
+ ]
+ for needle, replacement in replacements:
+ if needle:
+ sanitized = sanitized.replace(needle, replacement)
+ return sanitized.lstrip("\n")
+
+ def _offline_sanitize_message_content(self, processor, content: Any) -> Any:
+ if isinstance(content, str):
+ return self._offline_sanitize_prompt_text(processor, content)
+ if not isinstance(content, list):
+ return content
+
+ sanitized_items = []
+ for item in content:
+ if isinstance(item, dict):
+ item_copy = dict(item)
+ if "text" in item_copy:
+ item_copy["text"] = self._offline_sanitize_prompt_text(processor, item_copy.get("text"))
+ sanitized_items.append(item_copy)
+ elif isinstance(item, str):
+ sanitized_items.append(self._offline_sanitize_prompt_text(processor, item))
+ else:
+ sanitized_items.append(item)
+ return sanitized_items
+
+ def _offline_prepare_messages(self, processor, query: Dict[str, Any]) -> List[Dict[str, Any]]:
+ messages = query.get("messages")
+ if messages:
+ prepared_messages = []
+ for message in messages:
+ if not isinstance(message, dict):
+ continue
+ message_copy = dict(message)
+ message_copy["content"] = self._offline_sanitize_message_content(
+ processor,
+ message_copy.get("content", ""),
+ )
+ prepared_messages.append(message_copy)
+ if prepared_messages:
+ return prepared_messages
+
+ prompt = self._offline_sanitize_prompt_text(processor, query.get("prompt", ""))
+ images = list(query.get("images") or [])
+ videos = list(query.get("videos") or [])
+
+ content = []
+ for image in images:
+ content.append({"type": "image", "image": image})
+ for video in videos:
+ content.append({"type": "video", "video": video})
+ if prompt:
+ content.append({"type": "text", "text": prompt.lstrip("\n")})
+
+ if not content:
+ content = [{"type": "text", "text": ""}]
+
+ return [{"role": "user", "content": content}]
+
+ def _offline_prepare_input_text(self, processor, messages: List[Dict[str, Any]]) -> str:
+ processed_messages = []
+ for message in messages:
+ message_copy = dict(message)
+ message_copy["content"] = self._offline_flatten_content_with_vision_tokens(
+ message_copy.get("content", "")
+ )
+ processed_messages.append(message_copy)
+ return processor.apply_chat_template(
+ processed_messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ )
+
+ @staticmethod
+ def _offline_collect_media(messages: List[Dict[str, Any]]) -> tuple[List[Any], List[Any]]:
+ all_images: List[Any] = []
+ all_videos: List[Any] = []
+
+ for message in messages:
+ content = message.get("content")
+ if isinstance(content, list):
+ for item in content:
+ if not isinstance(item, dict):
+ continue
+ if item.get("type") == "image" or "image" in item:
+ image = item.get("image") or item.get("image_url")
+ if image is not None:
+ all_images.append(image)
+ elif item.get("type") == "video" or "video" in item:
+ video = item.get("video")
+ if video is not None:
+ all_videos.append(video)
+
+ return all_images, all_videos
+
+ def _offline_build_processor_kwargs(
+ self,
+ input_text: Union[str, List[str]],
+ all_images: List[Any],
+ all_videos: List[Any],
+ media_kwargs: Dict[str, Any],
+ ) -> Dict[str, Any]:
+ processor_kwargs: Dict[str, Any] = {
+ "text": input_text,
+ "images": all_images or None,
+ "videos": all_videos or None,
+ "return_tensors": "pt",
+ "padding": False,
+ }
+
+ if media_kwargs.get("min_pixels") is not None:
+ processor_kwargs["min_pixels"] = media_kwargs["min_pixels"]
+ if media_kwargs.get("max_pixels") is not None:
+ processor_kwargs["max_pixels"] = media_kwargs["max_pixels"]
+ if media_kwargs.get("video_fps") is not None:
+ processor_kwargs["video_fps"] = media_kwargs["video_fps"]
+
+ min_frames = media_kwargs.get("min_frames", media_kwargs.get("video_minlen"))
+ max_frames = media_kwargs.get("max_frames", media_kwargs.get("video_maxlen"))
+ if min_frames is not None:
+ processor_kwargs["min_frames"] = min_frames
+ if max_frames is not None:
+ processor_kwargs["max_frames"] = max_frames
+
+ return processor_kwargs
+
+ def _offline_prepare_inputs(self, processor, query: Dict[str, Any]):
+ messages = self._offline_prepare_messages(processor, query)
+ input_text = self._offline_prepare_input_text(processor, messages)
+ all_images, all_videos = self._offline_collect_media(messages)
+ media_kwargs = dict(query.get("media_kwargs") or {})
+ processor_kwargs = self._offline_build_processor_kwargs(
+ input_text,
+ all_images,
+ all_videos,
+ media_kwargs,
+ )
+
+ image_proc = getattr(processor, "image_processor", None)
+ video_proc = getattr(processor, "video_processor", None)
+ modified_multi_image = False
+ modified_video = False
+
+ try:
+ multi_image_max_pixels = media_kwargs.get("multi_image_max_pixels")
+ if multi_image_max_pixels is not None and image_proc is not None:
+ orig_multi_image_max_pixels = getattr(image_proc, "multi_image_max_pixels", None)
+ image_proc.multi_image_max_pixels = multi_image_max_pixels
+ modified_multi_image = True
+
+ video_max_pixels = media_kwargs.get("video_max_pixels")
+ if video_max_pixels is not None and video_proc is not None:
+ orig_video_max_pixels = getattr(video_proc, "video_max_pixels", None)
+ video_proc.video_max_pixels = video_max_pixels
+ modified_video = True
+
+ inputs = processor(**processor_kwargs)
+ finally:
+ if modified_multi_image and image_proc is not None:
+ image_proc.multi_image_max_pixels = orig_multi_image_max_pixels
+ if modified_video and video_proc is not None:
+ video_proc.video_max_pixels = orig_video_max_pixels
+
+ text_device = self.get_input_embeddings().weight.device
+ vision_device = self.visual.patch_embed.proj.weight.device
+ vision_input_keys = {"pixel_values", "grid_thw"}
+
+ for key, value in list(inputs.items()):
+ if not isinstance(value, torch.Tensor):
+ continue
+
+ target_device = vision_device if key in vision_input_keys else text_device
+ moved_value = value.to(target_device)
+ if moved_value.dtype == torch.float32:
+ moved_value = moved_value.to(torch.bfloat16)
+ inputs[key] = moved_value
+
+ return inputs, input_text
+
+ def _offline_build_session_messages(
+ self,
+ processor,
+ query: Dict[str, Any],
+ session_messages: List[Dict[str, Any]],
+ ) -> List[Dict[str, Any]]:
+ has_explicit_messages = bool(query.get("messages"))
+ if has_explicit_messages and not query.get("append_messages_to_session", False):
+ base_messages: List[Dict[str, Any]] = []
+ else:
+ base_messages = [dict(message) for message in session_messages]
+
+ turn_messages = self._offline_prepare_messages(processor, query)
+ has_system_message = any(
+ isinstance(message, dict) and message.get("role") == "system"
+ for message in (base_messages + turn_messages)
+ )
+
+ should_add_system_prompt = (
+ query.get("use_default_system_prompt", False)
+ or query.get("system_prompt") is not None
+ or query.get("system_prompt_type") is not None
+ or query.get("thinking_mode") is not None
+ )
+
+ if not base_messages and not has_system_message and should_add_system_prompt:
+ system_prompt = self._offline_resolve_system_prompt(query, turn_messages)
+ if system_prompt is not None:
+ base_messages.append({"role": "system", "content": system_prompt})
+
+ return base_messages + turn_messages
+
+ @staticmethod
+ def _offline_query_contains_video(query: Dict[str, Any], messages: List[Dict[str, Any]]) -> bool:
+ if query.get("videos"):
+ return True
+
+ for message in messages:
+ content = message.get("content") if isinstance(message, dict) else None
+ if isinstance(content, list) and any(
+ isinstance(item, dict) and (item.get("type") == "video" or "video" in item)
+ for item in content
+ ):
+ return True
+ return False
+
+ @staticmethod
+ def _offline_normalize_thinking_mode(value: Optional[str]) -> str:
+ if value is None:
+ return "no_thinking"
+
+ normalized = _OFFLINE_THINKING_MODE_ALIASES.get(str(value).strip().lower())
+ if normalized is None:
+ allowed = ", ".join(sorted(set(_OFFLINE_THINKING_MODE_ALIASES.values())))
+ raise ValueError(f"Unsupported thinking_mode: {value!r}. Supported values: {allowed}")
+ return normalized
+
+ @staticmethod
+ def _offline_normalize_system_prompt_type(value: Optional[str], has_video: bool) -> str:
+ if value is None:
+ return "video" if has_video else "text_image"
+
+ normalized_key = str(value).strip().lower().replace("/", "_").replace(" ", "_")
+ while "__" in normalized_key:
+ normalized_key = normalized_key.replace("__", "_")
+
+ normalized = _OFFLINE_SYSTEM_PROMPT_TYPE_ALIASES.get(normalized_key)
+ if normalized is None:
+ allowed = ", ".join(sorted(set(_OFFLINE_SYSTEM_PROMPT_TYPE_ALIASES.values())))
+ raise ValueError(f"Unsupported system_prompt_type: {value!r}. Supported values: {allowed}")
+ return normalized
+
+ def _offline_resolve_system_prompt(
+ self,
+ query: Dict[str, Any],
+ turn_messages: List[Dict[str, Any]],
+ ) -> Optional[str]:
+ explicit_system_prompt = query.get("system_prompt")
+ if explicit_system_prompt is not None:
+ return str(explicit_system_prompt)
+
+ has_video = self._offline_query_contains_video(query, turn_messages)
+ thinking_mode = self._offline_normalize_thinking_mode(query.get("thinking_mode"))
+ system_prompt_type = self._offline_normalize_system_prompt_type(
+ query.get("system_prompt_type"),
+ has_video=has_video,
+ )
+ return _OFFLINE_SYSTEM_PROMPTS[thinking_mode][system_prompt_type]
+
+ @staticmethod
+ def _offline_finalize_session_messages(
+ working_messages: List[Dict[str, Any]],
+ assistant_text: str,
+ ) -> List[Dict[str, Any]]:
+ next_messages = [dict(message) for message in working_messages]
+ next_messages.append({"role": "assistant", "content": assistant_text})
+ return next_messages
+
+ def _offline_prepare_generation(self, processor, query: Dict[str, Any]):
+ inputs, input_text = self._offline_prepare_inputs(processor, query)
+ generate_kwargs = dict(query.get("generate_kwargs") or {})
+
+ max_new_tokens = generate_kwargs.pop("max_new_tokens", 1024)
+ temperature = generate_kwargs.pop("temperature", 1.0)
+ top_k = generate_kwargs.pop("top_k", 50)
+ top_p = generate_kwargs.pop("top_p", 1.0)
+ repetition_penalty = generate_kwargs.pop("repetition_penalty", 1.0)
+ do_sample = generate_kwargs.pop("do_sample", False)
+ vision_chunked_length = generate_kwargs.pop("vision_chunked_length", None)
+
+ if temperature is None:
+ temperature = 1.0
+ if temperature <= 0:
+ temperature = 1.0
+ do_sample = False
+
+ call_kwargs = dict(
+ max_new_tokens=max_new_tokens,
+ temperature=temperature,
+ top_k=top_k,
+ top_p=top_p,
+ repetition_penalty=repetition_penalty,
+ do_sample=do_sample,
+ vision_chunked_length=vision_chunked_length,
+ **generate_kwargs,
+ )
+ return inputs, input_text, call_kwargs
+
+ @staticmethod
+ def _offline_normalize_shared_mapping(
+ values: List[Dict[str, Any]],
+ mapping_name: str,
+ ) -> Dict[str, Any]:
+ normalized_values = [dict(value or {}) for value in values]
+ if not normalized_values:
+ return {}
+
+ all_keys = set()
+ for value in normalized_values:
+ all_keys.update(value.keys())
+
+ merged: Dict[str, Any] = {}
+ mismatched_keys: List[str] = []
+ for key in sorted(all_keys):
+ unique_values = {repr(value.get(key)) for value in normalized_values}
+ if len(unique_values) > 1:
+ mismatched_keys.append(key)
+ else:
+ merged[key] = normalized_values[0].get(key)
+
+ if mismatched_keys:
+ mismatch_text = ", ".join(mismatched_keys)
+ raise ValueError(
+ f"All batch queries must share the same {mapping_name}. "
+ f"Mismatched keys: {mismatch_text}"
+ )
+ return merged
+
+ def _offline_prepare_batch_generation(
+ self,
+ processor,
+ queries: List[Dict[str, Any]],
+ session_states: Optional[List[List[Dict[str, Any]]]] = None,
+ ):
+ if not queries:
+ raise ValueError("`queries` must contain at least one query.")
+
+ if session_states is None:
+ session_states = [[] for _ in queries]
+ elif len(session_states) != len(queries):
+ raise ValueError("`session_states` must have the same length as `queries`.")
+
+ working_messages_list: List[List[Dict[str, Any]]] = []
+ input_texts: List[str] = []
+ all_images_per_query: List[List[Any]] = []
+ all_videos_per_query: List[List[Any]] = []
+
+ for query, session_state in zip(queries, session_states):
+ if not isinstance(query, dict):
+ raise TypeError("Each batch query must be a dict.")
+ if query.get("stop_offline_generate"):
+ raise ValueError("`stop_offline_generate` is not supported in offline_batch_generate.")
+ if query.get("stream_output", query.get("stream", False)):
+ raise ValueError("Streaming is not supported in offline_batch_generate.")
+ if query.get("cancel_current_generate") or query.get("stop_generation"):
+ raise ValueError("Cancel / stop controls are not supported in offline_batch_generate.")
+
+ current_session = [] if query.get("reset_session") or query.get("clear_history") else session_state
+ working_messages = self._offline_build_session_messages(
+ processor,
+ query,
+ current_session,
+ )
+ working_messages_list.append(working_messages)
+ input_texts.append(self._offline_prepare_input_text(processor, working_messages))
+
+ all_images, all_videos = self._offline_collect_media(working_messages)
+ all_images_per_query.append(all_images)
+ all_videos_per_query.append(all_videos)
+
+ media_kwargs = self._offline_normalize_shared_mapping(
+ [query.get("media_kwargs") or {} for query in queries],
+ mapping_name="media_kwargs",
+ )
+ processor_kwargs = self._offline_build_processor_kwargs(
+ input_text=input_texts,
+ all_images=[image for images in all_images_per_query for image in images],
+ all_videos=[video for videos in all_videos_per_query for video in videos],
+ media_kwargs=media_kwargs,
+ )
+ processor_kwargs["padding"] = True
+
+ image_proc = getattr(processor, "image_processor", None)
+ video_proc = getattr(processor, "video_processor", None)
+ tokenizer = getattr(processor, "tokenizer", None)
+ modified_multi_image = False
+ modified_video = False
+ orig_padding_side = None
+
+ try:
+ multi_image_max_pixels = media_kwargs.get("multi_image_max_pixels")
+ if multi_image_max_pixels is not None and image_proc is not None:
+ orig_multi_image_max_pixels = getattr(image_proc, "multi_image_max_pixels", None)
+ image_proc.multi_image_max_pixels = multi_image_max_pixels
+ modified_multi_image = True
+
+ video_max_pixels = media_kwargs.get("video_max_pixels")
+ if video_max_pixels is not None and video_proc is not None:
+ orig_video_max_pixels = getattr(video_proc, "video_max_pixels", None)
+ video_proc.video_max_pixels = video_max_pixels
+ modified_video = True
+
+ if tokenizer is not None and hasattr(tokenizer, "padding_side"):
+ orig_padding_side = tokenizer.padding_side
+ tokenizer.padding_side = "left"
+
+ inputs = processor(**processor_kwargs)
+ finally:
+ if modified_multi_image and image_proc is not None:
+ image_proc.multi_image_max_pixels = orig_multi_image_max_pixels
+ if modified_video and video_proc is not None:
+ video_proc.video_max_pixels = orig_video_max_pixels
+ if tokenizer is not None and orig_padding_side is not None:
+ tokenizer.padding_side = orig_padding_side
+
+ text_device = self.get_input_embeddings().weight.device
+ vision_device = self.visual.patch_embed.proj.weight.device
+ vision_input_keys = {"pixel_values", "grid_thw"}
+
+ for key, value in list(inputs.items()):
+ if not isinstance(value, torch.Tensor):
+ continue
+
+ target_device = vision_device if key in vision_input_keys else text_device
+ moved_value = value.to(target_device)
+ if moved_value.dtype == torch.float32:
+ moved_value = moved_value.to(torch.bfloat16)
+ inputs[key] = moved_value
+
+ generate_kwargs = self._offline_normalize_shared_mapping(
+ [query.get("generate_kwargs") or {} for query in queries],
+ mapping_name="generate_kwargs",
+ )
+ max_new_tokens = generate_kwargs.pop("max_new_tokens", 1024)
+ temperature = generate_kwargs.pop("temperature", 1.0)
+ top_k = generate_kwargs.pop("top_k", 50)
+ top_p = generate_kwargs.pop("top_p", 1.0)
+ repetition_penalty = generate_kwargs.pop("repetition_penalty", 1.0)
+ do_sample = generate_kwargs.pop("do_sample", False)
+ vision_chunked_length = generate_kwargs.pop("vision_chunked_length", None)
+
+ if temperature is None:
+ temperature = 1.0
+ if temperature <= 0:
+ temperature = 1.0
+ do_sample = False
+
+ call_kwargs = dict(
+ max_new_tokens=max_new_tokens,
+ temperature=temperature,
+ top_k=top_k,
+ top_p=top_p,
+ repetition_penalty=repetition_penalty,
+ do_sample=do_sample,
+ vision_chunked_length=vision_chunked_length,
+ **generate_kwargs,
+ )
+ return inputs, input_texts, working_messages_list, call_kwargs
+
+ def offline_batch_generate(
+ self,
+ processor,
+ queries: List[Dict[str, Any]],
+ session_states: Optional[List[List[Dict[str, Any]]]] = None,
+ vision_chunked_length: int = 64,
+ ) -> Dict[str, Any]:
+ """
+ Batch offline generation for multiple independent samples.
+
+ This method supports:
+ - batched single-turn generation
+ - batched multi-turn continuation through `session_states`
+
+ It intentionally does not support queue-style controls such as:
+ - `stream_output`
+ - `cancel_current_generate`
+ - `stop_generation`
+ - `stop_offline_generate`
+ """
+ if not queries:
+ return {"results": [], "session_states": []}
+
+ prepared_queries = [dict(query) for query in queries]
+ for query in prepared_queries:
+ generate_kwargs = query.setdefault("generate_kwargs", {})
+ generate_kwargs.setdefault("vision_chunked_length", vision_chunked_length)
+ if session_states is None:
+ session_states = [[] for _ in prepared_queries]
+ elif len(session_states) != len(prepared_queries):
+ raise ValueError("`session_states` must have the same length as `queries`.")
+
+ tokenizer = getattr(processor, "tokenizer", None)
+ bucketed_indices: Dict[Any, List[int]] = {}
+ for index, (query, session_state) in enumerate(zip(prepared_queries, session_states)):
+ current_session = [] if query.get("reset_session") or query.get("clear_history") else session_state
+ working_messages = self._offline_build_session_messages(processor, query, current_session)
+ input_text = self._offline_prepare_input_text(processor, working_messages)
+
+ if tokenizer is not None:
+ token_ids = tokenizer(input_text, add_special_tokens=False)["input_ids"]
+ bucket_key = len(token_ids)
+ else:
+ bucket_key = len(input_text)
+ bucketed_indices.setdefault(bucket_key, []).append(index)
+
+ results: List[Optional[Dict[str, Any]]] = [None] * len(prepared_queries)
+ next_session_states: List[Optional[List[Dict[str, Any]]]] = [None] * len(prepared_queries)
+
+ for bucket_indices in bucketed_indices.values():
+ bucket_queries = [prepared_queries[index] for index in bucket_indices]
+ bucket_session_states = [session_states[index] for index in bucket_indices]
+ inputs, input_texts, working_messages_list, call_kwargs = self._offline_prepare_batch_generation(
+ processor,
+ bucket_queries,
+ session_states=bucket_session_states,
+ )
+
+ with torch.no_grad():
+ outputs = self.generate(
+ **inputs,
+ **call_kwargs,
+ )
+
+ input_seq_len = inputs["input_ids"].shape[1]
+ generated_tokens = outputs[:, input_seq_len:]
+ decoded_texts = processor.batch_decode(generated_tokens, skip_special_tokens=True)
+
+ for local_index, (query, input_text, working_messages, text) in enumerate(
+ zip(bucket_queries, input_texts, working_messages_list, decoded_texts)
+ ):
+ original_index = bucket_indices[local_index]
+ if query.get("persist_session", True):
+ next_session_state = self._offline_finalize_session_messages(working_messages, text)
+ else:
+ next_session_state = working_messages
+ next_session_states[original_index] = next_session_state
+ results[original_index] = {
+ "index": original_index,
+ "text": text,
+ "input_text": input_text,
+ "messages": working_messages,
+ }
+
+ return {
+ "results": [item for item in results if item is not None],
+ "session_states": [item for item in next_session_states if item is not None],
+ }
+
+ def _offline_generate_one(self, processor, query: Dict[str, Any]) -> str:
+ working_messages = self._offline_build_session_messages(processor, query, [])
+ generation_query = dict(query)
+ generation_query["messages"] = working_messages
+ inputs, _, call_kwargs = self._offline_prepare_generation(processor, generation_query)
+
+ with torch.no_grad():
+ outputs = self.generate(
+ **inputs,
+ **call_kwargs,
+ )
+
+ new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+ return processor.decode(new_tokens, skip_special_tokens=True)
+
+ def offline_generate(
+ self,
+ processor,
+ new_queries: "queue.Queue[dict]",
+ output_text_queue: "queue.Queue[str]",
+ vision_chunked_length: int = 64,
+ ) -> None:
+ """
+ HF-style offline inference wrapper aligned with the previous backend output path.
+
+ This method intentionally reuses the checkpoint's existing processor and
+ `generate()` flow so that outputs stay consistent with the old external
+ backend inference implementation.
+
+ Supported query keys include:
+ - `prompt` / `messages`
+ - `images` / `videos`
+ - `media_kwargs` / `generate_kwargs`
+ - `thinking_mode` (`no_thinking` or `deep_thinking`, plus compatible aliases)
+ - `system_prompt_type` (`text_image` or `video`, plus compatible aliases)
+ - `system_prompt` for an explicit override
+ - `stream_output` / `stream`
+ - `reset_session` / `clear_history`
+ - `cancel_current_generate` / `stop_generation` / `stop_offline_generate`
+ """
+ buffered_queries: List[Dict[str, Any]] = []
+ session_messages: List[Dict[str, Any]] = []
+
+ while True:
+ if buffered_queries:
+ query = buffered_queries.pop(0)
+ else:
+ query = new_queries.get()
+ if not isinstance(query, dict):
+ continue
+
+ if query.get("stop_offline_generate"):
+ break
+
+ if query.get("reset_session") or query.get("clear_history"):
+ session_messages = []
+
+ try:
+ generate_kwargs = query.setdefault("generate_kwargs", {})
+ generate_kwargs.setdefault("vision_chunked_length", vision_chunked_length)
+ working_messages = self._offline_build_session_messages(
+ processor,
+ query,
+ session_messages,
+ )
+
+ generation_query = dict(query)
+ generation_query["messages"] = working_messages
+ inputs, input_text, call_kwargs = self._offline_prepare_generation(processor, generation_query)
+
+ stream_output = bool(query.get("stream_output", query.get("stream", False)))
+ cancel_event = threading.Event()
+ stopping_criteria = StoppingCriteriaList([_OfflineCancelStoppingCriteria(cancel_event)])
+ generation_state: Dict[str, Any] = {}
+
+ if stream_output:
+ output_text_queue.put("<|round_start|>")
+ streamer = _OfflineQueueStreamer(getattr(processor, "tokenizer", processor), output_text_queue)
+ else:
+ streamer = None
+
+ def _run_generation():
+ try:
+ with torch.no_grad():
+ generation_state["outputs"] = self.generate(
+ **inputs,
+ stopping_criteria=stopping_criteria,
+ streamer=streamer,
+ **call_kwargs,
+ )
+ except Exception as exc:
+ generation_state["exception"] = exc
+
+ worker = threading.Thread(target=_run_generation, daemon=True)
+ worker.start()
+
+ stop_conversation_after_turn = False
+ while worker.is_alive():
+ try:
+ control_query = new_queries.get(timeout=0.1)
+ except queue.Empty:
+ continue
+
+ if not isinstance(control_query, dict):
+ continue
+
+ if control_query.get("cancel_current_generate") or control_query.get("stop_generation"):
+ cancel_event.set()
+ stop_conversation_after_turn = stop_conversation_after_turn or control_query.get("stop_offline_generate", False)
+ continue
+
+ if control_query.get("stop_offline_generate"):
+ cancel_event.set()
+ stop_conversation_after_turn = True
+ continue
+
+ buffered_queries.append(control_query)
+
+ worker.join()
+ was_cancelled = cancel_event.is_set()
+
+ if "exception" in generation_state:
+ raise generation_state["exception"]
+
+ if stream_output and streamer is not None:
+ text = "".join(streamer.collected_chunks)
+ else:
+ outputs = generation_state["outputs"]
+ new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+ text = processor.decode(new_tokens, skip_special_tokens=True)
+ output_text_queue.put(text)
+
+ if query.get("persist_session", True) and (not was_cancelled or query.get("persist_cancelled_turn", False)):
+ session_messages = self._offline_finalize_session_messages(working_messages, text)
+
+ output_text_queue.put("<|round_end|>")
+
+ if stop_conversation_after_turn:
+ break
+ except Exception as exc:
+ output_text_queue.put(f"[ERROR] {exc}")
+ output_text_queue.put("<|round_end|>")
+
+
+__all__ = [
+ "MossVLVisionModel",
+ "MossVLForConditionalGeneration",
+ "MossVLModel",
+ "MossVLPreTrainedModel",
+ "MossVLTextModel",
+]