Unify modeling with the 0708 standard (transformers 4.57 compatible)

modeling_moss_vl.py

@@ -111,12 +111,12 @@ logger = logging.get_logger(__name__)
 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 
+            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).
@@ -129,7 +129,7 @@ class MossVLModelOutputWithPast(ModelOutput):
             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
@@ -142,7 +142,7 @@ class MossVLModelOutputWithPast(ModelOutput):
 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).
@@ -159,7 +159,7 @@ class MossVLCausalLMOutputWithPast(ModelOutput):
         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
@@ -225,7 +225,7 @@ class MossVLVisionPatchMerger(nn.Module):
         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
@@ -233,9 +233,9 @@ class MossVLVisionPatchMerger(nn.Module):
             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)
@@ -244,14 +244,14 @@ class MossVLVisionPatchMerger(nn.Module):
         # 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)
 
@@ -480,7 +480,7 @@ class MossVLTextRotaryEmbedding(nn.Module):
 
         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)
 
@@ -628,7 +628,7 @@ class MossVLTextCrossAttention(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        cross_attention_states: torch.Tensor,       
+        cross_attention_states: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         past_key_values: Optional[Cache] = None,
         use_cache: bool = None,
@@ -638,7 +638,7 @@ class MossVLTextCrossAttention(nn.Module):
         **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)
@@ -653,13 +653,13 @@ class MossVLTextCrossAttention(nn.Module):
             # 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)
 
             # Rotate new keys once; cached keys already include vision RoPE.
-            
+
             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)
@@ -670,13 +670,13 @@ class MossVLTextCrossAttention(nn.Module):
                 # 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!"
@@ -771,7 +771,7 @@ class MossVLSelfAttentionDecoderLayer(GradientCheckpointingLayer):
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
-        
+
         return hidden_states
 
 
@@ -785,10 +785,10 @@ class MossVLCrossAttentionDecoderLayer(GradientCheckpointingLayer):
 
         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))
@@ -815,7 +815,7 @@ class MossVLCrossAttentionDecoderLayer(GradientCheckpointingLayer):
         # 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,
@@ -839,7 +839,7 @@ class MossVLCrossAttentionDecoderLayer(GradientCheckpointingLayer):
             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
 
 
@@ -866,7 +866,7 @@ class MossVLPreTrainedModel(PreTrainedModel):
 
     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
         """
@@ -884,7 +884,7 @@ class MossVLPreTrainedModel(PreTrainedModel):
                 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:
@@ -915,14 +915,14 @@ class MossVLVisionModel(MossVLPreTrainedModel):
         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, 
+            config=config,
             num_deepstack_features=num_deepstack_features
         )
 
@@ -1027,9 +1027,9 @@ class MossVLVisionModel(MossVLPreTrainedModel):
         return patch_pos_embeds
 
     def forward(
-        self, 
-        hidden_states: torch.Tensor, 
-        grid_thw: torch.Tensor, 
+        self,
+        hidden_states: torch.Tensor,
+        grid_thw: torch.Tensor,
         **kwargs
     ) -> torch.Tensor:
         """
@@ -1097,10 +1097,10 @@ class MossVLTextModel(MossVLPreTrainedModel):
         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):
@@ -1114,7 +1114,7 @@ class MossVLTextModel(MossVLPreTrainedModel):
                 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
@@ -1136,7 +1136,7 @@ class MossVLTextModel(MossVLPreTrainedModel):
         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, 
+        vision_cache_position: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Union[tuple, BaseModelOutputWithPast]:
         """
@@ -1250,7 +1250,7 @@ class MossVLModel(MossVLPreTrainedModel):
         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 
+        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
@@ -1259,7 +1259,7 @@ class MossVLModel(MossVLPreTrainedModel):
         )
 
         self.post_init()
-    
+
 
 
     def convert_packed_to_batch(
@@ -1270,25 +1270,25 @@ class MossVLModel(MossVLPreTrainedModel):
     ) -> 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, 
+
+        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)
@@ -1297,21 +1297,21 @@ class MossVLModel(MossVLPreTrainedModel):
                 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': [],
@@ -1319,41 +1319,41 @@ class MossVLModel(MossVLPreTrainedModel):
                 '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,
@@ -1364,12 +1364,12 @@ class MossVLModel(MossVLPreTrainedModel):
                     '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
@@ -1385,29 +1385,29 @@ class MossVLModel(MossVLPreTrainedModel):
                 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
@@ -1415,45 +1415,45 @@ class MossVLModel(MossVLPreTrainedModel):
                 '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({
@@ -1466,12 +1466,12 @@ class MossVLModel(MossVLPreTrainedModel):
                     '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):
@@ -1494,7 +1494,7 @@ class MossVLModel(MossVLPreTrainedModel):
     ) -> 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.
@@ -1505,12 +1505,12 @@ class MossVLModel(MossVLPreTrainedModel):
                 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
@@ -1520,9 +1520,9 @@ class MossVLModel(MossVLPreTrainedModel):
                 "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:
@@ -1552,59 +1552,59 @@ class MossVLModel(MossVLPreTrainedModel):
             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(
@@ -1615,89 +1615,89 @@ class MossVLModel(MossVLPreTrainedModel):
     ) -> 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(
@@ -1711,21 +1711,21 @@ class MossVLModel(MossVLPreTrainedModel):
         """
         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
@@ -1738,17 +1738,17 @@ class MossVLModel(MossVLPreTrainedModel):
 
         # 1. Gather all frame metadata
         # We need to flatten the nested vision_token_info structure to align with image tokens in input_ids
-        
+
         # Find all image tokens in text: (num_occurrences, 2) -> [batch_idx, seq_idx]
         image_token_indices = (input_ids == image_token_id).nonzero()
-        
+
         # Flatten vision_token_info to parallel lists
         # We assume the order of medias in vision_token_info matches the appearance of image tokens in input_ids
         flat_eff_h = []
         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', [])
@@ -1759,7 +1759,7 @@ class MossVLModel(MossVLPreTrainedModel):
                 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)
@@ -1769,8 +1769,8 @@ class MossVLModel(MossVLPreTrainedModel):
 
         # Pre-allocate output
         vision_pos_ids = torch.zeros(
-            (3, batch_size, max_vision_seq_len), 
-            dtype=torch.long, 
+            (3, batch_size, max_vision_seq_len),
+            dtype=torch.long,
             device=device
         )
 
@@ -1781,152 +1781,152 @@ class MossVLModel(MossVLPreTrainedModel):
 
         # 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) 
+            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, 
+        self,
+        pixel_values: torch.FloatTensor,
         grid_thw: Optional[torch.LongTensor] = None,
         media_nums_per_sample: Optional[List[int]] = None
     ):
@@ -1941,12 +1941,12 @@ class MossVLModel(MossVLPreTrainedModel):
         """
         pixel_values = pixel_values.type(self.visual.dtype)
         hidden_states = self.visual(
-            pixel_values, 
+            pixel_values,
             grid_thw=grid_thw
         )
         vision_embeds, vision_token_info = self.convert_packed_to_batch(
-            hidden_states, 
-            grid_thw, 
+            hidden_states,
+            grid_thw,
             media_nums_per_sample
         )
         return vision_embeds, vision_token_info
@@ -2497,14 +2497,14 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
 
 
     def get_vision_features(
-        self, 
-        pixel_values: torch.FloatTensor, 
+        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
@@ -2558,8 +2558,12 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
             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_cache_position / full_vision_token_info: see `MossVLModel.forward`. Used by real-time generation
-                to incrementally append new frames to the cross-attention KV cache.
+            vision_chunked_length (`int`, *optional*):
+                Number of media items per visual-encoder chunk for offline prefill.
+            vision_cache_position (`torch.LongTensor`, *optional*):
+                Cache positions for incrementally appended realtime vision key/value states.
+            full_vision_token_info (`List[dict]`, *optional*):
+                Cumulative realtime vision metadata used to expand the cross-attention mask.
         """
         outputs = self.model(
             input_ids=input_ids,
@@ -2617,11 +2621,11 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
     ):
         """
         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 
+
+        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)
         """
@@ -2643,7 +2647,7 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         # 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)
@@ -2655,10 +2659,10 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
         # 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["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
@@ -2670,20 +2674,20 @@ class MossVLForConditionalGeneration(MossVLPreTrainedModel, GenerationMixin):
     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: