vLLM minimax_m3_tiny.py 源码精读

class MiniMaxM3TinyForConditionalGeneration(nn.Module, SupportsMultiModal,
                                          SupportsPP, HasInnerState):

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config

        # 缓存图片/视频 token 信息，用于动态分割 multimodal_embeddings
        self._cached_image_token_counts = []
        self._cached_video_token_counts = []

        # 四大核心组件
        self.vision_tower = MiniMaxVL7UVisionModel(
            config=config.vision_config, quant_config=quant_config,
            require_post_norm=False,
            prefix=maybe_prefix(prefix, "vision_tower"))

        self.multi_modal_projector = MiniMaxVL7UMultiModalProjector(
            vision_hidden_size=config.vision_config.hidden_size,
            text_hidden_size=config.text_config.hidden_size,
            projector_hidden_act=config.projector_hidden_act,
            multimodal_projector_bias=True, ...)

        self.patch_merge_mlp = MiniMaxVL7UPatchMerger(
            spatial_merge_size=config.img_token_compression_config.spatial_merge_size,
            text_hidden_size=config.text_config.hidden_size, ...)

        # Dense vs Sparse 语言模型后端
        self.is_sparse_attention_model = (
            getattr(config.text_config, "sparse_attention_config", None) is not None
            and config.text_config.sparse_attention_config["use_sparse_attention"])
        if self.is_sparse_attention_model:
            self.language_model = MiniMaxM2SparseForCausalLM(...)
        else:
            self.language_model = MiniMaxM2ForCausalLM(...)

class CLIPVisionEmbeddings(nn.Module):
    def __init__(self, config: CLIPVisionConfig):
        super().__init__()
        self.patch_size = config.patch_size
        self.embed_dim = config.hidden_size
        self.temporal_patch_size = config.img_token_compression_config.get(
            "temporal_patch_size", 2)
        self.spatial_merge_size = config.img_token_compression_config.get(
            "spatial_merge_size", 2)

        if self.temporal_patch_size == 1:
            # 图像路径：标准 2D Conv
            self.patch_embedding = nn.Conv2d(
                in_channels=config.num_channels,
                out_channels=self.embed_dim,
                kernel_size=self.patch_size,
                stride=self.patch_size, bias=False)
        else:
            # 视频路径：3D Conv，时间维度步长 = temporal_patch_size
            self.patch_embedding = nn.Conv3d(
                in_channels=config.num_channels,
                out_channels=self.embed_dim,
                kernel_size=(self.temporal_patch_size,
                             self.patch_size, self.patch_size),
                stride=(self.temporal_patch_size,
                        self.patch_size, self.patch_size),
                bias=False)

# 3D Conv 路径
# 输入: pixel_values [batch, channel, temporal, height, width]
if pixel_values.dim() == 4:
    pixel_values = pixel_values.unsqueeze(2)  # 添加时间维度

# 3D Conv: [B, C, T, H, W] -> [B, embed_dim, T', H', W']
# 其中 T' = T // temporal_patch_size, H' = H // patch_size, W' = W // patch_size
patch_embeds = self.patch_embedding(pixel_values)

_, embed_dim, n_temporal, n_height, n_width = patch_embeds.shape

# 关键：Patch 区域重排，与 3D RoPE 对齐
# [B, D, T', H', W'] -> [B, D, H'//m, W'//m, T', m, m]
patch_embeds = patch_embeds.reshape(
    batch_size, embed_dim, n_temporal,
    n_height // spatial_merge_size, spatial_merge_size,
    n_width // spatial_merge_size, spatial_merge_size
).permute(0, 1, 3, 5, 2, 4, 6)

# 展平为序列: [B, D, seq_len] -> [B, seq_len, D]
patch_embeds = patch_embeds.flatten(2).transpose(1, 2).contiguous()

def forward(self, hidden_states, cu_seq_len, max_seqlen=None,
            rotary_pos_emb=None, **kwargs):
    seq_length = hidden_states.shape[0]

    # QKV 投影（支持 Tensor Parallel）
    qkv_states, _ = self.qkv_proj(hidden_states)
    q, k, v = self.split_qkv(qkv_states)  # [seq, 1, tp_head, head_dim]
    q, k, v = q.squeeze(1), k.squeeze(1), v.squeeze(1)

    # 应用 RoPE 位置编码
    if self.config.position_embedding_type == "rope":
        q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
        k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)

    # 变长注意力：cu_seq_len 控制注意力边界
    attn_output = flash_attn_varlen_func(
        q, k, v,
        cu_seq_len, cu_seq_len,     # Q 和 KV 使用相同的边界
        max_seqlen, max_seqlen,
        causal=getattr(self.config, 'causal_attention', False))

    attn_output = attn_output.view(seq_length, -1)
    attn_output = self.out_proj(attn_output)[0].unsqueeze(1)
    return attn_output, None

def split_qkv(self, qkv):
    seq_len, bs, _ = qkv.shape

    if self.tp_size > 1:
        # 跨 TP rank 收集完整 QKV，交错排列
        qkv = all_gather_interleave(qkv, self.qkv_proj.hidden_size, self.tp_size)

    q, k, v = qkv.chunk(3, dim=2)

    if self.tp_size > 1:
        # 各 rank 取自己负责的 head 切片
        splitter = partial(dist_utils.split_tensor_along_last_dim,
                          num_partitions=self.tp_size)
        q = splitter(q)[self.tp_rank]
        k = splitter(k)[self.tp_rank]
        v = splitter(v)[self.tp_rank]

    # reshape 为 [seq, batch, num_heads_per_partition, head_dim]
    new_shape = (seq_len, bs,
                 self.num_attention_heads_per_partition,
                 self.hidden_size_per_attention_head)
    q, k, v = (x.view(*new_shape) for x in (q, k, v))
    return q, k, v

def _get_rope_embed_2d(self, grid_thw, spatial_merge_size):
    pos_ids = []
    max_grid_size = 0

    for h, w in grid_thw:
        w = w // self.config.patch_size
        h = h // self.config.patch_size
        max_grid_size = max(max_grid_size, w, h)

        # 高度位置 ID：按 spatial_merge_size 重排
        hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
        hpos_ids = hpos_ids.reshape(
            h // spatial_merge_size, spatial_merge_size,
            w // spatial_merge_size, spatial_merge_size,
        ).permute(0, 2, 1, 3).flatten()

        # 宽度位置 ID：同样重排
        wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
        wpos_ids = wpos_ids.reshape(
            h // spatial_merge_size, spatial_merge_size,
            w // spatial_merge_size, spatial_merge_size,
        ).permute(0, 2, 1, 3).flatten()

        pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1))

    pos_ids = torch.cat(pos_ids, dim=0)
    # 查频率表：rotary_pos_emb_full[pos_ids] → [seq_len, rope_dim]
    rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
    rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
    return rotary_pos_emb

def _get_3d_rope_embed(self, grid_h, grid_w, frame_index, spatial_merge_size):
    """单帧的 3D RoPE 位置嵌入"""
    tokens_per_frame = grid_h * grid_w

    # ===== 时间位置 ID =====
    # 当前帧的所有 token 使用相同的时间位置 frame_index
    tpos_ids = torch.full((tokens_per_frame,), frame_index, dtype=torch.long)

    # ===== 高度位置 ID（带 spatial_merge 重排）=====
    hpos_ids = torch.arange(grid_h).unsqueeze(1).expand(-1, grid_w)
    hpos_ids = hpos_ids.reshape(
        grid_h // spatial_merge_size, spatial_merge_size,
        grid_w // spatial_merge_size, spatial_merge_size,
    ).permute(0, 2, 1, 3).flatten()

    # ===== 宽度位置 ID（同样重排）=====
    wpos_ids = torch.arange(grid_w).unsqueeze(0).expand(grid_h, -1)
    wpos_ids = wpos_ids.reshape(
        grid_h // spatial_merge_size, spatial_merge_size,
        grid_w // spatial_merge_size, spatial_merge_size,
    ).permute(0, 2, 1, 3).flatten()

    # ===== 三维频率表 =====
    max_hw = max(grid_h, grid_w)
    freqs_t = torch.outer(torch.arange(max(frame_index + 1, 1)), self.inv_freq_t)
    freqs_h = torch.outer(torch.arange(max_hw), self.inv_freq_h)
    freqs_w = torch.outer(torch.arange(max_hw), self.inv_freq_w)

    # ===== 查表拼接 =====
    emb_t = freqs_t[tpos_ids]  # [seq_len, t_dim/2]
    emb_h = freqs_h[hpos_ids]  # [seq_len, h_dim/2]
    emb_w = freqs_w[wpos_ids]  # [seq_len, w_dim/2]

    rotary_pos_emb = torch.cat([emb_t, emb_h, emb_w], dim=-1)
    return rotary_pos_emb  # [seq_len, (t_dim + h_dim + w_dim) / 2]

def _compute_cu_seq_len(self, frame_token_counts, video_segments, device):
    """
    计算 flash_attn_varlen_func 的 cumulative sequence lengths。

    Args:
        frame_token_counts: 每个压缩帧的 token 数列表
        video_segments: 每个视频段的压缩帧数
            e.g., [8, 8] = 两段各 8 个压缩帧

    Examples:
        frame_token_counts = [576, 576, 576, 576]

        Case 1: video_segments = None（每帧独立）
            cu_seq_len = [0, 576, 1152, 1728, 2304]

        Case 2: video_segments = [4]（所有帧一个段 → VPA）
            cu_seq_len = [0, 2304]
            所有帧互相可见。

        Case 3: video_segments = [2, 2]（两段各 2 帧）
            cu_seq_len = [0, 1152, 2304]
            帧 0-1 互相可见，帧 2-3 互相可见。
    """
    num_frames = len(frame_token_counts)

    if video_segments is None:
        # 每帧独立
        cu_seq_len = [0] + frame_token_counts
        cu_seq_len = torch.tensor(cu_seq_len, device=device).to(torch.int32)
        return torch.cumsum(cu_seq_len, dim=0).to(torch.int32)

    # 应用 vision_segment_max_frames 限制
    effective_segments = self._apply_max_frames_limit(video_segments, num_frames)

    # 按段聚合 token 数
    cu_seq_len = [0]
    frame_idx = 0
    for segment_frame_count in effective_segments:
        segment_tokens = sum(frame_token_counts[frame_idx:frame_idx + segment_frame_count])
        cu_seq_len.append(segment_tokens)
        frame_idx += segment_frame_count

    cu_seq_len = torch.tensor(cu_seq_len, device=device).to(torch.int32)
    return torch.cumsum(cu_seq_len, dim=0).to(torch.int32)

def _compute_frame_indices(self, num_frames, video_segments):
    """
    根据 video_segments 计算每帧的 frame_index（段内索引）。

    video_segments=[3, 2] → frame_indices=[0, 1, 2, 0, 1]
    video_segments=None   → frame_indices=[0, 0, 0, ...]（全部视为独立图像）

    若设置 vision_segment_max_frames=4:
      video_segments=[10] → effective=[4, 4, 2]
      → frame_indices=[0,1,2,3, 0,1,2,3, 0,1]
    """
    if video_segments is None:
        return [0] * num_frames

    # 应用 max_frames 限制（内含 _normalize_video_segments 处理 [0] 特殊值）
    effective_segments = self._apply_max_frames_limit(video_segments, num_frames)

    frame_indices = []
    for segment_size in effective_segments:
        for i in range(segment_size):
            frame_indices.append(i)

    return frame_indices

def _apply_max_frames_limit(self, video_segments, num_frames):
    """
    将超长段拆分为多个小段。

    Example:
        temporal_patch_size=2, vision_segment_max_frames=4
        原始 20 帧 → 压缩 10 帧
        video_segments=[10] → effective_segments=[4, 4, 2]
    """
    normalized_segments = self._normalize_video_segments(video_segments, num_frames)

    if self.vision_segment_max_frames is None:
        return normalized_segments

    max_frames = self.vision_segment_max_frames
    effective_segments = []

    for segment_size in normalized_segments:
        if segment_size <= max_frames:
            effective_segments.append(segment_size)
        else:
            remaining = segment_size
            while remaining > 0:
                chunk_size = min(remaining, max_frames)
                effective_segments.append(chunk_size)
                remaining -= chunk_size

    return effective_segments

def forward(self, pixel_values, image_sizes,
            feature_sample_layers=None, video_segments=None):
    inputs = []
    frame_token_counts = []
    max_seqlen = 0

    for pixel_value in pixel_values:
        # 1. Patch Embedding（2D/3D Conv + 区域重排）
        hidden_states = self.embeddings(pixel_value.unsqueeze(0))
        hidden_states = self.pre_layrnorm(hidden_states)
        # [B, S, C] → [S, B, C]（flash-attn 输入格式）
        hidden_states = hidden_states.permute(1, 0, 2).contiguous()
        frame_token_counts.append(hidden_states.size(0))
        max_seqlen = max(max_seqlen, hidden_states.size(0))
        inputs.append(hidden_states)

    # 拼接所有帧的 token
    inputs = torch.cat(inputs, dim=0)  # [total_seq, 1, embed_dim]

    # 2. 计算 cu_seq_len（VPA 注意力分组）
    cu_seq_len = self._compute_cu_seq_len(
        frame_token_counts, video_segments, inputs.device)

    # 更新 max_seqlen 为段级别的最大值
    if video_segments is not None:
        effective_segments = self._apply_max_frames_limit(
            video_segments, len(frame_token_counts))
        frame_idx = 0
        segment_token_counts = []
        for seg_count in effective_segments:
            segment_tokens = sum(frame_token_counts[frame_idx:frame_idx + seg_count])
            segment_token_counts.append(segment_tokens)
            frame_idx += seg_count
        max_seqlen = max(segment_token_counts)

    # 3. 计算 RoPE 位置编码（2D 或 3D）
    rotary_pos_emb = self.rot_pos_emb(image_sizes, video_segments)

    # 4. N 层 CLIPEncoder（每层：LN → VisionFlashAttention2 → LN → CLIPMLP）
    encoder_outputs = self.encoder(
        inputs_embeds=inputs,
        cu_seq_len=cu_seq_len,
        max_seqlen=max_seqlen,
        rotary_pos_emb=rotary_pos_emb)

    # [S, B, C] → [total_seq, hidden_size]
    encoder_outputs = encoder_outputs.permute(1, 0, 2).contiguous()[0]
    return encoder_outputs

class MiniMaxVL7UMultiModalProjector(nn.Module):
    def __init__(self, vision_hidden_size, text_hidden_size,
                 projector_hidden_act, multimodal_projector_bias,
                 projector_hidden_size=None, quant_config=None, prefix=""):
        super().__init__()
        projector_intermediate = projector_hidden_size or text_hidden_size
        self.linear_1 = ColumnParallelLinear(
            vision_hidden_size, projector_intermediate,
            bias=multimodal_projector_bias, quant_config=quant_config)
        self.act = get_act_fn(projector_hidden_act)
        self.linear_2 = RowParallelLinear(
            projector_intermediate, text_hidden_size,
            bias=multimodal_projector_bias, quant_config=quant_config)

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

class MiniMaxVL7UPatchMerger(nn.Module):
    def __init__(self, spatial_merge_size, text_hidden_size, ...):
        super().__init__()
        # 输入维度 = text_hidden_size × spatial_merge_size²
        # 输出维度 = text_hidden_size
        self.linear_1 = ColumnParallelLinear(
            text_hidden_size * spatial_merge_size ** 2,
            projector_intermediate, ...)
        self.act = get_act_fn(projector_hidden_act)
        self.linear_2 = RowParallelLinear(
            projector_intermediate, text_hidden_size, ...)

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

class MiniMaxVL7UImagePixelInputs(TypedDict, total=False):
    type: Literal["pixel_values"]
    pixel_values: torch.Tensor
    # Shape: (batch_size * num_images, num_channels, height, width)
    # 或 list（当不同图片尺寸不同时）
    image_sizes: torch.Tensor
    # Shape: (batch_size * num_images, 2)  → (height, width)
    video_segments: Optional[List[int]]
    # 例如 [16, 16] = 两个视频段，每段 16 帧共享注意力

class MiniMaxVL7UVideoPixelInputs(TypedDict, total=False):
    type: Literal["pixel_values_videos"]
    pixel_values_videos: torch.Tensor
    # Shape: (batch_size * num_videos, num_channels, height, width)
    video_sizes: torch.Tensor
    # Shape: (batch_size * num_videos, 3)  → (temporal, height, width)
    video_segments: Optional[List[int]]

def _parse_and_validate_image_input(self, **kwargs):
    pixel_values = kwargs.pop("pixel_values", None)
    image_sizes = kwargs.pop("image_sizes", None)

    if pixel_values is not None:
        # 6D tensor: (batch, n_images, temporal_patch_size, C, H, W)
        if isinstance(pixel_values, torch.Tensor) and pixel_values.dim() == 6:
            new_pixel_values = []
            for i in range(pixel_values.shape[0]):
                for j in range(pixel_values.shape[1]):
                    # (T, C, H, W) → (C, T, H, W)  permute 通道到前面
                    new_pixel_values.append(
                        pixel_values[i, j].permute(1, 0, 2, 3))

        # 5D tensor: (batch, n_images, C, H, W) — 无时间维度
        elif isinstance(pixel_values, torch.Tensor) and pixel_values.dim() == 5:
            new_pixel_values = []
            for i in range(pixel_values.shape[0]):
                for j in range(pixel_values.shape[1]):
                    # (C, H, W) → (C, 1, H, W)  添加时间维度
                    new_pixel_values.append(
                        pixel_values[i, j].unsqueeze(1))

        # list of list：逐元素处理
        elif isinstance(pixel_values, list):
            new_pixel_values = []
            for i in range(len(pixel_values)):
                for j in range(len(pixel_values[i])):
                    item = pixel_values[i][j]
                    if item.dim() == 4:  # (T, C, H, W) → (C, T, H, W)
                        new_pixel_values.append(item.permute(1, 0, 2, 3))
                    elif item.dim() == 3:  # (C, H, W) → (C, 1, H, W)
                        new_pixel_values.append(item.unsqueeze(1))

    # 同步展平 image_sizes（3D tensor / nested list → flat list）
    return MiniMaxVL7UImagePixelInputs(
        type="pixel_values", pixel_values=new_pixel_values,
        image_sizes=new_image_sizes)

# 辅助函数 1：处理单个 tensor → (C, T, H, W)
def process_single_tensor(tensor: torch.Tensor) -> torch.Tensor:
    # 先 squeeze 掉所有前导的 dim=1
    while tensor.dim() > 4 and tensor.shape[0] == 1:
        tensor = tensor.squeeze(0)
    if tensor.dim() == 4:
        if tensor.shape[0] == 3:    # 已经是 (C, T, H, W)
            return tensor
        if tensor.shape[1] == 3:    # 是 (T, C, H, W)，需要 permute
            return tensor.permute(1, 0, 2, 3)
    if tensor.dim() == 3:           # (C, H, W) → 添加时间维度
        return tensor.unsqueeze(1)

# 辅助函数 2：递归展平嵌套的 tensor / list
def flatten_video_tensors(data) -> list[torch.Tensor]:
    if isinstance(data, torch.Tensor):
        if data.dim() <= 4:
            return [process_single_tensor(data)]
        flattened = []
        for child in data:
            flattened.extend(flatten_video_tensors(child))
        return flattened
    if isinstance(data, list):
        flattened = []
        for child in data:
            if child is None: continue
            flattened.extend(flatten_video_tensors(child))
        return flattened

# 辅助函数 3：递归展平 video_sizes
def flatten_video_sizes(data) -> list:
    # 1D tensor → [data], ≥2D tensor → 递归
    # tuple → [data], list of numbers → [data], nested list → 递归

# 辅助函数 4：递归展平 video_segments
def flatten_video_segments(data) -> Optional[list[int]]:
    # 0D tensor → [int(item)], ≥1D tensor → 递归
    # list/tuple → 递归, int/float → [int(data)]

class ImageTokensCalculator:
    def __init__(self):
        self.patch_size = 14
        # 28 种预定义分辨率（336×336 到 2016×1008）
        self.default_grid_pinpoints = [
            [336, 336], [336, 672], ..., [2016, 1008]
        ]

    def calculate_tokens(self, height, width):
        mode = self.config['process_image_mode']
        if mode == 'resize':
            # 最简单：直接按 patch 面积计算
            return int(height * width / self.patch_size ** 2)

        elif mode == 'dynamic_res':
            return self._get_dynamic_res_num_token(
                height, width, max_img_h, max_img_w)

        else:  # anyres
            return self._get_anyres_num_token(height, width)

def get_hw_multiple_of(image_size, multiple, max_size=None):
    w, h = image_size
    # 向上对齐到 multiple 的倍数
    new_w = w if w % multiple == 0 else w + (multiple - w % multiple)
    new_h = h if h % multiple == 0 else h + (multiple - h % multiple)

    if max_size is not None:
        max_w, max_h = max_size
        if new_w > max_w or new_h > max_h:
            # 按比例缩放到 max_size 内，再次对齐
            new_w_ = min((new_w * max_w) // new_w,
                         (new_w * max_h) // new_h)
            new_h_ = min((new_h * max_w) // new_w,
                         (new_h * max_h) // new_h)
            new_w = new_w_ + (multiple - new_w_ % multiple) % multiple
            new_h = new_h_ + (multiple - new_h_ % multiple) % multiple
    return new_w, new_h

def _get_prompt_updates(self, mm_items, hf_processor_mm_kwargs, out_mm_kwargs):
    image_token_id = hf_config.image_token_index  # 200025
    video_token_id = hf_config.video_token_index  # 200026
    START_TOKEN_ID = 200029   # ]<]start of image[>[
    END_TOKEN_ID   = 200030   # ]<]end of image[>[

    def get_image_replacement(item_idx):
        num_tokens = get_num_tokens_for_item(item_idx)
        token_ids = ([START_TOKEN_ID]
                   + [image_token_id] * num_tokens
                   + [END_TOKEN_ID])
        # 只有 image_token_id 位置会被视觉 embedding 替换
        return PromptUpdateDetails.select_token_id(
            token_ids, image_token_id)

    def get_video_replacement(item_idx):
        num_tokens = get_num_tokens_for_item(item_idx, modality="video")
        token_ids = ([START_TOKEN_ID]
                   + [video_token_id] * num_tokens
                   + [END_TOKEN_ID])
        return PromptUpdateDetails.select_token_id(
            token_ids, video_token_id)

    return [
        PromptReplacement(modality="image",
            target="]<]image[>[",
            replacement=get_image_replacement),
        PromptReplacement(modality="video",
            target="]<]video[>[",
            replacement=get_video_replacement),
    ]

def _process_image_input(self, image_input):
    self._cached_image_token_counts = []

    # Step 1: ViT 编码（无 video_segments → 每帧独立注意力）
    image_features = self._process_image_pixels(image_input)
    #   → self.vision_tower(pixel_values, image_sizes)

    # Step 2: 计算每张图片的 token 数量
    spatial_merge_size = config.img_token_compression_config.spatial_merge_size
    for image_size in image_input["image_sizes"]:
        height, width = image_size[0].item(), image_size[1].item()
        new_width, new_height = get_hw_multiple_of(
            (width, height), patch_size * spatial_merge_size, max_size=...)
        num_patches_w = new_width // patch_size // spatial_merge_size
        num_patches_h = new_height // patch_size // spatial_merge_size
        self._cached_image_token_counts.append(num_patches_w * num_patches_h)

    # Step 3: Projector（ViT dim → LLM dim）
    image_embeds = self.multi_modal_projector(image_features.unsqueeze(1))

    # Step 4: 按每张图片的 token 数分割，逐张 PatchMerger
    chunk_sizes = [t * spatial_merge_size**2
                   for t in self._cached_image_token_counts]
    image_embeds = torch.split(
        image_embeds.squeeze(1).unsqueeze(0), chunk_sizes, dim=1)

    final_image_features = []
    for chunk_size, image_feature in zip(chunk_sizes, image_embeds):
        # reshape: [1, chunk, hidden] → [chunk/s², hidden * s²]
        image_feature = self.patch_merge_mlp(
            image_feature.view(-1, hidden_size * spatial_merge_size**2))
        final_image_features.append(image_feature)

    return tuple(final_image_features)

def _process_video_input(self, video_input):
    self._cached_video_token_counts = []

    video_pixel_values = video_input["pixel_values_videos"]
    video_sizes = video_input["video_sizes"]
    video_segments = video_input.get("video_segments", None)

    # Step 1: ViT 编码 —— 传入 video_segments 启用 VPA
    video_features = self.vision_tower(
        video_pixel_values, video_sizes, video_segments=video_segments)

    # Step 2: 计算每个压缩帧的 token 数量
    for frame_size in video_sizes:
        height, width = frame_size[-2].item(), frame_size[-1].item()
        new_width, new_height = get_hw_multiple_of(
            (width, height), patch_size * spatial_merge_size, max_size=...)
        tokens = (new_width // patch_size // spatial_merge_size *
                  new_height // patch_size // spatial_merge_size)
        self._cached_video_token_counts.append(tokens)

    # Step 3: Projector
    video_embeds = self.multi_modal_projector(video_features.unsqueeze(1))

    # Step 4: 按帧分割 + PatchMerger
    chunk_sizes = [t * spatial_merge_size**2
                   for t in self._cached_video_token_counts]
    video_embeds = torch.split(
        video_embeds.squeeze(1).unsqueeze(0), chunk_sizes, dim=1)

    final_video_features = []
    for chunk_size, video_feature in zip(chunk_sizes, video_embeds):
        video_feature = self.patch_merge_mlp(
            video_feature.contiguous().view(
                -1, hidden_size * spatial_merge_size**2))
        final_video_features.append(video_feature)

    return tuple(final_video_features)

def get_input_embeddings(self, input_ids, multimodal_embeddings=None):
    inputs_embeds = self.language_model.embed_input_ids(input_ids)

    if multimodal_embeddings is not None and len(multimodal_embeddings) != 0:
        image_token_id = self.config.image_token_index  # 200025
        video_token_id = self.config.video_token_index  # 200026

        # 分别合并 image 和 video embedding
        if "image_embeddings" in multimodal_embeddings:
            inputs_embeds = merge_multimodal_embeddings(
                input_ids, inputs_embeds,
                multimodal_embeddings["image_embeddings"],
                [image_token_id])

        if "video_embeddings" in multimodal_embeddings:
            inputs_embeds = merge_multimodal_embeddings(
                input_ids, inputs_embeds,
                multimodal_embeddings["video_embeddings"],
                [video_token_id])

    return inputs_embeds

def merge_multimodal_embeddings(
    input_ids, inputs_embeds, multimodal_embeddings, placeholder_token_id):
    """
    文本序列结构示意：
    [text ... | 200029 | 200025×N | 200030 | text ... | 200029 | 200026×M | 200030 | ...]
              ↑ start   ↑ image tokens ↑ end          ↑ start   ↑ video tokens ↑ end

    placeholder_token_id 可以是列表（如 [200025]），
    通过 torch.isin 匹配所有占位位置，用视觉 embedding 逐个替换。
    """
    if isinstance(placeholder_token_id, list):
        placeholder_token_id = torch.tensor(placeholder_token_id, device=input_ids.device)
        return _merge_multimodal_embeddings(
            inputs_embeds,
            torch.isin(input_ids, placeholder_token_id),
            multimodal_embeddings)

def _merge_multimodal_embeddings(inputs_embeds, is_multimodal, multimodal_embeddings):
    flattened = _flatten_embeddings(multimodal_embeddings)
    # 直接原地替换
    inputs_embeds[is_multimodal] = flattened
    return inputs_embeds

# 3D Conv 时间戳处理：跟踪当前 video_seg 中的帧计数器
frame_counter_in_current_video_seg = 0
current_video_seg_index = 0
last_timestamp_text = None

# 计算时间压缩后的 video_segments
video_segments_after_temporal_merge = [
    (seg + self.temporal_patch_size - 1) // self.temporal_patch_size
    for seg in video_segments
]

# 遍历 prompt 中的每个 token
for i, _sample in enumerate(split_text):
    if _sample == self.video_token:
        if has_video_item and video_segments is not None:
            # 只保留每 temporal_patch_size 帧中的第一帧
            if frame_counter_in_current_video_seg % self.temporal_patch_size == 0:
                # 保留这个 video token
                final_text += _sample
                video_index += 1
            else:
                # 跳过这个 video token（不添加到 final_text）
                # 同时删除被跳过帧对应的时间戳
                if last_timestamp_text is not None:
                    final_text = final_text[:-len(last_timestamp_text)]
                    last_timestamp_text = None

            frame_counter_in_current_video_seg += 1
            if frame_counter_in_current_video_seg == video_segments[current_video_seg_index]:
                current_video_seg_index += 1
                frame_counter_in_current_video_seg = 0

def forward(self, input_ids, positions, intermediate_tensors=None,
            inputs_embeds=None, **kwargs):

    if intermediate_tensors is not None:
        inputs_embeds = None
    elif inputs_embeds is None:
        # Step 1: 获取视觉 embedding（分别处理 image 和 video）
        multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
        # Step 2: 合并到文本 embedding
        inputs_embeds = self.get_input_embeddings(
            input_ids, multimodal_embeddings)

    # Step 3: 语言模型前向传播
    hidden_states = self.language_model.model(
        input_ids, positions, intermediate_tensors,
        inputs_embeds=inputs_embeds, **kwargs)

    return hidden_states

def get_multimodal_embeddings(self, **kwargs):
    image_input = self._parse_and_validate_image_input(**kwargs)
    video_input = self._parse_and_validate_video_input(**kwargs)

    multimodal_embeddings = {}

    if image_input is not None:
        embeddings = self._process_image_input(image_input)
        multimodal_embeddings["image_embeddings"] = embeddings

    if video_input is not None:
        embeddings = self._process_video_input(video_input)
        multimodal_embeddings["video_embeddings"] = embeddings

    return multimodal_embeddings

def get_seqlen_agnostic_capture_inputs(self, batch_size):
    # 只有 Sparse 模型有 const_cache
    if self.is_sparse_attention_model:
        return self.language_model.model.const_cache \
            .get_seqlen_agnostic_capture_inputs(batch_size)
    return None

def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
    if self.is_sparse_attention_model:
        return self.language_model.model.const_cache \
            .copy_inputs_before_cuda_graphs(input_buffers, **kwargs)
    return None